PINN：物理信息神经网络求解微分方程

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

### 3.2 定义神经网络结构
```python
class PINN(nn.Module):
    def __init__(self, layers):
        super(PINN, self).__init__()
        self.net = nn.Sequential()
        for i in range(len(layers)-1):
            self.net.add_module(f'linear_{i}', nn.Linear(layers[i], layers[i+1]))
            if i < len(layers)-2:
                self.net.add_module(f'tanh_{i}', nn.Tanh())
        # 权重初始化（常用 Xavier 均匀分布）
        for m in self.net.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                nn.init.constant_(m.bias, 0.0)
    
    def forward(self, x, t):
        # 输入形状 (N, 1) 的 x 和 t，拼接后送入网络
        inputs = torch.cat([x, t], dim=1)
        return self.net(inputs)

这里我们设计网络层数为 [2, 20, 20, 20, 1]，采用 Tanh 激活函数。Xavier 初始化有助于加快训练。

3.3 定义损失函数

def compute_loss(model, x_f, t_f, x_b, t_b, u_b,
                 x_0, t_0, u_0, nu=0.01/torch.pi):
    # PDE 残差损失
    x_f.requires_grad = True
    t_f.requires_grad = True
    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 * u_x - nu * u_xx
    loss_pde = torch.mean(f ** 2)
    
    # 边界损失
    u_pred_b = model(x_b, t_b)
    loss_bc = torch.mean((u_pred_b - u_b) ** 2)
    
    # 初始条件损失
    u_pred_0 = model(x_0, t_0)
    loss_ic = torch.mean((u_pred_0 - u_0) ** 2)
    
    return loss_pde + loss_bc + loss_ic

3.4 生成训练数据

# 配置点（域内）
N_f = 10000
x_f = torch.rand(N_f, 1) * 2 - 1  # [-1,1]
t_f = torch.rand(N_f, 1)

# 边界点
N_b = 500
x_b_left = -1 + torch.zeros(N_b//2, 1)
t_b_left = torch.rand(N_b//2, 1)
x_b_right = 1 + torch.zeros(N_b//2, 1)
t_b_right = torch.rand(N_b//2, 1)
x_b = torch.cat([x_b_left, x_b_right], 0)
t_b = torch.cat([t_b_left, t_b_right], 0)
u_b = torch.zeros(N_b, 1)  # Dirichlet BC = 0

# 初始点
N_0 = 500
x_0 = torch.rand(N_0, 1) * 2 - 1
t_0 = torch.zeros(N_0, 1)
u_0 = -torch.sin(torch.pi * x_0)

3.5 训练模型

model = PINN([2, 20, 20, 20, 1])
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10001):
    optimizer.zero_grad()
    loss = compute_loss(model, x_f, t_f, x_b, t_b, u_b, x_0, t_0, u_0)
    loss.backward()
    optimizer.step()
    if epoch % 1000 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item():.3e}')

3.6 结果可视化

# 在均匀网格上预测
X = torch.linspace(-1, 1, 256)
T = torch.linspace(0, 1, 100)
X, T = torch.meshgrid(X, T, indexing='ij')
x_flat = X.reshape(-1, 1)
t_flat = T.reshape(-1, 1)
with torch.no_grad():
    u_pred = model(x_flat, t_flat).numpy().reshape(256, 100)

plt.pcolormesh(T.numpy(), X.numpy(), u_pred, shading='auto', cmap='jet')
plt.colorbar(label='u')
plt.xlabel('t')
plt.ylabel('x')
plt.title('PINN solution for Burgers equation')
plt.show()