Automatically get the dimensions or indices of gridspec axes

Question

Given a gridspec object in matplotlib, I want to automatically iterate through all its indices so I can add the corresponding Axes automatically, something like:

for i, j in gspec.indices:  # whatever those indices are
    axs[i,j] = fig.add_subplot(gspec[i][j])

How do I do that, without knowing how many rows or columns the gridspec has in advance?

Answer 1

gspec.get_geometry() returns the number of rows and of columns. Here is some example code:

import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure(constrained_layout=True)
gspec = fig.add_gridspec(3, 4)
nrows, ncols = gspec.get_geometry()
axs = np.array([[fig.add_subplot(gspec[i, j]) for j in range(ncols)] for i in range(nrows)])
t = np.linspace(0, 4 * np.pi, 1000)
for i in range(nrows):
    for j in range(ncols):
        axs[i, j].plot(np.sin((i + 1) * t), np.sin((j + 1) * t))
plt.show()

If axs isn't needed as numpy array, the conversion to numpy array can be left out.

Note that the code assumes you need a subplot in every possible grid position, which also can be obtained via fig, axs = plt.subplots(...). A gridspec is typically used when you want to combine grid positions to create custom layouts, as shown in the examples of the tutorial.