3. Burgers方程粘性解

任务说明

本案例使用物理信息神经网络（PINN）求解一维粘性Burgers方程，在给定初始条件和边界条件下，捕捉激波的形成与传播过程。通过将偏微分方程残差作为损失项加入训练，使神经网络在拟合初始/边界数据的同时满足控制方程，从而验证网络对非线性对流项与粘性扩散项耦合问题的拟合能力。

控制方程与定解条件

求解如下一维粘性Burgers方程（ν = 0.01/π）：

u_t + u·u_x = ν·u_xx,  x ∈ [-1, 1],  t ∈ [0, 1]

初始条件：u(x, 0) = -sin(πx)；边界条件：u(-1, t) = u(1, t) = 0。该设置在无粘极限下会形成激波，粘性项使解保持光滑，是PINN处理非线性对流项的经典基准问题。

环境依赖

需要 Python 3.9+、PyTorch 2.0+、NumPy、Matplotlib。无需额外数据文件，训练数据由程序内部采样生成。

完整实现

import torch
import torch.nn as nn
import numpy as np

# ---------- 网络定义 ----------
class PINN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 20), nn.Tanh(),
            nn.Linear(20, 20), nn.Tanh(),
            nn.Linear(20, 20), nn.Tanh(),
            nn.Linear(20, 20), nn.Tanh(),
            nn.Linear(20, 1)
        )
    def forward(self, x, t):
        return self.net(torch.cat([x, t], dim=1))

# ---------- 方程残差 ----------
def pde_residual(model, x, t):
    x.requires_grad_(True)
    t.requires_grad_(True)
    u = model(x, t)
    u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]
    nu = 0.01 / np.pi
    return u_t + u * u_x - nu * u_xx

# ---------- 数据采样 ----------
def sample_data(n_ic=100, n_bc=100, n_pde=10000, seed=0):
    np.random.seed(seed)
    torch.manual_seed(seed)
    # 初始条件点
    x_ic = torch.rand(n_ic, 1) * 2 - 1
    t_ic = torch.zeros_like(x_ic)
    u_ic = -torch.sin(np.pi * x_ic)
    # 边界点（x = ±1）
    t_bc = torch.rand(n_bc, 1)
    x_bc = torch.cat([torch.ones(n_bc//2, 1), -torch.ones(n_bc - n_bc//2, 1)], dim=0)
    u_bc = torch.zeros_like(x_bc)
    # PDE配点
    x_pde = torch.rand(n_pde, 1) * 2 - 1
    t_pde = torch.rand(n_pde, 1)
    return (x_ic, t_ic, u_ic), (x_bc, t_bc, u_bc), (x_pde, t_pde)

# ---------- 训练 ----------
def train(model, n_iter=20000, lr=1e-3):
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    (x_ic, t_ic, u_ic), (x_bc, t_bc, u_bc), (x_pde, t_pde) = sample_data()
    for it in range(n_iter):
        opt.zero_grad()
        # 初始条件损失
        u_pred_ic = model(x_ic, t_ic)
        loss_ic = torch.mean((u_pred_ic - u_ic) ** 2)
        # 边界条件损失
        u_pred_bc = model(x_bc, t_bc)
        loss_bc = torch.mean((u_pred_bc - u_bc) ** 2)
        # PDE残差损失
        r = pde_residual(model, x_pde, t_pde)
        loss_pde = torch.mean(r ** 2)
        # 总损失（权重可调）
        loss = loss_ic + loss_bc + loss_pde
        loss.backward()
        opt.step()
        if it % 5000 == 0:
            print(f'iter {it:5d}, loss={loss.item():.3e}, '
                  f'ic={loss_ic.item():.2e}, bc={loss_bc.item():.2e}, pde={loss_pde.item():.2e}')

# ---------- 运行 ----------
if __name__ == '__main__':
    model = PINN()
    train(model, n_iter=20000)
    torch.save(model.state_dict(), 'burgers_pinn.pt')
    print('训练完成，模型已保存为 burgers_pinn.pt')

运行方式

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

python burgers_pinn.py

训练约需2～5分钟（CPU即可）。程序每5000步打印一次总损失及各分量损失。

验证方法与判定标准

训练完成后，可通过以下方式验证模型是否成功捕捉激波：

1. 损失收敛检查：最终总损失应低于 1e-3，且PDE残差损失与边界/初始损失均呈下降趋势并趋于平稳。
2. 解剖面可视化：在训练后添加如下代码，绘制不同时刻的 u(x, t) 剖面，观察激波是否在 x=0 附近形成并随时间传播：

import matplotlib.pyplot as plt
model.eval()
x = torch.linspace(-1, 1, 200).reshape(-1, 1)
with torch.no_grad():
    for t_val in [0.0, 0.25, 0.5, 0.75, 1.0]:
        t = torch.full_like(x, t_val)
        u = model(x, t).numpy().flatten()
        plt.plot(x.numpy().flatten(), u, label=f't={t_val}')
plt.legend()
plt.xlabel('x'); plt.ylabel('u')
plt.title('Burgers 方程 PINN 解剖面')
plt.savefig('burgers_profile.png', dpi=150)
print('剖面图已保存为 burgers_profile.png')

通过条件：在 t=0 时解应接近 -sin(πx)；随着 t 增大，在 x=0 附近出现陡峭过渡（激波），且剖面保持光滑无振荡；t=1 时解在 x=0 处从负值陡升至正值，两侧趋于0。若剖面出现明显振荡或激波位置偏移超过0.1，可增大 n_iter 至40000或增加网络宽度至30再训练。