Stack Overflow is a community of 4.7 million programmers, just like you, helping each other.

Join them; it only takes a minute:

Sign up
Join the Stack Overflow community to:
  1. Ask programming questions
  2. Answer and help your peers
  3. Get recognized for your expertise

Is there a simple/clean way to iterate an array of axis returned by subplots like

 nrow = ncol = 2
 a = []
 fig, axs = plt.subplots(nrows=nrow, ncols=ncol)
 for i, row in enumerate(axs):
     for j, ax in enumerate(row):
         a.append(ax)

 for i, ax in enumerate(a):
    ax.set_ylabel(str(i))

which even works for nrow or ncol = 1.

I tried list comprehension like:

[element for tupl in tupleOfTuples for element in tupl]

but that fails if nrows or ncols == 1

share|improve this question
up vote 9 down vote accepted

The ax return value is a numpy array, which can be reshaped, I believe, without any copying of the data. If you use the following, you'll get a linear array that you can iterate over cleanly.

nrow = 1; ncol = 2;
fig, axs = plt.subplots(nrows=nrow, ncols=nrow)

for ax in axs.reshape(-1):
  ax.set_ylabel(str(i))

This doesn't hold when ncols and nrows are both 1, since the return value is not an array; you could turn the return value into an array with one element for consistency, though it feels a bit like a cludge:

nrow = 1; ncol = 1;
fig, axs = plt.subplots(nrows=nrow, ncols=nrow)
axs = np.array(axs)

for ax in axs.reshape(-1):
  ax.set_ylabel(str(i))

reshape docs. The argument -1 causes reshape to infer dimensions of the output.

share|improve this answer
    
For nrow=ncol=1 you can use squeeze=0. plt.subplots(nrows=nrow, ncols=nrow, squeeze=0) always returns a 2 dimensional array for the axes, even if both are one. – cronos Apr 6 at 11:26

The fig return value of plt.subplots has a list of all the axes. To iterate over all the subplots in a figure you can use:

nrow = 2
ncol = 2
fig, axs = plt.subplots(nrow, ncol)
for i, ax in enumerate(fig.axes):
    ax.set_ylabel(str(i))

This also works for nrow == ncol == 1.

share|improve this answer

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.