Stacked bar plot color by a diferent category-CodePudding

I want to make a plot where it shows the RFS by floor opened by the unit_id, but i want the color of the unit_id to be defined by the year. So far i make it work for a reduced set of data, but i think it will be difficult to scale the code.
What i do is, first i identify the order that each unit has at it's floor so i first plot all the units by floor in the first position, then the ones that are in second position and so on.
Thanks!

df_build = pd.DataFrame({'floor':[1,1,1,2,2,2,3,3,3],'unidad':[100,101,102,200,201,202,300,301,302],
                         'rsf':[2000,1000,1500,1500,2000,1000,1000,1500,2000],'order':[0,1,2,0,1,2,0,1,2],
                         'year':[2008,2009,2010,2009,2010,2011,2010,2011,2012]})

assign_colors = {2008:'tab:red',2009:'tab:blue',2010:'tab:green',2011:'tab:pink',2012:'tab:olive'}

labels = list(df_build.floor.unique())
order_0 = df_build[df_build.order==0].rsf.values
c1=list(df_build[df_build.order==0].year.replace(assign_colors).values)
order_1 = df_build[df_build.order==1].rsf.values
c2=list(df_build[df_build.order==1].year.replace(assign_colors).values)
order_2 = df_build[df_build.order==2].rsf.values
c3=list(df_build[df_build.order==2].year.replace(assign_colors).values)

width = 0.35 

fig, ax = plt.subplots()

ax.barh(labels, order_0, width,color=c1)
ax.barh(labels, order_1, width,left=order_0, color=c2)
ax.barh(labels, order_2, width,left=order_0 order_1, color=c3)

ax.set_ylabel('floor')
ax.set_title('Stacking Plan')
#ax.legend()

plt.show()

CodePudding user response：

Try pivoting the data and loop:

# map the color
df_build['color'] = df_build['year'].map(assign_colors)

# pivot the data
plot_df = df_build.pivot(index='floor', columns='order')

# plot by row
fig, ax = plt.subplots()

for i in df.index:
    rights = plot_df.loc[i,'rsf'].cumsum()
    lefts = rights.shift(fill_value=0)
    ax.barh(i, plot_df.loc[i,'rsf'], left=lefts, color=plot_df.loc[i,'color'])
    
    for j in range(len(rights)):
        label = plot_df.loc[i, 'unidad'].iloc[j]
        rsf = plot_df.loc[i, 'rsf'].iloc[j]
        x = (rights.iloc[j]   lefts.iloc[j]) / 2
        ax.text(x, i, f'{label}-{rsf}', ha='center')

Output: