pythonmatplotlib

dynamic subplot dimensions

I want to plot several subplots. Depending on the data it might by a N x M array of plots, but also only a single plot or a 1 x M or N x 1 array.

The code I'm using is as follows which works for M x N arrays. However if M or N or both are 1 I get the Error: matplotlib IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed

Sure I could do a switch case for the 4 different possibilities but there must be a better way?

fig, axs = plt.subplots(max_N, max_N)

    for N in range(max_N):
        for M in range(max_N):
        axs[N, M].plot(...)
Solution

  • You can set the squeeze kwarg option to False when using plt.subplots. From the docs:

    squeeze : bool, default: True

    • If True, extra dimensions are squeezed out from the returned array of Axes:
      • if only one subplot is constructed (nrows=ncols=1), the resulting single Axes object is returned as a scalar.
      • for Nx1 or 1xM subplots, the returned object is a 1D numpy object array of Axes objects.
      • for NxM, subplots with N>1 and M>1 are returned as a 2D array.
    • If False, no squeezing at all is done: the returned Axes object is always a 2D array containing Axes instances, even if it ends up being 1x1.

    So by setting squeeze=False, you will always have a 2D array, even if either or both axes only have 1 element. For your code:

    fig, axs = plt.subplots(max_N, max_N, squeeze=False)