Bar Chart Race with matplotlib: Bars changing colors

Question

I'm trying to write a Bar Chart Race with matplot. I don't use the "bar_chart_race" library because I need more options for customization later. But I use the basic explanations of the same autor: https://www.dunderdata.com/blog/create-a-bar-chart-race-animation-in-python-with-matplotlib

It works fine, but the bars of the countries are changing their colors. But each country must have it's own color. It should not change when changing it's position.

I think I know where the problem is: My dataset is much bigger than the dataset of the example (230 columns instead of 6) and I only want to show the highest ten values. For this I use “.nlargest(10)” and I think, this is the problem. I also tried to use ".sort_values(ascending=False).head(10)" but it didn't work either. If i don't use "nlargest(10)" i get the Bar Chart Race for all 230 Columns.

Furthermore I can't (and don't want to) manually define a color for each of the 230 columns in this dataset and over 400 columns in my next dataset. So this is not an option.

What can I do to keep the country colors the same?

After an advice of an user, here is a minimalistic code that shows the problem:

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

I found a possible solution. A random generator assigns a fixed color to all countries. Since the choice of colors isn't the best, I'll have to create a color palette by hand later. But for now it works. The solution:

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)