matplotlib условный цвет фона в python

Question

Как изменить цвет фона линейной диаграммы на основе переменной, которой нет в диаграмме? Например, если у меня есть следующий dataframe:

import numpy as np 
import pandas as pd 

dates = pd.date_range('20000101', periods=800) 
df = pd.DataFrame(index=dates) 
df['A'] = np.cumsum(np.random.randn(800)) 
df['B'] = np.random.randint(-1,2,size=800) 

Если бы я сделать линейный график df.A, как я могу изменить цвет фона на основе значений столбца «B» в тот момент?

Например, если B = 1 в эту дату, тогда фон в эту дату будет зеленым.

Если B = 0, то фон этой даты должен быть желтым.

Если B = -1, то фон этой даты должен быть красным.

Добавление обходного пути, которое я изначально собирался делать с помощью axvline, но ответ @jakevdp - это то, что именно искали, потому что нет необходимости в циклах: Сначала нужно добавить столбец «i» в качестве счетчика, а затем весь код выглядит следующим образом:

dates = pd.date_range('20000101', periods=800) 
df = pd.DataFrame(index=dates) 
df['A'] = np.cumsum(np.random.randn(800)) 
df['B'] = np.random.randint(-1,2,size=800) 
df['i'] = range(1,801) 

# getting the row where those values are true wit the 'i' value 
zeros = df[df['B']== 0]['i'] 
pos_1 = df[df['B']==1]['i'] 
neg_1 = df[df['B']==-1]['i'] 

ax = df.A.plot() 

for x in zeros: 
    ax.axvline(df.index[x], color='y',linewidth=5,alpha=0.03) 
for x in pos_1: 
    ax.axvline(df.index[x], color='g',linewidth=5,alpha=0.03) 
for x in neg_1: 
    ax.axvline(df.index[x], color='r',linewidth=5,alpha=0.03) 

Answer 1

Вы можете сделать это с помощью команды участка с последующим pcolor() или pcolorfast(). Например, с помощью данных можно определить выше:

ax = df['A'].plot() 
ax.pcolorfast(ax.get_xlim(), ax.get_ylim(), 
       df['B'].values[np.newaxis], 
       cmap='RdYlGn', alpha=0.3)