用DeepXDE搞定薛定谔方程：一个Python物理信息神经网络(PINN)实战教程
用DeepXDE求解薛定谔方程：物理信息神经网络的Python实践指南
当传统数值方法遇到高维偏微分方程时，计算成本往往呈指数级增长。物理信息神经网络（PINN）提供了一种全新的解决思路——将物理定律直接编码到神经网络中。本文将带你用DeepXDE库，从零开始构建一个求解非线性薛定谔方程的完整流程。
1. 环境准备与DeepXDE基础
在开始之前，确保你的Python环境已安装以下依赖：
pip install deepxde numpy matplotlib tensorflowDeepXDE的核心优势在于它将复杂的微分方程求解过程抽象为几个直观的步骤：
- 几何定义：用简洁的API描述问题域
- PDE定义：以自然数学表达式编写方程
- 边界条件：支持多种常见边界条件类型
- 网络架构：灵活配置神经网络结构
import deepxde as dde import numpy as np import matplotlib.pyplot as plt提示：使用GPU加速可以显著减少训练时间，推荐配置CUDA环境
2. 问题建模与方程拆分
我们考虑的非线性薛定谔方程形式为：
$$ i h_t + \frac{1}{2} h_{xx} + |h|^2 h = 0 $$
由于DeepXDE目前仅支持实数运算，我们需要将复数方程拆分为实部和虚部：
def pde(x, y): u = y[:, 0:1] # 实部 v = y[:, 1:2] # 虚部 u_t = dde.grad.jacobian(y, x, i=0, j=1) v_t = dde.grad.jacobian(y, x, i=1, j=1) u_x = dde.grad.jacobian(y, x, i=0, j=0) v_x = dde.grad.jacobian(y, x, i=1, j=0) u_xx = dde.grad.hessian(y, x, component=0, i=0, j=0) v_xx = dde.grad.hessian(y, x, component=1, i=0, j=0) f_u = u_t + 0.5 * v_xx + (u**2 + v**2) * v f_v = v_t - 0.5 * u_xx - (u**2 + v**2) * u return [f_u, f_v]3. 定义计算域与边界条件
设置时空域为x∈[-5,5]，t∈[0,π/2]，并配置周期性边界条件和初始条件：
# 定义几何 space_domain = dde.geometry.Interval(-5, 5) time_domain = dde.geometry.TimeDomain(0, np.pi/2) geomtime = dde.geometry.GeometryXTime(space_domain, time_domain) # 周期性边界条件 bc_u = dde.PeriodicBC(geomtime, 0, lambda _, on_boundary: on_boundary, derivative_order=0, component=0) bc_ux = dde.PeriodicBC(geomtime, 0, lambda _, on_boundary: on_boundary, derivative_order=1, component=0) # 初始条件 def init_cond_u(x): return 2 / np.cosh(x[:, 0:1]) ic_u = dde.IC(geomtime, init_cond_u, lambda _, on_initial: on_initial, component=0)4. 构建PINN模型与训练策略
DeepXDE提供了灵活的神经网络配置选项，下面是构建模型的完整流程：
# 配置训练数据 data = dde.data.TimePDE( geomtime, pde, [bc_u, bc_ux, ic_u], num_domain=10000, num_boundary=20, num_initial=200, train_distribution="pseudo", ) # 构建神经网络 net = dde.maps.FNN([2] + [100] * 4 + [2], "tanh", "Glorot normal") # 创建模型 model = dde.Model(data, net)训练过程采用两阶段优化策略：
| 阶段 | 优化器 | 学习率 | 迭代次数 | 用途 |
|---|---|---|---|---|
| 1 | Adam | 1e-3 | 1000 | 初步收敛 |
| 2 | L-BFGS | - | 1000 | 精细调优 |
# 第一阶段训练 model.compile("adam", lr=1e-3, loss="MSE") model.train(epochs=1000, display_every=100) # 第二阶段训练 dde.optimizers.config.set_LBFGS_options(maxiter=1000) model.compile("L-BFGS") model.train()5. 结果可视化与分析
训练完成后，我们可以提取预测结果并进行可视化：
# 创建测试网格 x = np.linspace(-5, 5, 256) t = np.linspace(0, np.pi/2, 201) X, T = np.meshgrid(x, t) X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None])) # 预测结果 prediction = model.predict(X_star) u = prediction[:, 0].reshape(X.shape) v = prediction[:, 1].reshape(X.shape) h = np.sqrt(u**2 + v**2) # 绘制结果 fig, ax = plt.subplots(3, figsize=(10, 12)) ax[0].imshow(u.T, cmap="viridis", aspect="auto") ax[1].imshow(v.T, cmap="viridis", aspect="auto") ax[2].imshow(h.T, cmap="viridis", aspect="auto") plt.tight_layout() plt.show()6. 性能优化与问题排查
在实际应用中，可能会遇到以下常见问题及解决方案：
收敛问题排查表
| 现象 | 可能原因 | 解决方案 |
|---|---|---|
| 损失震荡 | 学习率过高 | 降低学习率或使用自适应优化器 |
| 收敛缓慢 | 网络容量不足 | 增加隐藏层宽度或深度 |
| 过拟合 | 训练点不足 | 增加num_domain或调整采样策略 |
关键参数调优建议
- 网络深度：4-6层通常足够处理大多数PDE问题
- 激活函数：tanh在大多数情况下表现良好
- 采样策略：对于时空问题，伪随机采样通常优于均匀采样
# 示例：调整网络结构 net = dde.maps.FNN([2] + [150] * 5 + [2], "tanh", "Glorot normal") # 示例：改进采样策略 data = dde.data.TimePDE( ..., train_distribution="LHS", # 拉丁超立方采样 )7. 扩展应用与进阶技巧
掌握了基础求解流程后，可以尝试以下进阶应用：
- 多物理场耦合：在同一个网络中求解多个相互作用的PDE
- 不确定性量化：使用Dropout评估解的可靠性
- 自适应采样：根据解的特性动态调整训练点分布
# 示例：添加不确定性量化 net = dde.maps.FNN([2] + [100] * 4 + [2], "tanh", "Glorot normal", dropout_rate=0.1)对于更复杂的几何形状，DeepXDE支持通过CSG（构造实体几何）构建：
# 示例：复杂几何定义 circle = dde.geometry.Disk([0, 0], 1) rectangle = dde.geometry.Rectangle([-1, -1], [1, 1]) complex_geom = circle - rectangle # 布尔运算在实际项目中，我发现合理设置权重对多任务学习至关重要。特别是当初值条件与边界条件的重要性不同时，可以通过调整损失权重来平衡各项约束：
# 示例：自定义损失权重 loss_weights = [1, 1] + [100] * 4 # 加重边界条件权重 model.compile("adam", lr=1e-3, loss="MSE", loss_weights=loss_weights)