如何基于DataFrame为折线图分段添加背景颜色

问题描述 #

需要绘制带折线的图表，并根据color_var的值（1为红色、0为白色）为背景添加分段颜色，使用的DataFrame如下：

line_var color_var
datetime        
2023-04-20 13:45    3     1
2023-04-20 14:00    4     0
2023-04-20 15:00    5     0
2023-04-20 15:15    4     1
2023-04-20 15:45    3     1
2023-04-20 16:00    5     1
2023-04-20 16:45    6     0
2023-04-20 17:45    7     0
2023-04-20 18:00    5     0
2023-04-20 18:45    6     1
2023-04-20 19:45    8     1
2023-04-21 13:45    9     0

尝试代码后抛出维度不兼容错误：

TypeError: Dimensions of C (12, 1) are incompatible with X (12) and/or Y (2); see help(pcolor)

错误原因 #

pcolor要求输入的C数组维度必须与X、Y定义的网格匹配：

• X和Y是网格的节点坐标，长度分别为Nx和Ny
• C必须是形状为(Ny-1, Nx-1)的数组，对应每个网格单元的颜色

原代码中X是12个时间点，Y是2个y轴边界，C是(12,1)，不符合(1,12)的要求，导致维度不匹配。

解决代码 #

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

df = pd.DataFrame({
    'datetime':['2023-04-20 13:45','2023-04-20 14:00','2023-04-20 15:00',
                '2023-04-20 15:15','2023-04-20 15:45','2023-04-20 16:00',
                '2023-04-20 16:45','2023-04-20 17:45','2023-04-20 18:00',
                '2023-04-20 18:45','2023-04-20 19:45','2023-04-21 13:45'],
    'line_var':[3,4,5,4,3,5,6,7,5,6,8,9], 
    'color_var':[1,0,0,1,1,1,0,0,0,1,1,0]
})

df = df.assign(datetime=pd.to_datetime(df.datetime))
df = df.set_index('datetime')

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

fig = plt.figure(figsize=(12,6))
ax = fig.add_subplot()

# 绘制折线
ax.plot(df['line_var'], color='blue')
ax.set_xlabel('')

# 设置时间轴格式
plt.xticks(rotation=30)
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))

# 转换时间为matplotlib数值格式
x_num = mdates.date2num(df.index)
# 创建X边界：每个color_var对应两个时间点之间的区间，需要13个边界点
x_bounds = np.concatenate([x_num, [x_num[-1] + (x_num[-1] - x_num[-2])]])
# 获取Y轴边界
y_bounds = ax.get_ylim()
# 调整C的形状为 (Ny-1, Nx-1) = (1,12)
c = df['color_var'].values.reshape(1, -1)

# 绘制分段背景色
ax.pcolor(x_bounds, y_bounds, c, cmap=cmap, alpha=0.4, shading='flat')

plt.axhline(y=0, color='black')
plt.tight_layout()
plt.show()

关键修改说明 #

1. 时间转换：用mdates.date2num将datetime转换为matplotlib可识别的数值类型，方便生成边界
2. 边界生成：为每个color_var值创建对应的时间区间边界，确保X边界数量比数据点多1，对应12个背景区间
3. 数组形状调整：将color_var数组重塑为(1,12)，匹配Y边界（2个点）和X边界（13个点）对应的网格单元数量(2-1,13-1)

内容的提问来源于stack exchange，提问作者corianne1234