How to perform custom sort for both x- and y-axes on indexed DataFrame for heatmap?

Question

This question gives the solution to sort the y-axis: Data order in seaborn heatmap from pivot But how to perform custom sort for both x- and y-axes?

Without custom sort, we see the order:

• x-axis: phone, tv
• y-axis: apple, google, samsung

Code:

lol = [['apple', 'phone', 10], ['samsung', 'tv', 20], ['apple', 'tv', 5], ['google', 'tv', 8], ['google', 'phone', 9], ['samsung', 'phone', 3]]
df = pd.DataFrame(lol)
df = df.rename(columns={0:'brand', 1:'product', 2:'count'})
df = df.pivot('brand', 'product', 'count')
ax = sns.heatmap(df)
plt.show()

[out]:

If I need to sort the y-axis to show the order samsung, apple, google, I could do:

lol = [['apple', 'phone', 10], ['samsung', 'tv', 20], ['apple', 'tv', 5], ['google', 'tv', 8], ['google', 'phone', 9], ['samsung', 'phone', 3]]
df = pd.DataFrame(lol)
df = df.rename(columns={0:'brand', 1:'product', 2:'count'})
df = df.pivot('brand', 'product', 'count')

df.index = pd.CategoricalIndex(df.index, categories= ["samsung", "apple", "google"])
df.sortlevel(level=0, inplace=True)
ax = sns.heatmap(df)
plt.show()

[out]:

But how to perform custom sort for both x- and y-axes?, e.g.

• y-axis to show the order samsung, apple, google
• x-axis to show the order tv, phone (not just reversing the order)

Answer 1

I think you can use reindex:

a = ['samsung', 'apple', 'google']
b = ['tv','phone']

df = df.pivot('brand', 'product', 'count')
df = df.reindex(index=a, columns=b)
print (df)
product  tv  phone
brand             
samsung  20      3
apple     5     10
google    8      9

Or ordered categorical:

df['brand'] = df['brand'].astype('category', categories=a, ordered=True)
df['product'] = df['product'].astype('category', categories=b, ordered=True)

df = df.pivot('brand', 'product', 'count')
print (df)
product  tv  phone
brand             
samsung  20      3
apple     5     10
google    8      9

---

ax = sns.heatmap(df)
plt.show()