基于物理信息神经网络（PINN）求解Burgers方程(附PyTorch源代码)

基于物理信息神经网络（PINN）求解Burgers方程

在这篇博客中，我们将展示如何通过物理信息神经网络（Physics-Informed Neural Network, PINN）求解经典的 Burgers 方程。PINN 不仅依赖于训练数据的监督学习，还将偏微分方程的物理信息嵌入损失函数，从而在满足边界和初始条件的同时，约束神经网络解的物理一致性。

Burgers 方程在流体动力学、非线性波传播等领域有广泛的应用，是测试数值求解方法的经典模型之一。

1. 问题背景

Burgers 方程的表达式为：

u t + u u x = ν u x x , x ∈ [ − 1 , 1 ] , t ≥ 0 u_t + u u_x = \nu u_{xx}, \quad x \in [-1, 1], t \geq 0 ut+uux=νuxx,x∈[−1,1],t≥0

其中 ν = 0.01 / π \nu = 0.01/\pi ν=0.01/π 为粘性系数， u ( t , x ) u(t, x) u(t,x) 表示时间 t t t 和空间 x x x 上的解。

初始条件为：
u ( 0 , x ) = − sin ⁡ ( π x ) , x ∈ [ − 1 , 1 ] u(0, x) = -\sin(\pi x), \quad x \in [-1, 1] u(0,x)=−sin(πx),x∈[−1,1]

边界条件为：
u ( t , − 1 ) = u ( t , 1 ) = 0 , t ≥ 0 u(t, -1) = u(t, 1) = 0, \quad t \geq 0 u(t,−1)=u(t,1)=0,t≥0

Burgers 方程既包含非线性项 u u x u u_x uux，也包含扩散项 ν u x x \nu u_{xx} νuxx，为此类问题引入 PINN 可以有效地学习解的动态特性。

2. 神经网络架构

我们首先设计一个神经网络用于逼近 Burgers 方程的解。PINN 输入时间 t t t 和空间坐标 x x x，输出预测的解 u ( t , x ) u(t, x) u(t,x)。我们选用 7 层全连接神经网络，每层包含 100 个神经元，激活函数为 Tanh。该网络能够处理输入的非线性关系，并生成较为平滑的解。

import torch
import torch.nn as nn

class PINN(nn.Module):
    def __init__(self):
        super(PINN, self).__init__()
        self.layer = nn.Sequential(
            nn.Linear(2, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 1)
        )
    
    def forward(self, t, x):
        # 拼接时间和空间坐标作为输入
        u = self.layer(torch.cat([t, x], dim=1))
        return u

网络说明

• 输入：时间 t t t 和空间坐标 x x x（形状均为 ( n , 1 ) (n, 1) (n,1)）
• 输出：在 ( t , x ) (t, x) (t,x) 上的解 u ( t , x ) u(t, x) u(t,x)
• 使用 Tanh 激活函数能够有效处理非线性问题，同时可以使模型生成平滑的函数。

3. 损失函数设计

PINN 的核心在于将物理方程作为约束条件加入到网络的损失函数中。为此，我们设计了以下三类损失函数：

3.1 物理损失（Physics Loss）

物理损失是通过自动微分计算偏导数，并将其代入 Burgers 方程中进行约束。具体来说，通过 PyTorch 的自动求导功能，我们能够轻松计算模型输出关于 t t t 和 x x x 的偏导数，进一步计算残差。

def physics_loss(model, t, x):
    u = model(t, x)  # 预测解
    u_t = torch.autograd.grad(u, t, torch.ones_like(t), create_graph=True)[0]
    u_x = torch.autograd.grad(u, x, torch.ones_like(x), create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x, torch.ones_like(x), create_graph=True)[0]
    
    # Burgers 方程的残差
    f = (u_t + u * u_x - (0.01 / torch.pi) * u_xx).pow(2).mean()
    return f

3.2 边界条件损失（Boundary Loss）

边界条件要求 x = − 1 x = -1 x=−1 和 x = 1 x = 1 x=1 时，解为 0。因此，我们可以通过计算这些边界上的解值，定义边界损失。

def boundary_loss(model, t, x_left, x_right):
    u_left = model(t, x_left)  # x = -1 处的解
    u_right = model(t, x_right)  # x = 1 处的解
    return u_left.pow(2).mean() + u_right.pow(2).mean()

3.3 初始条件损失（Initial Condition Loss）

初始条件要求在 t = 0 t = 0 t=0 时， u ( 0 , x ) = − sin ⁡ ( π x ) u(0, x) = -\sin(\pi x) u(0,x)=−sin(πx)。因此，初始损失通过比较模型预测的解和精确初始条件来定义。

def initial_loss(model, x):
    t_0 = torch.zeros_like(x)  # t = 0
    u_init = model(t_0, x)  # 初始时刻模型预测的解
    u_exact = -torch.sin(torch.pi * x)  # 精确初始解
    return (u_init - u_exact).pow(2).mean()  # 初始损失

4. 训练过程

我们使用 Adam 优化器对 PINN 进行训练，训练过程结合了物理损失、边界损失和初始损失。

def train(model, optimizer, num_epochs):
    losses = []
    for epoch in range(num_epochs):
        optimizer.zero_grad()

        # 随机采样时间 t 和空间坐标 x
        t = torch.rand(10000, 1, requires_grad=True)
        x = torch.rand(10000, 1, requires_grad=True) * 2 - 1  # x ∈ [-1, 1]

        # 计算物理损失
        f_loss = physics_loss(model, t, x)

        # 计算边界条件损失
        t_bc = torch.rand(500, 1)
        x_left = -torch.ones(500, 1)
        x_right = torch.ones(500, 1)
        bc_loss = boundary_loss(model, t_bc, x_left, x_right)

        # 计算初始条件损失
        x_ic = torch.rand(1000, 1) * 2 - 1
        ic_loss = initial_loss(model, x_ic)

        # 总损失为三者之和
        loss = f_loss + bc_loss + ic_loss
        loss.backward()
        optimizer.step()

        losses.append(loss.item())

        if epoch % 100 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}')
    
    return losses

训练过程中，每个 epoch 随机采样 ( t , x ) (t, x) (t,x) 点，并计算各类损失。最终的总损失是物理损失、边界损失和初始损失的加权和。

5. 可视化

5.1 训练损失曲线

训练损失曲线展示了模型训练过程中的损失变化情况。通常，损失逐渐减小，表明模型在不断逼近精确解。

import matplotlib.pyplot as plt

def plot_loss(losses):
    plt.plot(losses)
    plt.yscale('log')  # 使用对数坐标轴
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss Curve')
    plt.show()

5.2 数值解二维等高线图

我们还可以在 ( x , t ) (x, t) (x,t) 平面上绘制解的二维等高线图。这里，我们通过采样 ( x , t ) (x, t) (x,t) 点并利用模型预测数值解 u ( t , x ) u(t, x) u(t,x)，绘制出解的空间-时间分布。

def plot_solution(model):
    x = torch.linspace(-1, 1, 100).unsqueeze(1)
    t = torch.linspace(0, 1, 100).unsqueeze(1)
    X, T = torch.meshgrid(x.squeeze(), t.squeeze())
    x_flat = X.reshape(-1, 1)
    t_flat = T.reshape(-1, 1)

    with torch.no_grad():
        u_pred = model(t_flat, x_flat).numpy().reshape(100, 100)

    plt.contourf(X.numpy(), T.numpy(), u_pred, 100, cmap='viridis')
    plt.colorbar(label='u(t, x)')
    plt.xlabel('x')
    plt.ylabel('t')
    plt.title('Burgers Equation Solution in (x, t)')
    plt.show()

通过以上步骤，我们可以得到 Burgers 方程在整个 ( x , t ) (x, t) (x,t) 平面上的解，并通过二维等高线图可视化展示解的变化趋势。

数值结果（未进行调参，仅供参考） ：

6. 结论

通过 PINN 结合自动微分和物理约束，能够高效求解 Burgers 方程这样的非线性偏微分方程。PINN 不仅能够逼近未知的数值解，还能通过物理一致性约束增强解的合理性和物理解释性。相比传统的数值方法，PINN 具有处理高维问题和复杂边界条件的潜力。

总结

本文介绍了如何通过物理信息神经网络（PINN）求解经典的 Burgers 方程，详细讲解了网络结构、损失函数设计、训练过程以及可视化解的步骤。PINN 是一种强大的工具，能够将物理方程与神经网络相结合，用于求解复杂的偏微分方程。

完整代码

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


# 绘制损失曲线
def plot_loss(losses):
    plt.figure(figsize=(8, 5))
    plt.plot(losses, color='blue', lw=2)
    plt.xlabel('Epoch', fontsize=14)
    plt.ylabel('Loss', fontsize=14)
    plt.title('Training Loss Curve', fontsize=16)
    plt.grid(True)
    plt.tight_layout()
    plt.show()

# 绘制数值解图像
def plot_solution(model):
    x = torch.linspace(-1, 1, 100).unsqueeze(1)
    t = torch.full((100, 1), 0.0)  # 在t=0时绘制解
    with torch.no_grad():
        u_pred = model(t, x).numpy()
    
    # 参考解 u(0,x) = -sin(πx)
    u_exact = -np.sin(np.pi * x.numpy())

    plt.figure(figsize=(8, 5))
    plt.plot(x.numpy(), u_pred, label='Predicted Solution', color='red', lw=2)
    plt.plot(x.numpy(), u_exact, label='Exact Solution (Initial)', color='blue', lw=2, linestyle='dashed')
    plt.xlabel('x', fontsize=14)
    plt.ylabel('u(t=0, x)', fontsize=14)
    plt.title('Burgers Equation Solution at t=0', fontsize=16)
    plt.legend(fontsize=12)
    plt.grid(True)
    plt.tight_layout()
    plt.show()
# 绘制整个 (x, t) 平面的解
def plot_solution_3d(model):
    # 创建 (x, t) 网格
    x = torch.linspace(-1, 1, 100).unsqueeze(1)
    t = torch.linspace(0, 1, 100).unsqueeze(1)
    X, T = torch.meshgrid(x.squeeze(), t.squeeze())
    
    # 将 X 和 T 拉平，方便模型预测
    x_flat = X.reshape(-1, 1)
    t_flat = T.reshape(-1, 1)

    with torch.no_grad():
        u_pred = model(t_flat, x_flat).numpy().reshape(16, 16)

    # 绘制三维曲面图
    fig = plt.figure(figsize=(10, 7))
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_surface(X.numpy(), T.numpy(), u_pred, cmap='viridis')

    ax.set_xlabel('x', fontsize=12)
    ax.set_ylabel('t', fontsize=12)
    ax.set_zlabel('u(t, x)', fontsize=12)
    ax.set_title('Solution of Burgers Equation on (x, t) Plane', fontsize=14)

    plt.show()

# 绘制二维等高线图
def plot_solution_contour(model):
    # 创建 (x, t) 网格
    x = torch.linspace(-1, 1, 100).unsqueeze(1)
    t = torch.linspace(0, 1, 100).unsqueeze(1)
    X, T = torch.meshgrid(x.squeeze(), t.squeeze())

    # 将 X 和 T 拉平，方便模型预测
    x_flat = X.reshape(-1, 1)
    t_flat = T.reshape(-1, 1)

    with torch.no_grad():
        u_pred = model(t_flat, x_flat).numpy().reshape(16, 16)

    # 绘制二维等高线图
    plt.figure(figsize=(8, 6))
    plt.contourf(X.numpy(), T.numpy(), u_pred, 100, cmap='viridis')
    plt.colorbar(label='u(t, x)')
    
    plt.xlabel('x', fontsize=12)
    plt.ylabel('t', fontsize=12)
    plt.title('Contour Plot of Burgers Equation Solution', fontsize=14)
    
    plt.tight_layout()
    plt.show()
    
    
# 定义网络结构
class PINN(nn.Module):
    def __init__(self):
        super(PINN, self).__init__()
        self.layer = nn.Sequential(
            nn.Linear(2, 16), nn.Tanh(),
            nn.Linear(16, 16), nn.Tanh(),
            nn.Linear(16, 16), nn.Tanh(),
            nn.Linear(16, 16), nn.Tanh(),
            nn.Linear(16, 16), nn.Tanh(),
            nn.Linear(16, 1)
        )
    
    def forward(self, t, x):
        u = self.layer(torch.cat([t, x], dim=1))
        return u

# 定义Burgers方程中的物理损失
def physics_loss(model, t, x):
    u = model(t, x)
    u_t = torch.autograd.grad(u, t, torch.ones_like(t), create_graph=True)[0]
    u_x = torch.autograd.grad(u, x, torch.ones_like(x), create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x, torch.ones_like(x), create_graph=True)[0]
    f = (u_t + u * u_x - (0.01 / torch.pi) * u_xx).pow(2).mean()
    return f

# 定义边界条件损失
def boundary_loss(model, t, x_left, x_right):
    u_left = model(t, x_left)
    u_right = model(t, x_right)
    return u_left.pow(2).mean() + u_right.pow(2).mean()

# 初始条件损失
def initial_loss(model, x):
    t_0 = torch.zeros_like(x)
    u_init = model(t_0, x)
    u_exact = -torch.sin(torch.pi * x)
    return (u_init - u_exact).pow(2).mean()

# 训练模型并记录损失
def train(model, optimizer, num_epochs):
    losses = []
    for epoch in range(num_epochs):
        optimizer.zero_grad()

        # 随机采样 t 和 x，并确保 requires_grad=True
        t = torch.rand(3000, 1, requires_grad=True)
        x = torch.rand(3000, 1, requires_grad=True) * 2 - 1  # x ∈ [-1, 1]

        # 物理损失
        f_loss = physics_loss(model, t, x)

        # 边界条件损失
        t_bc = torch.rand(500, 1)
        x_left = -torch.ones(500, 1)
        x_right = torch.ones(500, 1)
        bc_loss = boundary_loss(model, t_bc, x_left, x_right)

        # 初始条件损失
        x_ic = torch.rand(1000, 1) * 2 - 1
        ic_loss = initial_loss(model, x_ic)

        # 总损失
        loss = f_loss + bc_loss + ic_loss
        loss.backward()
        optimizer.step()

        # 记录损失
        losses.append(loss.item())

        if epoch % 100 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}')
    
    return losses

# 初始化模型和优化器
model = PINN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 训练模型
losses = train(model, optimizer, num_epochs = 100000)

# 绘制训练损失曲线
plot_loss(losses)

# 绘制数值解图像
plot_solution(model)
plot_solution_3d(model)   # 三维曲面图
plot_solution_contour(model)   # 二维等高线图