如何用VTK与Python生成场时间步快照及OpenFOAM结果后处理

阿华AIGC实验室

2026-5-14

解答：OpenFOAM后处理与VTK生成场快照

我来帮你搞定这两个需求，先从你关心的OpenFOAM cavity算例后处理说起，再讲VTK生成单时间步快照的方法~

一、处理OpenFOAM结果：生成压力场numpy数组与热图

针对你的cavity算例（v1812版本，时间步目录为0, 0.1, ..., 0.5），我们可以用VTK+NumPy+Matplotlib的组合来实现需求，不需要额外安装OpenFOAM专属的Python库（当然你也可以用pyFoam，但VTK的兼容性更强）。

1. 读取压力场并转换为NumPy数组

下面的代码会读取指定时间步（比如0.5）的压力场，转换成可被NumPy处理的二维数组：

import vtk
from vtk.util.numpy_support import vtk_to_numpy
import numpy as np

# 初始化OpenFOAM读取器
reader = vtk.vtkOpenFOAMReader()
reader.SetFileName("./system/controlDict")  # 指向算例的controlDict
reader.UpdateInformation()

# 设置要读取的时间步和场
reader.SetTimeValue(0.5)  # 指定时间步，也可以用SetTimeStepIndex(index)
reader.SetPatchArrayStatus("p", 1)  # 启用压力场"p"
reader.Update()

# 获取数据集并提取压力数据
data = reader.GetOutput()
pressure_vtk = data.GetPointData().GetArray("p")
pressure_np = vtk_to_numpy(pressure_vtk)

# 因为cavity是二维算例，把一维的点数据reshape成二维矩阵
# 先获取网格的维度（假设是结构化网格）
dims = data.GetDimensions()
pressure_matrix = pressure_np.reshape(dims[1], dims[0])  # 维度顺序可能需要根据实际情况调整

2. 生成压力场热图

有了NumPy数组后，用Matplotlib就能轻松画出类似示例的热图：

import matplotlib.pyplot as plt
import seaborn as sns

plt.figure(figsize=(8,6))
# 用seaborn画热图，样式更美观
sns.heatmap(pressure_matrix, cmap="viridis", annot=False, cbar=True)
plt.title("Pressure Field at Time Step 0.5")
plt.xlabel("X Direction")
plt.ylabel("Y Direction")
plt.savefig("pressure_heatmap_0.5.png", dpi=300)
plt.show()

如果是三维算例，只需要调整reshape的维度，或者用matplotlib的3D可视化工具即可。

二、用VTK+Python生成单时间步场的可视化快照

如果想要直接用VTK生成场的专业可视化快照（带网格、自定义颜色映射的渲染图），可以按照下面的步骤来：

完整示例代码

import vtk

# 1. 读取OpenFOAM数据（和之前的读取逻辑一致）
reader = vtk.vtkOpenFOAMReader()
reader.SetFileName("./system/controlDict")
reader.UpdateInformation()
reader.SetTimeValue(0.5)
reader.SetPatchArrayStatus("p", 1)
reader.Update()

# 2. 创建映射器和演员，设置颜色映射
mapper = vtk.vtkDataSetMapper()
mapper.SetInputConnection(reader.GetOutputPort())
mapper.SetScalarModeToUsePointData()
mapper.SelectColorArray("p")
mapper.SetScalarRange(reader.GetOutput().GetPointData().GetArray("p").GetRange())

# 自定义颜色映射表（可选，替换默认配色）
color_map = vtk.vtkColorTransferFunction()
color_map.AddRGBPoint(-0.005, 0.23, 0.30, 0.75)
color_map.AddRGBPoint(0, 0.86, 0.86, 0.86)
color_map.AddRGBPoint(0.005, 0.70, 0.02, 0.15)
mapper.SetLookupTable(color_map)

actor = vtk.vtkActor()
actor.SetMapper(mapper)

# 3. 设置渲染器、窗口
renderer = vtk.vtkRenderer()
render_window = vtk.vtkRenderWindow()
render_window.AddRenderer(renderer)
render_window.SetSize(800, 600)

# 添加演员到渲染器，设置背景
renderer.AddActor(actor)
renderer.SetBackground(1, 1, 1)  # 设置背景为白色

# 调整相机视角（针对二维cavity算例）
camera = renderer.GetActiveCamera()
camera.SetPosition(0, 0, 5)  # 从Z轴正方向俯视
camera.SetFocalPoint(0, 0, 0)
renderer.ResetCamera()

# 4. 生成快照并保存为PNG
window_to_image_filter = vtk.vtkWindowToImageFilter()
window_to_image_filter.SetInput(render_window)
window_to_image_filter.Update()

writer = vtk.vtkPNGWriter()
writer.SetFileName("pressure_snapshot_0.5.png")
writer.SetInputConnection(window_to_image_filter.GetOutputPort())
writer.Write()

# 可选：打开交互窗口手动调整视角
# interactor = vtk.vtkRenderWindowInteractor()
# interactor.SetRenderWindow(render_window)
# interactor.Start()

这段代码会直接生成一张PNG格式的压力场可视化快照，你可以根据需求调整颜色映射、窗口大小、相机位置等参数。