jax-pde

JAX - 微分可能な物理学と偏微分方程式

JAX は数値ソルバーを通して微分できるため、物理学の分野で独自の強みを発揮します。このガイドでは、「最適化に適した」従来型の偏微分方程式ソルバーの実装方法と、ニューラルネットワークと物理モデルを組み合わせたハイブリッドモデルの構築方法を説明します。

使用する場面

• GPU でナビエ・ストークス方程式、波動方程式、熱方程式を解く場合。
• 物理法則を組み込んだニューラルネットワーク（PINN）を実装する場合。
• 逆設計（観測結果から材料特性を特定すること）を行う場合。
• ロボット工学や気候モデリング向けに微分可能なシミュレーションを作成する場合。
• 物理システムの感度解析を行う場合。

基本原則

1. ソルバーを通した微分

JAX では、jax.numpy を使ってオイラー法やルンゲ・クッタ法の積分器を記述すると、∂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 では、境界条件を損失関数に追加する必要があります：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 による偏微分方程式の計算は、物理学を静的なシミュレーションから、動的で最適化可能な問題へと変えます。これにより研究者は「この結果を生み出した物理パラメーターは何か？」と問い、勾配の力を通じて答えを見つけられます。