matplotlib FuncAnimation - 確保圖例和 plot 行具有相同的 colors？

Question

考慮一個 Pandas dataframe 具有多列，每列一個國家名稱和多行，每行一個日期。 這些單元格是有關國家/地區的數據，這些數據會隨時間變化。 這是 CSV：

https://pastebin.com/bJbDz7ei

我想在 Jupyter 中制作動態 plot（動畫），以顯示數據如何隨時間演變。 在世界上所有國家中，我只想展示在任何給定時間排名前 10 的國家。 因此圖中顯示的國家/地區可能會不時發生變化（因為前 10 名正在演變）。

我還想在 colors 方面保持一致性。 任何時候只顯示 10 個國家，有些國家幾乎連續出現和消失，但任何國家的顏色在整個 animation 中都不應該改變。 任何國家的顏色都應該從頭到尾堅持下去。

這是我擁有的代碼（編輯：現在您可以將代碼復制/粘貼到 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()

它完成了這項工作 - 有點。 我將排名靠前的國家存儲在冠軍列表中，並根據每個國家在冠軍中的索引分配 colors。 但是根據 champs 中的索引，僅正確分配了繪制線的顏色。

傳說中的顏色是固定分配的，傳說中的第一個國家總是得到相同的顏色，傳說中的第二個國家總是得到某種顏色等等，基本上傳說中每個國家的顏色在整個 animation 中都是不同的當國家在傳說中上下移動時。

繪制線的 colors 服從以冠軍為單位的索引。 圖例中國家的colors是按照圖例中的順序排列的。 這不是我想要的。

如何以與 plot 線匹配的方式為圖例中的每個國家/地區分配顏色？

Answer 1

這是我的解決方案：

我刪除了您生成 colors 的代碼並設置了一個新的工作代碼：

首先，我在字典中用自己獨特的顏色初始化了每個國家：

# 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 圖例 label 和 colors 的代碼

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 的回答很棒。 我只改變了一件事——我希望每次運行都確定性地分配 colors，所以我將 colors 分配給這樣的國家，而不是隨機方法：

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