Setting same frame width in matplotlib subplots with external colorbar element

Question:

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?

enter image description here

Asked By: gioarma

||

Answers:

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

enter image description here

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:

enter image description here

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.

Answered By: gboffi

Reference: https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.colorbar.html

fig.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap), ax=ax)

ax: one or more parent
axes from which space for a new colorbar axes will be stolen, if cax (axes into which the colorbar will be drawn)
is None. This has no effect if cax is set.

By setting ax=axs, the space for the colorbar is stolen by the figure which it is applied to. However, plt.tight_layout() is not supported in this mode.

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, ax=axs, **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, 2, figsize=(10, 5))
df1.plot(ax=axs[0], legend=False, cmap='turbo')
df2.plot(ax=axs[1], legend=False, cmap='turbo')

colorbar(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')

# not compatible with tight_layout
# plt.tight_layout()
plt.show()

Result:

enter image description here

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:

enter image description here

Answered By: コリン
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.