Looping through a function to plot several subplots, Python

Question

I have variables x and y

def function(a,b):
    x = x[(x>a)*(x<b)]
    y = y[(y<a)*(y>b)]

    # perform some fitting routine using curve_fit on x and y

    fig = plt.figure()
    ax = fig.add_subplot(111)
    phist,xedge,yedge,img = ax.hist2d(x,y,bins=20,norm=LogNorm())
    im = ax.imshow(phist,cmap=plt.cm.jet,norm=LogNorm(),aspect='auto')
    fig.colorbar(im,ax=ax)
    fig.show()

All works fine. But I have 6 pairs of different input parameters a and b. I would like to somehow call function(a,b) using a loop and plot the six different x and y (corresponding to the 6 input pairs) as 6 subplots.

like we do

ax1 = fig.add_subplot(231) # x vs y for a1,b1
ax2 = fig.add_subplot(232) # x vs y for a2,b2
....
ax6 = fig.add_subplot(236) # x vs y for a6,b6

I would like to get an idea of how to proceed to get the final subplot!

I know that it can be done manually by specifying different variables, like x1 and y1 for the first input pair a and b and so on for the other 6 pairs (x2,y2...,x6,y6). But it will be a very lengthy and confusing code.

chthonicdaemon · Accepted Answer · 2015-12-03 14:12:32Z

3

The key is using the three parameter form of subplot:

import matplotlib.pyplot as plt

# Build a list of pairs for a, b
ab = zip(range(6), range(6))

#iterate through them
for i, (a, b) in enumerate(ab):
    plt.subplot(2, 3, i+1)
    #function(a, b)
    plt.plot(a, b)

plt.show()

You'll just have to take the call to figure out of the function.

edited Dec 3, 2015 at 14:12

answered Dec 3, 2015 at 13:35

chthonicdaemon

19.9k2 gold badges55 silver badges70 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

user3397243 Over a year ago

can you elaborate further?

chthonicdaemon Over a year ago

I've changed the answer to be completely self-contained and to operate on a list of pairs ab. This code runs fine on my machine. You can comment the plt.plot and substitute the function call where I've put it commented.

TheBlackCat · Accepted Answer · 2015-12-03 14:58:45Z

Use plt.subplots instead of plt.subplot (note the "s" at the end). fig, axs = plt.subplots(2, 3) will create a figure with 2x3 group of subplots, where fig is the figure, and axs is a 2x3 numpy array where each element is the axis object corresponding to the axis in the same position in the figure (so axs[1, 2] is the bottom-right axis).

You can then either use a pair of loops to loop over each row then each axis in that row:

fig, axs = plt.subplots(2, 3)
for i, row in enumerate(axs):
   for j, ax in enumerate(row):
       ax.imshow(foo[i, j])
fig.show()

Or you can use ravel to flatten the rows and whatever you want to get the data from:

fig, axs = plt.subplots(2, 3)
foor = foo.ravel()
for i, ax in enumerate(axs.ravel()):
    ax.imshow(foor[i])
fig.show()

Note that ravel is a view, not a copy, so this won't take any additional memory.

Collectives™ on Stack Overflow

Looping through a function to plot several subplots, Python

2 Answers 2

2 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related