matplotlib 條形圖競賽：條形變化 colors

Question

我正在嘗試用 matplot 編寫條形圖競賽。 我不使用“bar_chart_race”庫，因為我以后需要更多的自定義選項。 但我使用同一作者的基本解釋： https://www.dunderdata.com/blog/create-a-bar-chart-race-animation-in-python-with-matplotlib

它工作正常，但國家的酒吧正在改變他們的 colors。但每個國家必須有自己的顏色。 更改為 position 時不應更改。

我想我知道問題出在哪里：我的數據集比示例的數據集大得多（230 列而不是 6 列），我只想顯示最高的十個值。 為此，我使用“.nlargest(10)”，我想，這就是問題所在。 我還嘗試使用“.sort_values(ascending=False).head(10)”，但它也沒有用。 如果我不使用“nlargest(10)”，我會得到所有 230 列的條形圖競賽。

此外，我不能（也不想）為這個數據集中的 230 列和我的下一個數據集中的 400 多列手動定義顏色。 所以這不是一個選擇。

我該怎么做才能使國家 colors 保持不變？

在用戶的建議之后，這里是一個顯示問題的簡約代碼：

import pandas as pd
from matplotlib.animation import FuncAnimation
import numpy as np

data = {"year": [1950,1960,1970,1980,1990,2000,2010,2020,2030],
"USA" : [10,20,30,40,50,50,50,50,55],
"GB" : [5,10,15,45,60,70,80,90,95],
"FR" : [5,15,16,17,18,25,50,60,65],
"BEL" : [3,34,11,23,34,23,12,22,27],
"GER" : [5,15,16,23,34,40,23,50,55],
"POL" : [5,14,19,20,23,45,50,70,75],
"KAN" : [1,5,18,22,34,45,46,60,65],
"ISR" : [2,15,25,32,43,57,66,67,70],
"IND" : [3,12,16,17,23,25,45,50,55],
"CH" : [2,19,21,22,22,22,25,26,30],
"AUS" : [4,4,14,17,22,25,30,34,37],
}
df = pd.DataFrame(data).set_index("year")

def nice_axes(ax):
    ax.set_facecolor('.8')
    ax.tick_params(labelsize=8, length=0)
    ax.grid(True, axis='x', color='white')
    ax.set_axisbelow(True)
    [spine.set_visible(False) for spine in ax.spines.values()]

def prepare_data(df, steps=5):
    df = df.reset_index()
    df.index = df.index * steps
    last_idx = df.index[-1] + 1
    df_expanded = df.reindex(range(last_idx))
    df_expanded['year'] = df_expanded['year'].fillna(method='ffill')
    df_expanded = df_expanded.set_index('year')
    df_rank_expanded = df_expanded.rank(axis=1, method='first')
    df_expanded = df_expanded.interpolate()
    df_rank_expanded = df_rank_expanded.interpolate()
    return df_expanded, df_rank_expanded

df_expanded, df_rank_expanded = prepare_data(df)

colors = plt.cm.viridis(np.linspace(0, 1, 10))

def init():
    ax.clear()
    nice_axes(ax)

def update(i):
    for bar in ax.containers:
        bar.remove()
   
    y = df_rank_expanded.iloc[i].nlargest(10)
    width = df_expanded.iloc[i].nlargest(10)

    ax.barh(y=y, width=width, color = colors, tick_label=y.index)
  
fig = plt.Figure(figsize=(8, 3), dpi=144)
ax = fig.add_subplot()
anim = FuncAnimation(fig=fig, func=update, init_func=init, frames=len(df_expanded), 
                     interval=100, repeat=False)

from IPython.display import HTML
html = anim.to_html5_video()
HTML(html)

Answer 1

我找到了一個可能的解決方案。 隨機生成器為所有國家分配固定顏色。 由於 colors 的選擇不是最好的，我將不得不稍后手動創建一個調色板。 但現在它有效。 解決方案：

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.animation import FuncAnimation
import numpy as np
import random

data = {"year": [1950,1960,1970,1980,1990,2000,2010,2020,2030],
"USA" : [10,20,30,40,50,50,50,50,55],
"GB" : [5,10,15,45,60,70,80,90,95],
"FR" : [5,15,16,17,18,25,50,60,65],
"BEL" : [3,34,11,23,34,23,12,22,27],
"GER" : [5,15,16,23,34,40,23,50,55],
"POL" : [5,14,19,20,23,45,50,70,75],
"KAN" : [1,5,18,22,34,45,46,60,65],
"ISR" : [2,15,25,32,43,57,66,67,70],
"IND" : [3,12,16,17,23,25,45,50,55],
"CH" : [2,19,21,22,22,22,25,26,30],
"AUS" : [4,4,14,17,22,25,30,34,37],
}
df = pd.DataFrame(data).set_index("year")

def nice_axes(ax):
    ax.set_facecolor('.8')
    ax.tick_params(labelsize=8, length=0)
    ax.grid(True, axis='x', color='white')
    ax.set_axisbelow(True)
    [spine.set_visible(False) for spine in ax.spines.values()]

#Prepare Data (expand the dataframe for better animation)
def prepare_data(df, steps=10):
    df = df.reset_index()
    df.index = df.index * steps
    last_idx = df.index[-1] + 1
    df_expanded = df.reindex(range(last_idx))
    df_expanded['year'] = df_expanded['year'].fillna(method='ffill')
    df_expanded = df_expanded.set_index('year')
    df_rank_expanded = df_expanded.rank(axis=1, method='first')
    df_expanded = df_expanded.interpolate()
    df_rank_expanded = df_rank_expanded.interpolate()
    return df_expanded, df_rank_expanded

df_expanded, df_rank_expanded = prepare_data(df)

# RGB to RGBA 
def get_color(r, g, b):
    return (r, g, b, 1.0)

#Randomized colors
color_dict = {}
for idx in range(df_expanded.shape[1]):
    r = random.random()
    b = random.random()
    g = random.random()
    color = get_color(r, g, b)
    color_dict[df_expanded.columns[idx]] = color 

def init():
    ax.clear()
    nice_axes(ax)

def update(i):
    for bar in ax.containers:
        bar.remove()

    ax.clear()
    nice_axes(ax)
    y = df_rank_expanded.iloc[i].nlargest(10)
    width = df_expanded.iloc[i].nlargest(10)
    
    # Color for each country
    bar_colors = [color_dict.get(country) for country in y.index]

    
    #Plot
    ax.barh(y=y, width=width, color = bar_colors, tick_label=y.index, alpha=1.0, align='center')
 
fig = plt.Figure(figsize=(8, 3), dpi=144)
ax = fig.add_subplot()
anim = FuncAnimation(fig=fig, func=update, init_func=init, frames=len(df_expanded), 
                     interval=100, repeat=False)

from IPython.display import HTML
html = anim.to_html5_video()
HTML(html)