Matplotlib contourf输入数组维度错误：CFD压力云图绘制问题

我尝试绘制CFD导出数据的压力云图，步骤为：创建矩形网格→根据导入坐标过滤网格→在过滤后的网格上插值数据→调用.contourf绘制云图。最初报错TypeError: Input z must be 2D, not 1D，排查后确认原数据存在问题，修复后meshgrid+插值功能正常，但现在新问题是如何将过滤后的数据整理为2D形状。

示例代码 #

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

# Read data
input_filename = 'period_rot_shadow_vel-mag_p-stat.txt'
df_results = pd.read_csv(input_filename,
                         sep =',', skiprows = 1,
                         names = ['nodenumber', 'X', 'Y', 'Z','p_stat', 'vel_mag'])

#get the size of the data to generate the same size array
x_data_min = df_results['X'].min()
x_data_max = df_results['X'].max()
y_data_min = df_results['Y'].min()
y_data_max = df_results['Y'].max()

#from max and min values calculate the range (size) for both axis
x_size = np.abs(x_data_max) + np.abs(x_data_min)
y_size = np.abs(y_data_max) + np.abs(y_data_min)

rect_mesh_spacing = 0.001

#number of points on each edge is calculated from the x_size and spacing, +1 is added so that we always
#cover the whole area covered in exported data
N_x_steps = int(x_size/rect_mesh_spacing) + 1
N_y_steps = int(y_size/rect_mesh_spacing) + 1

#create an empty list to fill with the values
rect_mesh_x_axis = np.array([])
rect_mesh_y_axis = np.array([])

#for loop to fill the list of evenly spaced value in the x axis
for i in range(0, N_x_steps):
    x = x_data_min + i * rect_mesh_spacing
    rect_mesh_x_axis = np.append(rect_mesh_x_axis, x)

for i in range(0, N_y_steps):
    y = y_data_min + i * rect_mesh_spacing
    rect_mesh_y_axis = np.append(rect_mesh_y_axis, y)

#use meshgrid function to create a rectangular grid, meshgrid returns a tuple of ndarrays
xx, yy = np.meshgrid(rect_mesh_x_axis, rect_mesh_y_axis)

#create an array from meshgrid points
xy_mesh = np.array((xx,yy)).T

#define input coordinates (imporeted data point coordinates) as numpy arrays
input_coords = np.array(list(zip(df_results['X'].to_numpy(), df_results['Y'].to_numpy())))

# scipy.spatial.KDTree function is used to lookup nearest-neighbour for defined points
tree = scipy.spatial.KDTree(input_coords)

dist, idx = tree.query(xy_mesh)
#define the radius of lookup of vicinity of points
radius = (rect_mesh_spacing + rect_mesh_spacing) / 2
distance_filter = dist <= radius

#filter out the points in rectangular grid that are close to the exported data points
xx_filtered, yy_filtered = xy_mesh[distance_filter].T

# scipy.interpolate.griddata(points, values, xi)
p_stat_interpol= scipy.interpolate.griddata((df_results['X'], df_results['Y']), df_results['p_stat'],
                                               (xx_filtered, yy_filtered), method = 'nearest')
p_stat_interpol2= scipy.interpolate.griddata((df_results['X'], df_results['Y']), df_results['p_stat'],
                                               (xx, yy), method = 'nearest')

# Contours are plotted in matplotlib with contourf function .contourf(X, Y, Z, levels...)
fig, axs = plt.subplots(1,2)
fig.set_size_inches(13, 4)

# Define number of contour levels
contour_levels = 20

# Plot non-interpolated data
axs[0].contourf(df_results['X'], df_results['Y'], df_results['p_stat'], contour_levels,  cmap='RdGy')

# Plot interpolated data
axs[1].contourf(xx, yy, p_stat_interpol2, contour_levels, cmap='RdGy')

plt.show()

数据样本 #

