生成用PINN求解薛定谔方程的pytorch代码

下面是一个用 PINN（Physics-Informed Neural Networks）求解一维薛定谔方程的 PyTorch 代码示例：

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# 设置随机数种子
np.random.seed(1234)
torch.manual_seed(1234)

# 定义神经网络模型
class SchrodingerPINN(nn.Module):
    def __init__(self):
        super(SchrodingerPINN, self).__init__()
        self.fc1 = nn.Linear(1, 50)
        self.fc2 = nn.Linear(50, 50)
        self.fc3 = nn.Linear(50, 50)
        self.fc4 = nn.Linear(50, 1)

    def forward(self, x):
        u = torch.sin(np.pi * x)
        x = torch.cat([x, u], 1)
        x = torch.tanh(self.fc1(x))
        x = torch.tanh(self.fc2(x))
        x = torch.tanh(self.fc3(x))
        x = self.fc4(x)
        return x

# 定义损失函数
def loss_fn(model, x_bc, u_bc, x_ic):
    # 边界条件损失
    u_bc_pred = model(x_bc)
    loss_bc = torch.mean(torch.abs(u_bc - u_bc_pred))

    # 初始条件损失
    u_ic_pred = model(x_ic)
    loss_ic = torch.mean(torch.square(u_ic_pred))

    # 总损失
    loss = loss_bc + loss_ic
    return loss

# 定义训练函数
def train(model, x_bc, u_bc, x_ic, epochs):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    for epoch in range(epochs):
        optimizer.zero_grad()
        loss = loss_fn(model, x_bc, u_bc, x_ic)
        loss.backward()
        optimizer.step()
        if epoch % 1000 == 0:
            print(f"Epoch {epoch}, Loss {loss.item()}")

# 定义边界条件和初始条件
x_bc = torch.tensor([[0.0], [1.0]], requires_grad=True)
u_bc = torch.tensor([[0.0], [0.0]], requires_grad=True)
x_ic = torch.tensor([[0.5]], requires_grad=True)

# 创建神经网络模型
model = SchrodingerPINN()

# 训练模型
train(model, x_bc, u_bc, x_ic, epochs=5000)

# 可视化结果
x = torch.linspace(0.0, 1.0, 100).reshape(-1, 1)
u_pred = model(x).detach().numpy()

plt.plot(x, np.sin(np.pi * x), label="True solution")
plt.plot(x, u_pred, label="PINN solution")
plt.legend()
plt.show()

在上面的代码中，我们定义了一个 `SchrodingerPINN` 类作为神经网络模型，其中包含了 4 个全连接层。在训练函数中，我们使用 Adam 优化器对模型进行训练。

损失函数包含了边界条件损失和初始条件损失。在边界条件损失中，我们计算了边界点处的函数值和模型预测值之间的平均绝对误差；在初始条件损失中，我们计算了初始点处的函数值和模型预测值之间的平均平方误差。

最后，我们使用 matplotlib 可视化了真实解和模型预测解的比较。