如何基于文件读取的二元函数数据实现简洁优雅的可视化方案

原有代码存在几个可优化的问题：

• 用set获取唯一坐标会打乱顺序，容易出现网格坐标和值不匹配的bug
• 自定义匹配函数循环查找对应值，时间复杂度高，数据量较大时性能极差
• 存在大量多余的类型转换操作，代码冗余
• 绘图时交换了X和Y的传入顺序，会导致坐标轴和实际数据不匹配，优化后的代码已默认修正该问题，如需保持原有显示逻辑可自行调整传入顺序

场景1：数据为规整网格（每个(x,y)组合都有对应值） #

纯Numpy实现（依赖最少） #

import numpy as np
import matplotlib.pyplot as plt

# 直接读取文件为numpy数组，替代原列表推导式写法
data = np.loadtxt('data.txt')
x_col, y_col, z_col = data.T

# 去重并排序坐标，避免顺序混乱
x = np.unique(x_col)
y = np.unique(y_col)

# 直接重塑z列为对应维度的Z矩阵，无需循环查找
Z = z_col.reshape(len(y), len(x))

X, Y = np.meshgrid(x, y)
cp = plt.contourf(X, Y, Z)
plt.colorbar(cp)
plt.show()

Pandas实现（兼容性最好，不受数据存储顺序影响） #

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

# 直接读取分隔符分隔的文本为数据表
df = pd.read_csv('data.txt', sep='\s+', names=['x', 'y', 'z'])
# 透视表一键生成x-y对应的Z矩阵，自动匹配坐标值
Z = df.pivot_table(index='y', columns='x', values='z').values

x = df['x'].unique()
y = df['y'].unique()
X, Y = np.meshgrid(x, y)

cp = plt.contourf(X, Y, Z)
plt.colorbar(cp)
plt.show()

场景2：数据为非规整散点（存在坐标缺失） #

如果你的数据不是刚好覆盖所有网格点，可以用插值方式生成平滑的热力图：

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import griddata

data = np.loadtxt('data.txt')
x_col, y_col, z_col = data.T

# 自定义生成指定精度的网格
x = np.linspace(x_col.min(), x_col.max(), 100)
y = np.linspace(y_col.min(), y_col.max(), 100)
X, Y = np.meshgrid(x, y)

# 插值生成Z矩阵，支持linear、cubic、nearest三种插值方式
Z = griddata((x_col, y_col), z_col, (X, Y), method='cubic')

cp = plt.contourf(X, Y, Z)
plt.colorbar(cp)
plt.show()

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