jax-pde

JAX - 可微分物理与 PDE

JAX 特别适用于物理计算，因为它可以对数值求解器进行求导。本指南介绍如何实现适合优化的传统 PDE 求解器，以及如何构建神经网络混合物理模型。

适用场景

• 在 GPU 上求解 Navier-Stokes、Wave 或 Heat 方程。
• 实现物理信息神经网络（PINN）。
• 执行逆向设计（根据观测结果求出材料属性）。
• 为机器人或气候建模创建可微分模拟。
• 对物理系统进行灵敏度分析。

核心原则

1. 对求解器求导

在 JAX 中，如果使用 jax.numpy 编写 Euler 或 Runge-Kutta 积分器，就可以自动计算 ∂Result/∂InitialCondition 或 ∂Result/∂Viscosity。

2. 交错网格与 Vmap

物理场（速度、压力）通常存储在网格上。JAX 的 vmap 可以让你立即针对不同边界条件或参数集并行运行求解器。

3. 伴随方法

对于非常大的系统，JAX 的反向模式自动微分有效地实现了传统 CFD/地球物理学中用于梯度计算的“伴随状态法”。

实现模式

1. PINN（物理信息神经网络）

import jax.numpy as jnp
from jax import grad, vmap

# A simple MLP representing the solution u(x, t)
def model(params, x, t):
    # standard neural net logic...
    return result

# Residual of the PDE: u_t + u*u_x - nu*u_xx = 0 (Burgers Equation)
def pde_loss(params, x, t, nu):
    u = lambda x, t: model(params, x, t)
    
    # Automatic derivatives of the MODEL
    u_t = grad(u, argnums=1)(x, t)
    u_x = grad(u, argnums=0)(x, t)
    u_xx = grad(grad(u, argnums=0), argnums=0)(x, t)
    
    return jnp.mean((u_t + u * u_x - nu * u_xx)**2)

2. 可微分有限差分求解器

@jit
def update_step(u, dt, dx, nu):
    """One step of a diffusion solver."""
    # Vectorized Laplacian using shifts (Zero-copy views)
    u_left = jnp.roll(u, -1)
    u_right = jnp.roll(u, 1)
    laplacian = (u_left + u_right - 2*u) / (dx**2)
    return u + dt * nu * laplacian

# We can now differentiate this solver!
def loss(initial_u, target_u):
    final_u = integrate_pde(initial_u) # Loop of update_step
    return jnp.sum((final_u - target_u)**2)

grad_initial_condition = grad(loss)(initial_u, target_u)

关键规则

✅ 应该做

• 使用 jax.lax.scan 编写时间循环 - 标准 Python for 循环会创建庞大的 XLA 计算图。scan 会将循环编译成单个高效内核。
• 对网格进行归一化 - 与机器学习类似，如果将 x、t 缩放到 [0,1] 或 [-1,1]，PINN 的收敛速度会更快。
• 结合数据与物理规律 - 在有一些传感器数据和物理定律的情况下使用 PINN，以“填补空缺”。
• 物理计算使用双精度 - 对于对精度敏感的数值求解器，使用 jax.config.update("jax_enable_x64", True)。

❌ 不应该做

• 不要什么问题都用 PINN - 对于“正向”问题，传统求解器（FDM/FEM）速度快得多。PINN 更适合“逆向”问题。
• 不要忽略边界条件（BC） - 在 PINN 中，必须将 BC 加入损失函数：Loss = PDE_loss + BC_loss。
• 不要忘记“幽灵单元” - 实现 FDM 时要谨慎处理边界，避免产生伪影。

实用工作流：逆问题

根据流体视频求出黏度

def objective(nu_guess):
    # 1. Run simulation with nu_guess
    final_state = run_simulation(initial_state, nu_guess)
    # 2. Compare with experimental data
    return jnp.mean((final_state - experimental_frame)**2)

# Gradient descent to find the real physical property
optimal_nu = optimize(grad(objective))

JAX PDE 将物理计算从静态模拟转变为动态、可优化的空间。研究人员可以借此询问“哪些物理参数产生了这个结果？”，并通过梯度的力量找到答案。