如何用Matplotlib自动识别数组维度绘制三列数据的密度等高线图

问题背景 #

拥有一份三列格式（x坐标、y坐标、密度值）的数据集，示例数据如下：

0 0 1.9
0 1 1.92
0 2 1.93
0 3 1.98
0 4 1.89

1 0 1.8
1 1 1.83
1 2 1.79
1 3 1.8
1 4 1.86

2 0 1.5
2 1 1.6
2 2 1.4
2 3 1.55
2 4 1.62

3 0 1.4
3 1 1.35
3 2 1.36
3 3 1.34
3 4 1.3

4 0 1.1
4 1 1.2
4 2 1.15
4 3 1.05
4 4 1.06

现有Gnuplot脚本可自动识别数据集维度（支持5x5、128x128等任意规模），但Matplotlib脚本需手动指定reshape的尺寸，无法通用。

通用化解决方案 #

修改Matplotlib脚本，通过提取x、y的唯一值数量自动计算数据集维度，无需手动指定：

import matplotlib.pyplot as plt
import numpy as np

# 加载三列数据
data = np.loadtxt("n2d.dat", usecols=range(3), dtype=np.float32)
x = data[:, 0]
y = data[:, 1]
z = data[:, 2]

# 自动获取数据集的维度：x的唯一值数量为行数，y的唯一值数量为列数
nx = len(np.unique(x))
ny = len(np.unique(y))

# 重塑密度数组为2D格式，并转置以匹配坐标映射
datamod = z.reshape(nx, ny).T

# 绘制密度图，保持原脚本的视觉配置
density = plt.imshow(datamod, cmap='viridis', origin='lower', interpolation='none')
plt.colorbar(density)
plt.savefig("density.png", dpi=300)

关键逻辑说明 #

• 自动计算维度：利用np.unique()获取x轴和y轴的唯一值数量，这两个值就是数据集的行列维度（nx × ny），适配任意规模的规整数据集。
• 数组重塑与转置：原始数据按「固定x、遍历y」的顺序排列，z.reshape(nx, ny)将一维密度数组按x分组，转置.T后匹配imshow(origin='lower')的坐标显示逻辑，确保x对应横轴、y对应纵轴，和Gnuplot的显示效果一致。
• 兼容性：只要数据集保持「x固定时y连续递增、x按顺序遍历」的格式，脚本就能自动适配5x5、128x128等任意规模的输入。

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