Adjusting width of subplots on seaborn heatmap

Question

After coming across this: Combining two heat maps in seaborn I did the following:

df = pd.DataFrame(np.random.rand(8,1), columns=list("A"))
df2 = pd.DataFrame(np.random.rand(8,1), columns=list("B"))
df3 = pd.DataFrame(np.random.rand(8,4), columns=list("CDEF"))
df4 = pd.DataFrame(np.random.rand(8,1), columns=list("G"))
df5 = pd.DataFrame(np.random.rand(8,1), columns=list("H"))

fig, (ax, ax2, ax3, ax4, ax5) = plt.subplots(figsize=(30, 10), ncols=5)
fig.subplots_adjust(wspace=0.01)

sb.heatmap(df,  ax=ax, cbar=False, cmap = 'coolwarm_r', vmax = 30, vmin = 10, annot = True, fmt = '.2f', annot_kws = {'size':12})
sb.heatmap(df2, ax=ax2, cbar=False, cmap = 'coolwarm_r', vmax = 20, vmin = -10, annot = True, fmt = '.2f', annot_kws = {'size':12})
sb.heatmap(df3, ax=ax3, cbar=False, cmap = 'coolwarm_r', vmax = 50, vmin = 30, annot = True, fmt = '.2f', annot_kws = {'size':12})
sb.heatmap(df4, ax=ax4, cbar=False, cmap = 'coolwarm_r', vmax = 10, vmin = 5, annot = True, fmt = '.2f', annot_kws = {'size':12})
sb.heatmap(df5, ax=ax5, cbar=False, cmap = 'coolwarm_r', vmax = 8, vmin = 2, annot = True, fmt = '.2f', annot_kws = {'size':12})

ax2.set_yticks([])
ax3.set_yticks([])
ax4.set_yticks([])
ax5.set_yticks([])

ax.xaxis.tick_top()  # Put x axis on top
ax2.xaxis.tick_top()  # Put x axis on top
ax3.xaxis.tick_top()  # Put x axis on top
ax4.xaxis.tick_top()  # Put x axis on top
ax5.xaxis.tick_top()  # Put x axis on top

fig.colorbar(ax5.collections[0], ax=ax5, location="right", use_gridspec=False, pad=0.2)
ax.tick_params(rotation=0)  # Do not rotate y tick labels

plt.show()

But I get the following output where the width of each column of the heatmap is not the same:

Is there any way to resize each column (including those 4 columns from that one df in the middle) so that they're all of the same width?

Answer 1

Maybe I don't fully understand your question, but what about combining all the data frames and drawing a heat map?

df_new = pd.concat([df,df2,df3,df4,df5],axis=1)

sb.heatmap(df_new,  ax=ax, cbar=True, cmap = 'coolwarm_r', annot = True, fmt = '.2f', annot_kws = {'size':12})

Answer 2

The 'width_ratios' of plt.subplot's gridspec_kw sets the ratios between the axes. Note that adding the colorbar to ax5 will take away some of its space which would make it smaller. Therefore, it's easier to create an additional ax for the colorbar. As the axes are very close to each other, a dummy ax can be added to create some padding. As the colorbar ticks only refer to ax5, an idea is to remove them with fig.colorbar(..., ticks=[]).

The code can be a bit easier to maintain using loops.

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

df = pd.DataFrame(np.random.rand(8, 1) * 50, columns=list("A"))
df2 = pd.DataFrame(np.random.rand(8, 1) * 50, columns=list("B"))
df3 = pd.DataFrame(np.random.rand(8, 4) * 50, columns=list("CDEF"))
df4 = pd.DataFrame(np.random.rand(8, 1) * 50, columns=list("G"))
df5 = pd.DataFrame(np.random.rand(8, 1) * 50, columns=list("H"))

dfs = [df, df2, df3, df4, df5]
widths = [len(d.columns) for d in dfs]
fig, axes = plt.subplots(figsize=(20, 7), ncols=7,
                         gridspec_kw={'wspace': 0.01, 'width_ratios': widths + [0.2, 0.1]})

for ax, dfi, vmin, vmax in zip(axes[:len(dfs)], dfs, [10, -10, 30, 5, 2], [30, 20, 50, 10, 8]):
    sns.heatmap(dfi, ax=ax, cbar=False, cmap='coolwarm_r', vmax=vmax, vmin=vmin, annot=True, fmt='.2f',
                annot_kws={'size': 12})
    ax.xaxis.tick_top()  # Put x axis on top
    ax.tick_params(rotation=0)  # Do not rotate y tick labels
    if ax != axes[0]:
        ax.set_yticks([])

fig.colorbar(ax5.collections[0], cax=axes[-1])
axes[-2].axis('off')  # turn dummy as off

plt.show()