ML: PINN代码实例-耦合ODE

本文尝试使用PINN求解Lotka-Volterra方程。在阅读本文之前，请阅读dyfluid.com - ML: PINN代码实例-1D槽道流，防止出现不理解跟不上的情况。

Lotka-Volterra方程是2个耦合的ODE，其中不同的参数，可以具有不同的解。在这个链接中，可以自行更改ODE的参数，实时的获得解析解。Lotka-Volterra方程可以表示为

其中。用户可以自行调试这些参数，可以发现Lotka-Volterra方程会预测非常不同的结果。

PINN网格划分

在OpenFOAM libtorch环境下，可以通过下面几行来生成网格：

auto mesh = torch::arange(0, 1.9, 0.1).reshape({-1,1});
mesh.set_requires_grad(true);

其中的arrange函数，请参考dyfluid.com - ML: libtorch张量（二）。

神经网络搭建

在OpenFOAM libtorch环境下，可以通过下述代码搭建神经网络：

class NN
:
    public torch::nn::Module 
{
    torch::nn::Sequential net_;

public:

    NN()
    {
        net_ = 
            register_module
            (
                "net", 
                torch::nn::Sequential
                (
                    torch::nn::Linear(1,10),
                    torch::nn::Tanh(),
                    torch::nn::Linear(10,10),
                    torch::nn::Tanh(),
                    torch::nn::Linear(10,2)
                )
            );
    }

    auto forward(torch::Tensor x)
    {
        return net_->forward(x);
    }
};

其为一个常规的全连接网络，采用Tanh激活函数，每次10个神经元，输入层1个参数，输出层2个参数。

超参数以及自动微分

PINN在训练网络的时候，需要指定梯度下降的方法。这部分理论请参考dyfluid.com - ML: 手算反向传播与自动微分。具体的，本算例训练50000次，采用Adam梯度更新方法。学习率为0.001。损失计算的方法采用均方差最小化。在OpenFOAM libtorch环境下的代码可以写为

torch::manual_seed(0); 
int iter = 1;
int IterationNum = 50000;
double lr = 1e-3;
auto crit = torch::nn::MSELoss();
auto model = std::make_shared<NN>();
auto opti = 
    std::make_shared<torch::optim::Adam>
    (
        model->parameters(), 
        torch::optim::AdamOptions(lr)
    );

上述内容均属于常规操作。对于本文的神经网络，可以通过下面的代码进行输出：

auto xy = model->forward(mesh);
auto x = xy.index({torch::indexing::Slice(), 0}).reshape({-1,1});
auto y = xy.index({torch::indexing::Slice(), 1}).reshape({-1,1});

需要注意的是，本文的神经网络输出的为2个参数，因此xytensor的形状为{19,2}，其中的19表示网格单元数。其中的index操作，将xy的列元素提取出来，形成单独的x以及y。在此基础上，可以进行自动微分：

auto dxdt = 
     torch::autograd::grad
     (
         {x}, {mesh}, {torch::ones_like(x)}, true, true
     )[0];

auto dydt = 
     torch::autograd::grad
     (
         {y}, {mesh}, {torch::ones_like(y)}, true, true
     )[0];

损失

本文的耦合的ODE具有2个方程，为了更好地训练，本文定义3个损失：

• 方程损失；
• 边界损失；
• 数据损失；

方程的损失可以定义为

auto f1 = dxdt - a1*x + c12*x*y;
auto f2 = dydt + a2*y - c21*x*y;
auto loss1 = crit(f1, 0.0*f1);
auto loss2 = crit(f2, 0.0*f2);

初次之外，边界损失可以定义为

auto loss3 = crit(x0, x.index({0}));
auto loss4 = crit(y0, y.index({0}));

其中的index函数可以参考dyfluid.com - ML: libtorch张量（二）。下面我们定义数据损失：

#include "data.H"

数据损失即为神经网络预测的点与提取的解析解的点的差异。最终有总损失：

auto loss = loss1 + loss3 + loss2 + loss4 + loss5 + loss6;

随后进行训练即可。

运行算例

请首先按照cfd-china.com - OpenFOAM libtorch tutorial step by step来配置OpenFOAM libtorch环境。也可以在dyfluid.com下载配置好的虚拟机。本算例的源代码请参考本文文末。

将源代码直接创立一个.C文件后，即可通过wmake来编译。编译后可以进行训练。训练过程的log如下：

dyfluid@dyfluid-virtual-machine:~/source/code$ lotkaVolterra 
500 loss = 229.306
1000 loss = 61.6972
1500 loss = 47.0121
2000 loss = 41.3949
....
49000 loss = 0.0256056
50000 loss = 0.0149014
Done!

在这里我们将PINN计算的结果（线条）与解析解（圆点）对比，如下图（通过gnuplot绘制）。

源代码，阅读原文，免费下载！去找【数据驱动-ML: PINN代码实例-耦合ODE】。