基于物理信息神经网络(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。该网络能够处理输入的非线性关系,并生成较为平滑的解。
python
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 的偏导数,进一步计算残差。
python
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。因此,我们可以通过计算这些边界上的解值,定义边界损失。
python
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)。因此,初始损失通过比较模型预测的解和精确初始条件来定义。
python
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 进行训练,训练过程结合了物理损失、边界损失和初始损失。
python
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 训练损失曲线
训练损失曲线展示了模型训练过程中的损失变化情况。通常,损失逐渐减小,表明模型在不断逼近精确解。
python
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),绘制出解的空间-时间分布。
python
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 是一种强大的工具,能够将物理方程与神经网络相结合,用于求解复杂的偏微分方程。
完整代码
python
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) # 二维等高线图