matplotlib FuncAnimation - 确保图例和绘图线具有相同的颜色？

Question

考虑一个包含多列的 Pandas 数据框，每列一个国家名称，多行，每行一个日期。这些单元格是有关国家/地区的数据，随时间变化。这是 CSV：

我想在 Jupyter 中制作一个动态图（动画）来显示数据如何随时间演变。在世界上所有国家中，我只想在任何给定时间显示前 10 个国家。因此图表中显示的国家可能会不时发生变化（因为前 10 名在不断变化）。

我也想在颜色方面保持一致性。任何时候只显示 10 个国家，有些国家几乎连续出现和消失，但任何国家的颜色在整个动画中都不应改变。任何国家的颜色都应该从头到尾保持一致。

这是我的代码（编辑：现在你可以 copy/paste 把代码放到 Jupyter 中，它开箱即用，所以你可以很容易地看到我说的错误）：

import pandas as pd
import requests
import os
from matplotlib import pyplot as plt
import matplotlib.animation as ani

rel_big_file = 'rel_big.csv'
rel_big_url = 'https://pastebin.com/raw/bJbDz7ei'

if not os.path.exists(rel_big_file):
    r = requests.get(rel_big_url)
    with open(rel_big_file, 'wb') as f:
        f.write(r.content)

rel_big = pd.read_csv(rel_big_file, index_col='Date')

# history of top N countries
champs = []
# frame draw function
def animate_graph(i=int):
    N = 10
    # get current values for each country
    last_index = rel_big.index[i]
    # which countries are top N in last_index?
    topN = rel_big.loc[last_index].sort_values(ascending=False).head(N).index.tolist()
    # if country not already in champs, add it
    for c in topN:
        if c not in champs:
            champs.append(c)
    # pull a standard color map from matplotlib
    cmap = plt.get_cmap("tab20")
    # draw legend
    plt.legend(topN)
    # make a temporary dataframe with only top N countries
    rel_plot = rel_big[topN].copy(deep=True)
    # plot temporary dataframe
    p = plt.plot(rel_plot[:i].index, rel_plot[:i].values)
    # set color for each country based on index in champs
    for i in range(0, N):
        p[i].set_color(cmap(champs.index(topN[i]) % 20))

%matplotlib notebook
fig = plt.figure(figsize=(10, 6))
plt.xticks(rotation=45, ha="right", rotation_mode="anchor")
# x ticks get too crowded, limit their number
plt.gca().xaxis.set_major_locator(plt.MaxNLocator(nbins=10))
animator = ani.FuncAnimation(fig, animate_graph, interval = 333)
plt.show()

它在某种程度上起到了作用。我将排名靠前的国家存储在冠军列表中，并根据冠军中每个国家的索引分配颜色。但是根据 champs 中的索引，只有绘制线条的颜色被正确分配。

图例中的颜色分配很死板，图例中的第一个国家总是相同的颜色，图例中的第二个国家总是特定的颜色等等，基本上每个国家的颜色都是图例中的当国家在图例中上下移动时，图例在整个动画中会发生变化。

绘制线条的颜色遵循 champs 中的索引。传说中国家的颜色是根据传说中的顺序排列的。这不是我想要的。

如何以匹配情节线的方式为图例中的每个国家分配颜色？

Answer 1

这是我的解决方案：

我删除了生成颜色的代码并设置了一个新的工作代码：

首先，我在字典中用自己独特的颜色初始化了每个国家：

# initializing fixed color to all countries
colorsCountries = {}
for country in rel_big.columns:
    colorsCountries[country] = random.choice(list(mcd.CSS4_COLORS.keys()))

然后我替换了这个：

# plot temporary dataframe
p = plt.plot(rel_plot[:i].index, rel_plot[:i].values)

用这个：

# plot temporary dataframe
for keyIndex in rel_plot[:i].keys() :
    p = plt.plot(rel_plot[:i].index,rel_plot[:i][keyIndex].values,color=colorsCountries[keyIndex])

然后添加了更新 matplotlib 图例标签和颜色的代码

leg = plt.legend(topN)
for line, text in zip(leg.get_lines(), leg.get_texts()):
    line.set_color(colorsCountries[text.get_text()])

别忘了添加导入：

import matplotlib._color_data as mcd
import random

这是完整的建议解决方案：

import pandas as pd
import requests
import os
from matplotlib import pyplot as plt
import matplotlib.animation as ani
import matplotlib._color_data as mcd
import random

rel_big_file = 'rel_big.csv'
rel_big_url = 'https://pastebin.com/raw/bJbDz7ei'

if not os.path.exists(rel_big_file):
    r = requests.get(rel_big_url)
    with open(rel_big_file, 'wb') as f:
        f.write(r.content)

rel_big = pd.read_csv(rel_big_file, index_col='Date')

# history of top N countries
champs = []
# initializing fixed color to all countries
colorsCountries = {}
for country in rel_big.columns:
    colorsCountries[country] = random.choice(list(mcd.CSS4_COLORS.keys()))
# frame draw function
def animate_graph(i=int):
    N = 10
    # get current values for each country
    last_index = rel_big.index[i]
    # which countries are top N in last_index?
    topN = rel_big.loc[last_index].sort_values(ascending=False).head(N).index.tolist()
    # if country not already in champs, add it
    for c in topN:
        if c not in champs:
            champs.append(c)
    # pull a standard color map from matplotlib
    cmap = plt.get_cmap("tab20")
    # draw legend
    plt.legend(topN)
    # make a temporary dataframe with only top N countries
    rel_plot = rel_big[topN].copy(deep=True)
    # plot temporary dataframe
    #### Removed Code
    #p = plt.plot(rel_plot[:i].index, rel_plot[:i].values)
    #### Removed Code
    for keyIndex in rel_plot[:i].keys() :
        p = plt.plot(rel_plot[:i].index,rel_plot[:i][keyIndex].values,color=colorsCountries[keyIndex])
    # set color for each country based on index in champs
    #### Removed Code
    #for i in range(0, N):
        #p[i].set_color(cmap(champs.index(topN[i]) % 20))
    #### Removed Code
    leg = plt.legend(topN)
    for line, text in zip(leg.get_lines(), leg.get_texts()):
        line.set_color(colorsCountries[text.get_text()])

%matplotlib notebook
fig = plt.figure(figsize=(10, 6))
plt.xticks(rotation=45, ha="right", rotation_mode="anchor")
# x ticks get too crowded, limit their number
plt.gca().xaxis.set_major_locator(plt.MaxNLocator(nbins=10))
animator = ani.FuncAnimation(fig, animate_graph, interval = 333)
plt.show()

Answer 2

ZINE Mahmoud 的回答很棒。我只改变了一件事——我想要每个运行的颜色确定性分配，所以我没有使用随机方法为国家/地区分配颜色，如下所示：

colorsCountries = {}
colorPalette = list(mcd.CSS4_COLORS.keys())
for country in rel_big.columns:
    colorsCountries[country] = colorPalette[rel_big.columns.tolist().index(country) % len(colorPalette)]

matplotlib FuncAnimation - 确保图例和绘图线具有相同的颜色？

matplotlib FuncAnimation - make sure legend and plot lines have the same colors?

python

matplotlib

data-science

jupyter-notebook

matplotlib-animation