pytorch pinn求解满足区域[0,4]×[0,4],连续初始条件为u0(x) = 0.5+sin(π(x+y)/2)的二维无黏 Burgers 方程的代码

以下是使用 PyTorch 实现求解无黏 Burgers 方程的代码：

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

# 定义模型
class PINN(nn.Module):
    def __init__(self):
        super(PINN, self).__init__()
        self.fc1 = nn.Linear(2, 50)
        self.fc2 = nn.Linear(50, 50)
        self.fc3 = nn.Linear(50, 50)
        self.fc4 = nn.Linear(50, 1)

    def forward(self, x):
        a = torch.sin(torch.pi * (x[:, 0] + x[:, 1]) / 2)
        h = torch.cat([x, a.unsqueeze(1)], dim=1)
        h = torch.tanh(self.fc1(h))
        h = torch.tanh(self.fc2(h))
        h = torch.tanh(self.fc3(h))
        out = self.fc4(h)
        return out

# 定义损失函数
def loss_fn(u_pred, f_pred):
    return torch.mean(torch.pow(u_pred, 2)) + torch.mean(torch.pow(f_pred, 2))

# 定义边界条件
def bc_fn(x):
    return torch.cat([model(x), torch.zeros(x.shape[0], 1)], dim=1)

# 定义初始条件
x = torch.linspace(0, 4, 100)
y = torch.linspace(0, 4, 100)
X, Y = torch.meshgrid(x, y)
x0 = torch.cat([X.reshape(-1, 1), Y.reshape(-1, 1)], dim=1)
u0 = 0.5 + torch.sin(torch.pi * (X + Y) / 2)

# 转换为张量
x0 = x0.float()
u0 = u0.float()

# 定义模型和优化器
model = PINN()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 训练模型
for epoch in range(10000):
    optimizer.zero_grad()

    # 计算模型预测值和损失函数
    u_pred = model(x0)
    x0.requires_grad = True
    u_grad = torch.autograd.grad(u_pred, x0, grad_outputs=torch.ones_like(u_pred), create_graph=True)[0]
    u_x = u_grad[:, 0]
    u_y = u_grad[:, 1]
    u_xx = torch.autograd.grad(u_x, x0, grad_outputs=torch.ones_like(u_x), create_graph=True)[0][:, 0]
    u_yy = torch.autograd.grad(u_y, x0, grad_outputs=torch.ones_like(u_y), create_graph=True)[0][:, 1]
    f_pred = u_grad[:, 2] + u_pred * (u_xx + u_yy)

    # 计算边界条件
    x_lb = x0[x0[:, 0] == 0.0]
    x_ub = x0[x0[:, 0] == 4.0]
    y_lb = x0[x0[:, 1] == 0.0]
    y_ub = x0[x0[:, 1] == 4.0]
    u_lb = bc_fn(x_lb)
    u_ub = bc_fn(x_ub)
    u_lb_y = bc_fn(y_lb)
    u_ub_y = bc_fn(y_ub)

    # 计算损失函数
    loss = loss_fn(u_pred, f_pred) + loss_fn(u_lb, torch.zeros_like(u_lb)) + \
           loss_fn(u_ub, torch.zeros_like(u_ub)) + loss_fn(u_lb_y, torch.zeros_like(u_lb_y)) + \
           loss_fn(u_ub_y, torch.zeros_like(u_ub_y))

    # 反向传播和优化
    loss.backward()
    optimizer.step()

    # 打印损失函数值
    if epoch % 1000 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, 10000, loss.item()))

# 绘制结果
u_pred = model(x0).detach().numpy().reshape(100, 100)
plt.contourf(X.numpy(), Y.numpy(), u_pred, cmap='jet')
plt.colorbar()
plt.xlabel('x')
plt.ylabel('y')
plt.show()

代码中首先定义了一个名为 `PINN` 的类，它继承自 `nn.Module` 并包含了一个四层的全连接神经网络。该神经网络的输入是二维坐标 $x$，输出是 $u(x)$ 的预测值。在神经网络中，我们将 $x$ 和 $\sin(\pi(x+y)/2)$ 按列拼接起来，并通过 `tanh` 激活函数进行非线性变换，最终输出一个标量。

接着定义了损失函数 `loss_fn`，它由两部分组成，分别对应于 $u(x)$ 的平方和 $f(x)$ 的平方和。其中 $f(x)$ 是无黏 Burgers 方程的右侧项，可以通过对 $u(x)$ 求导得到。

然后定义了边界条件函数 `bc_fn`，它将模型预测值和 0 拼接在一起作为边界条件。在本例中，我们选择将 $x=0$，$x=4$，$y=0$ 和 $y=4$ 作为边界条件。

接着定义了初始条件，它由网格点 $(x,y)$ 和对应的 $u(x,y)$ 值组成。

在模型训练过程中，我们首先计算模型预测值和损失函数，然后计算边界条件的损失函数，最后将两部分损失函数加权求和。在反向传播和优化过程中，我们使用 PyTorch 内置的自动求导功能计算损失函数对模型参数的导数，并使用 Adam 优化算法更新模型参数。

最后，我们将模型预测值的等高线图绘制出来。可以看到，模型能够很好地拟合出初始条件，并在整个区域内得到了比较准确的解。