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

Question

I have a dataframe with a datetime index:

            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

I want to make a lineplot to track A over time and colour the background of the plot red when the values of B are 1. I have implemented this code to make the graph:

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

This gives me this graph:

But the red region incorrectly starts from 2020-05-21 rather than 2020-05-24 and it doesn't end at the end date in the dataframe. How can I alter my code to fix this?

Answer 1

If you change ax.pcolorfast(ax.get_xlim(), ... by ax.pcolor(data.index, ... you get what you want. The problem with the current code is that by using ax.get_xlim() , it creates a uniform rectangular grid while your index is not uniform (dates are missing), so the coloredmeshed is not like expected. The whole thing is:

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

and you get

Answer 2

I prefer axvspan in this case, see here for more information.

This adaptation will color the areas where data.B==1 , including the potential where data.B might not be a continuous block .

With a modified dataframe data from data1.csv (added some more points that are 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()

If data.B==1 will always be "one block" you can do away with the for loop and just use something like this in its place:

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)

Regarding "why" your data does not align, @Ben.T has this solution .

UPDATE : as pointed out, the for loop could be too crude for large datasets. The following uses numpy to find the falling and rising edges of data.B and then loops on those results:

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)