优化NumPy双层循环实现自定义热力图颜色计算的性能

核心思路 #

放弃Python双层循环，改用NumPy的向量化运算——底层基于C实现的批量计算，能让10000×10000级别的数组处理速度提升几个数量级，同时严格保留你需要的配色逻辑：

• 负值区间（[-1, 0]）：红(255,0,0) → 黄(255,255,0)，即G通道从0线性升至255，R保持255，B保持0
• 正值区间（[0, 1]）：黄(255,255,0) → 绿(0,255,0)，即R通道从255线性降至0，G保持255，B保持0
• 最终输出**[B, G, R]**格式的uint8类型数组（适配OpenCV的图像格式）

实现代码 #

import numpy as np

def generate_custom_heatmap(arr):
    # 确保输入是[-1, 1]范围的二维NumPy数组
    assert arr.min() >= -1 and arr.max() <= 1, "输入数组需在[-1,1]范围内"
    
    # 初始化三个通道的数组
    height, width = arr.shape
    b = np.zeros((height, width), dtype=np.uint8)
    g = np.full((height, width), 255, dtype=np.uint8)
    r = np.full((height, width), 255, dtype=np.uint8)
    
    # 处理负值区间：G从0→255
    neg_mask = arr < 0
    # 将[-1,0]映射到[0,1]，再乘255得到G通道值
    g[neg_mask] = ((arr[neg_mask] + 1) * 255).astype(np.uint8)
    
    # 处理正值区间：R从255→0
    pos_mask = arr > 0
    # 将[0,1]映射到[1,0]，乘255得到R通道值
    r[pos_mask] = ((1 - arr[pos_mask]) * 255).astype(np.uint8)
    
    # 组合成[B, G, R]格式的三维数组（OpenCV图像格式）
    heatmap = np.stack([b, g, r], axis=-1)
    return heatmap

关键优化点 #

1. 避免Python循环：所有计算都是NumPy批量操作，跳过了Python解释器的循环开销
2. 掩码筛选：用布尔掩码直接定位正负区间，不用逐元素判断
3. 类型提前指定：直接初始化uint8类型数组，避免中途类型转换的额外开销

测试示例 #

# 生成测试用10000×10000数组
test_arr = np.random.uniform(-1, 1, (10000, 10000))
# 生成热力图
heatmap = generate_custom_heatmap(test_arr)
# 可直接用OpenCV显示/保存
# import cv2
# cv2.imwrite("custom_heatmap.png", heatmap)

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