Python: Create matplotlib subplots using a for loop

Question

I would like to know how can I make my code more efficient with a for loop. I'm interested in making multiple subplots. In this example I'm using 4 but in reality I have 14.

So far, I've been copying/pasting the same block of code

df_A = df.loc[df['category'] == 'A'].copy()
df_B = df.loc[df['category'] == 'B'].copy()
df_C = df.loc[df['category'] == 'C'].copy()
df_D = df.loc[df['category'] == 'D'].copy()

dataframes = [df_A, df_B, df_C, df_D]

fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 12), dpi=100)

fig.suptitle('Distributions per category in minutes', fontweight="bold", fontsize=15)

# category A
ax1.hist(df_A['time_spent'], color="darkcyan", edgecolor='black', bins=20)
ax1.set_title('Category A', fontsize=10)

# category B
ax2.hist(df_B['time_spent'], color="darkcyan", edgecolor='black', bins=20)
ax2.set_title('Category B', fontsize=10)

# category C
ax3.hist(df_C['time_spent'], color="darkcyan", edgecolor='black', bins=20)
ax3.set_title('Category C', fontsize=10)

# category D
ax4.hist(df_D['time_spent'], color="darkcyan", edgecolor='black', bins=20)
ax4.set_title('Category D', fontsize=10)

fig.tight_layout()
plt.show()

Joehat · Accepted Answer · 2022-09-02 11:19:24Z

2

You could store the axes and conditions within lists and iterate over them to create your plots. Something like this would do the job:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from itertools import chain

df = pd.DataFrame(
    {
        "category": ["A", "B", "C", "D"] * 1000,
        "time_spent": np.random.rand(4000)
    }
)

# categories = ["A", "B", "C", "D"]
categories = sorted(list(df["category"].unique()))
fig, axes = plt.subplots(2, 2, figsize=(12, 12), dpi=100)
axes = list(chain.from_iterable(axes))

fig.suptitle('Distributions per category in minutes', fontweight="bold", fontsize=15)
for i in range(len(axes)):
  axes[i].hist(df.loc[df['category'] == categories[i], 'time_spent'], color="darkcyan", edgecolor='black', bins=20)
  axes[i].set_title(f'Category {categories[i]}', fontsize=10)
  
fig.tight_layout()
plt.show()

edited Sep 2, 2022 at 11:19

Joehat

1,1391 gold badge17 silver badges53 bronze badges

answered Jul 7, 2022 at 15:20

ko3

1,8117 silver badges14 bronze badges

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

Collectives™ on Stack Overflow

Python: Create matplotlib subplots using a for loop

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related