Skip to content
Advertisement

Multiple Seaborn Heatmaps from Pandas Dataframe

I have a Pandas dataframe that looks like this:

store_id    days       times                     rating
    100     monday     '1:00pm - 3:00pm'         0
    100     monday     '3:00pm - 6:00pm'         1
    100     monday     '6:00pm - 9:00pm'         2
    ...
    store n      

Where there are ~60 stores and the ratings range from 0 – 2. I would like to create a 6×5 grid Seaborn of heatmaps, with one heatmap per store. I would like for the x-axis to be the days and for the y-axis to be the times.

I tried this:

f, axes = plt.subplots(5,6)
i=0
for store in df['store_id']:
    sns.heatmap(data=df[df['store_id']==store]['rating'], ax=axes[i])
    i+=1

This creates the 5×6 grid, but generates an error (‘Inconsistent shape between the condition and the input…’). What’s the best way to do this?

Advertisement

Answer

For heat map, you need to transpose/pivot your data so as the days becomes columns (x-axis) and times becomes index:

f, axes = plt.subplots(5,6)

# flatten axes for looping
axes = axes.ravel()

# use groupby to extract data faster
for ax, (store, data) in zip(axes, df.groupby('store_id')):
    pivot = data.pivot_table(index='times', columns='days', values='rating')
    sns.heatmap(data=pivot, ax=ax)
User contributions licensed under: CC BY-SA
3 People found this is helpful
Advertisement