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(...)
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 1×1.
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)