Matplotlib: given a figure, how to efficiently handle a variable number of axis?-CodePudding

Consider the following example

from matplotlib import pyplot as plt
import numpy as np
from random import randrange

N = 100
x = np.arange(0,100,1)
y = np.random.rand(len(x),1)

p = randrange(4) 1
q = randrange(4) 1

fig, ax = plt.subplots(p, q, sharex=True)
for ii in range(p):
    for jj in range(q):
        if p == 1 and q == 1:
            ax.plot(x,y)
            ax.grid(True)
            ax.set_xlim([-3,5])
        elif p == 1 or q == 1:
            ax[max(ii, jj)].plot(x,y)
            ax[max(ii, jj)].grid(True)
            ax[max(ii, jj)].set_xlim([-3,5])
        else:
            ax[ii, jj].plot(x,y)
            ax[ii, jj].grid(True)
            ax[ii, jj].set_xlim([-3,5])
plt.suptitle('Test')
plt.show()

It works fairly well, but due to that p and q are varying, I am not allowed to write something more elegant like

from matplotlib import pyplot as plt
import numpy as np
from random import randrange

N = 100
x = np.arange(0,100,1)
y = np.random.rand(len(x),1)

p = randrange(4) 1
q = randrange(4) 1

fig, ax = plt.subplots(p, q, sharex=True)
for ii in range(p):
    for jj in range(q):
            ax[ii, jj].plot(x,y)
            ax[ii, jj].grid(True)
            ax[ii, jj].set_xlim([-3,5])
plt.suptitle('Test')
plt.show()

because, for example, when p=q=1 I would get an error. However, my solution is rather ugly and I am wondering how that can be made nicer.

CodePudding user response：

You can use the argument squeeze=False of plt.subplots:

from matplotlib import pyplot as plt
import numpy as np
from random import randrange

N = 100
x = np.arange(0,100,1)
y = np.random.rand(len(x),1)

p = randrange(4) 1
q = randrange(4) 1

fig, ax = plt.subplots(p, q, sharex=True, squeeze=False)
for ii in range(p):
    for jj in range(q):
            ax[ii, jj].plot(x,y)
            ax[ii, jj].grid(True)
            ax[ii, jj].set_xlim([-3,5])
plt.suptitle('Test')
plt.show()

Source