pytorch实现：PINN寻求一维非线性薛定谔方程数值解
pytorch实现：PINN寻求一维非线性薛定谔方程数值解
1. 非线性薛定谔方程
非线性薛定谔方程（Nonlinear Schrödinger Equation, NLS）是对经典薛定谔方程的一种推广，用于描述那些包含显著非线性效应的系统，尤其是在波动力学领域，如非线性光学、Bose-Einstein凝聚体、水波、等离子体物理学以及某些类型的流体力学系统中。
典型的非线性薛定谔方程形式为：
i
ℏ
∂
ψ
∂
t
+
ℏ
2
2
m
∂
2
ψ
∂
x
2
+
g
∣
ψ
∣
2
ψ
=
0
i\hbar \frac{\partial \psi}{\partial t}+ \frac{\hbar^2}{2m}\frac{\partial^2\psi}{\partial x^2}+ g|\psi|^2\psi=0
iℏ∂t∂ψ+2mℏ2∂x2∂2ψ+g∣ψ∣2ψ=0
其中：
- ψ ( x , t ) \psi(x,t) ψ(x,t) 是复值波函数，描述系统的量子态（在经典物理中，可以被解释为场的幅度）。
- i i i 是虚数单位。
- ℏ \hbar ℏ 是约化普朗克常数。
- m m m 是粒子质量。
- g g g 是非线性系数，其正负决定了非线性项的性质。
- ∣ ψ ∣ 2 |\psi|^2 ∣ψ∣2 是波函数模的平方，代表波的强度或粒子数密度。
2. PINN实例
2.1 偏微分方程条件
考虑如下偏微分方程：
i
h
t
+
0.5
h
x
x
+
∣
h
∣
2
h
=
0
，
x
∈
[
−
5
,
5
]
,
t
∈
[
0
,
π
2
]
ih_t+0.5h_{xx}+|h|^2h=0，x\in[-5,5],t\in[0,\frac{\pi}{2}]
iht+0.5hxx+∣h∣2h=0，x∈[−5,5],t∈[0,2π]
具有初始条件：
h
(
0
,
x
)
=
2
s
e
c
h
(
x
)
h(0, x)=2sech(x)
h(0,x)=2sech(x)
具有周期边界条件：
h
(
t
,
−
5
)
=
h
(
t
,
5
)
h(t, -5)=h(t, 5)
h(t,−5)=h(t,5)
h x ( t , − 5 ) = h x ( t , 5 ) h_x(t, -5)=h_x(t, 5) hx(t,−5)=hx(t,5)
2.2 损失函数推导
h ( t , x ) h(t, x) h(t,x) 是一个复数波函数，可以写为 h ( t , x ) = u ( t , x ) + i v ( t , x ) h(t, x)=u(t, x)+iv(t,x) h(t,x)=u(t,x)+iv(t,x)，其中 u ( t , x ) u(t, x) u(t,x) 和 v ( t , x ) v(t,x) v(t,x) 分别是 h ( t , x ) h(t, x) h(t,x) 的实部和虚部。为了在实数计算框架中处理这个上述偏微分方程，将偏微分方程 i h t + 0.5 h x x + ∣ h ∣ 2 h = 0 ih_t+0.5h_{xx}+|h|^2h=0 iht+0.5hxx+∣h∣2h=0分解成实部和虚部。 可以分解为两个实数方程：
实部方程：
−
v
t
+
0.5
u
x
x
+
(
u
2
+
v
2
)
u
=
0
-v_t+0.5u_{xx}+(u^2+v^2)u=0
−vt+0.5uxx+(u2+v2)u=0
虚部方程：
u
t
+
0.5
v
x
x
+
(
u
2
+
v
2
)
v
=
0
u_t+0.5v_{xx}+(u^2+v^2)v=0
ut+0.5vxx+(u2+v2)v=0
2.3 损失函数定义
损失函数可定义为以下几部分：
M
S
E
=
M
S
E
0
+
M
S
E
b
+
M
S
E
f
MSE=MSE_0+MSE_b+MSE_f
MSE=MSE0+MSEb+MSEf
其中
M
S
E
0
MSE_0
MSE0 为初始条件损失函数：
M
S
E
0
=
1
N
0
∑
i
=
1
N
0
∣
h
(
0
,
x
0
i
)
−
h
0
i
∣
MSE_0 = \frac{1}{N_0}\sum_{i=1}^{N_0}|h(0, x_0^i) − h_0^i|
MSE0=N01i=1∑N0∣h(0,x0i)−h0i∣
M
S
E
b
MSE_b
MSEb 为周期边界条件损失函数：
M
S
E
b
=
1
N
b
∑
i
=
1
N
b
(
∣
h
i
(
t
b
i
,
−
5
)
−
h
i
(
t
b
i
,
5
)
∣
2
+
∣
h
x
i
(
t
b
i
,
−
5
)
−
h
x
i
(
t
b
i
,
5
)
∣
2
)
MSE_b = \frac{1}{N_b}\sum_{i=1}^{N_b} (|h^i(t_b^i, -5) − h^i(t_b^i, 5)|^2+|h_x^i(t_b^i, -5) − h_x^i(t_b^i, 5)|^2)
MSEb=Nb1i=1∑Nb(∣hi(tbi,−5)−hi(tbi,5)∣2+∣hxi(tbi,−5)−hxi(tbi,5)∣2)
M
S
E
f
MSE_f
MSEf 为偏微分方程损失函数：
M
S
E
f
=
1
N
f
∑
i
=
1
N
f
∣
f
(
t
f
i
,
x
f
i
)
∣
MSE_f = \frac{1}{N_f}\sum_{i=1}^{N_f}|f(t_f^i, x_f^i)|
MSEf=Nf1i=1∑Nf∣f(tfi,xfi)∣
其中
N
0
=
N
b
=
50
N_0=N_b=50
N0=Nb=50,
N
f
=
20000
N_f=20000
Nf=20000，
u
0
i
u_0^i
u0i，
v
0
i
v_0^i
v0i 为真实解。代码中损失函数的具体形式如下：
l
=
l
1
+
l
2
+
l
3
+
l
4
+
l
5
+
l
6
+
l
7
+
l
8
l=l_1+l_2+l_3+l_4+l_5+l_6+l_7+l_8
l=l1+l2+l3+l4+l5+l6+l7+l8
其中，
l
1
=
1
N
0
∑
i
=
1
N
0
∣
u
(
0
,
x
0
i
)
−
u
0
i
∣
2
l_1=\frac{1}{N_0}\sum_{i=1}^{N_0}|u(0,x_0^i)-u_0^i|^2
l1=N01i=1∑N0∣u(0,x0i)−u0i∣2
l 2 = 1 N 0 ∑ i = 1 N 0 ∣ v ( 0 , x 0 i ) − v 0 i ∣ 2 l_2=\frac{1}{N_0}\sum_{i=1}^{N_0}|v(0,x_0^i)-v_0^i|^2 l2=N01i=1∑N0∣v(0,x0i)−v0i∣2
l 3 = 1 N b ∑ i = 1 N b ∣ u i ( t b i , − 5 ) − u i ( t b i , 5 ) ∣ 2 l_3=\frac{1}{N_b}\sum_{i=1}^{N_b}|u^i(t_b^i,-5)-u^i(t_b^i,5)|^2 l3=Nb1i=1∑Nb∣ui(tbi,−5)−ui(tbi,5)∣2
l 4 = 1 N b ∑ i = 1 N b ∣ v i ( t b i , − 5 ) − v i ( t b i , 5 ) ∣ 2 l_4=\frac{1}{N_b}\sum_{i=1}^{N_b}|v^i(t_b^i,-5)-v^i(t_b^i,5)|^2 l4=Nb1i=1∑Nb∣vi(tbi,−5)−vi(tbi,5)∣2
l 5 = 1 N b ∑ i = 1 N b ∣ u x i ( t b i , − 5 ) − u x i ( t b i , 5 ) ∣ 2 l_5=\frac{1}{N_b}\sum_{i=1}^{N_b}|u_x^i(t_b^i,-5)-u_x^i(t_b^i,5)|^2 l5=Nb1i=1∑Nb∣uxi(tbi,−5)−uxi(tbi,5)∣2
l 6 = 1 N b ∑ i = 1 N b ∣ v x i ( t b i , − 5 ) − v x i ( t b i , 5 ) ∣ 2 l_6=\frac{1}{N_b}\sum_{i=1}^{N_b}|v_x^i(t_b^i,-5)-v_x^i(t_b^i,5)|^2 l6=Nb1i=1∑Nb∣vxi(tbi,−5)−vxi(tbi,5)∣2
l 7 = 1 N f ∑ i = 1 N f ∣ u t + 0.5 v x x + ( u 2 + v 2 ) v ∣ 2 l_7=\frac{1}{N_f}\sum_{i=1}^{N_f}|u_t+0.5v_{xx}+(u^2+v^2)v|^2 l7=Nf1i=1∑Nf∣ut+0.5vxx+(u2+v2)v∣2
l 8 = 1 N f ∑ i = 1 N f ∣ − v t + 0.5 u x x + ( u 2 + v 2 ) u ∣ 2 l_8=\frac{1}{N_f}\sum_{i=1}^{N_f}|-v_t+0.5u_{xx}+(u^2+v^2)u|^2 l8=Nf1i=1∑Nf∣−vt+0.5uxx+(u2+v2)u∣2
3. 代码实现
原代码GitHub链接，原代码使用TensorFlow 1.0框架实现，本文使用pytorch进行复现。
定义 PINN 类继承 nn.Module，实现网络的基本结构。网络由4个隐藏层组成，均为线性层，采用 Tanh 激活函数，在 PINN 类中定义 init_bias 方法对网络的参数进行初始化，具体为对权重 weight 采用 Xavier 正态初始化，偏置 bias 置零。
# 定义 PINN 类: 网络结构
class PINN(nn.Module):
def __init__(self, layers):
super().__init__()
# 定义网络
self.net = nn.Sequential(
nn.Linear(layers[0], layers[1]),
nn.Tanh(),
nn.Linear(layers[1], layers[2]),
nn.Tanh(),
nn.Linear(layers[2], layers[3]),
nn.Tanh(),
nn.Linear(layers[3], layers[4]),
nn.Tanh(),
nn.Linear(layers[4], layers[5]),
)
# 初始化偏置
self.init_bias()
# 初始化偏置
def init_bias(self):
for layer in self.net.children():
if isinstance(layer, nn.Linear):
# 权重采用Xavier正态初始化
nn.init.xavier_normal_(layer.weight, gain=5 / 3) # gain 根据激活函数选择, torch.nn.init.calculate_gain()查看
# 偏置置0
nn.init.constant_(layer.bias, 0.)
# 前向传播
def forward(self, x, t):
xt = torch.cat((x, t), dim=1)
xt = self.net(xt)
u = xt[:, 0].unsqueeze(-1)
v = xt[:, 1].unsqueeze(-1)
return u, v
封装 PINNSolver 类，该类实现对方程进行求解，并且验证。其中 get_grad 方法实现对网络进行前向传播得到预测值，并对相关叶子节点进行梯度计算，包括一阶导数和二阶导数；loss_fn 方法定义了该网络的损失函数。
在 train 方法中，使用Adam优化器进行梯度下降，初始学习率为0.001，并采用学习率衰减策略，具体策略为 ReduceLROnPlateau，这个调度器会根据监控的指标动态地调整学习率。当监测到性能指标在一段时间内停止改善（即达到了一个性能平台期，plateau），它就会按照预先设定的规则降低学习率，这种方法有助于模型跳出局部最优。
class PINNSolver:
def __init__(self, model, x0, u0, v0, tb, X_f, lb, ub):
self.lb = torch.from_numpy(lb.astype(np.float32)).to(device)
self.ub = torch.from_numpy(ub.astype(np.float32)).to(device)
self.u0 = torch.from_numpy(u0.astype(np.float32)).to(device)
self.v0 = torch.from_numpy(v0.astype(np.float32)).to(device)
# 左边界
X0 = np.concatenate((x0, 0 * x0), 1) # (x0, 0)
self.x0 = torch.from_numpy(X0[:, 0:1].astype(np.float32)).to(device)
self.t0 = torch.from_numpy(X0[:, 1:2].astype(np.float32)).to(device)
#
X_lb = np.concatenate((0 * tb + lb[0], tb), 1) # (lb[0], tb)
X_ub = np.concatenate((0 * tb + ub[0], tb), 1) # (ub[0], tb)
self.x_lb = torch.from_numpy(X_lb[:, 0:1].astype(np.float32)).to(device)
self.t_lb = torch.from_numpy(X_lb[:, 1:2].astype(np.float32)).to(device)
self.x_ub = torch.from_numpy(X_ub[:, 0:1].astype(np.float32)).to(device)
self.t_ub = torch.from_numpy(X_ub[:, 1:2].astype(np.float32)).to(device)
self.x_f = torch.from_numpy(X_f[:, 0:1].astype(np.float32)).to(device)
self.t_f = torch.from_numpy(X_f[:, 1:2].astype(np.float32)).to(device)
self.model = model
# 前向传播并获取梯度信息
def get_grad(self, model, x, t, fist_gras=False, second_grad=False):
# 前向传播
x.requires_grad, t.requires_grad = True, True
u, v = model(x, t)
if not (fist_gras or second_grad):
return u, v
if fist_gras:
# 计算一阶导数
u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]
v_x = torch.autograd.grad(v, x, grad_outputs=torch.ones_like(v), create_graph=True)[0]
u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]
v_t = torch.autograd.grad(v, t, grad_outputs=torch.ones_like(v), create_graph=True)[0]
if second_grad:
# 计算二阶导数
u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]
v_xx = torch.autograd.grad(v_x, x, grad_outputs=torch.ones_like(v_x), create_graph=True)[0]
return u, v, u_x, v_x, u_t, v_t, u_xx, v_xx
return u, v, u_x, v_x, u_t, v_t
# 损失函数定义
def loss_fn(self, model):
# 计算预测值和梯度
u0_pred, v0_pred = self.get_grad(model, self.x0, self.t0,
fist_gras=False, second_grad=False)
u_lb_pred, v_lb_pred, u_lb_x, v_lb_x, _, _ = self.get_grad(model, self.x_lb, self.t_lb,
fist_gras=True, second_grad=False)
u_ub_pred, v_ub_pred, u_ub_x, v_ub_x, _, _ = self.get_grad(model, self.x_ub, self.t_ub,
fist_gras=True, second_grad=False)
u_f_pred, v_f_pred, _, _, u_f_t, v_f_t, u_f_xx, v_f_xx = self.get_grad(model, self.x_f, self.t_f,
fist_gras=True, second_grad=True)
# 损失函数
l1 = torch.mean(torch.square(u0_pred - self.u0))
l2 = torch.mean(torch.square(v0_pred - self.v0))
l3 = torch.mean(torch.square(u_lb_pred - u_ub_pred))
l4 = torch.mean(torch.square(v_lb_pred - v_ub_pred))
l5 = torch.mean(torch.square(u_lb_x - u_ub_x))
l6 = torch.mean(torch.square(v_lb_x - v_ub_x))
l7 = torch.mean(torch.square(u_f_t+0.5*v_f_xx+(torch.square(u_f_pred)+torch.square(v_f_pred))*v_f_pred))
l8 = torch.mean(torch.square((-1)*v_f_t+0.5*u_f_xx+(torch.square(u_f_pred)+torch.square(v_f_pred))*u_f_pred))
loss = l1 + l2 + l3 + l4 + l5 + l6 + l7 + l8
return loss
# 训练
def train(self, nIter):
# 优化器
optim = torch.optim.Adam(self.model.parameters(), lr=1e-3)
# 学习率衰减策略
scheduler = ReduceLROnPlateau(optim, factor=0.1, patience=100, verbose=True)
# 开始训练时间
t0 = time()
# 开始训练
for epoch in tqdm(range(nIter)):
loss = self.loss_fn(self.model)
# 清零梯度
optim.zero_grad()
# 反向传播
loss.backward()
# 更新参数
optim.step()
# 更新学习率
scheduler.step(loss)
# print(f"Iteration={epoch+1}, loss={loss.cpu().detach().numpy()}")
if epoch % 100 == 0:
tqdm.write(f"Iteration={epoch:04d}, loss={loss.cpu().detach().numpy():.5f}")
# 训练时间
print(f"TrainTime: {time()-t0:0.4f} seconds")
# 预测
def predict(self, X_start, u_start, v_start, h_start):
print("开始预测")
x_start = torch.from_numpy(X_start[:, 0].astype(np.float32)).reshape(-1, 1).to(device)
t_start = torch.from_numpy(X_start[:, 1].astype(np.float32)).reshape(-1, 1).to(device)
u, v, u_x, v_x, u_t, v_t, u_xx, v_xx = self.get_grad(self.model, x_start, t_start, fist_gras=True, second_grad=True)
h = torch.sqrt(u ** 2 + v ** 2)
f_u_star = torch.square(u_t+0.5*v_xx+(torch.square(u)+torch.square(v))*v)
f_v_star = torch.square((-1)*v_t+0.5*u_xx+(torch.square(u)+torch.square(v))*u)
# 误差(二范数)
print("计算误差")
error_u = np.linalg.norm(u_start - u.cpu().detach().numpy(), 2) / np.linalg.norm(u_start, 2)
error_v = np.linalg.norm(v_start - v.cpu().detach().numpy(), 2) / np.linalg.norm(v_start, 2)
error_h = np.linalg.norm(h_start - h.cpu().detach().numpy(), 2) / np.linalg.norm(h_start, 2)
print(f'Error u: {error_u}')
print(f'Error v: {error_v}')
print(f'Error h: {error_h}')
print("预测完毕")
return u, v, h, f_u_star, f_v_star
4. 训练结果
运行结果如图所示：
原代码运行结果：
5. 源代码
完整源代码GitHub链接，带有详细注释：https://github.com/freshman-ip/PINN
若有不足之处，欢迎批评指正！
标签：PINN,pred,self,torch,0.5,pytorch,薛定谔,np,grad From： https://blog.csdn.net/m0_57543713/article/details/139869154