Saving plots created in a for loop

Viewed 61

The code below produces a set of plots. However, when exporting the plots using PdfPages only 1 plot is exported for the fig2 plots. Is it possible to export the other plots created - also to the pdf?

Dataframe

for c in df3.columns:
    if df3[c].value_counts().count() < 5:
        vc = df3[c].value_counts().sort_values()
        
        # Visualization
        fig1 = plt.figure();
        ax = vc.plot(kind='barh',color=(0.2, 0.4, 0.6, 0.6));
        text   = c +' (Total)'
        ax.set_title(text)
# If 
    elif df3[c].value_counts().count() / df3.shape[0]  < 0.9:
        vc = df3[c].value_counts().head().sort_values()

        # Visualization
        fig2 = plt.figure();
        ax = vc.plot(kind='barh',color=(0.2, 0.4, 0.6, 0.6));
        text   = c +' (Sample)'
        ax.set_title(text)
        
    else:
        pass

from matplotlib.backends.backend_pdf import PdfPages
pp = PdfPages('foo.pdf')
pp.savefig(fig1)
pp.savefig(fig2)
pp.close()
1 Answers

One way is to store all the figure handles in a list. Then the list will contain all the figures, not just the last iteration:

figs = []

for c in df3.columns:
    if df3[c].value_counts().count() < 5:
        ...
        figs.append(fig1)
    elif df3[c].value_counts().count() / df3.shape[0] < 0.9:
        ...
        figs.append(fig2)
    else:
        pass

with PdfPages('foo.pdf') as pp:
    for fig in figs:
        pp.savefig(fig)

Or another option is to open the PdfPages context around the whole plotting loop and save each figure immediately:

with PdfPages('foo.pdf') as pp:
    for c in df3.columns:
        if df3[c].value_counts().count() < 5:
            ...
            pp.savefig(fig1)
        elif df3[c].value_counts().count() / df3.shape[0] < 0.9:
            ...
            pp.savefig(fig2)
        else:
            pass
Related