[sharonyue]

Hello there,

I would like to share my PINN-libtorch code for cavity flow. It can be compiled in the environment of OpenFOAM-8,9,10,11, etc. Please see the following:

#include <torch/torch.h>

using namespace std;

//创建神经网络
struct NN
:
    torch::nn::Cloneable<NN>
{
    torch::nn::Sequential net_;
public:
    NN()
    {
        reset();
    }
    void reset() override
    {
        int neu = 100;
        net_ = register_module
        (
            "net", 
            torch::nn::Sequential
            (
                torch::nn::Linear(2,neu),
                torch::nn::Tanh(),
                torch::nn::Linear(neu,neu),
                torch::nn::Tanh(),
                torch::nn::Linear(neu,neu),
                torch::nn::Tanh(),
                torch::nn::Linear(neu,3)
            )
        );
    }
    auto forward(torch::Tensor x)
    {
        return net_->forward(x);
    }
};

int main()
{
    int iter = 1;
    double learningRate = 0.0005;
    int IterationNum = 10000;
    auto crit = torch::nn::MSELoss();
    auto net = std::make_shared<NN>();

    auto opti = 
       torch::optim::Adam
       (
           net->parameters(),
           torch::optim::AdamOptions(learningRate)
       );   

    torch::manual_seed(0); 
    int pixel = 70;//网格数
    double nu = 0.1;//粘度
    double vel = 3.0;//顶盖速度
    
    auto bottom = torch::full({pixel, 2}, 0.0);
    auto top = torch::full({pixel, 2}, 1.0);
    auto left = torch::full({pixel, 2}, 0.0);
    auto right = torch::full({pixel, 2}, 1.0);
    auto u_top = torch::full({pixel}, vel);
    auto v_top = torch::full({pixel}, 0.0);
    
    for (int i = 0; i < pixel; i++)
    {
        int min = 1;
        int max = 100;
        int numb = min + rand() % (max - min + 1);
        bottom[i][0] = numb/100.0;
    }
    
    for (int i = 0; i < pixel; i++)
    {
        int min = 1;
        int max = 100;
        int numb = min + rand() % (max - min + 1);
        top[i][0] = numb/100.0;
    }
    
    for (int i = 0; i < pixel; i++)
    {
        int min = 1;
        int max = 100;
        int numb = min + rand() % (max - min + 1);
        left[i][1] = numb/100.0;
    }
    
    for (int i = 0; i < pixel; i++)
    {
        int min = 1;
        int max = 100;
        int numb = min + rand() % (max - min + 1);
        right[i][1] = numb/100.0;
    }
    
    auto mesh = torch::rand({pixel*pixel, 2});
    
    mesh.requires_grad_(true);
    top.requires_grad_(true);
    left.requires_grad_(true);
    right.requires_grad_(true);
    bottom.requires_grad_(true);

    for (int i = 0; i < IterationNum; i++) 
    {
        //#include "toCUDA.H" 

        opti.zero_grad();

        auto UP = net->forward(mesh);
        auto UPtop = net->forward(top);
        auto UPright = net->forward(right);
        auto UPbottom = net->forward(bottom);
        auto UPleft = net->forward(left);

        auto u_pred = UP.index({torch::indexing::Slice(), 0});
        auto v_pred = UP.index({torch::indexing::Slice(), 1});
        auto p_pred = UP.index({torch::indexing::Slice(), 2});

        auto u_pred_top = UPtop.index({torch::indexing::Slice(), 0});
        auto v_pred_top = UPtop.index({torch::indexing::Slice(), 1}); 
        auto p_pred_top = UPtop.index({torch::indexing::Slice(), 2});  

        auto u_pred_left = UPleft.index({torch::indexing::Slice(), 0});
        auto v_pred_left = UPleft.index({torch::indexing::Slice(), 1}); 
        auto p_pred_left = UPleft.index({torch::indexing::Slice(), 2});  
        
        auto u_pred_bottom = UPbottom.index({torch::indexing::Slice(), 0});
        auto v_pred_bottom = UPbottom.index({torch::indexing::Slice(), 1}); 
        auto p_pred_bottom = UPbottom.index({torch::indexing::Slice(), 2});  

        auto u_pred_right = UPright.index({torch::indexing::Slice(), 0});
        auto v_pred_right = UPright.index({torch::indexing::Slice(), 1}); 
        auto p_pred_right = UPright.index({torch::indexing::Slice(), 2});  

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

        auto dpdx = dpdMesh.index({torch::indexing::Slice(), 0});
        auto dpdy = dpdMesh.index({torch::indexing::Slice(), 1});

        auto dpdxMesh = 
            torch::autograd::grad
            (
                {dpdx},
                {mesh},
                {torch::ones_like(dpdx)},
                true,
                true
            )[0];
        auto dpdyMesh = 
            torch::autograd::grad
            (
                {dpdy},
                {mesh},
                {torch::ones_like(dpdy)},
                true,
                true
            )[0];
        auto dpdxx = dpdxMesh.index({torch::indexing::Slice(), 0});
        auto dpdyy = dpdyMesh.index({torch::indexing::Slice(), 1});

        auto dpdLeft = 
            torch::autograd::grad
            (
                {p_pred_left},
                {left},
                {torch::ones_like(p_pred_left)},
                true,
                true
            )[0];

        auto dpdx_left = dpdLeft.index({torch::indexing::Slice(), 0});
        auto dpdy_left = dpdLeft.index({torch::indexing::Slice(), 1});

        auto dpdBottom = 
            torch::autograd::grad
            (
                {p_pred_bottom},
                {bottom},
                {torch::ones_like(p_pred_bottom)},
                true,
                true
            )[0];

        auto dpdx_bottom = dpdBottom.index({torch::indexing::Slice(), 0});
        auto dpdy_bottom = dpdBottom.index({torch::indexing::Slice(), 1});

        auto dpdRight = 
            torch::autograd::grad
            (
                {p_pred_right},
                {right},
                {torch::ones_like(p_pred_right)},
                true,
                true
            )[0];

        auto dpdx_right = dpdRight.index({torch::indexing::Slice(), 0});
        auto dpdy_right = dpdRight.index({torch::indexing::Slice(), 1});

        auto dpdTop = 
            torch::autograd::grad
            (
                {p_pred_top},
                {top},
                {torch::ones_like(p_pred_top)},
                true,
                true
            )[0];

        auto dpdx_top = dpdTop.index({torch::indexing::Slice(), 0});
        auto dpdy_top = dpdTop.index({torch::indexing::Slice(), 1});

        auto dudMesh = 
            torch::autograd::grad
            (
                {u_pred},
                {mesh},
                {torch::ones_like(u_pred)},
                true,
                true
            )[0];
        auto dudx = dudMesh.index({torch::indexing::Slice(), 0});
        auto dudy = dudMesh.index({torch::indexing::Slice(), 1});
 
        auto dvdMesh = 
            torch::autograd::grad
            (
                {v_pred},
                {mesh},
                {torch::ones_like(v_pred)},
                true,
                true
            )[0];
        auto dvdx = dvdMesh.index({torch::indexing::Slice(), 0});
        auto dvdy = dvdMesh.index({torch::indexing::Slice(), 1});

        auto dudxMesh = 
            torch::autograd::grad
            (
                {dudx},
                {mesh},
                {torch::ones_like(dudx)},
                true,
                true
            )[0];
        auto dudxx = dudxMesh.index({torch::indexing::Slice(), 0});

        auto dudyMesh = 
            torch::autograd::grad
            (
                {dudy},
                {mesh},
                {torch::ones_like(dudy)},
                true,
                true
            )[0];
        auto dudyy = dudyMesh.index({torch::indexing::Slice(), 1});
        auto dudxy = dudyMesh.index({torch::indexing::Slice(), 0});

        auto dvdxMesh = 
            torch::autograd::grad
            (
                {dvdx},
                {mesh},
                {torch::ones_like(dvdx)},
                true,
                true
            )[0];
        auto dvdyMesh = 
            torch::autograd::grad
            (
                {dvdy},
                {mesh},
                {torch::ones_like(dvdy)},
                true,
                true
            )[0];
        auto dvdxx = dvdxMesh.index({torch::indexing::Slice(), 0});
        auto dvdxy = dvdxMesh.index({torch::indexing::Slice(), 1});
        auto dvdyy = dvdyMesh.index({torch::indexing::Slice(), 1});
    
        auto cont1 = dudx + dvdy;
        auto mom1 = 2.0*u_pred*dudx + v_pred*dudy + u_pred*dvdy 
          - nu*(dudxx + dudyy) - nu*(dudxx + dvdxy) + dpdx;
        auto mom2 = u_pred*dvdx + v_pred*dudx + 2.0*v_pred*dvdy 
          - nu*(dvdxx + dvdyy) - nu*(dudxy + dvdyy) + dpdy;
        auto loss_cont1 = crit(cont1, 0*cont1);
        auto loss_mom1 = crit(mom1, 0*mom1);
        auto loss_mom2 = crit(mom2, 0*mom2);
    
        auto loss_top1 = crit(u_pred_top, u_top); 
        auto loss_top2 = crit(v_pred_top, v_top); 
        auto loss_top3 = crit(dpdy_top, 0*dpdy_top);
        auto loss_top = loss_top1 + loss_top2 + loss_top3;
        auto loss_left1 = crit(u_pred_left, 0*u_pred_left);
        auto loss_left2 = crit(v_pred_left, 0*v_pred_left);
        auto loss_left3 = crit(dpdx_left, 0*dpdx_left);
        auto loss_left = loss_left1 + loss_left2 + loss_left3;
        auto loss_right1 = crit(u_pred_right, 0*u_pred_right);
        auto loss_right2 = crit(v_pred_right, 0*v_pred_right);
        auto loss_right3 = crit(dpdx_right, 0*dpdx_right);
        auto loss_right = loss_right1 + loss_right2 + loss_right3;
        auto loss_bottom1 = crit(u_pred_bottom, 0*u_pred_bottom);
        auto loss_bottom2 = crit(v_pred_bottom, 0*v_pred_bottom);
        auto loss_bottom3 = crit(dpdy_bottom, 0*dpdy_bottom);
        auto loss_bottom = loss_bottom1 + loss_bottom2 + loss_bottom3;

        auto loss_dataP = crit(p_pred[0], 0*p_pred[0]);
	    auto loss_data = 100*(loss_dataP);
        auto loss_boundary = 
            10.0*(loss_top + loss_left + loss_right + loss_bottom);
        auto loss_pde = 1.0*(loss_cont1 + loss_mom1 + loss_mom2); 

        auto loss = loss_pde + loss_boundary + loss_data;
        
        loss.backward();
        opti.step();
        iter++;
        
        if (iter % 50 == 0)
        {
            cout<< iter<< ": cont: " << loss_cont1.item<double>()
                << ", pde: " << loss_pde.item<double>()
                << ", boundary: " << loss_boundary.item<double>()
                << ", All: " << loss.item<double>() << endl;

            std::ofstream meshfile("meshfile");
            meshfile << mesh << "\n";
            std::ofstream ufile("U");
            ufile << UP << "\n";
            
            auto device = torch::kCPU;
            net->to(device);
            mesh = mesh.to(device);
            top = top.to(device);
            bottom = bottom.to(device);
            right = right.to(device);
            left = left.to(device);
            u_top = u_top.to(device);
            v_top = v_top.to(device);
            torch::save(net, "net.pth");
        }
    }

    return 0;
}

