使用pcolormesh绘制热图时高斯曲线拟合的维度不匹配问题

问题根源 #

报错显示(300,11658)和(12000,)形状无法广播，原因是：

• X、Y是二维网格（形状(300,11658)），对应整个热图的坐标网格
• z_non_zero是一维数组（形状(12000,)），存储所有非零点的振幅值
• 直接将一维振幅与二维高斯指数项相乘，不符合numpy广播规则

修正方案 #

核心思路是为每个非零点生成独立的高斯核，再将所有高斯核累加得到最终热图数据。通过扩展维度让振幅与网格坐标支持广播计算，具体步骤如下：

1. 重写高斯计算函数 #

替换原gaussian函数为以下实现：

import numpy as np

def gaussian(X, Y, z_amps, x_centers, y_centers, sigma):
    # 将中心坐标和振幅扩展为三维数组，适配广播规则
    x_centers = x_centers[:, np.newaxis, np.newaxis]
    y_centers = y_centers[:, np.newaxis, np.newaxis]
    z_amps = z_amps[:, np.newaxis, np.newaxis]
    
    # 计算每个非零点对应的高斯核
    gaussian_kernels = z_amps * np.exp(-((X - x_centers)**2 + (Y - y_centers)**2) / (2 * sigma**2))
    
    # 累加所有高斯核，得到与原网格同形状的结果
    return np.sum(gaussian_kernels, axis=0)

2. 验证维度匹配 #

调用该函数后，gaussian_curves的形状会与X、Y完全一致（(300,11658)），可直接传入pcolormesh绘制。

3. 额外注意事项 #

• 确保non_zero_indices正确提取所有非零点坐标：np.nonzero(Z_cleaned)返回二维数组的行、列索引，对应提取的x_non_zero、y_non_zero、z_non_zero均为一维数组且长度一致
• 若高斯核叠加后结果过强，可调整sigma值或对z_non_zero做归一化处理

完整修正后代码片段 #

# 保留原数据加载、网格构建、非零点筛选代码

# 重写后的高斯函数
def gaussian(X, Y, z_amps, x_centers, y_centers, sigma):
    x_centers = x_centers[:, np.newaxis, np.newaxis]
    y_centers = y_centers[:, np.newaxis, np.newaxis]
    z_amps = z_amps[:, np.newaxis, np.newaxis]
    gaussian_kernels = z_amps * np.exp(-((X - x_centers)**2 + (Y - y_centers)**2) / (2 * sigma**2))
    return np.sum(gaussian_kernels, axis=0)

sigma = 0.2
gaussian_curves = gaussian(X, Y, z_non_zero, x_non_zero, y_non_zero, sigma)

# 绘图部分
fig, ax = plt.subplots(figsize=(5, 8))
shw = ax.pcolormesh(X, Y, gaussian_curves, shading='gouraud', norm=mcolors.PowerNorm(gamma=0.2), cmap=plt.get_cmap(cmap))
bar = plt.colorbar(shw, ticks=np.linspace(vmin, vmax, 5))
plt.show()

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