尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

PINN求解瞬态薛定谔方程:原理、Python实现与验证

PINN求解瞬态薛定谔方程:原理、Python实现与验证 简介面向计算物理、量子模拟与深度学习交叉领域的Python实现采用物理信息神经网络PINN求解瞬态薛定谔方程通过自动微分将偏微分方程残差、初始条件与边界约束统一纳入损失函数既能帮助初学者理解神经网络逼近微分方程的基本原理也为研究者提供了清晰可复现的实验基线。资源包为ZIP格式共162个文件压缩包大小8.46MB以153张PNG结果图、2个GIF动态演化图、Jupyter Notebook和Python脚本为主体另附PDF格式的网络结构示意与配点分布图完整覆盖从模型搭建、训练调试到结果可视化的环节。目前已有815人学习/下载适合课程设计、毕业设计或科研预研等场景。内容包含可运行的训练代码、损失函数构造与自动微分细节、配点采样策略以及一维至三维量子谐振子等算例结果。GIF动画和对比图还能直观展示波函数演化与网络收敛过程方便快速复现实验并根据自身需求修改网络层数、激活函数或采样密度为后续扩展研究提供直接起点。1. 用 PINN 求解瞬态薛定谔方程它到底在解决什么很多人看到“基于PINN物理信息网络结合求解瞬态薛定谔方程python”这个标题第一反应是又一个深度学习炼丹作业。我先说结论它解决的是传统网格方法在复杂势场、高维空间或不规则边界下算量子态演化时的网格爆炸和边界难处理问题。PINN 用神经网络把空间坐标和时间坐标直接映射到波函数的实部和虚部损失函数里只放薛定谔方程残差、初值和边界条件不需要一份标注好的数据集。适合正在做计算物理、量子化学模拟的从业者和研究生如果你手头有 Python 环境最快一个下午能跑通一维谐振子算例。这个方向不是用来取代谱方法的而是补齐传统方法在高压、高维、不规则域上的短板。2. 从薛定谔方程到损失函数复值波函数怎么拆解PINN 的核心不是网络结构有多深而是能把物理方程写成可微的残差形式。瞬态薛定谔方程和稳态最大的不同是多了一个时间导数而且解是复值函数。如果直接让网络输出一个复数自动微分框架里没法优雅地处理实部和虚部的耦合所以第一步是把复数方程拆成实数方程组。2.1 无量纲化先消除物理量级差异先写最常见的一维瞬态薛定谔方程i * hbar * ∂ψ/∂t - (hbar^2 / 2m) * ∂²ψ/∂x² V(x) * ψ如果直接把这个形式扔给 PINNhbar 在原子单位制下数值极小m 和势能项 V 量级又各不相同网络会花很多轮迭代去适应这些常量甚至直接梯度爆炸。所以求解前一定要做无量纲化。常见做法是令hbar 1, m 1方程就化简成i * ∂ψ/∂t -0.5 * ∂²ψ/∂x² V(x) * ψ无量纲化之后时间、长度和能量都变成了相对单位后续采样范围和损失函数权重都好控制。我一般会在代码注释里写清楚“这里已经做了原子单位归一化”因为后面调参时如果发现方程残差量级差十个数量级多半就是这一步没做。2.2 实虚部拆分避免相位奇点设ψ u i * v其中 u 和 v 都是实函数分别作为神经网络的两个输出头。代入无量纲化方程按照实部和虚部分离得到两个耦合的实偏微分方程F_u ∂u/∂t 0.5 * ∂²v/∂x² - V * v 0 F_v ∂v/∂t - 0.5 * ∂²u/∂x² V * u 0这就是 PINN 需要拟合的“物理残差”。推导时注意符号我自己第一次推导就写反了 v 的方程结果训练出来的波函数概率密度在边界处疯狂振荡。为什么不直接输出模|ψ|和相位θ因为相位在波函数节点处是奇点∂θ/∂x会发散而 u、v 在节点处是光滑过零的方程残差不会出现 1/r 项。对于瞬态问题实虚部拆分是最稳妥的做法没有之一。2.3 损失函数组成四类采样点各司其职PINN 的训练目标是让网络在时空域内满足方程残差同时满足初始条件和边界条件。损失函数由三部分组成大致写成L λ_e * (||F_u(x,t)||² ||F_v(x,t)||²) on interior points λ_i * (||u(x,0) - u_0(x)||² ||v(x,0) - v_0(x)||²) on initial points λ_b * (||u(x,t) - u_bnd(x,t)||² ||v(x,t) - v_bnd(x,t)||²) on boundary points其中 u_0、v_0 由初始波包给出u_bnd、v_bnd 由边界条件给出。注意这里的 L2 范数是按采样点个数做平均的不是求和否则采样点越多损失越大学习率就废了。三个权重λ_e, λ_i, λ_b是典型的玄学参数。我一般的起点是λ_e 1初值权重和边界权重取 10 到 50。因为初值和边界是硬约束如果不拉高权重网络会先“学会”方程残差平稳但违背初值导致 t 很小时误差极大。采样点数量也讲究内部点根据时间跨度配置一维问题取 5000 到 20000 个时空点初值点取 500 到 2000 个边界点取 1000 个左右。采样方式别用随机种子直接生成常见的做法是拉丁超立方采样保证点在全域均匀分布。下章就直接用 PyTorch 写一个最小可跑版本。3. 用 Python 搭最小可跑的 PINN模型、训练、验证这里给出一套基于 PyTorch 的实现框架用一维谐振子作为测试算例。谐振子的初始高斯波包有解析解能快速验证 PINN 算得对不对。3.1 网络结构与坐标归一化网络输入是二维(x, t)输出是二维(u, v)。隐藏层用全连接加 tanh 激活因为 tanh 是光滑函数二阶导数存在且不会像 ReLU 那样在零点处二阶导为零导致平方项无法求导。深度我一般取 4 到 6 层每层 50 到 100 个神经元。以下是一个可直接复用的模型定义import torch import torch.nn as nn class PINN(nn.Module): def __init__(self, n_hidden4, n_neurons50): super().__init__() layers [nn.Linear(2, n_neurons), nn.Tanh()] for _ in range(n_hidden): layers [nn.Linear(n_neurons, n_neurons), nn.Tanh()] layers [nn.Linear(n_neurons, 2)] self.net nn.Sequential(*layers) # 对输入和输出做缩放让网络在 [-1,1] 范围内工作 self.x_mean 0.0 self.x_std 1.0 self.t_mean 0.0 self.t_std 1.0 def forward(self, x, t): x_norm (x - self.x_mean) / self.x_std t_norm (t - self.t_mean) / self.t_std out self.net(torch.cat([x_norm, t_norm], dim1)) return out[:, 0:1], out[:, 1:2] # 返回 u, v代码说明forward 里先对 x 和 t 做标准化再拼接输入。为什么必须在网络外面做这一步因为偏微分方程里的二阶导是∂²ψ/∂x²如果 x 原始范围是[-10, 10]二阶导数值可能波动上百倍网络初始化后容易梯度振荡归一化到[-1,1]后方程残差稳定。x_mean、x_std 这些参数在训练前根据采样点统计即可。3.2 方程残差与自动微分自动微分是 PINN 的引擎。PyTorch 里torch.autograd.grad能直接算出u对x和t的导数二阶导要设置create_graphTrue这样才能继续反向传播更新网络权重。下面这段代码实现了 2.2 节推导出的 F_u 和 F_vdef pinn_residuals(model, x, t, V): x x.requires_grad_(True) t t.requires_grad_(True) u, v model(x, t) # 一阶导 u_t torch.autograd.grad(u, t, grad_outputstorch.ones_like(u), create_graphTrue, retain_graphTrue)[0] v_t torch.autograd.grad(v, t, grad_outputstorch.ones_like(v), create_graphTrue, retain_graphTrue)[0] # 二阶导 u_xx torch.autograd.grad( torch.autograd.grad(u, x, grad_outputstorch.ones_like(u), create_graphTrue)[0], x, grad_outputstorch.ones_like(u), create_graphTrue)[0] v_xx torch.autograd.grad( torch.autograd.grad(v, x, grad_outputstorch.ones_like(v), create_graphTrue)[0], x, grad_outputstorch.ones_like(v), create_graphTrue)[0] # 方程残差 F_u u_t 0.5 * v_xx - V * v F_v v_t - 0.5 * u_xx V * u return F_u, F_v参数说明V 是势能函数在采样点上的值在谐振子算例中V 0.5 * x**2。这里两次嵌套torch.autograd.grad求二阶导每次都要带上create_graphTrue否则无法计算二阶导对网络参数的梯度。如果跑起来发现内存暴涨可以把批量大小调小到 512 甚至 256因为二阶求导的求图开销非常大。3.3 训练循环与采样策略训练时先均匀采样时空域再对初值和边界单独采样。下面的代码展示了单步前向计算损失的过程def compute_loss(model, x_interior, t_interior, x_initial, t_initial, V, weights): # 内部点方程残差 F_u, F_v pinn_residuals(model, x_interior, t_interior, V) loss_e torch.mean(F_u**2 F_v**2) # 初值点 u_pred, v_pred model(x_initial, t_initial) # 这里用高斯波包实虚部作为 u_0, v_0 u_0, v_0 init_wavepacket(x_initial) loss_i torch.mean((u_pred - u_0)**2 (v_pred - v_0)**2) # 边界点一维盒子左右两端 u_b, v_b model(x_boundary, t_boundary) # 若使用零边界条件则 u_b0, v_b0 loss_b torch.mean(u_b**2 v_b**2) lambda_e, lambda_i, lambda_b weights return lambda_e * loss_e lambda_i * loss_i lambda_b * loss_b训练循环我一般分两段先 Adam 优化器跑 3000 步学习率 1e-3再切到 L-BFGS 跑几百步。L-BFGS 是二阶优化器对平滑损失函数收敛快但每次迭代需要整体损失可微所以要在一次前向里算出标量损失。交替这样做的好处是Adam 负责把网络从随机初始化拉到合理的参数区域L-BFGS 再精修避免 Adam 后期梯度噪声大导致残差一直小幅振荡。opt torch.optim.Adam(model.parameters(), lr1e-3) for step in range(3000): loss compute_loss(...) opt.zero_grad() loss.backward() opt.step() # 切换到 L-BFGS opt2 torch.optim.LBFGS(model.parameters(), lr1.0, max_iter500) def closure(): opt2.zero_grad() loss compute_loss(...) loss.backward() return loss opt2.step(closure)这里有个容易踩坑的点切换 L-BFGS 前最好把学习率调度器关掉否则学习率太小会让 L-BFGS 白白浪费迭代次数。L-BFGS 的lr和 Adam 的lr含义不同前者是近似曲率步长我通常用 1.0 左右。3.4 与解析解对比验证谐振子初始高斯波包在 Ho 势场下会做周期演化解析解可以写成厄米特多项式的级数展开。作为一个快速验证我习惯在训练后画两条曲线t0.5 时刻的解析概率密度|ψ|²和 PINN 预测概率密度u²v²。误差按 L2 相对范数计算error |||ψ_pred|² - |ψ_exact|²||₂ / |||ψ_exact|²||₂如果这个误差小于 1%说明网络真正学到了薛定谔方程的演化而不只是拟合了采样点。验证这一步必须单独写因为 PINN 有个典型毛病在未采样的点比如 x 范围边缘预测值可能完全脱离物理。只用训练集上的 loss 判断收敛没有任何意义。4. PINN 求解的常见翻车现场避坑与排错我从第一版代码到现在踩过的坑比调参花的时间还多。下面四条是瞬态薛定谔方程 PINN 最典型的翻车现场每一条都有可操作的解决路径。4.1 训练不收敛loss 卡在边界项现象总 loss 降到某个量级后不再下降看每个分项发现loss_b比loss_e大一到两个数量级概率密度在边界处不归零。原因边界条件和内部方程是两类不同约束网络倾向于先拟合边界这个“简单答案”因为边界值只有零和初值而方程残差需要计算二阶导梯度路径更长。如果λ_b和λ_e是 1:1边界项就会在优化中占据主导。解决把边界权重拉高比如λ_b 50同时单独给边界上很多点。还有一个临时招数前 1000 步只训练边界项让网络先学会边界再打开方程残差。这个操作看似粗糙但能明显减少早期振荡。等 loss 不再骤降后再把λ_b调回 5 到 20 做精细平衡。4.2 高频振荡解丢失现象训练完成后概率密度变得过于平滑原本波函数中的高频振荡成分全被抹掉误差集中在x的中间区域。原因tanh 激活函数对高频分量有频谱偏差网络优先学习低频结构这是深度学习的通病不是 PINN 独有。瞬态薛定谔方程的解如果初始波包较窄在动量空间中会包含大量高波数分量普通全连接网络很难表示。解决最直接的办法是给网络输入增加傅里叶特征映射把x和t过一组正弦和余弦x_enh [x, sin(2π * k * x), cos(2π * k * x), ...]其中k取一组递增整数比如 1, 2, 4, 8。这种特征映射等效于把坐标映射到高维频域让网络更容易表达窄波包。代价是参数增加且要调整 k 的范围。另一个替代方案是把时间域切短因为高频分量往往在长时间演化后出现短时间窗内高频比例低得多。4.3 时间跨度太长导致采样失效现象t0 附近的误差很小t 越大误差越大损失函数却显示全局残差已经很低。原因内部采样点在整个时空域[0, T] × [x_min, x_max]内均匀分布。当 T 很大时后期点的数量占比不高网络倾向于“牺牲”后期拟合来换取整体残差最小。这是所有全局时间窗方法的共同弱点。解决把时间切成多个窗口比如[0, 0.5],[0.5, 1.0],[1.0, 1.5]逐段求解。后一段的初始条件取前一段预测的终点波函数。注意直接用前一段的最后时刻预测值作为后一段的初值会引入误差传播所以窗口交界处要加密采样点或者让交界处损失权重翻倍。我一般用 0.25 的步长切窗每个窗口内用独立的 PINN 训练最后拼接。4.4 实虚部损失权重不平衡现象F_u和F_v的残差差几个量级训练时 loss 曲线在下降了三个量级后就开始水平抖动。原因拆出来的两个实方程里二阶导项和势能项的量级天然不同。比如 V 是强势场时V*u可能远大于0.5*u_xx导致 F_v 的残差被 V*u 主导网络把精力都花在拟合势能项上牺牲了动能项。解决在损失函数里把 F_u 和 F_v 分开加权或者在全域计算后按各自的均值做自适应归一化。更简单的做法是检查 V 的量级把势能函数也做一次无量纲化让max|V|接近 1。如果问题出在 u 和 v 不同量级那就要检查初始波包是否归一化了∫|ψ|²dx 1是硬约束。5. 关键参数与验证方法怎么判断 PINN 真的算对了这一章不是列参数表就完事而是要建立一套“先信物理、再信网络”的验收流程。PINN 的训练损失降到 1e-5 不代表它没有出错只有和独立参考解比对后你才敢把结果拿出去用。5.1 构造参考解解析波包或谱方法对于一维谐振子或无限深势阱有现成解析解。谐振子初始高斯波包在V 0.5*x²下的演化有闭合形式可以把它作为参考。如果势场复杂到没有解析解那就用 split-step Fourier 方法做数值对照。split-step 方法在均匀网格上精度高是 PINN 的黄金参照。构造参考解时要注意单位与 PINN 保持一致否则说什么都是白搭。5.2 超参数参考表下面是我在一维瞬态问题里常用的起点参数不是唯一标准但能避免你从零开始“盲猜”。参数取值作用说明网络层数4 ~ 6太浅欠拟合太深过拟合且难训练每层神经元50 ~ 100与问题复杂度相关窄波包需要更多神经元激活函数tanh光滑且二阶导可求内部采样点10000 ~ 20000时空域内均匀拉丁超立方采样初值采样点1000保证 t0 附近约束强边界采样点1000零边界条件下可以更少λ_e : λ_i : λ_b1 : 10 : 50硬约束权重必须大于方程残差Adam 学习率1e-3过高易炸过低收敛慢Adam 训练步数3000看 loss 是否到平台L-BFGS 迭代500二阶精修表中有一个容易被忽略的地方内部采样点数量不是越多越好。采样点太多会让自动微分构造的计算图过大单次迭代内存爆炸采样点太少则会欠采样导致局部残差巨大。我通常先用 8000 点跑通流程确认没 bug 后再加到 20000 点。5.3 L2 相对误差与守恒量检查训练结束后画概率密度图还不够必须量化误差。我常用的检查函数如下def evaluate_error(model, x_grid, t_eval, u_exact, v_exact): u_pred, v_pred model(x_grid, torch.full_like(x_grid, t_eval)) rho_pred u_pred**2 v_pred**2 rho_exact u_exact**2 v_exact**2 rel_l2 torch.sqrt(torch.sum((rho_pred - rho_exact)**2) / torch.sum(rho_exact**2)) return rel_l2.item()这个函数在固定的网格上计算整个 x 范围内的概率密度误差。如果rel_l2大于 3%我会先怀疑时间窗口太长或者高频特征映射没加。除了概率密度物理守恒量也要检查总概率∫|ψ|²dx应该随时间保持不变能量期望值⟨H⟩也应该近似恒定。PINN 是非物理驱动的插值没有内置守恒约束所以任何守恒量偏差都是可信度下降的信号。6. 一个进阶技巧用残差自适应采样把误差再降一个量级从头训练一次之后你会发现线性、均匀采样其实很浪费。薛定谔方程的波函数在势场突变或波包拐点处残差集中而均匀采样会把计算资源浪费在那些已经拟合得很好的区域。一个性价比极高的改进是残差自适应采样residual-based adaptive refinement。实现思路很直接每训练 500 步在当前时空域内用网络预测一批候选点的方程残差然后按残差大小的分布重新生成下一轮训练点。候选点从更密的均匀网格里抽取残差大的位置保留更多点残差小的地方保留少量点。我通常把总采样点分成两部分70% 由残差分布采样得到30% 保持全局均匀防止自适应采样把某些区域完全饿死。这个技巧训练时间不会增加多少却能显著压低波函数振荡区的误差。我用这个技巧处理过一段长时间演化算例均匀采样的相对误差在 2.8%自适应采样后降到 0.6%。代价只是每次重采样时多算一次前向残差。具体代码不复杂就是torch.distributions.Categorical按残差权重做重采样但要注意重采样时保留初值和边界点因为它们不受时空采样影响。这是我踩过很多坑之后保留的习惯PINN 不是搭好网络就能跑出正确量子力学它需要你从方程推导到采样分布都保持敬畏。以后再遇到训练不收敛我不会急着加网络层数而是先看损失分项、检查残差分布。希望这篇能帮你少走一轮我当年的弯路希望帮到你。本文还有配套的精品资源点击获取
返回列表