sample_data.py

# Load in sample dataset
df = pd.read_csv('../data/segmented_bar_graph_example_data.csv').set_index('country').fillna(0)

# Extract the log(deaths)
# Use the sorted index from weeklies to sort deaths in the same way.
deaths = df.loc[weeklies.index]['log_deaths']

# Separate the weekly mobility drops and sort them in descending order of the sum of the
# the weekly drops
weeklies = df.drop(columns=['log_deaths'])\
    .assign(s=df.sum(axis=1))\
    .sort_values('s', ascending=False)\
    .drop('s', axis=1)

weeklies.head()