參考文獻(xiàn)

PINN(Physics-informed Neural Networks)的原理部分可參見https://maziarraissi.github.io/PINNs/

考慮Burgers方程，如下圖所示，初始時(shí)刻u符合sin分布，隨著時(shí)間推移在x=0處發(fā)生間斷.
這是一個(gè)經(jīng)典問題，可使用pytorch通過PINN實(shí)現(xiàn)對Burgers方程的求解。

源代碼與注釋

源代碼共含有三個(gè)文件，來源于Github https://github.com/jayroxis/PINNs

network.py文件用于定義神經(jīng)網(wǎng)絡(luò)的結(jié)構(gòu)
train.py文件用于訓(xùn)練神經(jīng)網(wǎng)絡(luò)
evaluate.py文件用于測試訓(xùn)練好的模型繪制結(jié)果圖

建議使用Anaconda構(gòu)建運(yùn)行環(huán)境，需要安裝pytorch和一些輔助包。

1、network.py 文件

import torch
import torch.nn as nn
from collections import OrderedDict# 定義神經(jīng)網(wǎng)絡(luò)的架構(gòu)
class Network(nn.Module):# 構(gòu)造函數(shù)def __init__(self,input_size, # 輸入層神經(jīng)元數(shù)hidden_size, # 隱藏層神經(jīng)元數(shù)output_size, # 輸出層神經(jīng)元數(shù)depth, # 隱藏層數(shù)act=torch.nn.Tanh, # 輸入層和隱藏層的激活函數(shù)):super(Network, self).__init__()#調(diào)用父類的構(gòu)造函數(shù)# 輸入層layers = [('input', torch.nn.Linear(input_size, hidden_size))]layers.append(('input_activation', act()))# 隱藏層for i in range(depth):layers.append(('hidden_%d' % i, torch.nn.Linear(hidden_size, hidden_size)))layers.append(('activation_%d' % i, act()))# 輸出層layers.append(('output', torch.nn.Linear(hidden_size, output_size)))#將這些層組裝為神經(jīng)網(wǎng)絡(luò)self.layers = torch.nn.Sequential(OrderedDict(layers))# 前向計(jì)算方法def forward(self, x):return self.layers(x)

2、train.py 文件

