Setting same frame width in matplotlib subplots with external colorbar element-CodePudding

I want to produce two subplots that contain a lot of curves, so I defined a function that produces a colorbar, to avoid having a super long legend that is not readable. This is the function to create the colorbar:

import matplotlib as mpl, matplotlib.pyplot as plt


def colorbar (cmap, vmin, vmax, label, ax=None, **cbar_opts):
    
    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax, clip=False)
    cbar = plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
                        label=label, ax=ax, **cbar_opts)
    return cbar

I want to show just one colorbar, since its values are the same for the two plots. So, I place it only on the right of the second axis. Here is the code.

import pandas as pd, numpy as np

df1 = pd.DataFrame({i: np.linspace(i*i, 10, 10) for i in range(50)})
df2 = pd.DataFrame({i: np.linspace(2*i*i, 10, 10) for i in range(50)})

fig, axs = plt.subplots(1,2, figsize=(5,3))
df1.plot(ax=axs[0], legend=False, cmap='turbo')
df2.plot(ax=axs[1], legend=False, cmap='turbo')

colorbar(ax = axs[1], cmap='turbo', vmin=0, vmax=49, label='My title')

axs[1].set_title('I want this frame \n as large as \n the first one')

plt.tight_layout()

My problem is that now the two plots have different width, because the colorbar is considered in the measurement of the width of the second axis. How can I get the two frames to have the same width?

CodePudding user response：

You can pass to colorbar a list of axes, from which space is stolen in equal way.

import matplotlib as M
f, a = M.pyplot.subplots(1, 2)
#                                 ↓↓↓↓
f.colorbar(M.cm.ScalarMappable(), ax=a)
#                                 ↑↑↑↑
f.show()

If you want a "tight layout" you can specify layout='constrained' when instantiating the Figure and the Axes:

import matplotlib as M
f, a = M.pyplot.subplots(1, 2, layout='constrained')
f.colorbar(M.cm.ScalarMappable(), ax=a)
f.show()

constrained layout is, so to say, the next gen of tight layout and, as you can see, it's way smarter than its predecessor, that remains supported because of the tons of preexisting code that uses it. I recommend that you use constrained layout in all the new code.

ps The syntax layout='constrained' is recent (Matplotlib 3.6?), previously one had to use the more verbose constrained_layout=True, that is still supported. If you want to use your code on an old installation, it's hence safer to use the second format.

CodePudding user response：

Reference:

Alternatively, you can create a third new axs[2] just for the colorbar. In this case, plt.tight_layout() is supported.

import matplotlib as mpl, matplotlib.pyplot as plt
import pandas as pd, numpy as np


def colorbar(cmap, vmin, vmax, label, **cbar_opts):
    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax, clip=False)
    cbar = plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
                        label=label, **cbar_opts)
    return cbar


df1 = pd.DataFrame({i: np.linspace(i * i, 10, 10) for i in range(50)})
df2 = pd.DataFrame({i: np.linspace(2 * i * i, 10, 10) for i in range(50)})

fig, axs = plt.subplots(1, 3, figsize=(10, 5))
df1.plot(ax=axs[0], legend=False, cmap='turbo')
df2.plot(ax=axs[1], legend=False, cmap='turbo')

colorbar(ax=axs[2], cmap='turbo', vmin=0, vmax=49, label='My title', cax=axs[2])  # make a third axs for colorbar

axs[2].set_aspect(0.5)  # adjust colorbar's ratio

axs[1].set_title('I want this frame \n as large as \n the first one')

plt.tight_layout()
plt.show()

Result: