3. 波动方程初边值问题

任务目标

使用物理信息神经网络（PINN）求解一维波动方程初边值问题，模拟波的传播过程，并与经典有限差分法（FDM）的数值解进行对比验证。

问题定义

求解如下一维波动方程：

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

初始条件：u(x, 0) = sin(πx)，∂u/∂t(x, 0) = 0
边界条件：u(0, t) = u(1, t) = 0
波速：c = 1

该问题的解析解为 u(x, t) = sin(πx)cos(πt)，可用于验证PINN与FDM的精度。

实现方案

采用全连接神经网络近似解 u_θ(x, t)，损失函数包含三部分：

  1. 方程残差：L_PDE = (∂²u/∂t² − c²∂²u/∂x²)²
  2. 初始条件残差：L_IC = (u(x,0) − sin(πx))² + (∂u/∂t(x,0))²
  3. 边界条件残差：L_BC = u(0,t)² + u(1,t)²

使用自动微分计算各阶导数，通过梯度下降最小化总损失。FDM采用中心差分格式，时间步进使用蛙跳法。

环境依赖

pip install torch numpy matplotlib

完整代码


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

# ============ PINN 模型 ============
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):
        return self.net(torch.cat([x, t], dim=1))

# ============ 损失函数 ============
def pinn_loss(model, x, t):
    x.requires_grad_(True)
    t.requires_grad_(True)
    u = model(x, t)
    
    # 一阶导数
    u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_t = torch.autograd.grad(u, t, 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]
    u_tt = torch.autograd.grad(u_t, t, grad_outputs=torch.ones_like(u_t), create_graph=True)[0]
    
    # PDE残差 (c=1)
    f_pde = u_tt - u_xx
    
    # 初始条件: u(x,0)=sin(πx), u_t(x,0)=0
    x_ic = torch.linspace(0, 1, 100).reshape(-1, 1).requires_grad_(True)
    t_ic = torch.zeros_like(x_ic)
    u_ic = model(x_ic, t_ic)
    u_ic_t = torch.autograd.grad(u_ic, t_ic, grad_outputs=torch.ones_like(u_ic), create_graph=True)[0]
    loss_ic = torch.mean((u_ic - torch.sin(np.pi * x_ic))**2) + torch.mean(u_ic_t**2)
    
    # 边界条件: u(0,t)=u(1,t)=0
    t_bc = torch.linspace(0, 1, 100).reshape(-1, 1)
    x_bc0 = torch.zeros_like(t_bc)
    x_bc1 = torch.ones_like(t_bc)
    loss_bc = torch.mean(model(x_bc0, t_bc)**2) + torch.mean(model(x_bc1, t_bc)**2)
    
    # 内部点PDE残差
    loss_pde = torch.mean(f_pde**2)
    
    return loss_pde + loss_ic + loss_bc

# ============ 有限差分法 (FDM) ============
def fdm_solve(Nx=101, Nt=1001, c=1.0):
    dx = 1.0 / (Nx - 1)
    dt = 1.0 / (Nt - 1)
    r = (c * dt / dx)**2
    
    x = np.linspace(0, 1, Nx)
    u = np.zeros((Nt, Nx))
    u[0, :] = np.sin(np.pi * x)
    
    # 蛙跳法第一步
    u[1, 1:-1] = u[0, 1:-1] + 0.5 * r * (u[0, :-2] - 2*u[0, 1:-1] + u[0, 2:])
    
    # 时间推进
    for n in range(1, Nt-1):
        u[n+1, 1:-1] = 2*u[n, 1:-1] - u[n-1, 1:-1] + r * (u[n, :-2] - 2*u[n, 1:-1] + u[n, 2:])
    
    return x, u

# ============ 训练与对比 ============
def train_and_compare():
    torch.manual_seed(42)
    model = PINN()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    
    # 内部采样点
    x = torch.rand(2000, 1) * 1.0
    t = torch.rand(2000, 1) * 1.0
    
    # 训练
    for epoch in range(3000):
        optimizer.zero_grad()
        loss = pinn_loss(model, x, t)
        loss.backward()
        optimizer.step()
        if epoch % 500 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.6f}")
    
    # PINN预测
    x_test = torch.linspace(0, 1, 101).reshape(-1, 1)
    t_test = torch.linspace(0, 1, 101).reshape(-1, 1)
    X, T = torch.meshgrid(x_test.squeeze(), t_test.squeeze(), indexing='ij')
    u_pinn = model(X.reshape(-1,1), T.reshape(-1,1)).detach().numpy().reshape(101, 101)
    
    # FDM求解
    x_fdm, u_fdm = fdm_solve()
    
    # 解析解
    Xn, Tn = np.meshgrid(x_fdm, np.linspace(0, 1, 1001))
    u_exact = np.sin(np.pi * Xn) * np.cos(np.pi * Tn)
    
    # 对比误差（取t=0.5时刻）
    t_idx = 500
    pinn_t05 = u_pinn[:, 50]  # t=0.5
    fdm_t05 = u_fdm[t_idx, :]
    exact_t05 = np.sin(np.pi * x_fdm) * np.cos(np.pi * 0.5)
    
    print(f"\nPINN 最大误差 (t=0.5): {np.max(np.abs(pinn_t05 - exact_t05)):.6f}")
    print(f"FDM  最大误差 (t=0.5): {np.max(np.abs(fdm_t05 - exact_t05)):.6f}")
    
    # 绘制对比图
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    
    im1 = axes[0].imshow(u_pinn.T, extent=[0,1,0,1], origin='lower', aspect='auto', cmap='RdBu')
    axes[0].set_title('PINN Solution')
    axes[0].set_xlabel('x'); axes[0].set_ylabel('t')
    plt.colorbar(im1, ax=axes[0])
    
    im2 = axes[1].imshow(u_fdm[::10, :].T, extent=[0,1,0,1], origin='lower', aspect='auto', cmap='RdBu')
    axes[1].set_title('FDM Solution')
    axes[1].set_xlabel('x'); axes[1].set_ylabel('t')
    plt.colorbar(im2, ax=axes[1])
    
    axes[2].plot(x_fdm, exact_t05, 'k-', label='Exact', linewidth=2)
    axes[2].plot(x_fdm, pinn_t05, 'r--', label='PINN', linewidth=1.5)
    axes[2].plot(x_fdm, fdm_t05, 'b:', label='FDM', linewidth=1.5)
    axes[2].set_title('t=0.5 Cross-section')
    axes[2].set_xlabel('x'); axes[2].set_ylabel('u')
    axes[2].legend()
    axes[2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('wave_comparison.png', dpi=150)
    plt.show()

if __name__ == "__main__":
    train_and_compare()

运行与验证

  1. 直接运行脚本：python wave_pinn.py
  2. 训练过程中损失应持续下降，最终收敛至约1e-4～1e-5量级
  3. 程序输出t=0.5时刻PINN与FDM相对解析解的最大误差
  4. 生成wave_comparison.png，包含三个子图：PINN解云图、FDM解云图、t=0.5时刻的横截面对比曲线

通过标准：PINN与FDM的最大误差均在0.05以内，且对比图中PINN与FDM曲线均与解析解曲线基本重合，说明PINN成功模拟了波的传播过程。

常见问题

  • 若训练不收敛，可增大内部采样点数量或增加网络层数/宽度
  • FDM的稳定性条件要求 r = (c·dt/dx)² ≤ 1，本实现中 r ≈ 0.01，满足稳定性要求
  • PINN训练时间约2～5分钟（CPU），若使用GPU可显著加速
← 上一章 📖 返回目录 下一章 →