用PyTorch和JAX复现PINN：手把手教你用物理信息神经网络求解薛定谔方程

用PyTorch和JAX实现物理信息神经网络：从薛定谔方程实战看PINN技术内核

在深度学习与科学计算的交叉领域，物理信息神经网络（PINN）正掀起一场方法论革命。不同于传统数值模拟的黑箱特性，PINN将物理定律直接编码到神经网络架构中，实现了"物理规律学习"与"数据驱动建模"的有机融合。本文将以量子力学中的薛定谔方程为研究对象，通过PyTorch和JAX双框架对比实现，揭示PINN在微分方程求解中的独特优势。无论您是计算物理研究者还是AI工程师，这份包含完整代码实现的指南都将帮助您快速掌握这一前沿技术。

1. 环境配置与工具链选择

1.1 框架选型：PyTorch vs JAX

现代深度学习框架的自动微分（AD）能力是PINN实现的核心支柱。我们选择PyTorch和JAX进行对比实现，二者在自动微分机制上有着显著差异：

# PyTorch环境安装 pip install torch torchdiffeq torchphysics # JAX环境安装（根据CUDA版本选择） pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

1.2 薛定谔方程数学表述

我们考虑一维非线性薛定谔方程(NLS)：

$$ i\frac{\partial h}{\partial t} + \frac{1}{2}\frac{\partial^2 h}{\partial x^2} + |h|^2 h = 0,\quad x\in[-5,5], t\in[0,\pi/2] $$

其中$h(t,x)$是复值波函数。边界条件设为周期性边界，初始条件为：

$$ h(0,x) = 2,\text{sech}(x) $$

提示：在实现时需将复值函数分解为实部和虚部两个通道处理

2. PyTorch实现详解

2.1 网络架构设计

采用具有傅里特征映射(Feature Mapping)的全连接网络：

import torch import torch.nn as nn class PINN(nn.Module): def __init__(self, layers): super().__init__() self.activation = nn.Tanh() self.layers = nn.ModuleList() # 输入层傅里特征映射 self.fourier = nn.Linear(2, layers[0], bias=False) nn.init.normal_(self.fourier.weight, mean=0, std=10.0) # 隐藏层 for i in range(len(layers)-1): self.layers.append(nn.Linear(layers[i], layers[i+1])) def forward(self, t, x): X = torch.cat([t,x], dim=1) H = torch.sin(self.fourier(X)) # 傅里基函数 for layer in self.layers[:-1]: H = self.activation(layer(H)) # 输出实部和虚部 out = self.layers[-1](H) return out[:, 0:1], out[:, 1:2]

关键设计考量：

• 傅里特征映射增强网络对高频信号的捕捉能力
• 双通道输出分别对应波函数的实部(Re)和虚部(Im)
• Tanh激活函数保证二阶导数稳定性

2.2 损失函数构建

PINN的核心创新在于将物理方程融入损失函数：

def physics_loss(model, t, x): # 启用梯度追踪 t.requires_grad_(True) x.requires_grad_(True) # 网络预测 h_real, h_imag = model(t, x) # 一阶导数 dh_real = torch.autograd.grad(h_real.sum(), [t,x], create_graph=True) dh_imag = torch.autograd.grad(h_imag.sum(), [t,x], create_graph=True) # 二阶导数 d2h_real = torch.autograd.grad(dh_real[1].sum(), [x], create_graph=True)[0] d2h_imag = torch.autograd.grad(dh_imag[1].sum(), [x], create_graph=True)[0] # 薛定谔方程残差 f_real = dh_imag[0] + 0.5*d2h_real + (h_real**2 + h_imag**2)*h_real f_imag = -dh_real[0] + 0.5*d2h_imag + (h_real**2 + h_imag**2)*h_imag return torch.mean(f_real**2 + f_imag**2)

2.3 训练策略优化

针对PINN训练不稳定的问题，采用分阶段训练策略：

