如何更改感兴趣区域的 pyplot 背景颜色？

Question

我有一个带有日期时间索引的数据框：

            A  B
date            
2020-05-04  0  0
2020-05-05  5  0
2020-05-07  2  0
2020-05-09  2  0
2020-05-18 -5  0
2020-05-19 -1  0
2020-05-20  0  0
2020-05-21  1  0
2020-05-22  0  0
2020-05-23  3  0
2020-05-24  1  1
2020-05-25  0  1
2020-05-26  4  1
2020-05-27  3  1

我想制作一个线图来随着时间的推移跟踪 A，并在 B 的值为 1 时将绘图的背景涂成红色。我已经实现了以下代码来制作图表：

from matplotlib import dates as mdates
from matplotlib.colors import ListedColormap
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

cmap = ListedColormap(['white','red'])

ax.plot(data['A'])
ax.set_xlabel('')
plt.xticks(rotation = 30)
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
ax.pcolorfast(ax.get_xlim(), ax.get_ylim(),
              data['B'].values[np.newaxis],
              cmap = cmap, alpha = 0.4)
plt.axhline(y = 0, color = 'black')
plt.tight_layout()

这给了我这个图表：

但红色区域错误地从 2020-05-21 而不是 2020-05-24 开始，并且它没有在数据框中的结束日期结束。如何更改我的代码来解决这个问题？

Answer 1

如果您将ax.pcolorfast(ax.get_xlim(), ... 更改为ax.pcolor(data.index, ...，您将得到您想要的。当前代码的问题在于，通过使用ax.get_xlim()，它会创建一个统一的矩形网格，而您的索引不是统一的（缺少日期），因此彩色网格并不像预期的那样。整件事是：

from matplotlib import dates as mdates
from matplotlib.colors import ListedColormap
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

cmap = ListedColormap(['white','red'])

fig = plt.figure()
ax = fig.add_subplot()

ax.plot(data['A'])
ax.set_xlabel('')

plt.xticks(rotation = 30)
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
#here are the two changes use pcolor
ax.pcolor(data.index, #use data.index to create the proper grid
          ax.get_ylim(),
          data['B'].values[np.newaxis], 
          cmap = cmap, alpha = 0.4, 
          linewidth=0, antialiased=True)
plt.axhline(y = 0, color = 'black')
plt.tight_layout()

然后你得到

Answer 2

在这种情况下，我更喜欢axvspan，有关详细信息，请参阅here。

此调整将为data.B==1 所在的区域着色，包括data.B 可能不是连续块的潜在区域。

使用来自 data1.csv 的修改后的数据框 data（添加了更多点，即 1）：

date        A   B
5/4/2020    0   0
5/5/2020    5   0
5/7/2020    2   1
5/9/2020    2   1
5/18/2020   -5  0
5/19/2020   -1  0
5/20/2020   0   0
5/21/2020   1   0
5/22/2020   0   0
5/23/2020   3   0
5/24/2020   1   1
5/25/2020   0   1
5/26/2020   4   1
5/27/2020   3   1

from matplotlib import dates as mdates
import pandas as pd
import matplotlib.pyplot as plt

data = pd.read_csv('data1.csv',index_col='date')
data.index = pd.to_datetime(data.index)

fig = plt.figure()
ax = fig.add_subplot()

ax.plot(data['A'])
plt.xticks(rotation = 30)
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
plt.axhline(y = 0, color = 'black')

# in this case I'm looking for a pair of ones to determine where to color
for i in range(1,len(data.B)):
    if data.B[i]==True and data.B[i-1]==True:
        plt.axvspan(data.index[i-1], data.index[i], color='r', alpha=0.4, lw=0)

plt.tight_layout()

如果data.B==1 将始终是“一个块”，您可以取消for 循环并使用类似的东西代替它：

first = min(idx for idx, val in enumerate(data.B) if val == 1) 
last = max(idx for idx, val in enumerate(data.B) if val == 1) 

plt.axvspan(data.index[first], data.index[last], color='r', alpha=0.4, lw=0)

关于“为什么”您的数据不对齐，@Ben.T has this solution。

更新：正如所指出的，for 循环对于大型数据集可能过于粗糙。下面使用 numpy 查找data.B 的下降沿和上升沿，然后循环这些结果：

import numpy as np
diffB = np.append([0], np.diff(data.B))
up = np.where(diffB == 1)[0]
dn = np.where(diffB == -1)[0]

if diffB[np.argmax(diffB!=0)]==-1:
    # we have a falling edge before rising edge, must have started 'up'
    up = np.append([0], up)

if diffB[len(diffB) - np.argmax(diffB[::-1]) - 1]==1:
    # we have a rising edge that never fell, force it 'dn'
    dn = np.append(dn, [len(data.B)-1])

for i in range(len(up)):
    plt.axvspan(data.index[up[i]], data.index[dn[i]], color='r', alpha=0.4, lw=0)