Create subplot, for every group/id of a dataframe, in Python

June 18, 2023

I have the below dataframe:

#Load the required libraries
import pandas as pd
import matplotlib.pyplot as plt

#Create dataset
data = {'id': [1, 1, 1, 1, 1, 1,1, 1, 1, 1, 1, 1,
               2, 2, 2, 2, 2, 2, 2,
               3, 3, 3, 3, 3, 3, 3,3,
               4, 4, 4, 4, 4,4,],
        'cycle': [0.0, 0.2,0.4, 0.6, 0.8, 1,1.2,1.4,1.6,1.8,2.0,2.2,
                  0.0, 0.2,0.4, 0.6,0.8,1.0,1.2,
                  0.0, 0.2,0.4, 0.6, 0.8,1.0,1.2,1.4,
                  0.0, 0.2,0.4, 0.6, 0.8,1.0,],
        'Salary': [6, 7, 7, 7,8,9,10,11,12,13,14,15,
                   3, 4, 4, 4,4,5,6,
                   2, 8,9,10,11,12,13,14,
                   1, 8,9,10,11,12,],
        'Children': ['Yes', 'No', 'Yes', 'Yes', 'Yes', 'Yes', 'No','No', 'Yes', 'Yes', 'Yes', 'No',
                     'Yes', 'Yes', 'Yes', 'No', 'Yes', 'Yes', 'Yes', 
                     'Yes', 'No','Yes', 'Yes', 'No','No', 'Yes','Yes',
                     'Yes', 'Yes', 'No','Yes', 'Yes','Yes',],
        'Days': [141, 123, 128, 66, 66, 120, 141, 52,96, 120, 141, 52,
                 141, 96, 120,120, 141, 52,96,
                 141,  15,123, 128, 66, 120, 141, 141,
                 141, 141,123, 128, 66,67,],
        }

#Convert to dataframe
df = pd.DataFrame(data)
print("\n df = \n",df)

Now, here I wish to plot the cycle vs Salary, for all id‘s of the dataframe in one single plot. Thus I need to use subplot function as such:

## Plot for all id's
plt_fig_verify = plt.figure(figsize=(10,8))

## id1: 
plt.subplot(4,1,1)
plt.plot(df.groupby(by="id").get_group(1)['cycle'], df.groupby(by="id").get_group(1)['Salary'], 'b',  linewidth = '1', label ='id1')
plt.xlabel('cycle')
plt.ylabel('Salary')
plt.legend()

## id2: 
plt.subplot(4,1,2)
plt.plot(df.groupby(by="id").get_group(2)['cycle'], df.groupby(by="id").get_group(2)['Salary'], 'b',  linewidth = '1', label ='id2')
plt.xlabel('cycle')
plt.ylabel('Salary')
plt.legend()

## id3: 
plt.subplot(4,1,3)
plt.plot(df.groupby(by="id").get_group(3)['cycle'], df.groupby(by="id").get_group(3)['Salary'], 'b',  linewidth = '1', label ='id3')
plt.xlabel('cycle')
plt.ylabel('Salary')
plt.legend()

## id4: 
plt.subplot(4,1,4)
plt.plot(df.groupby(by="id").get_group(4)['cycle'], df.groupby(by="id").get_group(4)['Salary'], 'b',  linewidth = '1', label ='id4')
plt.xlabel('cycle')
plt.ylabel('Salary')
plt.legend()

plt.show()

However, here I need to write the codes for the subplot function four times, i.e. for all four id’s of the dataframe.

Is there any way out, by which we can have some iterative function and write the subplot function only once and get all four subplots.

Can somebody please let me know how to achieve this task in Python?

>Solution :

One option is a for-loop to plot each group of ids :

fig, ax = plt.subplots(figsize=(10, 8))

for n, g in df.groupby("id"):
    g.plot(
        x="cycle", y="Salary",
        xlabel="Cycle", ylabel="Salary",
        label=f"id {n}",
        ax=ax,
)

plt.show();

Or, if you need to create a subplot for each id, you can use :

fig, axs = plt.subplots(figsize=(10, 8), nrows=2, ncols=2)

for (n, g), ax in zip(df.groupby("id"), axs.flatten()):
    g.plot(
        x="cycle", y="Salary",
        xlabel="Cycle", ylabel="Salary",
        label=f"id {n}",
        ax=ax,
    )

plt.tight_layout()

plt.show();

Output :