大家好，我是讯享网，很高兴认识大家。这里提供最前沿的Ai技术和互联网信息。

# 用PyTorch实战PINN：从零构建求解薛定谔方程的物理信息神经网络

在深度学习与科学计算的交叉领域，物理信息神经网络（PINN）正掀起一场革命。想象一下，当你训练神经网络时，不再仅仅依赖数据标签，而是将牛顿定律、量子力学方程直接编码到损失函数中——这就是PINN的魔力所在。本文将带你用PyTorch实现一个完整的PINN框架，求解量子力学中的经典方程：薛定谔方程。不同于理论讲解，我们会聚焦于可运行的代码实现，让你在Google Colab或本地环境中能立即复现结果。

1. 环境准备与问题定义

1.1 安装必要库

确保你的Python环境包含以下核心库（推荐使用conda虚拟环境）：

pip install torch numpy matplotlib scipy 

1.2 薛定谔方程数学表述

我们考虑一维非线性薛定谔方程（NLSE）：

\[ ifrac{partial h}{partial t} + frac{1}{2}frac{partial^2 h}{partial x^2} + |h|^2 h = 0 \]

其中\(h(t,x)\)是复值波函数，边界条件为周期性边界。我们的目标是构建PINN来学习该方程的动力学演化。

2. PINN网络架构设计

2.1 神经网络结构

采用全连接网络作为基础架构，输入是时空坐标\((t,x)\)，输出是波函数的实部和虚部：

import torch import torch.nn as nn class PINN(nn.Module): def __init__(self, layers): super().__init__() self.activation = nn.Tanh() self.loss_function = nn.MSELoss() # 构建全连接层 self.linears = nn.ModuleList() for i in range(len(layers)-1): self.linears.append(nn.Linear(layers[i], layers[i+1])) def forward(self, x): if not isinstance(x, torch.Tensor): x = torch.tensor(x, dtype=torch.float32) for i in range(len(self.linears)-1): x = self.activation(self.linears[i](x)) # 最后一层不使用激活函数 x = self.linears[-1](x) return x 

2.2 物理信息编码关键

PINN的核心在于将物理方程融入损失函数。我们需要计算：

1. 网络输出的实部\(u(t,x)\)和虚部\(v(t,x)\)
2. 通过自动微分计算各阶偏导数
3. 构建物理残差项

def compute_physics_loss(model, coords): # 启用梯度追踪 coords.requires_grad_(True) # 网络预测 output = model(coords) u = output[:, 0:1] # 实部 v = output[:, 1:2] # 虚部 # 一阶导数计算 du_dt = torch.autograd.grad(u.sum(), coords, create_graph=True)[0][:, 0:1] dv_dt = torch.autograd.grad(v.sum(), coords, create_graph=True)[0][:, 0:1] # 二阶空间导数 du_dx = torch.autograd.grad(u.sum(), coords, create_graph=True)[0][:, 1:2] dv_dx = torch.autograd.grad(v.sum(), coords, create_graph=True)[0][:, 1:2] d2u_dx2 = torch.autograd.grad(du_dx.sum(), coords, create_graph=True)[0][:, 1:2] d2v_dx2 = torch.autograd.grad(dv_dx.sum(), coords, create_graph=True)[0][:, 1:2] # 物理方程残差 f_u = dv_dt + 0.5*d2u_dx2 + (u2 + v2)*v f_v = -du_dt + 0.5*d2v_dx2 + (u2 + v2)*u return torch.mean(f_u2 + f_v2) 

3. 训练策略与技巧

3.1 数据采样策略

成功的PINN实现高度依赖训练点的分布：

采样区域	采样方法	比例	作用
初始条件	均匀网格	20%	满足t=0的初始状态
边界条件	随机采样	20%	强制周期性边界
内部区域	LHS采样	60%	覆盖整个解空间

def generate_training_data(t_domain, x_domain, n_points): # 初始条件采样 t_init = torch.zeros(n_points//5, 1) x_init = torch.rand(n_points//5, 1)*(x_domain[1]-x_domain[0]) + x_domain[0] # 边界条件采样 t_bound = torch.rand(2*n_points//5, 1)*(t_domain[1]-t_domain[0]) + t_domain[0] x_bound = torch.cat([ torch.ones(n_points//5, 1)*x_domain[0], torch.ones(n_points//5, 1)*x_domain[1] ]) # 内部点采样 t_inner = torch.rand(2*n_points//5, 1)*(t_domain[1]-t_domain[0]) + t_domain[0] x_inner = torch.rand(2*n_points//5, 1)*(x_domain[1]-x_domain[0]) + x_domain[0] return torch.cat([ torch.cat([t_init, x_init], dim=1), torch.cat([t_bound, x_bound], dim=1), torch.cat([t_inner, x_inner], dim=1) ]) 

