Plotting on multiple figures with subplots in a single loop

Question

I\'m plotting on two figures and each of these figures have multiple subplots. I need to do this inside a single loop. Here is what I do when I have only one figure:

<

Answer 1

Each of the pyplot function has its corresponding method in the object oriented API. If you really want to loop over both figures' axes at the same time, this would look like this:

import numpy as np
import matplotlib.pyplot as plt

x1 = x2 = np.arange(10)
y1 = y2 = c = np.random.rand(10,6)

fig1, axes1 = plt.subplots(nrows=2,ncols=3)
fig1.subplots_adjust(hspace=.5,wspace=0.4)

fig2, axes2 = plt.subplots(nrows=2,ncols=3)
fig2.subplots_adjust(hspace=.5,wspace=0.4)

for i, (ax1,ax2) in enumerate(zip(axes1.flatten(), axes2.flatten())):
    ax1.set_title('day='+str(i))
    ax2.set_title('day='+str(i))
    sc1 = ax1.scatter(x1,y1[:,i], c=c[:,i])
    sc2 = ax2.scatter(x2,y2[:,i], c=c[:,i])
    fig1.colorbar(sc1, ax=ax1)
    fig2.colorbar(sc2, ax=ax2)

plt.savefig("plot.png") 
plt.show()   
plt.close()

Here you loop over the two flattened axes arrays, such that ax1 and ax2 are the matplotlib axes to plot to. fig1 and fig2 are matplotlib figures (matplotlib.figure.Figure).

In order to obtain an index as well, enumerate is used. So the line

for i, (ax1,ax2) in enumerate(zip(axes1.flatten(), axes2.flatten())):
    # loop code

is equivalent here to

for i in range(6):
    ax1 = axes1.flatten()[i]
    ax2 = axes2.flatten()[i]
    # loop code

or

i = 0
for ax1,ax2 in zip(axes1.flatten(), axes2.flatten()):
    # loop code
    i += 1

which are both longer to write.

At this point you may be interested in the fact that althought the above solution using the object oriented API is surely more versatile and preferable, a pure pyplot solution still is possible. This would look like

import numpy as np
import matplotlib.pyplot as plt

x1 = x2 = np.arange(10)
y1 = y2 = c = np.random.rand(10,6)

plt.figure(1)
plt.subplots_adjust(hspace=.5,wspace=0.4)

plt.figure(2)
plt.subplots_adjust(hspace=.5,wspace=0.4)

for i in range(6):
    plt.figure(1)
    plt.subplot(2,3,i+1)
    sc1 = plt.scatter(x1,y1[:,i], c=c[:,i])
    plt.colorbar(sc1)

    plt.figure(2)
    plt.subplot(2,3,i+1)
    sc2 = plt.scatter(x1,y1[:,i], c=c[:,i])
    plt.colorbar(sc2)

plt.savefig("plot.png") 
plt.show()   
plt.close()