用DeepXDE搞定薛定谔方程：一个Python代码示例带你入门物理信息神经网络

物理信息神经网络实战：用DeepXDE求解非线性薛定谔方程

在科学计算的广阔天地里，物理信息神经网络（PINN）正掀起一场静默的革命。不同于传统数值方法的繁琐网格划分和离散化处理，PINN将微分方程直接编码进神经网络的损失函数，让偏微分方程的求解过程变得像训练图像分类模型一样直观。本文将带你用Python库DeepXDE，从零开始构建一个求解非线性薛定谔方程的PINN模型，无需深厚的数学物理背景，只需基础的Python知识就能上手。

1. 环境准备与DeepXDE基础

开始前需要确保已安装最新版的DeepXDE和TensorFlow后端。推荐使用Python 3.8+环境，通过以下命令安装：

pip install deepxde tensorflow

DeepXDE的核心优势在于它将复杂的微分方程求解抽象为几个直观的组件：

• 几何对象：定义问题的空间和时间域
• 边界条件：封装各类数学约束
• PDE定义：以自然数学形式表达方程
• 神经网络：作为近似解的万能函数逼近器

让我们先看一个最简单的示例，了解DeepXDE的基本工作流程：

import deepxde as dde # 定义区间几何 geom = dde.geometry.Interval(-1, 1) # 定义微分方程：u'' + u = 0 def pde(x, y): dy_xx = dde.grad.hessian(y, x) return dy_xx + y # 边界条件 bc = dde.DirichletBC(geom, lambda x: 0, lambda _, on_boundary: on_boundary) # 组装问题 data = dde.data.PDE(geom, pde, bc, num_domain=100, num_boundary=2) # 构建神经网络 net = dde.maps.FNN([1] + [50] * 3 + [1], "tanh", "Glorot normal") # 定义模型并训练 model = dde.Model(data, net) model.compile("adam", lr=0.001) model.train(epochs=1000)

这个简单例子展示了DeepXDE处理常微分方程的典型流程。接下来我们将处理更复杂的非线性薛定谔方程。

2. 非线性薛定谔方程的问题建模

非线性薛定谔方程(NLS)是量子力学中的基本方程，描述光脉冲在非线性介质中的传播。其标准形式为：

i∂h/∂t + (1/2)∂²h/∂x² + |h|²h = 0

由于DeepXDE仅支持实数运算，我们需要将复数方程拆分为实部和虚部：

∂u/∂t + (1/2)∂²v/∂x² + (u²+v²)v = 0 ∂v/∂t - (1/2)∂²u/∂x² - (u²+v²)u = 0

其中u和v分别表示h的实部和虚部。在DeepXDE中，这转化为两个耦合的PDE残差。

2.1 定义时空几何与初始条件

首先设置计算域和初始条件：

import numpy as np import deepxde as dde # 定义空间和时间范围 x_lower, x_upper = -5, 5 t_lower, t_upper = 0, np.pi/2 # 创建时空几何 space_domain = dde.geometry.Interval(x_lower, x_upper) time_domain = dde.geometry.TimeDomain(t_lower, t_upper) geomtime = dde.geometry.GeometryXTime(space_domain, time_domain) # 初始条件函数 def init_cond_u(x): """2 sech(x)""" return 2 / np.cosh(x[:, 0:1]) def init_cond_v(x): """虚部初始为0""" return 0

2.2 构建PDE残差函数

PDE残差函数的定义是PINN的核心，需要正确计算各阶导数：

def pde(x, y): """定义PDE残差""" u = y[:, 0:1] # 实部 v = y[:, 1:2] # 虚部 # 一阶导数 u_t = dde.grad.jacobian(y, x, i=0, j=1) # ∂u/∂t v_t = dde.grad.jacobian(y, x, i=1, j=1) # ∂v/∂t u_x = dde.grad.jacobian(y, x, i=0, j=0) # ∂u/∂x v_x = dde.grad.jacobian(y, x, i=1, j=0) # ∂v/∂x # 二阶导数 u_xx = dde.grad.hessian(y, x, component=0, i=0, j=0) # ∂²u/∂x² v_xx = dde.grad.hessian(y, x, component=1, i=0, j=0) # ∂²v/∂x² # 两个残差方程 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. 边界条件与模型配置

非线性薛定谔方程通常需要周期性边界条件，这在DeepXDE中很容易实现：