nodenumber,    x-coordinate,    y-coordinate,    z-coordinate,        pressure,velocity-magnitude
1,-2.745435307E-02, 2.104994697E-03,-2.185633883E-02,-8.094133146E+03, 4.686428765E+01
2,-2.738908254E-02, 3.158549336E-03,-2.185629997E-02,-8.104864467E+03, 4.606212337E+01
3,-2.757460451E-02, 5.167262860E-03,-2.185623337E-02,-8.093051020E+03, 4.495193692E+01
4,-2.733931890E-02, 5.656880572E-03,-2.185622490E-02,-8.116433300E+03, 4.405860916E+01
5,-2.759126430E-02, 6.070293390E-03,-2.185618260E-02,-8.084872243E+03, 4.428866265E+01
6,-2.737704207E-02, 6.507508463E-03,-2.185627796E-02,-8.106663765E+03, 4.361527022E+01
7,-2.762655064E-02, 0.000000000E+00,-2.185611682E-02,-8.057108514E+03, 7.720657592E+01
8,-2.762655053E-02, 1.920716437E-05,-2.185609461E-02,-8.057029042E+03, 7.591235710E+01
9,-2.762654491E-02, 4.228439350E-05,-2.185606795E-02,-8.057013286E+03, 7.368296305E+01
10,-2.762653069E-02, 6.997596792E-05,-2.185603598E-02,-8.057068875E+03, 7.113144845E+01
11,-2.762650289E-02, 1.032044066E-04,-2.185599766E-02,-8.057152704E+03, 6.841504839E+01
12,-2.762645410E-02, 1.430766767E-04,-2.185595174E-02,-8.057232638E+03, 6.565721945E+01
13,-2.740227795E-02, 1.197447086E-03,-2.185604704E-02,-8.086002417E+03, 4.769658095E+01
14,-2.729995772E-02, 4.555508762E-03,-2.185613769E-02,-8.113811400E+03, 4.473595941E+01
15,-2.762637329E-02, 1.909744910E-04,-2.185589666E-02,-8.057315245E+03, 6.283751408E+01
16,-2.762624634E-02, 2.484489809E-04,-2.185583068E-02,-8.057450247E+03, 5.993119264E+01
17,-2.762465252E-02, 6.184882281E-04,-2.185591455E-02,-8.059810059E+03, 5.089516358E+01
18,-2.762605091E-02, 3.174151315E-04,-2.185575167E-02,-8.057686244E+03, 5.707278224E+01
19,-2.762531127E-02, 4.994699682E-04,-2.185577109E-02,-8.058658000E+03, 5.248349769E+01
20,-2.762366891E-02, 7.612957389E-04,-2.185583175E-02,-8.062030940E+03, 4.956215663E+01

解决方案：恢复过滤后数据的2D网格结构 #

过滤操作会把原始2D网格的点打散成1D数组，导致无法直接用于contourf。解决核心是利用过滤时生成的布尔掩码，重新构建带无效区域标记的2D数组：

具体实现 #

在现有代码的插值步骤之后，添加以下代码：

# 初始化和原始网格同形状的2D数组，用nan填充无效区域
p_stat_interpol_2d = np.full(xx.shape, np.nan)
# 将1D插值数据填充到掩码标记的有效位置
p_stat_interpol_2d[distance_filter] = p_stat_interpol

# 绘制过滤后的压力云图
fig, ax = plt.subplots(figsize=(8,6))
ax.contourf(xx, yy, p_stat_interpol_2d, contour_levels, cmap='RdGy')
ax.set_xlabel('X坐标')
ax.set_ylabel('Y坐标')
plt.colorbar(ax.collections[0], label='静压')
plt.show()

原理说明 #

• distance_filter是和原始2D网格xx、yy完全同形状的布尔数组，每个元素标记对应网格点是否有效
• 用np.full创建和网格同形状的数组，nan值会被Matplotlib自动忽略，不会在云图中显示
• 通过掩码索引直接将1D插值数据放回对应的2D网格位置，完美匹配contourf对2D输入的要求

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