基于黑色曲线值为Matplotlib子图分区域设置背景色

需求说明 #

我编写了一段创建两个双轴子图的Python代码，希望根据黑色曲线（event序列）的数值为子图背景分区域上色，规则如下：

• 黑色曲线值为1的区间前1个索引段：蓝色
• 黑色曲线值为1的区间本身：红色
• 黑色曲线值为1的区间后2个索引段：绿色

示例区间对应：

• 索引[9:10]（蓝色）→ [10:14]（红色）→ [14:16]（绿色）
• 索引[23:24]（蓝色）→ [24:35]（红色）→ [35:37]（绿色）

修改后的完整代码 #

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# 生成event序列
event = pd.DataFrame(np.random.randint(1, size=(56, 1)))
event.iloc[10:14, :] = 1
event.iloc[24:36, :] = 1

# 提取event中连续为1的区间（起始和结束索引）
def get_event_intervals(event_series):
    arr = event_series.values.flatten()
    # 计算差分定位区间起止点
    diff = np.diff(np.concatenate([[0], arr, [0]]))
    starts = np.where(diff == 1)[0]
    ends = np.where(diff == -1)[0]
    return list(zip(starts, ends))

event_intervals = get_event_intervals(event)

# 给子图添加背景色的工具函数
def add_event_background(ax, intervals, total_length):
    color_map = {
        'pre': '#87CEEB',    # 浅蓝色
        'event': '#FF6347',  # 浅红色
        'post': '#90EE90'    # 浅绿色
    }
    for start, end in intervals:
        # 处理前导1段区间，避免索引越界
        pre_start = max(start - 1, 0)
        ax.axvspan(pre_start, start, facecolor=color_map['pre'], alpha=0.3, zorder=0)
        # 处理event本身区间
        ax.axvspan(start, end, facecolor=color_map['event'], alpha=0.3, zorder=0)
        # 处理后续2段区间，避免索引越界
        post_end = min(end + 2, total_length)
        ax.axvspan(end, post_end, facecolor=color_map['post'], alpha=0.3, zorder=0)

plt.figure(figsize=(38, 16))

# 第一个双轴子图
ax1 = plt.subplot(2, 2, 1)
ax2 = ax1.twinx()  

# 绘制数据曲线
ax1.plot(pd.DataFrame(np.random.randint(300, size=(56,1))), label='1', color='g')
ax1.plot(pd.DataFrame(np.random.randint(3, size=(56,1))), label='2', color='r')

# 添加背景色块
add_event_background(ax1, event_intervals, len(event))

# 子图样式设置
ax1.set_ylabel('m', fontsize=20)
ax1.legend(loc=1, prop={'size': 20})
ax1.set_title('0', fontsize=30)
ax1.tick_params(axis='y', labelsize=20)
ax1.tick_params(axis='x', labelsize=20) 

# 绘制event曲线
ax2.plot(event, color='k')
ax2.set_ylabel('t', color='k', fontsize=20)
ax2.tick_params(axis='y', labelsize=20)

# 第二个双轴子图
ax3 = plt.subplot(2, 2, 2)
ax4 = ax3.twinx()  

# 绘制数据曲线
ax3.plot(pd.DataFrame(np.random.randint(300, size=(56,1))), label='1', color='g')
ax3.plot(pd.DataFrame(np.random.randint(400, size=(56,1))), label='2', color='r')

# 添加背景色块
add_event_background(ax3, event_intervals, len(event))

# 子图样式设置
ax3.set_ylabel('m', fontsize=30)
ax3.legend(loc=1, prop={'size': 20})
ax3.set_title('1', fontsize=30)
ax3.tick_params(axis='y', labelsize=20) 
ax3.tick_params(axis='x', labelsize=20) 

# 绘制event曲线
ax4.plot(event, color='k')
ax4.set_ylabel('t', color='k', fontsize=20)
ax4.tick_params(axis='y', labelsize=20)

plt.show()

关键逻辑说明 #

1. 提取连续区间：通过numpy差分计算，自动定位event序列中所有连续为1的区间起止索引，无需手动指定区间
2. 背景色绘制：使用axvspan绘制垂直背景块，通过zorder=0确保色块位于曲线下方，不会遮挡数据；同时对区间做边界检查，避免索引越界
3. 复用性：将背景色绘制逻辑封装为函数，可快速应用到多个子图上

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