pythonmatplotlibmatplotlib-gridspec

Automatically get the dimensions or indices of gridspec axes


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?


Solution

  • 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.

    creating the axes for gridspec

    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.