.py

# Create 2-row gridspec to hold the colour bar (cax) and the main graph (ax)
fig, (cax, ax) = plt.subplots(nrows=2,figsize=(15,10), gridspec_kw={"height_ratios":[.05, 1]})

# --- Creating the  colour map
# Use the coolwarm colour map, which comes with Matplotlib.
cmap = plt.cm.coolwarm

# Normalise the log-deaths into a 0, 1 scale
norm = colors.Normalize(vmin=deaths.min(), vmax=deaths.max())

# Now the ith colour in colours corresponds to the ith country in deaths.
colours = cmap(norm(deaths))

# Build a simple colour dictionary indexed by country name.
colour_dict = dict(zip(deaths.index, colours))


# --- Plotting the segments
# Plot the segments as a stacked bar at the x coordinate.
def plot_segments(ax, x, segments, width=0.7, min_y=0, colour_dict=colour_dict):
    current_y = min_y
    for segment in segments:
        
        # Plot a single bar at x with height=segment and its base at current_y
        # Colour th ebar based on the colour code for x
        ax.bar([x], [segment], width, bottom=current_y, color=colour_dict[x])
        
        # Update current_y so that the next segment is plotted on top.
        current_y += segment
        
# Construct the bar chart - plot each of the country's segments using apply.
weeklies.apply(lambda segments: plot_segments(ax, segments.name, segments), axis=1)


# --- Setup the axes (labels and limits)
# Rotate the x-axis labels for a better fit.
[label.set_rotation(90) for label in ax.get_xticklabels()]

ax.set_xlim(-1, len(weeklies))
ax.set_ylabel('Total Lockdown Mobility Drop')


# --- Adding the colour bar as a key
colour_scale = plt.cm.ScalarMappable(cmap=cmap, norm=norm)  # Create the colour scale
cbar = plt.colorbar(colour_scale, cax=cax, orientation="horizontal")  # Add the bar

# Add the colour bar title
cbar.ax.set_title('Deaths/1M')

# Transform the colour bar ticks: since the deaths are in log base 10 right now, 
# let's transform them back into deaths per million.
ticks = [int(10**n) for n in cbar.get_ticks()]
cbar.ax.set_xticklabels(ticks)

fig.tight_layout()