Update legend when adding data to existing plot (Pandas)

user_17663921 0
Viewed 54

I've written a small Python code to read Covid statistics from ourworldindata.org and plot a certain data series for a certain country.

from pandas import read_csv
import pandas as pd
import matplotlib.pyplot as plt

filename = "https://covid.ourworldindata.org/data/owid-covid-data.csv"

dataset = read_csv(filename)
dataset["date"] = pd.to_datetime(dataset["date"])

country = "Norway"
data = "new_cases"  

mask = dataset["location"] == country
dataset.loc[mask].set_index("date")[data].plot()

plt.ylabel(data)
plt.legend([country])
plt.show()

It works as intended and plots the number of new cases in Norway as a function of date in the example above. If I change "country" and rerun it, it will plot a new curve for the new country with a different color in the same plot, which is what I want. But there's a problem with the legend. It shows the name of the last plotted country but the color of the first plotted country. I would like it to show both with the correct name and color. How can I do that?

The link shows a figure with the result when first plotting Norway (blue curve) and then Denmark (yellow curve): Plot of new cases in Norway and Denmark

pandas plot python
1 Answers

I'm not quite sure how exactly you "rerun" the code but you can define your countries in a list and print them in a loop:

import pandas as pd
import matplotlib.pyplot as plt

filename = "owid-covid-data.csv"

dataset = pd.read_csv(filename)
dataset["date"] = pd.to_datetime(dataset["date"])

countries = ["Denmark", "Norway"]
data = "new_cases"  

for country in countries:
    mask = dataset["location"] == country
    dataset.loc[mask].set_index("date")[data].plot()

plt.ylabel(data)
plt.legend(countries)
plt.show()

Or you can use seaborn instead of the loop:

import seaborn as sns
df = dataset[dataset["location"].isin(countries)][["date", "location", data]]
sns.lineplot(data=df, x="date", y=data, hue="location")
plt.show()
user_135255120
Related