Creating Subplots with a Function and Loop

Question

I have a nice function that makes a plot with two lines ... it works totally fine on it's own. I'd like to run it however 4 times to make 2row x 2col subplots. I can't find a good way to run the function multiple times and add each one to a subplot.

import matplotlib.pyplot as plt

patient = [['P1', [1,6,12,18,24], [15,17,16,19,15]],
           ['P2', [1,6,12,18,24], [12,13,17,18,18]],
           ['P3', [1,6,12,18,24], [19,19,12,11,9]],
           ['P4', [1,6,12,18,24], [8,7,3,12,15]]]

def plot_graph(patient):
    name = patient[0]
    X = patient[1]
    Y = patient[2]

    fig, ax = plt.subplots()

    #first line
    ax.plot([X[0],X[-1]],[Y[0],Y[-1]+20], marker='o', label='Line 1')

    #second line
    ax.plot(X,Y,marker='o',label='Line 2')

    #axis settings
    ax.set_ylim([0,80])
    ax.invert_yaxis()
    ax.set_title('My Graph')
    for x,y in zip(X,Y): ax.annotate(y,(x,y-4))
    ax.legend(loc=8)


plot_graph(patient[0])

plt.show()

Answer 1

• Remove fig, ax = plt.subplots() and add an ax parameter to the function
• Iterate through subplots outside of the function.
  • .flatten converts axes from a (2, 2) to a (4, ) array, which is easier to iterate though.

def plot_graph(data, ax):
    name = data[0]
    X = data[1]
    Y = data[2]

    #first line
    ax.plot([X[0],X[-1]],[Y[0],Y[-1]+20], marker='o', label='Line 1')

    #second line
    ax.plot(X,Y,marker='o',label='Line 2')

    #axis settings
    ax.set_ylim([0,80])
    ax.invert_yaxis()
    ax.set_title('My Graph')
    for x,y in zip(X,Y): ax.annotate(y,(x,y-4))
    ax.legend(loc=8)


fig, axes = plt.subplots(2, 2, figsize=(10, 8))
axes = axes.flatten()

for i, axe in enumerate(axes):
    plot_graph(data=patient[i], ax=axe)