Python流体模拟实战:从Navier-Stokes求解到GPU加速
流体模拟在计算机图形学里一直是个很有“话题感”的领域。不管做游戏特效、影视视效还是交互装置大家最终都希望能让屏幕上出现一团像样的烟、一汪能泛起涟漪的水或者一场不穿帮的火焰。我这次用 Python 从一个最朴素的想法出发把完整流程跑通了——从 Navier-Stokes 方程求解到密度场演化再到 GPU 加速渲染算是把“物理仿真从原型到落地”这条路从头走了一遍。这篇东西适合两类人看一类是刚接触计算机图形学和物理仿真、想搞明白流体到底怎么算的同学另一类是 Python 用得挺熟但还没想好怎么把数值求解和可视化串起来的开发者和创意工作者。1. 先想清楚流体模拟要算的到底是什么1.1 Navier-Stokes 方程没那么吓人很多人一听到 Navier-Stokes 就本能地后退一步觉得这是流体力学里最难啃的骨头。实际上如果你只是想做计算机图形学里“看起来对”的流体而不是去发论文那你要面对的方程比教科书里写的要友好得多。不可压缩流体的 Navier-Stokes 方程在动量守恒上长这样[ \frac{\partial \mathbf{u}}{\partial t} -(\mathbf{u} \cdot \nabla)\mathbf{u} - \frac{1}{\rho}\nabla p \nu \nabla^2 \mathbf{u} \mathbf{f} ]配合不可压缩条件 (\nabla \cdot \mathbf{u} 0)合起来就是流体模拟的全部物理核心。拆开看这个方程描述的就四件事速度场随时间的改变等于“自己跟着自己跑”平流加上压力带来的“推力”加上粘性造成的“摩擦”再加上外力重力、鼠标扰动、风的输入。在图形学里我们一般不会花时间算真实的物理参数。你要的是一团烟的流动感、涡旋感、被阻挡后的绕流感。所以最常走的路线是别追求真实数值追求运动规律“看起来合理”。这就是为什么 1999 年 Jos Stam 提出的 Stable Fluids 方法成了行业基石后来的实时流体模拟、游戏场景烟雾、甚至不少商业软件里的基础求解器都脱胎于这套思路。1.2 为什么第一步是“简化”而不是“求解”真正的一手 CFD计算流体力学工程里NS 方程会被丢给有限体积法、有限元法或者谱方法配合复杂的湍流模型和高精度网格去求解。但在计算机图形学里我们的目标是每秒生成几十到几千帧对精度的要求完全不同。稳定流体方法的核心思想是“算子分裂”不试图一次性地直接求解整个偏微分方程而是把方程按物理过程拆成独立步骤按顺序逐个处理。每个步骤都是成熟且数值稳定的小算子把它们拼起来整个求解器就稳定了。这在工程上有个很大的好处你可以单独替换任何一个步骤。比如把平流算子从一阶迎风改成更高精度的插值或者把压力求解器从雅可比迭代换成多重网格而不需要重构整个系统。这也是我推荐所有人入门流体模拟时从这套方法下手的原因——它把复杂的物理按逻辑边界切开了理解成本和实现成本都低得多。1.3 方案选型CPU 原型 GPU 加速的组合拳项目开始前我定了两条原则先用 Python NumPy 做 CPU 原型把每个算子的正确性验证清楚。这个阶段不追求性能追求逻辑透明、好调试。等效果稳定之后再做 GPU 加速迁移。这一步的目标是把计算压到实时或者接近实时。选择 Python 做原型的原因很朴素NumPy 向量化能省掉大半的循环书写matplotlib 一画就能看到流体形态调试时随时打印字段数据也不会心疼。更重要的是Python 的生态里有好几个能无缝衔接 GPU 的方案Numba CUDA、PyOpenCL、Taichi原型代码很容易平移过去而不是推倒重写。2. 核心细节解析稳定求解器的四个轮子2.1 空间离散把连续方程搬到网格上做网格法流体模拟第一步是把连续的物理量离散化。我用的方案是规则网格也就是把二维空间切成均匀的 (N \times N) 个小格子在每个格点上存储标量如密度、压力和矢量如速度。网格布局有两种常见选择布局方式物理量位置优点缺点共置网格所有量都在格心实现简单适合快速原型容易产生棋盘格伪影交错网格MAC压力在格心速度在面心天然避免压力-速度耦合伪影索引处理繁琐内存布局复杂我这次在 CPU 原型里选了共置网格因为重点是跑通流程在 GPU 版本里转向了 Taichi 自带的结构化数据布局本质上也接近交错网格的写法和存储。对初学者来说直接从共置网格入手完全没问题只要记住两个关键点每个格点上的速度向量代表这个位置流体的大致运动趋势标量场密度、压力、温度和速度场必须采用同样的网格尺寸和边界策略否则插值和投影阶段会错位。网格大小 N 的选择直接影响效果和性能。我实测下来二维场景 N128 时 CPU 原型还算能交互N256 以上就开始吃力。N 也不是越大越好网格细化到一定程度后如果数值耗散numerical diffusion没控制好小尺度涡旋反而会被“抹平”。2.2 平流Advection的半拉格朗日取巧平流是流体模拟里最能体现“运动的灵魂”的步骤。物理含义很简单流体里的物质密度、速度本身、温度标量顺着流场移动上一时刻在某个位置的东西下一时刻会跑到新的位置。半拉格朗日方法的思路非常反直觉但是极其有效它不做“向前追踪”——因为向前追踪粒子的位置容易算过头导致数值发散——而是反过来从当前格点出发沿着速度场回溯时间步长 (dt)找到这个格点的流体在上一时刻来自哪个位置然后把那个位置的值取回来。用一维的方式解释假设你在网格点 i 上速度为 u时间步长为 dt格间距为 dx。回溯距离就是[ x_{back} x_{current} - u \cdot dt ]然后找到 (x_{back}) 所在的网格区间用双线性插值把密度或者速度分量取回来直接作为当前格点的新值。这里有两个关键点决定了这个方法好不好用插值方法双线性插值是默认选择比最近邻插值平滑很多想保留更多涡旋细节可以尝试双三次插值但计算量会上去。越界处理回溯位置可能超出计算域。边界处一般做钳制clamping或者环绕wrapping具体取决于你是想让流体撞墙还是穿墙。半拉格朗日法的最大好处是无条件稳定无论 dt 取多大它都不会因为这一步直接数值爆炸。这句话听起来很爽但不要误会——稳定性不代表正确性dt 太大时回溯距离跨过太多格子插值会严重平滑画面就会变得“糊成一团”。2.3 扩散和投影让流体“不塌陷”的关键扩散项对应的是 Navier-Stokes 方程里的粘性项[ \nu \nabla^2 \mathbf{u} ]它的作用是让速度场里的尖锐差异比如剪切层逐渐趋缓是流体看起来“有质感和黏稠感”的关键。但显式求解扩散方程在数值上是个坑时间步长必须小于某个阈值否则高频分量会放大。Stam 的做法是用隐式离散把扩散变成一个线性方程组求解问题联合雅可比迭代来处理这样它也是无条件稳定的。扩散的实现里迭代次数是个实用指标。我一般默认设置 4 到 8 次雅可比迭代粘性系数 (\nu) 取 0.0001 附近。粘性调高到 0.001 以上时会明显感觉到流体“变稠”动两下就停住了调太低则容易出现局部速度锯齿。投影Projection步骤是整个求解器里最抽象也最关键的环节。它的任务是把速度场变成“无散度”的也就是让流体不可压缩。直观理解如果压力不介入速度场里会有局部“聚集”或“发散”的趋势流体会在某个点堆积或抽空画面表现为不自然的“蜂窝状”。为了修正这一点要解一个压力泊松方程[ \nabla^2 p \frac{\nabla \cdot \mathbf{u}}{dt} ]解出压力场 p 之后把压力梯度从速度场里减掉[ \mathbf{u} \mathbf{u} - dt \cdot \nabla p ]结果就是每个格点的速度满足 (\nabla \cdot \mathbf{u} \approx 0)。这个求解器我用的是 Gauss-Seidel 迭代或者雅可比迭代迭代次数从 20 到 50 不等。注意投影迭代如果次数不足会表现出“软绵绵”的压缩效应迭代次数太足当然好但计算时间线性增长。批量测试时记住验证一条铁律投影做完之后速度场的散度应该接近零。打印一下 max(abs(div))如果比平流前的值小了 2 到 3 个数量级说明投影是健康的。2.4 边界条件与数值稳定性边界条件在流体模拟里的影响远比很多人想象的要大。对于二维烟模拟最常用的边界是自由边界流体可以自由离开计算域和墙边界流体在墙的垂直方向速度为 0。在 Stable Fluids 的标准实现里边界条件的处理方式是在最外层格点之外“幻影”一格把边界处的分量对称地复制过去保证速度场的法向分量为 0标量的梯度在边界处为 0。这样做的好处是自洽且容易实现。稳定性方面有一条特别实际的经验法则——CFL 条件[ dt \le \frac{dx}{u_{max}} \cdot C_{cfl} ]其中 (u_{max}) 是当前速度场最大值(C_{cfl}) 通常取 0.5 到 1.0。实际项目中我不会手算 (u_{max})而是每帧统计速度场最大值动态调整 dt然后在界面上把当前 dt 显示出来。这样既能保证稳定又能让系统在外力剧烈变化时自动放慢步子不至于直接爆成一片雪花。3. 实操记录用 NumPy 把 CPU 版跑起来3.1 核心求解器的代码骨架下面的代码是我实际项目里的简化版去掉了工程封装只保留求解逻辑。它完整展示了四个核心步骤如何按顺序拼接。import numpy as np N 128 dx 1.0 / N dt 0.01 visc 0.0001 GaussSeidel_iters 30 u np.zeros((2, N, N), dtypenp.float32) # 速度场 u[0]x分量, u[1]y分量 p np.zeros((N, N), dtypenp.float32) # 压力场 d np.zeros((N, N), dtypenp.float32) # 密度场 force np.zeros((2, N, N), dtypenp.float32) def advect(f, u, dt, dx): 半拉格朗日平流回溯-插值 N f.shape[0] f_new np.zeros_like(f) x_src np.zeros((N, N)) y_src np.zeros((N, N)) idx np.arange(N) # 格点中心坐标 回溯 x_src (idx 0.5) * dx - dt * u[0] y_src (idx 0.5) * dx - dt * u[1] # 钳制与双线性插值向量化版本 # 具体实现可拆成 i0,j0,i1,j1 双线性系数这里略 return f_new def apply_force(u, f, dt): return u f * dt def diffuse(f, visc, dt, dx, iters8): 隐式扩散 雅可比迭代 f_new f.copy() a visc * dt / (dx * dx) for _ in range(iters): f_new[1:-1, 1:-1] ( f[1:-1, 1:-1] a * ( f_new[:-2, 1:-1] f_new[2:, 1:-1] f_new[1:-1, :-2] f_new[1:-1, 2:] ) ) / (1 4 * a) return f_new def project(u, iters30): 投影解压力泊松方程去掉速度发散 N u.shape[1] div np.zeros((N, N), dtypenp.float32) p np.zeros((N, N), dtypenp.float32) # 计算散度中心差分 div[1:-1, 1:-1] ( (u[0][1:-1, 2:] - u[0][1:-1, :-2]) (u[1][2:, 1:-1] - u[1][:-2, 1:-1]) ) / (2 * dx) # Gauss-Seidel 迭代解 -div 泊松方程 for _ in range(iters): p[1:-1, 1:-1] ( p[:-2, 1:-1] p[2:, 1:-1] p[1:-1, :-2] p[1:-1, 2:] - div[1:-1, 1:-1] * dx * dx ) / 4.0 # 减压力梯度 u[0][1:-1, 1:-1] - 0.5 * (p[1:-1, 2:] - p[1:-1, :-2]) / dx u[1][1:-1, 1:-1] - 0.5 * (p[2:, 1:-1] - p[:-2, 1:-1]) / dx return u, p def step(u, d, force, visc, dt, dx): # 1. 外力 u apply_force(u, force, dt) # 2. 平流速度场先平流密度场再平流 u advect(u, u, dt, dx) d advect(d, u, dt, dx) # 3. 扩散 u diffuse(u, visc, dt, dx) # 4. 投影关键 u, _ project(u, GaussSeidel_iters) return u, d实际工程里我不会每个算子都单独写一份核心逻辑然后让读者自己拼——但这个骨架已经够你理解顺序和依赖关系了。特别注意平流的对象包括速度场自身。速度场不先平流后面的所有步骤都会在错误的速度场上进行模拟会变得僵硬。3.2 可视化与交互CPU 原型阶段我用 matplotlib 做实时显示。每次 step 之后把密度场 d 扔给imshow配合plt.pause(0.001)就能基本达到 20 到 30 FPS 的 Visual 交互效果N128 时。交互方面最常见的输入方式是在鼠标位置注入力和密度# 鼠标拖拽处加扰动def add_source_at_xy(d, u, x, y, strength): i int(x / dx) j int(y / dx) d[j-2:j3, i-2:i3] strength u[0, j-1:j2, i-1:i2] strength * 0.5 u[1, j-1:j2, i-1:i2] strength * 0.5这种“在固定半径内注入”的方式特别直观能立刻看到涡旋生成。密度注入的 strength 我常用 5000 到 20000速度注入则小得多大约 0.5 到 2.0。这里没有标准答案跟你用的速度尺度密切相关测几组数据就熟了。matplotlib 虽然方便但它是 CPU 光栅化的N 稍微上去就拖后腿。如果你追求更顺滑的观看体验可以用PySide6或GLFW做一个简单的窗口把密度场转成纹理贴到 Quad 上性能会好很多。不过这步属于渲染优化后面 GPU 阶段我会用 Taichi 自带的 GUI 直接接管。3.3 参数怎么定CFL、迭代次数与粘性数值模拟参数是最容易“玄学化”的部分很多新手上来随便设结果不是爆炸就是糊成一团。我整理了一份自己在项目中实际使用的参数表基本能覆盖大多数入门场景参数推荐范围我的实测默认值说明网格 N64 ~ 512128越高细节越多CPU 原型不要超 256时间步长 dt由 CFL 决定0.005 ~ 0.01画面发散就减半粘性 visc0.00001 ~ 0.0010.0001太小会出锯齿太大流不动外力强度自定义5000 密度 / 1.0 速度大速度输入需要更小 dt雅可比迭代20 ~ 6030迭代不足会有压缩感扩散迭代2 ~ 106默认 6 次足够这里有个心得初学时最容易盯着物理参数调但实际上一旦模拟“看起来不对”先检查的不是粘性系数而是dt 是否满足 CFL。超过一半的“模拟爆炸”案例都是 dt 太大速度场局部异常放大。用一句话总结先保证稳定再追求细节最后才调“观感”。4. GPU 加速从“能跑”到“能玩”4.1 CPU 版本的瓶颈到底在哪Numpy 版本在 N128 时跑出 30 FPS 不算难但分辨率一旦变成 256 甚至 512性能就断崖式下跌。原因很直接平流阶段的回溯和双线性插值涉及大量的非连续内存访问NumPy 向量化在这里帮不上大忙。投影迭代是全局耦合计算每次迭代都要访问整个网格的相邻数据内存带宽很快成为瓶颈。中间产生的临时数组div、p、f_new频繁分配和回收Python 层的开销也不可忽视。CPU 版适合验证正确性但如果你想用鼠标实时搅动一团 512x512 的烟雾并且要求画面不太卡顿CPU 是撑不住的。这就是 GPU 加速登场的时机。4.2 工具选择PyOpenCL / Numba CUDA / TaichiPyOpenCL上手门槛最高要自己写 OpenCL C 代码但可控性最强适合有底层经验的开发者。Numba CUDA可以用 Python 风格写 kernel但对 CUDA 的 thread/block 模型仍然要求较深理解且迁移现有 CPU 代码并不轻松。Taichi太极我最推荐给工作流偏图形图像、想迅速出效果的人。它的编程模型完全贴合网格型模拟写起来像带类型标注的 Python底层的并行调度和内存管理交给编译器处理。我用 Taichi 的体验是原来 NumPy 版本 30 行核心逻辑迁移过去差不多 30 行核心逻辑只是把np换成ti循环部分变成ti.kernel里的显式循环。差别在于每个格子上的计算会被自动并行化不用你手动分配 block。4.3 Taichi 改造的要点示例下面是我 Taichi 版本的一个核心片段展示平流算子的 GPU 版本长什么样import taichi as ti ti.init(archti.gpu) # 自动选择 CUDA / Metal / Vulkan N 256 dx 1.0 / N u ti.vector.field(2, dtypeti.f32, shape(N, N)) # 速度场 d ti.field(dtypeti.f32, shape(N, N)) # 密度场 p ti.field(dtypeti.f32, shape(N, N)) div ti.field(dtypeti.f32, shape(N, N)) ti.kernel def advect_kernel(dst: ti.template(), src: ti.template(), u: ti.template(), dt: f32): for I in ti.grouped(src): # 回溯坐标 x_src (I.x 0.5) * dx - dt * u[I].x y_src (I.y 0.5) * dx - dt * u[I].y # 双线性插值做了边界钳制 i0 ti.max(0, ti.min(int(x_src / dx), N - 1)) i1 ti.max(0, ti.min(i0 1, N - 1)) j0 ti.max(0, ti.min(int(y_src / dx), N - 1)) j1 ti.max(0, ti.min(j0 1, N - 1)) s x_src / dx - (i0 0.5) t y_src / dx - (j0 0.5) dst[I] (1-s)*(1-t)*src[j0, i0] s*(1-t)*src[j0, i1] \ (1-s)*t*src[j1, i0] s*t*src[j1, i1]注意到ti.grouped(src)这句本身就是暗含“把 src 的所有元素打包成一组并行循环”的语法糖。Taichi 会在底层编译成 GPU kernel循环体里没有明显的 Python 函数调用开销。投影部分同理写起来几乎和 NumPy 时代的 Gauss-Seidel 一模一样ti.kernel def project_kernel(p: ti.template(), div: ti.template()): for I in ti.grouped(p): p[I] (p[I ti.Vector([-1, 0])] p[I ti.Vector([1, 0])] p[I ti.Vector([0, -1])] p[I ti.Vector([0, 1])] - div[I] * dx * dx) * 0.25这只是单次 Gauss-Seidel 迭代主循环里按固定次数调用它就行。GPU 上循环次数多时尽量把整个迭代过程放进同一个 kernel减少 kernel 启动的开销。4.4 性能对比与调参经验同一个求解器我在两张设备之间做了对比。配置如下CPU 是 Intel i7-12700HGPU 是 NVIDIA RTX 3060 Laptop网格 256x256。阶段CPU NumPyFPSGPU TaichiFPS提升倍数N12825180约 7 倍N2566110约 18 倍N5121.540约 27 倍分辨率越高GPU 的并行优势越明显。不过要注意Taichi 在分野后的加速效果跟迭代次数强相关。Projection 迭代次数从 30 改成 80GPU 版本只掉十几帧CPU 版本直接掉到 1 FPS。所以 GPU 加速的另一个隐藏利好是——你可以放心加大迭代次数来提高物理质量而不必牺牲实时性。视觉观感方面我的调参顺序是先把 dt 设大一点比如 0.02让流体动起来明显再根据是否发散往回收。加外力时在鼠标位置同时注入速度扰动和密度比例大约是速度密度 1 : 20000这样既能看到烟又被拖着走的感觉。粘性保持在 0.00005 到 0.0002 之间视觉观感最像空气偶尔故意调高到 0.002能做出水下或油状特效。这里再强调一次我的大坑教训调试 GPU 版本时不要一上来就开 512 分辨率。先跑 128确认逻辑正确再拉高网格否则一旦出问题你在像素海洋里排查窗口根本盯不住。5. 常见问题与避坑实录5.1 数值爆炸最常见的翻车方式现象运行几秒后屏幕上出现白色噪点、彩色雪花甚至直接变成 NaN整个窗口变白。排查路径检查 dt。这是首选原因CFL 超标时几乎必炸。把 dt 除以 2 看是否缓解。检查速度场里是否出现了超大值。如果某格速度是几万方向还是乱的多半是平流回溯时索引出 bug或者外力 injection 时把模型撑爆了。检查投影是否被跳过或迭代次数为零。投影一旦失效速度散度会持续积累每帧都在“压气”迟早爆掉。我自己的经验是一旦炸了立刻打印速度场最大值、散度最大值、密度最大值三个值基本能定位问题在哪一层。5.2 边界穿越与插值异常现象流体“漏”到计算域外或者边界处出现奇怪的条纹。排查路径半拉格朗日回溯可能跳到边界之外。如果钳制逻辑写错了边界格点的 e 值会引用未初始化的内存区域。处理边界时仔细检查速度的法向分量是否在边界上被强制清零。墙边界最容易犯的错误是只处理了 x 方向或只处理了 y 方向导致流体在角落“斜着穿墙”。双线性插值的 i0, i1, j0, j1 需要保证落在 [0, N-1] 范围内通常配合钳制一起做。如果源坐标超出范围太多不要直接索引宁可利用边界幻影网格复制值。5.3 性能上不去时的排查思路如果你发现 GPU 版本没有想象中的快或者 CPU 版本在某个 N 值下卡顿明显试试这几个排查角度迭代次数是不是设太高。投影 60 次和 30 次性能差一倍视觉差可能很难察觉。是不是被显示刷新拖累了。matplotlib 的 imshow 在 N256 时本身就慢和数据模拟无关。GPU 版本有没有真正用 GPU。检查 Taichi 的ti.init的日志输出确认带的是CUDA / Vulkan而不是 CPU backend——如果archti.gpu失败Taichi 会静默退化到 CPU导致性能没有提升。临时数组分配。CPU 版最容易踩的坑是在主循环里用np.zeros_like创建新数组每帧都会触发垃圾回收改用预分配缓存数组。一句话避坑总结别在没确认 backend 是 GPU 的情况下就去优化“模拟”本身。我见过太多人卡在ti.init(archti.gpu)配错了环境跑了一个小时的 CPU 代码还以为是 GPU 的。最后再分享一个小技巧项目做到后期你会发现最影响观感的不是分辨率也不是 solver 精度而是“初始条件”。很多入门教程喜欢从静止场开始然后让用户拖动鼠标产生烟。但实际上一团烟的初始形状决定了后面的涡旋结构。我的经验是在 simulation 开始时先在中间注入一个静止的圆形密度区域半径大约 15 个格子然后在边缘施加一圈垂直于半径方向的切向速度这样会立刻产生一对对称的涡旋对画面瞬间有了“流体感”。这个技巧在演示时特别好用也适合用来验证你的求解器是否正确——如果连一对对称涡旋都画不出来说明压力投影或者平流实现肯定有问题。这套从 Navier-Stokes 方程到 GPU 加速渲染的流程跑完之后我自己回头看觉得最有价值的不是最终那片烟而是中间踩过的每一个坑。它们让我彻底理解了什么叫“数值离散”什么叫“稳定性”什么是“渲染瓶颈”。如果你也想在计算机图形学里做更多有意思的尝试我强烈建议从流体模拟开始。它既不要求复杂的几何处理又能让人直观看到算法与物理的交汇是进入物理仿真世界里性价比最高的一条路。