Построить временные ряды с цветовой шкалой в пандах + матплотлиб

Question

Я пытаюсь построить цветную полосу под этой диаграммой, где цвет зависит от того, когда начинается каждый из временных рядов:

Код, сгенерированный для создания графика, выглядит так:

import pandas as pd
import numpy as np

import matplotlib.pyplot as plt

import seaborn as sns
sns.set()

def partial_cum_returns(start, cum_returns):
    return cum_returns.loc[start:].div(cum_returns.loc[start])

index = pd.DatetimeIndex(pd.date_range('20170101', '20190101', freq='W'))
np.random.seed(5)
returns = pd.Series(np.exp(np.random.normal(loc=0, scale=0.05, size=len(index))), index=index)
cum_returns = returns.cumprod()
df = pd.DataFrame(index=index)
for date in index:
    df[date] = partial_cum_returns(date, cum_returns)

df.plot(legend=False, colormap='viridis');
plt.colorbar();

Но при выполнении появляется эта ошибка:

RuntimeError: Не найдено ни одного сопоставимого элемента для создания цветовой панели. Сначала определите сопоставляемую информацию, такую ​​как изображение (с imshow) или набор контуров (с контуром).

Я пытался добавить цветовую панель по-разному, например fig, ax = plt.figure()..., но пока не мог заставить ее работать. Есть идеи? Спасибо!

Answer 1

Во-первых, вам нужно создать ScalarMappable для вашей цветовой панели. Вам необходимо определить цветовую карту, которая в вашем случае равна 'viridis', и указать максимальное и минимальное значения, которые вы хотите для цветовой панели. Тогда, потому что он использует числовые значения времени, вы хотите переформатировать их.

import matplotlib.pyplot as plt
import pandas as pd

# Define your mappable for colorbar creation
sm = plt.cm.ScalarMappable(cmap='viridis', 
                           norm=plt.Normalize(vmin=df.index.min().value,
                                              vmax=df.index.max().value))
sm._A = []  

df.plot(legend=False, colormap='viridis', figsize=(12,7));

cbar = plt.colorbar(sm);
# Change the numeric ticks into ones that match the x-axis
cbar.ax.set_yticklabels(pd.to_datetime(cbar.get_ticks()).strftime(date_format='%b %Y'))