为多个图表重新使用轴对象

Question

I have a pandas DataFrame with values and an indicator for season. I want to create a density plot for the whole set and then one for each season, which should include the overall density plus the season's density.
Since the overall density estimation takes some time (I have over 40000 data points), I want to avoid repeating it for each season.

So far, the closest I got is this code, based on this answer :

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

n = 25
seasons = ['winter', 'spring', 'summer', 'autumn']
df = pd.DataFrame({'value': np.random.randn(n), 'season': np.random.choice(seasons, n)})

ax = df['value'].plot.density(label="all", legend=True)
plt.savefig("test_density_all.png")
for s in seasons:
    sdf = df[df['season'] == s]
    sdf['value'].plot.density(label=s, legend=True)
    plt.savefig("test_density_" + s + ".png")
    del ax.lines[-1]

There, instead of saving the overall plot, I am deleting the season-line after save. The problem is that it does not delete the legend, so while the summer-plot has the correct two densities, its legend includes four items (all, winter, spring, summer).

So, what I need is either to make the above code work with legend, or to find a way to store the overall plot in such a way that I can use it as a start point for each of the season-plots...

Answer 1

Use ax.legend() to get a legend in the plot.
Note that this is different from using the legend=True argument, because it creates a new legend based on the currently present artists in the plot.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

n = 25
seasons = ['winter', 'spring', 'summer', 'autumn']
df = pd.DataFrame({'value': np.random.randn(n), 'season': np.random.choice(seasons, n)})

ax = df['value'].plot.density(label="all", legend=True)
plt.savefig("test_density_all.png")
for s in seasons:
    sdf = df[df['season'] == s]
    sdf['value'].plot.density(label=s)
    ax.legend()
    ax.figure.savefig("test_density_" + s + ".png")
    ax.lines[-1].remove()