How to make group colorbar xticks yticks in Seaborn or Matplotplib

Question

Given an array normal_data with a label mgroup, the following heatmap is generated.

But, I wonder whether it is possible using Matplotlib or Seaborn to have a xtick-ytick color bar (The rectangle shape of color blue, orange, light green, dark blue, dark green) to represent a group membership, as shown above?

import seaborn as sns
from matplotlib import pyplot as plt
import numpy as np
mgroup=np.array([1,1,1,2,2,3,3,3,3,4,5,5])
normal_data = np.random.randn(12, 12)

ax = sns.heatmap(normal_data, center=0)
plt.show()

user_12046409 · Answer

You can use plt.Rectangle to draw colored rectangles. Setting clip_on=False allows drawing outside the main axes. The xaxis transform uses x in the horizontal direction, and axes coordinates in the y direction. The yaxis transform works similar in the other orientation.

import seaborn as sns
from matplotlib import pyplot as plt
import numpy as np

mgroup = np.array([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 5, 5])
normal_data = np.random.randn(12, 12)

grouping = [3, 2, 4, 1, 2]
colors = ['dodgerblue', 'orange', 'lime', 'darkblue', 'green']
labels = [f'Group {i + 1}' for i in range(len(colors))]
ax = sns.heatmap(normal_data, center=0, cbar_kws={'orientation': 'horizontal', 'pad': 0.08})

for width, pos, color, label in zip(grouping, np.cumsum([0] + grouping), colors, labels):
    ax.add_patch(plt.Rectangle((pos, 1.02), width, 0.1, color=color,
                               clip_on=False, transform=ax.get_xaxis_transform()))
    ax.add_patch(plt.Rectangle((1.02, pos), 0.1, width, color=color,
                               clip_on=False, transform=ax.get_yaxis_transform()))
    ax.text(1.14, pos + width/2, label, ha='left', va='center', clip_on=False, transform=ax.get_yaxis_transform())
plt.tight_layout()
plt.show()

user_13107804 · Answer

For the group values, we use np.unique() to get the array and frequency, and use that to create a stacked graph. The created graphs are placed on the right and top of the heat map, with new axes added to them. As for the labels, I decided their positions manually. I am getting there by hand to complete this, so it may not be the best way. Also, some of the code may be redundant.

import seaborn as sns
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd

mgroup=np.array([1,1,1,2,2,3,3,3,3,4,5,5])
normal_data = np.random.randn(12, 12)
u,counts = np.unique(mgroup, return_counts=True)
df = pd.DataFrame({'value':counts},index=u)

fig, axs = plt.subplots(1, 1, figsize=(4,4), sharey=True)# , sharey=True
cbar_ax = fig.add_axes([0.1,0.0,0.8,0.05])
g = sns.heatmap(normal_data, center=0, cbar_ax=cbar_ax, cbar_kws={"orientation": "horizontal"}, ax=axs)
#g.set_aspect('auto')
df_ax = fig.add_axes([0.9,0.12,0.05,0.80])
d = df.T.plot(kind='bar', stacked=True, width=0.7, legend=False, ax=df_ax)
d.spines['top'].set_visible(False)
d.spines['bottom'].set_visible(False)
d.spines['right'].set_visible(False)
d.spines['left'].set_visible(False)
df_ax.set_xticks([])
df_ax.set_yticks([])

df_ax2 = fig.add_axes([0.12,0.85,0.82,0.05])
d = df.T.plot(kind='barh', stacked=True, width=0.7, legend=False, ax=df_ax2)
d.spines['top'].set_visible(False)
d.spines['bottom'].set_visible(False)
d.spines['right'].set_visible(False)
d.spines['left'].set_visible(False)
df_ax2.set_xticks([])
df_ax2.set_yticks([])

df_ax.text(1.0, 0.20, 'Group5', transform=fig.transFigure)
df_ax.text(1.0, 0.35, 'Group4', transform=fig.transFigure)
df_ax.text(1.0, 0.55, 'Group3', transform=fig.transFigure)
df_ax.text(1.0, 0.70, 'Group2', transform=fig.transFigure)
df_ax.text(1.0, 0.80, 'Group1', transform=fig.transFigure)
plt.show()

How to make group colorbar xticks yticks in Seaborn or Matplotplib

2 Answers