ML: PINN代码实例-顶盖驱动流
本文尝试使用PINN求解2D顶盖驱动流。据岳子所知，在2024年10月，本代码是国际上首次开源的基于OpenFOAM/libtorch的顶盖驱动流开源代码。在阅读本文之前，请先阅读：
要不然不会理解代码的含义。相应的代码如下。由于代码量过大，不对其进行介绍。建议系统学习OpenFOAM以及libtorch后再来理解。
#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;
}
该代码运行后，会自动输出结果文件。可以通过下面的代码来出图：
#!/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
相应的结果如下：