3.2 损失函数组合

总损失函数由三部分组成：

1. 初始条件损失：确保t=0时满足初始波形
2. 边界条件损失：强制周期性边界
3. 物理残差损失：方程本身的约束

def total_loss(model, coords, init_cond_fn): # 初始条件损失 init_coords = coords[coords[:,0]==0] # t=0的点 pred_init = model(init_coords) true_init = init_cond_fn(init_coords[:,1:]) init_loss = torch.mean((pred_init - true_init)2) # 边界条件损失 boundary_coords = coords[(coords[:,1]==0) | (coords[:,1]==L)] # x=0或x=L的点 pred_left = model(boundary_coords[boundary_coords[:,1]==0]) pred_right = model(boundary_coords[boundary_coords[:,1]==L]) boundary_loss = torch.mean((pred_left - pred_right)2) # 物理残差损失 physics_loss = compute_physics_loss(model, coords) return 10.0*init_loss + 5.0*boundary_loss + physics_loss 

4. 完整训练流程与可视化

4.1 训练循环实现

采用Adam优化器配合学习率衰减：

def train(model, coords, epochs, lr): optimizer = torch.optim.Adam(model.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2000, gamma=0.5) for epoch in range(epochs): optimizer.zero_grad() loss = total_loss(model, coords, initial_condition) loss.backward() optimizer.step() scheduler.step() if epoch % 100 == 0: print(f"Epoch {epoch}: Loss = {loss.item():.4e}") 

4.2 结果可视化

训练完成后，我们可以将预测结果与解析解或数值解比较：

import matplotlib.pyplot as plt def visualize_results(model, t_range, x_range): # 生成测试网格 t = np.linspace(t_range[0], t_range[1], 100) x = np.linspace(x_range[0], x_range[1], 100) tt, xx = np.meshgrid(t, x) # 模型预测 with torch.no_grad(): coords = torch.tensor(np.vstack([tt.ravel(), xx.ravel()]).T, dtype=torch.float32) pred = model(coords).numpy() # 绘制振幅 plt.figure(figsize=(12,4)) plt.subplot(121) plt.pcolormesh(tt, xx, np.sqrt(pred[:,0]2 + pred[:,1]2).reshape(tt.shape)) plt.colorbar(label='|h(t,x)|') # 绘制相位 plt.subplot(122) plt.pcolormesh(tt, xx, np.arctan2(pred[:,1], pred[:,0]).reshape(tt.shape)) plt.colorbar(label='Phase') 

5. 实战技巧与常见问题

5.1 梯度不稳定解决方案

PINN训练中常遇到的梯度爆炸问题可通过以下方法缓解：

• 输入归一化：将时空坐标归一化到[-1,1]范围
• 损失加权：不同损失项采用自适应权重
• 梯度裁剪：限制梯度最大值

# 在训练循环中添加梯度裁剪 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) 

5.2 网络深度与宽度选择

根据经验，对于薛定谔方程这类复杂PDE：

• 隐藏层数：4-8层为宜
• 每层神经元：20-50个足够
• 激活函数：Tanh通常优于ReLU

> 注意：过大的网络会增加训练难度，而太小的网络可能无法捕捉解的快速振荡特性

5.3 采样点动态调整

采用自适应采样策略可显著提升精度：

1. 首轮训练使用均匀采样
2. 计算残差分布，在残差大的区域增加采样密度
3. 迭代优化采样分布

def adaptive_sampling(model, old_coords, n_new_points): # 计算旧坐标点的残差 with torch.no_grad(): residuals = compute_residuals(model, old_coords) # 按残差分布采样新点 prob = residuals / residuals.sum() new_indices = torch.multinomial(prob, n_new_points) return torch.cat([old_coords, old_coords[new_indices]]) 

6. 扩展应用与性能提升

6.1 多GPU并行训练

对于大规模问题，可采用数据并行加速：

model = nn.DataParallel(PINN(layers=[2,50,50,50,50,2])).cuda() 

6.2 混合精度训练

减少内存占用并提升速度：

scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): loss = total_loss(model, coords, initial_condition) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() 

6.3 与其他数值方法对比

PINN与传统方法的比较优势：

方法	计算复杂度	并行性	适用场景
有限差分	O(N^d)	低	规则网格问题
谱方法	O(N logN)	中	周期性问题
PINN	O(网络参数)	高	复杂几何/高维问题

在实际项目中，我通常会先用传统方法获得基准解，再用PINN作为补充。特别是在需要快速探索参数空间时，训练好的PINN能实现实时预测，比重复运行数值模拟高效得多。