After you run it, you can postProcess the data by gnuplot:

#!/bin/bash

paste meshfile U > data

gnuplot<<EOF
#set terminal pngcairo size 2800,700 enhanced font 'Verdana,40' transparent
#set output 'cavity.png'
set terminal jpeg size 2800,700 enhanced font 'Verdana,40' 
set output 'cavity.jpeg'

set title 'Cavity flow predicted by PINN in the env of OpenFOAM + libtorch'

set pm3d map
set palette rgb 33,13,10
set colorbox
set dgrid3d 20,20
set xtics scale 0
set ytics scale 0
set format x ""
set format y ""
set style data points
set pointsize 2

# Start multiplot mode
set multiplot layout 1,4 title 'Cavity flow predicted by PINN in the env of OpenFOAM + libtorch'

# First plot
set title 'Ux'
#set cbrange [-0.2:2]
plot 'data' using 1:2:3 with points pointtype 7  palette notitle
unset cbrange 

# Second plot
set title 'Uy'
plot 'data' using 1:2:4 with points pointtype 7  palette notitle

# Third plot
set title 'Umag'
#set cbrange [0:2]
plot 'data' using 1:2:(sqrt(column(3)*column(3)+column(4)*column(4))) with points pointtype 7  palette notitle
unset cbrange 

# Fourth plot
set title 'p'
plot 'data' using 1:2:5 with points pointtype 7  palette notitle

# End multiplot mode
unset multiplot
EOF

Enjoy!



Regarding how to use libtorch in the environment of OpenFOAM. Please see this: https://cfd-china.com/topic/6967