python中的matplotlib条件背景颜色

Question

如何根据图表中没有的变量更改折线图的背景颜色？例如,如果我有以下数据帧:

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回答正是因为不需要for循环,所以首先需要添加一个'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)