将数据添加到现有绘图时更新图例（熊猫）

Question

我编写了一个小的 Python 代码来从 ourworldindata.org 读取 Covid 统计数据并绘制某个国家/地区的某个数据系列。

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()  

它按预期工作，并在上面的示例中将挪威的新病例数绘制为日期的函数。如果我更改“国家”并重新运行它，它会在同一个图中为新国家绘制一条不同颜色的新曲线，这就是我想要的。但是传说有问题。它显示最后绘制的国家/地区的名称，但显示第一个绘制的国家/地区的颜色。我希望它同时显示正确的名称和颜色。我该怎么做？

该链接显示了先绘制挪威（蓝色曲线）然后绘制丹麦（黄色曲线）时的结果图： Plot of new cases in Norway and Denmark

Answer 1

我不太确定您是如何“重新运行”代码，但您可以在列表中定义您的国家并循环打印它们：

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()

或者您可以使用seaborn 代替循环：

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()