求助：拆分大型3D structured array并绘制多组colormap的实现方法

问题说明 #

现有一个包含48400000个元素的3D NumPy数组，实际结构为1000组独立的484×10二维数组，需要为每组数据单独绘制热力图（colormap）。

原代码问题分析 #

• np.split(x, [4840])仅将数组分割为前后两部分，无法得到1000组目标子数组
• reshape(y, 4840*10)破坏了每组484×10的二维结构，无法用于热力图绘制

正确实现方案 #

步骤1：数据加载与形状校准 #

先加载数据并将原始数组重塑为(1000, 484, 10)的三维结构，确保每组对应一个484×10的二维数组。

步骤2：批量绘制热力图 #

通过循环遍历每组数据，为每组单独创建热力图，可选择逐个展示或批量保存。

完整代码 #

import numpy as np
import matplotlib.pyplot as plt

# 加载数据
data = np.load('C:/Users/**/***.npz', allow_pickle=True)
weights = data['weights']

# 重塑为1000组484×10的结构
# 先确认原始形状是否匹配：weights.size 应该等于 1000*484*10=48400000
assert weights.size == 48400000, "数据规模不匹配目标结构"
grouped_data = weights.reshape(1000, 484, 10)

# 循环绘制每组的热力图
for i in range(1000):
    plt.figure(figsize=(8, 12))  # 根据484×10的比例调整画布大小
    plt.imshow(grouped_data[i], cmap='viridis')  # 可替换其他colormap，比如'jet'
    plt.title(f"第{i+1}组数据热力图")
    plt.colorbar(label='数值')
    plt.xlabel("X轴（10列）")
    plt.ylabel("Y轴（484行）")
    
    # 可选：保存图片到本地
    # plt.savefig(f"group_{i+1}_heatmap.png", dpi=100, bbox_inches='tight')
    
    plt.show()
    plt.close()  # 关闭画布释放内存

补充说明 #

• 若不需要逐个展示，可注释掉plt.show()，仅保留plt.savefig()实现批量导出
• cmap参数可替换为Matplotlib支持的任意配色方案，比如'plasma'、'coolwarm'等
• 画布大小figsize可根据需求调整，适配484行×10列的长条形结构

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