ML: PINN代码实例-耦合ODE

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

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

\[\begin{split} \frac{\rd x}{\rd t} - a_1x + c_{12}xy=0 \\ \frac{\rd y}{\rd t} + a_2y - c_{21}xy=0 \end{split}\]

其中\(a_1=5,a_2=6,c_{12}=c_{21}=1\)。用户可以在这个链接中自行调试这些参数，可以发现Lotka-Volterra方程会预测非常不同的结果。

PINN网格划分

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

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

其中的arrange函数，请参考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在训练网络的时候，需要指定梯度下降的方法。这部分理论请参考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函数可以参考ML: libtorch张量（二）。下面我们定义数据损失：

#include "data.H"

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

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

随后进行训练即可。

运行算例

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

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

dyfluid@dyfluid-virtual-machine:~/source/code/ML/4-lotkaVolterra-coupled$ lotkaVolterra 
loss = 229.306
loss = 61.6972
loss = 47.0121
loss = 41.3949
loss = 37.382
loss = 34.4413
loss = 31.5408
loss = 15.2306
loss = 10.4132
loss = 8.61934
loss = 7.34001
loss = 6.29833
loss = 5.37021
loss = 4.4909
loss = 3.69553
loss = 3.01945
loss = 2.4769
loss = 2.06165
loss = 1.7428
loss = 1.48779
loss = 1.27579
loss = 1.08665
loss = 0.912517
loss = 0.777866
loss = 0.670766
loss = 0.583666
loss = 0.512146
loss = 0.471097
loss = 0.411536
loss = 0.366413
loss = 0.334259
loss = 0.307333
loss = 0.284608
loss = 0.26518
loss = 0.248345
loss = 0.233543
loss = 0.220394
loss = 0.208376
loss = 0.197417
loss = 0.187142
loss = 0.177603
loss = 0.168636
loss = 0.159863
loss = 0.152558
loss = 0.143973
loss = 0.1415
loss = 0.129767
loss = 0.123372
loss = 0.138523
loss = 0.111974
loss = 0.106436
loss = 0.10646
loss = 0.105635
loss = 0.0929342
loss = 0.0890507
loss = 0.0854735
loss = 0.0819435
loss = 0.078699
loss = 0.0756077
loss = 0.0841585
loss = 0.0698074
loss = 0.0687944
loss = 0.0667595
loss = 0.0620816
loss = 0.0595187
loss = 0.0570337
loss = 0.0547324
loss = 0.0525469
loss = 0.0503508
loss = 0.0484192
loss = 0.0463215
loss = 0.0444001
loss = 0.0425323
loss = 0.0408346
loss = 0.0390203
loss = 0.0374206
loss = 0.0358038
loss = 0.0343059
loss = 0.0348738
loss = 0.0315029
loss = 0.0302128
loss = 0.0289717
loss = 0.027824
loss = 0.0266658
loss = 0.0256056
loss = 0.024596
loss = 0.0237517
loss = 0.0227224
loss = 0.0218631
loss = 0.0210535
loss = 0.0202801
loss = 0.0196323
loss = 0.0188489
loss = 0.0181904
loss = 0.0175674
loss = 0.016975
loss = 0.0165304
loss = 0.0158832
loss = 0.0153789
loss = 0.0149014
Done!

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

源代码：

C文件：

#include <torch/torch.h>
using namespace std;

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);
    }
};


int main()
{
    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)
        );

    double a1 = 5.0;
    double a2 = 6.0;
    double c12 = 1.0;
    double c21 = 1.0;
    auto x0 = torch::full({1}, 2.0);
    auto y0 = torch::full({1}, 2.0);

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

    for (int i = 0; i < IterationNum; i++) 
    {
        opti->zero_grad();
        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});

        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];
        
        auto f1 = dxdt - a1*x + c12*x*y;
        auto f2 = dydt + a2*y - c21*x*y;

        #include "data.H"

        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}));

        auto loss = loss1 + loss3 + loss2 + loss4 + loss5 + loss6;
        loss.backward();
        opti->step();

        if (iter % 500 == 0) 
        {
            cout<< iter << " loss = " << loss.item<double>() << endl;

            auto meshxy = torch::cat({mesh,x,y}, 1);
            std::ofstream filex("results");
            filex << meshxy << "\n";
            filex.close();
        }
        iter++;
    }

    std::cout<< "Done!" << std::endl;
    return 0;
}

data.H:

auto dataPox = torch::tensor({{0.55991}, {0.77715}});
auto train1 = model->forward(dataPox);
auto xtrain = train1.index({torch::indexing::Slice(), 0}).unsqueeze(1);
auto dataX = torch::tensor({{15.35796}, {6.48939}}); 
auto dataPoy = torch::tensor({{0.65818}, {0.88577}});
auto dataY = torch::tensor({{7.91881}, {11.90934}});
auto train2 = model->forward(dataPoy);
auto ytrain = train2.index({torch::indexing::Slice(), 1}).unsqueeze(1);
auto loss5 = crit(xtrain, dataX);
auto loss6 = crit(ytrain, dataY);