1. 预训练阶段：先用少量边界数据训练网络满足初始/边界条件

   # 边界数据采样 t_initial = torch.zeros((N,1), device=device) # t=0 x_boundary = torch.rand((N,1), device=device)*10-5 # x∈[-5,5]

2. 物理约束阶段：逐步增加方程残差项的权重

   def weighted_loss(epoch): alpha = min(1.0, epoch/1000) # 渐进加权 return alpha*physics_loss + (1-alpha)*boundary_loss

3. 自适应采样：在残差较大区域增加采样点密度

3. JAX实现对比

3.1 函数式编程范式

JAX的实现展现截然不同的编程风格：

import jax import jax.numpy as jnp from flax import linen as nn class PINN(nn.Module): layer_sizes: list @nn.compact def __call__(self, inputs): t, x = inputs[:, 0:1], inputs[:, 1:2] X = jnp.concatenate([t,x], axis=1) # 傅里特征映射 W = self.param('W', jax.nn.initializers.normal(10.0), (64, 2)) H = jnp.sin(jnp.dot(X, W.T)) # 隐藏层 for size in self.layer_sizes[1:-1]: H = nn.tanh(nn.Dense(size)(H)) # 双通道输出 out = nn.Dense(self.layer_sizes[-1])(H) return out[:, 0:1], out[:, 1:2]

JAX优势体现：

• 纯函数式设计更利于微分运算
• jit编译大幅提升计算效率
• vmap实现自动批量处理

3.2 自动微分实现

JAX的grad函数实现更简洁的物理约束：

@jax.jit def physics_loss(params, batch): t, x = batch[:, 0:1], batch[:, 1:2] # 定义内部函数用于微分 def h_fn(tx): h_real, h_imag = model.apply(params, tx) return h_real, h_imag # 一阶导数 dh_real = jax.grad(lambda tx: h_fn(tx)[0].sum(), argnums=0)(batch) dh_imag = jax.grad(lambda tx: h_fn(tx)[1].sum(), argnums=0)(batch) # 二阶导数 d2h_real = jax.grad(lambda tx: dh_real(tx)[1].sum(), argnums=0)(batch)[1] d2h_imag = jax.grad(lambda tx: dh_imag(tx)[1].sum(), argnums=0)(batch)[1] # 方程残差 f_real = dh_imag[0] + 0.5*d2h_real + (h_real**2 + h_imag**2)*h_real f_imag = -dh_real[0] + 0.5*d2h_imag + (h_real**2 + h_imag**2)*h_imag return jnp.mean(f_real**2 + f_imag**2)

4. 结果分析与工程实践

4.1 性能对比测试

在NVIDIA V100 GPU上的基准测试结果：

注意：实际性能会随超参数配置和硬件环境变化

4.2 常见问题排查

梯度爆炸问题：

• 现象：损失值出现NaN
• 解决方案：
  1. 调整网络初始化尺度
  2. 添加梯度裁剪

  torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

模式崩溃(Mode Collapse)：

• 现象：网络退化为平凡解
• 解决方案：
  1. 增加傅里特征映射的随机频率
  2. 采用课程学习策略逐步扩大计算域

训练效率优化：

• 使用torch.compile()(PyTorch 2.0+)或@jax.jit加速计算图
• 对周期性边界条件采用硬约束编码：

  def forward(self, t, x): # 将x映射到[-π,π]周期 x_mapped = torch.pi * torch.sin(x/5 * torch.pi) return super().forward(t, x_mapped)

4.3 扩展应用方向

基于相同框架可扩展的物理系统：

1. 非线性波动方程：

   # KdV方程残差 f = h_t + 6*h*h_x + h_xxx

2. Navier-Stokes方程：

   # 不可压缩流体约束 f_continuity = u_x + v_y

3. Maxwell方程组：

   # 法拉第定律残差 f_faraday = E_y - B_t

在实际项目中，PINN常与传统数值方法（如有限元）结合使用，形成混合求解器。例如用PINN求解边界层区域，而用有限元处理主体区域，这种协同方式往往能获得意想不到的效果。