Как изменить цвет фона pyplot в интересующей вас области?

Question

У меня есть фрейм данных с индексом datetime:

            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

Я предпочитаю axvspan в этом случае, см. здесь для получения дополнительной информации.

Эта адаптация закрасит области, где data.B==1, , включая потенциал, где data.B может быть не непрерывным блоком .

С измененным фреймом данных data из data1.csv (добавлено еще несколько точек, равных 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 l oop и просто использовать вместо него что-то вроде этого:

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)

Что касается «почему» ваши данные not align, @Ben.T имеет это решение .

UPDATE : как уже указывалось, for l oop может быть слишком грубым для больших наборов данных . Следующее использует 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)

Answer 2

Если вы замените ax.pcolorfast(ax.get_xlim(), ... на ax.pcolor(data.index, ..., вы получите то, что хотите. Проблема с текущим кодом заключается в том, что при использовании ax.get_xlim() он создает однородный прямоугольник angular сетку , в то время как ваш индекс не является однородным (даты отсутствуют), поэтому цветная сетка не такая, как ожидалось. Все это:

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()

и вы получите