用 PINN（Physics-Informed Neural Networks）求解二阶Diffusion方程的PyTorch 代码

好的，以下是用 PINN 求解二阶 Diffusion 方程的 PyTorch 代码：

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

# 定义求解区域的左右边界
x_left = 0.0
x_right = 1.0

# 定义时间和空间步长
delta_t = 0.001
delta_x = 0.05

# 定义神经网络的输入维度和隐藏层维度
input_dim = 2
hidden_dim = 10

# 定义神经网络的输出维度
output_dim = 1

# 定义 Diffusion 方程的初始条件和边界条件
def u_init(x):
    return np.sin(np.pi * x)

def u_left(t):
    return 0.0

def u_right(t):
    return 0.0

# 定义神经网络模型
class PINN(torch.nn.Module):
    def __init__(self):
        super(PINN, self).__init__()
        self.fc1 = torch.nn.Linear(input_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = torch.nn.Linear(hidden_dim, output_dim)

    def forward(self, x, t):
        # 拼接时间和空间坐标作为神经网络的输入
        xt = torch.cat([x, t], dim=1)
        h1 = torch.tanh(self.fc1(xt))
        h2 = torch.tanh(self.fc2(h1))
        out = self.fc3(h2)
        return out

# 定义 PINN 的损失函数
def pinn_loss(model, x_left, x_right, delta_t, delta_x):
    # 定义空间和时间的网格
    x = torch.linspace(x_left, x_right, int((x_right - x_left) / delta_x) + 1).unsqueeze(1)
    t = torch.linspace(0, 1, int(1 / delta_t) + 1).unsqueeze(1)

    # 计算模型在内部网格点上的预测值
    x_internal = x[1:-1]
    t_internal = t[1:-1]
    u_internal = model(x_internal, t_internal)

    # 计算模型在左右边界上的预测值
    u_left_boundary = model(x_left, t_internal)
    u_right_boundary = model(x_right, t_internal)

    # 计算初始条件和边界条件的损失
    init_loss = torch.mean((model(x, torch.zeros_like(x)) - u_init(x)) ** 2)
    left_boundary_loss = torch.mean((u_left_boundary - u_left(t_internal)) ** 2)
    right_boundary_loss = torch.mean((u_right_boundary - u_right(t_internal)) ** 2)

    # 计算 Diffusion 方程的残差
    u_x = torch.autograd.grad(u_internal, x_internal, grad_outputs=torch.ones_like(u_internal), create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x_internal, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]
    f = u_xx - (1 / delta_t) * (u_internal - u_init(x_internal))

    # 计算残差的损失
    residual_loss = torch.mean(f ** 2)

    # 计算总的损失
    loss = init_loss + left_boundary_loss + right_boundary_loss + residual_loss
    return loss

# 初始化神经网络模型和优化器
model = PINN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练神经网络模型
for i in range(10000):
    optimizer.zero_grad()
    loss = pinn_loss(model, x_left, x_right, delta_t, delta_x)
    loss.backward()
    optimizer.step()

    if i % 1000 == 0:
        print("Iteration {}, Loss = {}".format(i, loss.item()))

# 使用训练好的神经网络模型进行预测
x_test = torch.linspace(x_left, x_right, 101).unsqueeze(1)
t_test = torch.linspace(0, 1, 1001).unsqueeze(1)
u_test = model(x_test, t_test).detach().numpy()

# 绘制预测结果
X, T = np.meshgrid(x_test, t_test)
plt.pcolormesh(X, T, u_test, cmap='coolwarm')
plt.colorbar()
plt.xlabel("x")
plt.ylabel("t")
plt.title("PINN Solution to Diffusion Equation")
plt.show()

这段代码中，首先定义了求解区域的左右边界、时间和空间步长等参数，然后定义了 Diffusion 方程的初始条件和边界条件，以及神经网络模型和损失函数。在训练神经网络模型时，使用了 Adam 优化器，最后使用训练好的模型进行预测，并绘制了预测结果。