1. 一维热方程PINN求解

任务目标

使用物理信息神经网络（PINN）求解一维热方程，获得与解析解一致的数值解，并可视化预测误差。本案例聚焦于PINN的核心实现路径：构建神经网络、定义损失函数（含控制方程残差与边界/初始条件）、训练并评估。

问题定义

求解如下一维热方程（扩散方程）：

∂u/∂t = α · ∂²u/∂x²,  x ∈ [0, 1], t ∈ [0, 1]

取扩散系数 α = 1，初始条件 u(x, 0) = sin(πx)，边界条件 u(0, t) = u(1, t) = 0。该问题的解析解为：

u_exact(x, t) = sin(πx) · exp(−π²t)

环境依赖

需要 Python 3.9+，以及以下库：

pip install torch numpy matplotlib

完整实现

以下代码完整实现了PINN求解过程，包含数据生成、网络定义、损失计算、训练循环与可视化：

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

# 设置随机种子保证可复现
torch.manual_seed(42)
np.random.seed(42)

# 解析解
def exact_solution(x, t):
    return torch.sin(np.pi * x) * torch.exp(-np.pi**2 * t)

# 神经网络模型
class PINN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 50), nn.Tanh(),
            nn.Linear(50, 50), nn.Tanh(),
            nn.Linear(50, 50), nn.Tanh(),
            nn.Linear(50, 1)
        )
    
    def forward(self, x, t):
        # 输入拼接为 (x, t)
        xt = torch.cat([x, t], dim=1)
        return self.net(xt)

# 生成训练数据
N_f = 10000  # 内部配点数
N_b = 200    # 边界点数
N_i = 200    # 初始点数

# 内部点（控制方程残差点）
x_f = torch.rand(N_f, 1, requires_grad=True)
t_f = torch.rand(N_f, 1, requires_grad=True)

# 边界点 x=0 和 x=1
t_b = torch.rand(N_b, 1)
x_b0 = torch.zeros_like(t_b)
x_b1 = torch.ones_like(t_b)

# 初始点 t=0
x_i = torch.rand(N_i, 1)
t_i = torch.zeros_like(x_i)

# 模型与优化器
model = PINN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 训练循环
epochs = 3000
for epoch in range(epochs):
    optimizer.zero_grad()
    
    # 1. 控制方程残差损失
    u = model(x_f, t_f)
    u_t = torch.autograd.grad(u, t_f, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_x = torch.autograd.grad(u, x_f, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x_f, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]
    f = u_t - u_xx  # 热方程残差
    loss_f = torch.mean(f**2)
    
    # 2. 边界条件损失
    u_b0 = model(x_b0, t_b)
    u_b1 = model(x_b1, t_b)
    loss_b = torch.mean(u_b0**2) + torch.mean(u_b1**2)
    
    # 3. 初始条件损失
    u_i = model(x_i, t_i)
    u_exact_i = exact_solution(x_i, t_i)
    loss_i = torch.mean((u_i - u_exact_i)**2)
    
    # 总损失
    loss = loss_f + loss_b + loss_i
    loss.backward()
    optimizer.step()
    
    if (epoch + 1) % 500 == 0:
        print(f'Epoch {epoch+1}, Loss: {loss.item():.2e}')

# 在网格上评估
x_grid = torch.linspace(0, 1, 100).reshape(-1, 1)
t_grid = torch.linspace(0, 1, 100).reshape(-1, 1)
X, T = torch.meshgrid(x_grid.squeeze(), t_grid.squeeze(), indexing='ij')
X_flat = X.reshape(-1, 1)
T_flat = T.reshape(-1, 1)

# 预测解
u_pred = model(X_flat, T_flat).detach().reshape(100, 100)
# 解析解
u_exact = exact_solution(X_flat, T_flat).reshape(100, 100)
# 误差
error = torch.abs(u_pred - u_exact)

# 可视化
fig, axes = plt.subplots(1, 3, figsize=(15, 4))

im0 = axes[0].imshow(u_pred.numpy(), extent=[0, 1, 0, 1], origin='lower', aspect='auto', cmap='viridis')
axes[0].set_title('PINN Prediction')
axes[0].set_xlabel('t')
axes[0].set_ylabel('x')
plt.colorbar(im0, ax=axes[0])

im1 = axes[1].imshow(u_exact.numpy(), extent=[0, 1, 0, 1], origin='lower', aspect='auto', cmap='viridis')
axes[1].set_title('Exact Solution')
axes[1].set_xlabel('t')
axes[1].set_ylabel('x')
plt.colorbar(im1, ax=axes[1])

im2 = axes[2].imshow(error.numpy(), extent=[0, 1, 0, 1], origin='lower', aspect='auto', cmap='hot')
axes[2].set_title('Absolute Error')
axes[2].set_xlabel('t')
axes[2].set_ylabel('x')
plt.colorbar(im2, ax=axes[2])

plt.tight_layout()
plt.savefig('pinn_heat_equation.png', dpi=150)
plt.show()

# 输出最大误差
print(f'Max absolute error: {error.max().item():.4e}')

运行方式

将上述代码保存为 pinn_heat.py，在终端执行：

python pinn_heat.py

验证方法与判定标准

运行后应观察到：

1. 训练损失随迭代次数下降，最终损失值通常低于 1e-3（具体值取决于随机种子与训练时长）；
2. 生成三幅并排热力图：PINN预测解、解析解、误差分布；
3. 预测解与解析解在视觉上几乎一致，误差热力图整体颜色偏暗（数值小）；
4. 终端输出的最大误差通常小于 0.05（3000轮训练后典型值约 1e-2 量级）。

若误差过大（如最大误差超过0.1），可增加训练轮数至5000或增加内部配点数至20000后重新运行。

关键说明

本实现采用Adam优化器与固定学习率，未使用学习率衰减或L-BFGS二阶优化器。对于更复杂方程，可引入学习率调度或混合优化策略以加速收敛。损失函数中三项权重均为1，实际应用中可根据各项收敛速度调整权重系数。