# 周期性边界条件 bc_u_0 = dde.PeriodicBC( geomtime, 0, lambda _, on_boundary: on_boundary, derivative_order=0, component=0 ) bc_u_1 = dde.PeriodicBC( geomtime, 0, lambda _, on_boundary: on_boundary, derivative_order=1, component=0 ) bc_v_0 = dde.PeriodicBC( geomtime, 0, lambda _, on_boundary: on_boundary, derivative_order=0, component=1 ) bc_v_1 = dde.PeriodicBC( geomtime, 0, lambda _, on_boundary: on_boundary, derivative_order=1, component=1 ) # 初始条件 ic_u = dde.IC(geomtime, init_cond_u, lambda _, on_initial: on_initial, component=0) ic_v = dde.IC(geomtime, init_cond_v, lambda _, on_initial: on_initial, component=1)

现在我们可以组装完整的PDE问题：

data = dde.data.TimePDE( geomtime, pde, [bc_u_0, bc_u_1, bc_v_0, bc_v_1, ic_u, ic_v], num_domain=10000, num_boundary=20, num_initial=200, train_distribution="pseudo", )

4. 神经网络架构与训练策略

4.1 构建深度神经网络

DeepXDE提供了多种网络架构选择。对于这个问题，我们使用4层全连接网络：

net = dde.maps.FNN([2] + [100] * 4 + [2], "tanh", "Glorot normal") model = dde.Model(data, net)

这里输入层有2个节点(对应x和t坐标)，输出层有2个节点(对应u和v)，中间是4层每层100个节点的隐藏层。

4.2 两阶段训练策略

实践证明，结合Adam和L-BFGS的混合优化策略效果最佳：

# 第一阶段：Adam优化 model.compile("adam", lr=1e-3, loss="MSE") model.train(epochs=1000, display_every=100) # 第二阶段：L-BFGS优化 dde.optimizers.config.set_LBFGS_options( maxcor=50, ftol=1.0 * np.finfo(float).eps, gtol=1e-08, maxiter=1000, maxfun=1000, maxls=50, ) model.compile("L-BFGS") model.train()

两阶段训练的关键参数对比如下：

5. 结果可视化与分析

训练完成后，我们可以提取预测结果并进行可视化：

# 创建测试网格 x = np.linspace(x_lower, x_upper, 256) t = np.linspace(t_lower, t_upper, 201) X, T = np.meshgrid(x, t) X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None])) # 预测 prediction = model.predict(X_star) u = griddata(X_star, prediction[:, 0], (X, T), method="cubic") v = griddata(X_star, prediction[:, 1], (X, T), method="cubic") h = np.sqrt(u**2 + v**2) # 振幅 # 绘图 import matplotlib.pyplot as plt fig, ax = plt.subplots(3, figsize=(10, 12)) ax[0].imshow(u.T, cmap="viridis", extent=[t_lower, t_upper, x_lower, x_upper]) ax[0].set_title("Real Part") ax[1].imshow(v.T, cmap="viridis", extent=[t_lower, t_upper, x_lower, x_upper]) ax[1].set_title("Imaginary Part") ax[2].imshow(h.T, cmap="viridis", extent=[t_lower, t_upper, x_lower, x_upper]) ax[2].set_title("Amplitude") plt.tight_layout() plt.show()

典型训练过程中损失值的变化趋势如下：

6. 常见问题与调优技巧

在实际应用中，可能会遇到以下典型问题及解决方案：

1. 训练不收敛

   • 检查PDE定义是否正确，特别是导数项
   • 尝试减小学习率
   • 增加网络容量或调整激活函数

2. 过拟合

   • 增加域内采样点数量
   • 添加正则化项
   • 使用早停策略

3. 梯度爆炸

   • 使用梯度裁剪
   • 尝试不同的权重初始化方法
   • 调整网络深度

一个实用的调优检查清单：

• [ ] PDE残差函数验证
• [ ] 边界条件正确实现
• [ ] 采样点分布合理
• [ ] 网络架构足够表达解
• [ ] 优化器参数适当

提示：对于复杂问题，可以先用少量训练点快速验证模型基本功能，再逐步增加采样点提高精度。

7. 扩展应用与进阶方向

掌握了基础PINN方法后，可以考虑以下进阶应用：

• 多物理场耦合问题：DeepXDE支持同时求解多个相互作用的PDE
• 不确定性量化：结合Dropout等技术评估解的可靠性
• 参数反演：从观测数据中推断PDE的未知参数
• 复杂几何：利用CSG构建非规则计算域

以下是一个多物理场问题的示例框架：

def coupled_pde(x, y): y1, y2 = y[:, 0:1], y[:, 1:2] # 定义第一个物理场的PDE pde1 = ... # 定义第二个物理场的PDE pde2 = ... # 耦合项 coupling = ... return [pde1 + coupling, pde2 + coupling]

物理信息神经网络为科学计算提供了全新的范式，将深度学习与物理定律完美融合。通过这个完整的非线性薛定谔方程求解示例，你应该已经掌握了PINN的核心概念和实现技巧。实际应用中，可能需要针对具体问题调整网络架构和训练策略，但基本框架是相通的。