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()

Answer 2

Here's a version that shows how to run scatter plots on two different figures. Basically you reference the axes that are created with plt.subplots.

import matplotlib.pyplot as plt
import numpy as np

x1 = y1 = range(10)
x2 = y2 = range(5)

nRows = nCols = 6
fig1, axesArray1 = plt.subplots(nrows=nRows,ncols=nCols,figsize=(20, 20))
fig1.subplots_adjust(hspace=.5,wspace=0.4)
fig1.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=None, hspace=None)

fig2, axesArray2 = plt.subplots(nrows=nRows,ncols=nCols,figsize=(20, 20))
fig2.subplots_adjust(hspace=.5,wspace=0.4)
fig2.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=None, hspace=None)

days = range(1, 32)
dayRowCol = np.array([i + 1 for i in range(nRows * nCols)]).reshape(nRows, nCols)
for day in days:
    rowIdx, colIdx = np.argwhere(dayRowCol == day)[0]

    axis1 = axesArray1[rowIdx, colIdx]
    axis1.set_title('day=' + str(day))
    axis1.scatter(x1, y1)

    axis2 = axesArray2[rowIdx, colIdx]
    axis2.set_title('day=' + str(day))
    axis2.scatter(x2, y2)

    # This didn't run in the original script, so I left it as is
    # plt.colorbar().set_label('Distance from ocean',rotation=270)

fig1.savefig('plots/everyday_D1_color.png')
fig2.savefig('plots/everyday_D2_color.png')
plt.close('all')

When I took the original code from the post plt.colorbar() raised an error, so I left it out in the answer. If you have an example of how colorbar was intended to work we could look at how to make that happen for two figures, but the rest of the code should work as intended!

Note that if day every does not appear in dayRolCol numpy will raise an error, it's up to you to decide how you want to handle that case. Also, using numpy is definitely not the only way to do it, just a way I'm comfortable with - all you really need to do is find a way to link a certain day/plot with the (x, y) indices of the axis you want to plot on.