使用 AI 编写 Python 流体仿真代码
前言
流体仿真代码通常涉及复杂的偏微分方程、数值方法和边界条件处理,对开发者的数学功底和编程能力都有较高要求。AI 编程助手的出现改变了这一局面——你不需要从零推导离散化格式,也不需要独自排查数值不稳定的 bug,AI 可以帮你生成代码框架、解释物理原理、定位错误并优化性能。
Claude Code 作为终端原生的 AI 编程工具,特别适合科学计算项目的开发。它能直接读取你的项目文件、运行代码、查看输出结果,形成一个完整的"编写-运行-调试"闭环。本文将分享使用 Claude Code 编写 Python 流体仿真代码的完整工作流,并通过一个 2D 流体求解器的实战案例演示具体操作。
使用 Claude Code 编写流体仿真代码的工作流
创建项目结构
在开始编码前,先用 CLAUDE.md 文件描述项目背景和目标。Claude Code 启动时会自动读取这个文件,从而理解你的项目上下文。
mkdir fluid-sim && cd fluid-sim创建 CLAUDE.md:
# 流体仿真项目
## 目标
使用 Python 实现基于格子玻尔兹曼方法(LBM)的 2D 流体仿真求解器。
## 技术栈
- Python 3.10+
- NumPy(数值计算)
- Matplotlib(可视化)
## 物理模型
- D2Q9 格子玻尔兹曼方法
- BGK 碰撞算子
- 顶盖驱动腔体流动(lid-driven cavity)
## 项目结构
- solver.py:核心求解器
- visualize.py:可视化模块
- main.py:入口脚本这样 Claude Code 在后续交互中就能理解你的物理模型选择和代码组织方式,生成的代码会更有针对性。
逐步构建仿真求解器
不要一次性让 AI 生成全部代码。科学计算项目适合分步构建:先实现核心算法,再添加边界条件,最后做可视化和优化。每一步都运行验证,确保正确后再进入下一步。
使用 Claude Code 调试和优化代码
当仿真结果不符合预期时,直接把错误信息或异常现象描述给 Claude Code。它能分析数值不稳定的原因(如松弛参数超出范围、边界条件实现错误),并给出修复建议。性能优化方面,Claude Code 可以识别 NumPy 代码中的瓶颈,建议向量化改写或使用 Numba 加速。
实战示例:使用 Claude Code 从零开始编写一个 2D 流体求解器
下面演示用 Claude Code 逐步构建一个基于 D2Q9 格子玻尔兹曼方法的顶盖驱动腔体流动求解器。
步骤1:描述需求让 Claude Code 生成初始代码框架
在 Claude Code 中输入:
请帮我创建 solver.py,实现 D2Q9 格子玻尔兹曼方法的核心数据结构:
- 定义 D2Q9 的 9 个离散速度方向和对应权重
- 实现分布函数的初始化(平衡态分布)
- 实现 BGK 碰撞步骤
- 实现流动步骤(streaming)Claude Code 会生成类似如下的代码:
import numpy as np
# D2Q9 离散速度方向
# 6 2 5
# 3 0 1
# 7 4 8
e = np.array([
[0, 0], # 0: 静止
[1, 0], # 1: 右
[0, 1], # 2: 上
[-1, 0], # 3: 左
[0, -1], # 4: 下
[1, 1], # 5: 右上
[-1, 1], # 6: 左上
[-1, -1], # 7: 左下
[1, -1], # 8: 右下
])
# D2Q9 权重
w = np.array([
4/9, # 0: 静止
1/9, 1/9, 1/9, 1/9, # 1-4: 轴向
1/36, 1/36, 1/36, 1/36 # 5-8: 对角
])
# 反方向索引(用于 bounce-back 边界条件)
opp = np.array([0, 3, 4, 1, 2, 7, 8, 5, 6])
class LBMSolver:
def __init__(self, nx: int, ny: int, tau: float):
"""
D2Q9 LBM 求解器
Args:
nx: x 方向格子数
ny: y 方向格子数
tau: 松弛时间(控制粘度,nu = (tau - 0.5) / 3)
"""
self.nx = nx
self.ny = ny
self.tau = tau
self.nu = (tau - 0.5) / 3.0
# 分布函数:shape = (9, ny, nx)
self.f = np.zeros((9, ny, nx))
# 宏观量
self.rho = np.ones((ny, nx))
self.ux = np.zeros((ny, nx))
self.uy = np.zeros((ny, nx))
# 初始化为平衡态
self._init_equilibrium()
def _init_equilibrium(self):
"""将分布函数初始化为平衡态(速度为零、密度为1)"""
for i in range(9):
self.f[i] = self._feq(i, self.rho, self.ux, self.uy)
def _feq(self, i: int, rho: np.ndarray, ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
"""计算第 i 个方向的平衡态分布函数"""
eu = e[i, 0] * ux + e[i, 1] * uy # e_i · u
usq = ux ** 2 + uy ** 2 # |u|^2
return w[i] * rho * (1.0 + 3.0 * eu + 4.5 * eu ** 2 - 1.5 * usq)
def compute_macroscopic(self):
"""从分布函数计算宏观密度和速度"""
self.rho = np.sum(self.f, axis=0)
self.ux = np.sum(self.f * e[:, 0, np.newaxis, np.newaxis], axis=0) / self.rho
self.uy = np.sum(self.f * e[:, 1, np.newaxis, np.newaxis], axis=0) / self.rho
def collide(self):
"""BGK 碰撞步骤"""
feq = np.zeros_like(self.f)
for i in range(9):
feq[i] = self._feq(i, self.rho, self.ux, self.uy)
self.f += -(1.0 / self.tau) * (self.f - feq)
def stream(self):
"""流动步骤:将分布函数沿各自方向移动一格"""
for i in range(9):
self.f[i] = np.roll(self.f[i], e[i, 0], axis=1)
self.f[i] = np.roll(self.f[i], e[i, 1], axis=0)步骤2:逐步完善求解器
初始框架缺少边界条件。继续让 Claude Code 补充:
请为 solver.py 添加顶盖驱动腔体的边界条件:
- 四面墙壁使用 bounce-back(无滑移)边界条件
- 顶盖(y=ny-1)以速度 u_lid 向右移动
- 添加 step() 方法整合碰撞、流动和边界条件Claude Code 会在类中添加边界条件处理和步进方法:
def apply_boundary(self, u_lid: float):
"""应用顶盖驱动腔体边界条件"""
# 保存流动后的分布函数(用于 bounce-back)
f_pre = self.f.copy()
# 四面墙壁 bounce-back
# 底壁 (y=0)
for i in [2, 5, 6]: # 指向y正方向的方向
self.f[i, 0, :] = f_pre[opp[i], 0, :]
# 顶壁 (y=ny-1):顶盖移动,使用修正的 bounce-back
for i in [4, 7, 8]: # 指向y负方向的方向
self.f[i, -1, :] = f_pre[opp[i], -1, :] - 2.0 * w[opp[i]] * 3.0 * e[opp[i], 0] * u_lid * self.rho[-1, :]
# 左壁 (x=0)
for i in [1, 5, 8]: # 指向x正方向的方向
self.f[i, :, 0] = f_pre[opp[i], :, 0]
# 右壁 (x=nx-1)
for i in [3, 6, 7]: # 指向x负方向的方向
self.f[i, :, -1] = f_pre[opp[i], :, -1]
def step(self, u_lid: float):
"""执行一个完整的 LBM 时间步"""
self.compute_macroscopic()
self.collide()
self.stream()
self.apply_boundary(u_lid)步骤3:可视化结果
让 Claude Code 生成可视化脚本:
请创建 visualize.py,实现以下可视化功能:
- 绘制速度场的矢量图和大小云图
- 绘制中心线速度剖面(与 Ghia 等人的基准数据对比)
- 保存为 PNG 文件生成的可视化代码:
import numpy as np
import matplotlib.pyplot as plt
def plot_velocity_magnitude(solver, save_path: str = "velocity_magnitude.png"):
"""绘制速度大小云图和矢量图"""
fig, ax = plt.subplots(figsize=(8, 8))
speed = np.sqrt(solver.ux ** 2 + solver.uy ** 2)
im = ax.imshow(speed, origin="lower", cmap="jet", extent=[0, solver.nx, 0, solver.ny])
plt.colorbar(im, ax=ax, label="|u|")
# 矢量图(稀疏采样)
skip = max(solver.nx // 20, 1)
x = np.arange(0, solver.nx, skip)
y = np.arange(0, solver.ny, skip)
X, Y = np.meshgrid(x, y)
ax.quiver(X, Y, solver.ux[::skip, ::skip], solver.uy[::skip, ::skip],
color="white", scale=None, alpha=0.7)
ax.set_title("Velocity Magnitude & Vectors")
ax.set_xlabel("x")
ax.set_ylabel("y")
plt.savefig(save_path, dpi=150, bbox_inches="tight")
plt.close()
def plot_centerline_profile(solver, u_lid: float, save_path: str = "centerline_profile.png"):
"""绘制竖直中心线上的 ux 剖面"""
fig, ax = plt.subplots(figsize=(6, 8))
mid_x = solver.nx // 2
ux_profile = solver.ux[:, mid_x] / u_lid
y_norm = np.linspace(0, 1, solver.ny)
# Ghia et al. (1982) 基准数据 (Re=100)
y_ghia = [0.0, 0.0547, 0.0625, 0.1016, 0.1719, 0.2813,
0.4531, 0.5, 0.6172, 0.7344, 0.8281, 0.8906,
0.9375, 0.9453, 1.0]
u_ghia = [0.0, -0.03717, -0.04119, -0.06436, -0.08081, -0.08420,
0.01374, 0.05702, 0.14563, 0.25812, 0.38569, 0.50267,
0.63356, 0.65590, 1.0]
ax.plot(ux_profile, y_norm, "b-", linewidth=2, label="LBM")
ax.plot(u_ghia, y_ghia, "ro", markersize=4, label="Ghia et al. (Re=100)")
ax.set_xlabel("u_x / u_lid")
ax.set_ylabel("y / L")
ax.set_title("Centerline Velocity Profile")
ax.legend()
ax.grid(True, alpha=0.3)
plt.savefig(save_path, dpi=150, bbox_inches="tight")
plt.close()然后创建主入口脚本 main.py:
from solver import LBMSolver
from visualize import plot_velocity_magnitude, plot_centerline_profile
def main():
# 参数设置
nx, ny = 100, 100
u_lid = 0.1
Re = 100
nu = u_lid * ny / Re
tau = 3.0 * nu + 0.5
print(f"网格: {nx}x{ny}, Re={Re}, tau={tau:.4f}, nu={nu:.6f}")
solver = LBMSolver(nx, ny, tau)
# 运行仿真
n_steps = 10000
for step in range(n_steps):
solver.step(u_lid)
if step % 1000 == 0:
speed = (solver.ux ** 2 + solver.uy ** 2).max() ** 0.5
print(f"Step {step}/{n_steps}, max speed = {speed:.6f}")
# 可视化
plot_velocity_magnitude(solver, "velocity_magnitude.png")
plot_centerline_profile(solver, u_lid, "centerline_profile.png")
print("结果已保存")
if __name__ == "__main__":
main()运行仿真:
python main.py步骤4:性能优化
原始实现使用 Python 循环处理碰撞和流动,速度较慢。让 Claude Code 进行优化:
solver.py 中的 collide() 和 stream() 使用了 Python 循环,性能较差。
请用 NumPy 向量化重写这两个方法,消除所有 for 循环。优化后的关键方法:
def collide(self):
"""BGK 碰撞步骤(向量化版本)"""
feq = np.zeros_like(self.f)
for i in range(9):
eu = e[i, 0] * self.ux + e[i, 1] * self.uy
usq = self.ux ** 2 + self.uy ** 2
feq[i] = w[i] * self.rho * (1.0 + 3.0 * eu + 4.5 * eu ** 2 - 1.5 * usq)
self.f += -(1.0 / self.tau) * (self.f - feq)
def stream(self):
"""流动步骤(向量化版本)"""
for i in range(9):
self.f[i] = np.roll(self.f[i], e[i, 0], axis=1)
self.f[i] = np.roll(self.f[i], e[i, 1], axis=0)如果需要进一步加速,可以让 Claude Code 添加 Numba JIT 编译:
from numba import njit
@njit
def collide_numba(f, rho, ux, uy, tau):
"""Numba 加速的碰撞步骤"""
ny, nx = rho.shape
for j in range(ny):
for i in range(nx):
for k in range(9):
eu = e[k, 0] * ux[j, i] + e[k, 1] * uy[j, i]
usq = ux[j, i] ** 2 + uy[j, i] ** 2
feq = w[k] * rho[j, i] * (1.0 + 3.0 * eu + 4.5 * eu ** 2 - 1.5 * usq)
f[k, j, i] += -(1.0 / tau) * (f[k, j, i] - feq)
return fNumba 版本通常能带来 10-50 倍的加速,使得 100x100 网格的 10000 步仿真在几秒内完成。
使用 Claude Code 的技巧
精确描述物理方程和边界条件
AI 不是万能的,模糊的描述会产生模糊的代码。与其说"帮我写一个流体仿真",不如明确指定:
- 使用什么方法(如"格子玻尔兹曼 D2Q9"而非"某种数值方法")
- 碰撞算子类型(BGK、MRT 等)
- 边界条件的具体实现方式(bounce-back、Zou-He 等)
- 雷诺数和松弛参数的关系
精确的描述让 Claude Code 一次就能生成正确的代码,减少反复修改的次数。
分步骤开发而非一次性生成全部代码
科学计算代码的正确性需要逐步验证。推荐的开发顺序:
- 先实现核心算法并验证守恒律(质量和动量守恒)
- 添加边界条件并检查边界处的物理行为
- 加入可视化并与已知基准解对比
- 最后做性能优化
每一步都运行代码、检查结果,确认无误后再进入下一步。一次性生成全部代码往往会在某一步引入错误,而定位错误的代价远大于分步验证。
让 Claude Code 解释代码原理
当 Claude Code 生成的代码你不完全理解时,直接问它。例如:
请解释顶盖边界条件中这行代码的物理含义:
self.f[i, -1, :] = f_pre[opp[i], -1, :] - 2.0 * w[opp[i]] * 3.0 * e[opp[i], 0] * u_lid * self.rho[-1, :]Claude Code 会解释 bounce-back 边界条件的推导过程、修正项的来源以及为什么使用 opp[i] 方向的权重和速度分量。理解原理比盲目使用代码更重要,尤其是当仿真结果异常时,理解原理才能有效排查问题。
从 Claude Code 辅助开发到自主开发的过渡
AI 辅助开发的目标不是永远依赖 AI,而是加速学习过程。建议的过渡路径:
- 初期:让 Claude Code 生成完整代码,逐行理解每一部分的含义
- 中期:自己写出代码框架和关键算法,让 Claude Code 补充细节和边界条件
- 后期:独立编写代码,仅在遇到疑难问题时咨询 Claude Code
当你能独立判断 Claude Code 生成的代码是否正确时,说明你已经掌握了相关知识。一个实用的检验方法:在 Claude Code 生成代码后,先不看代码,自己思考实现方案,然后对比两者。如果思路一致,说明理解到位;如果有差异,深入分析哪种方案更优以及为什么。
拓展方向
掌握 2D LBM 求解器后,可以继续探索:
- 3D 仿真:将 D2Q9 扩展为 D3Q19 或 D3Q27,模拟三维流动问题
- 多松弛时间(MRT)模型:替代 BGK 碰撞算子,提高数值稳定性
- 多相流:使用 Shan-Chen 模型或自由能方法模拟气液界面
- 热对流:耦合温度场,模拟 Rayleigh-Benard 对流
- 湍流模拟:结合大涡模拟(LES)方法处理高雷诺数流动
- GPU 加速:使用 CuPy 或 Taichi 将计算迁移到 GPU,处理更大规模的网格
每个方向都可以用同样的工作流:先在 CLAUDE.md 中描述目标和物理模型,然后分步让 Claude Code 辅助实现。
总结
AI 编程工具让流体仿真代码的开发门槛大幅降低。通过 Claude Code,你可以:
- 用自然语言描述物理模型,快速生成代码框架
- 分步构建和验证,确保每一步的正确性
- 让 AI 解释代码原理,加速对数值方法的理解
- 识别性能瓶颈并自动优化
关键在于:精确描述需求、分步验证结果、理解代码原理。AI 是加速器而非替代品,真正的工程判断力仍然需要通过实践积累。从简单的 2D 问题开始,逐步挑战更复杂的物理模型,AI 辅助开发会让你的学习曲线更加平滑。