import math
import torch
import numpy as np
from network import Network# 定義一個(gè)類，用于實(shí)現(xiàn)PINN(Physics-informed Neural Networks)
class PINN:# 構(gòu)造函數(shù)def __init__(self):# 選擇使用GPU還是CPUdevice = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")# 定義神經(jīng)網(wǎng)絡(luò)self.model = Network(input_size=2,  # 輸入層神經(jīng)元數(shù)hidden_size=16,  # 隱藏層神經(jīng)元數(shù)output_size=1,  # 輸出層神經(jīng)元數(shù)depth=8,  # 隱藏層數(shù)act=torch.nn.Tanh  # 輸入層和隱藏層的激活函數(shù)).to(device)  # 將這個(gè)神經(jīng)網(wǎng)絡(luò)存儲(chǔ)在GPU上（若GPU可用）self.h = 0.1  # 設(shè)置空間步長self.k = 0.1  # 設(shè)置時(shí)間步長x = torch.arange(-1, 1 + self.h, self.h)  # 在[-1,1]區(qū)間上均勻取值，記為xt = torch.arange(0, 1 + self.k, self.k)  # 在[0,1]區(qū)間上均勻取值，記為t# 將x和t組合，形成時(shí)間空間網(wǎng)格，記錄在張量X_inside中self.X_inside = torch.stack(torch.meshgrid(x, t)).reshape(2, -1).T# 邊界處的時(shí)空坐標(biāo)bc1 = torch.stack(torch.meshgrid(x[0], t)).reshape(2, -1).T  # x=-1邊界bc2 = torch.stack(torch.meshgrid(x[-1], t)).reshape(2, -1).T  # x=+1邊界ic = torch.stack(torch.meshgrid(x, t[0])).reshape(2, -1).T  # t=0邊界self.X_boundary = torch.cat([bc1, bc2, ic])  # 將所有邊界處的時(shí)空坐標(biāo)點(diǎn)整合為一個(gè)張量# 邊界處的u值u_bc1 = torch.zeros(len(bc1))  # x=-1邊界處采用第一類邊界條件u=0u_bc2 = torch.zeros(len(bc2))  # x=+1邊界處采用第一類邊界條件u=0u_ic = -torch.sin(math.pi * ic[:, 0])  # t=0邊界處采用第一類邊界條件u=-sin(pi*x)self.U_boundary = torch.cat([u_bc1, u_bc2, u_ic])  # 將所有邊界處的u值整合為一個(gè)張量self.U_boundary = self.U_boundary.unsqueeze(1)# 將數(shù)據(jù)拷貝到GPUself.X_inside = self.X_inside.to(device)self.X_boundary = self.X_boundary.to(device)self.U_boundary = self.U_boundary.to(device)self.X_inside.requires_grad = True  # 設(shè)置：需要計(jì)算對X的梯度# 設(shè)置準(zhǔn)則函數(shù)為MSE，方便后續(xù)計(jì)算MSEself.criterion = torch.nn.MSELoss()# 定義迭代序號，記錄調(diào)用了多少次lossself.iter = 1# 設(shè)置lbfgs優(yōu)化器self.lbfgs = torch.optim.LBFGS(self.model.parameters(),lr=1.0,max_iter=50000,max_eval=50000,history_size=50,tolerance_grad=1e-7,tolerance_change=1.0 * np.finfo(float).eps,line_search_fn="strong_wolfe",)# 設(shè)置adam優(yōu)化器self.adam = torch.optim.Adam(self.model.parameters())# 損失函數(shù)def loss_func(self):# 將導(dǎo)數(shù)清零self.adam.zero_grad()self.lbfgs.zero_grad()# 第一部分loss: 邊界條件不吻合產(chǎn)生的lossU_pred_boundary = self.model(self.X_boundary)  # 使用當(dāng)前模型計(jì)算u在邊界處的預(yù)測值loss_boundary = self.criterion(U_pred_boundary, self.U_boundary)  # 計(jì)算邊界處的MSE# 第二部分loss:內(nèi)點(diǎn)非物理產(chǎn)生的lossU_inside = self.model(self.X_inside)  # 使用當(dāng)前模型計(jì)算內(nèi)點(diǎn)處的預(yù)測值# 使用自動(dòng)求導(dǎo)方法得到U對X的導(dǎo)數(shù)du_dX = torch.autograd.grad(inputs=self.X_inside,outputs=U_inside,grad_outputs=torch.ones_like(U_inside),retain_graph=True,create_graph=True)[0]du_dx = du_dX[:, 0]  # 提取對第x的導(dǎo)數(shù)du_dt = du_dX[:, 1]  # 提取對第t的導(dǎo)數(shù)# 使用自動(dòng)求導(dǎo)方法得到U對X的二階導(dǎo)數(shù)du_dxx = torch.autograd.grad(inputs=self.X_inside,outputs=du_dX,grad_outputs=torch.ones_like(du_dX),retain_graph=True,create_graph=True)[0][:, 0]loss_equation = self.criterion(du_dt + U_inside.squeeze() * du_dx, 0.01 / math.pi * du_dxx)  # 計(jì)算物理方程的MSE# 最終的loss由兩項(xiàng)組成loss = loss_equation + loss_boundary# loss反向傳播，用于給優(yōu)化器提供梯度信息loss.backward()# 每計(jì)算100次loss在控制臺(tái)上輸出消息if self.iter % 100 == 0:print(self.iter, loss.item())self.iter = self.iter + 1return loss# 訓(xùn)練def train(self):self.model.train()  # 設(shè)置模型為訓(xùn)練模式# 首先運(yùn)行5000步Adam優(yōu)化器print("采用Adam優(yōu)化器")for i in range(5000):self.adam.step(self.loss_func)# 然后運(yùn)行l(wèi)bfgs優(yōu)化器print("采用L-BFGS優(yōu)化器")self.lbfgs.step(self.loss_func)# 實(shí)例化PINN
pinn = PINN()# 開始訓(xùn)練
pinn.train()# 將模型保存到文件
torch.save(pinn.model, 'model.pth')

運(yùn)行該文件后模型結(jié)果保存在model.pth文件中

3、evaluate.py 文件

import torch
import seaborn as sns
import matplotlib.pyplot as plt# 選擇GPU或CPU
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")# 從文件加載已經(jīng)訓(xùn)練完成的模型
model_loaded = torch.load('model.pth', map_location=device)
model_loaded.eval()  # 設(shè)置模型為evaluation狀態(tài)# 生成時(shí)空網(wǎng)格
h = 0.01
k = 0.01
x = torch.arange(-1, 1, h)
t = torch.arange(0, 1, k)
X = torch.stack(torch.meshgrid(x, t)).reshape(2, -1).T
X = X.to(device)# 計(jì)算該時(shí)空網(wǎng)格對應(yīng)的預(yù)測值
with torch.no_grad():U_pred = model_loaded(X).reshape(len(x), len(t)).cpu().numpy()# 繪制計(jì)算結(jié)果
plt.figure(figsize=(5, 3), dpi=300)
xnumpy = x.numpy()
plt.plot(xnumpy, U_pred[:, 0], 'o', markersize=1)
plt.plot(xnumpy, U_pred[:, 20], 'o', markersize=1)
plt.plot(xnumpy, U_pred[:, 40], 'o', markersize=1)
plt.figure(figsize=(5, 3), dpi=300)
sns.heatmap(U_pred, cmap='jet')
plt.show()

運(yùn)行該文件后，可繪制u場的結(jié)果