波动方程的PINN方法求解

最新推荐文章于 2025-03-14 09:54:54 发布

邓肯黑火

最新推荐文章于 2025-03-14 09:54:54 发布

阅读量3.8k

点赞数 44

文章标签： python pytorch 深度学习

本文链接：https://blog.csdn.net/qq_58509609/article/details/138239469

版权

基于Pytorch的无震源二维波动方程PINN方法求解

文章目录

基于Pytorch的无震源二维波动方程PINN方法求解
前言
一、二维波动方程的建立
二、PINN结构
三、代码部分
总结

前言

波动方程的求解一直是地球物理领域中的一大重要问题。波动方程的求解有多种方法，传统的方法有差分法、谱元法、伪谱法等，此外利用深度学习算法求解偏微分方程的方法在近几年也出了很多文章。这段时间刚好在做毕业设计，内容是基于深度学习算法求解Biot方程，其中简单二维波动方程的求解是至关重要的一步，因此将计算的流程和代码分享出来，供给大家参考，也请指正。

一、二维波动方程的建立

本文求解的二维波动方程的参数参考Wave Propagation with Physics Informed Neural Networks（DOI:10.1190/segam2020-3425406.1）一文给出的参数模型。
以下给出波动方程的参数设置：
（最后一行公式是真实解）

二、PINN结构

PINN，全称物理信息神经网络，是一种经过训练来求解偏微分方程的神经网络。该方法是无需设置网格求解的，因此可以相对容易地扩展到任意复杂的计算边界。以下是PINN的结构图。
此图系本作者绘制，转载请注明出处

三、代码部分

此代码建议使用GPU加速，否则会运行较慢。并在D盘建立相应的文件夹存储损失函数图像、网络结构、真实解图像、预测解图像、误差图像。

1.引入库

import torch
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import torch.nn as nn
from pyDOE import lhs

2.初步设置

#自动更换CPU,GPU，并设置浮点型位数
default_device = "cuda" if torch.cuda.is_available() else "cpu"
dtype=torch.float32
torch.set_default_dtype(dtype)
#####################################################################
# 设置随机数种子
def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True

setup_seed(888888)

该处设置数据的容器，便于使用GPU和CPU计算的切换。并设置随机种子，让每次结果一致。

3.基本子函数的设置

#####################################################################
#转换张量函数，存储损失的时候用到
def Tran_fun(tensor):
    a = []
    for i in tensor:
        #i = i.to("cpu")
        i = i.detach().numpy()
        a.append(i)
    return(np.array(a))
#####################################################################
#设置一种采样策略，这里用到lhs函数
def LHS_Sampling(N,be,en):
    return be + (en-be)*lhs(3, N)
#####################################################################
 #求导函数
def gradients(u, x, order=1):
    u.requires_grad_(True)
    x.requires_grad_(True)
    if order == 1:
        return torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u),
                                   create_graph=True,
                                   only_inputs=True, )[0]
    else:
        return gradients(gradients(u, x), x, order=order - 1)
#####################################################################

这里设置了三个重要子函数，包括转换张量函数、求导函数、采样函数。

4.PINN网络设置

######################################################################
#网络结构的编译
class PINN(nn.Module):
    def __init__(self, in_features, out_features, num_neurons, activation=torch.sin):
        super(PINN, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.num_neurons = num_neurons
        self.act_func = activation
        
        self.block1_layer1 = nn.Linear(self.in_features, self.num_neurons)
        self.block1_layer2 = nn.Linear(self.num_neurons, self.num_neurons)
        self.block1 = [self.block1_layer1, self.block1_layer2]
        
        self.block2_layer1 = nn.Linear(self.in_features + self.num_neurons, self.num_neurons)
        self.block2_layer2 = nn.Linear(self.num_neurons, self.num_neurons)
        self.block2 = [self.block2_layer1, self.block2_layer2]
        
        self.layer_after_block = nn.Linear(self.num_neurons + self.in_features, self.num_neurons)
        self.layer_output = nn.Linear(self.num_neurons, self.out_features)

    def forward(self, x):
        
        x_temp = x
        
        for dense in self.block1:
            x_temp = self.act_func(dense(x_temp))
        x_temp = torch.cat([x_temp, x], dim=-1)
        
        for dense in self.block2:
            x_temp = self.act_func(dense(x_temp))
        x_temp = torch.cat([x_temp, x], dim=-1)
        
        x_temp = self.act_func(self.layer_after_block(x_temp))
        x_temp = self.layer_output(x_temp)
        return x_temp
#####################################################################

网络是最简单的全连接层。重要在于之后的损失函数的设置。

5.损失函数，网络，优化器对象化

######################################################################
#损失函数，网络，优化器对象化
loss = torch.nn.MSELoss()
loss = loss.to(device = default_device)
pinn = PINN(3, 1, 8)
pinn = pinn.to(device = default_device)
#opt = torch.optim.Adam(params=pinn.parameters())
opt = torch.optim.Adam(pinn.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=5000, gamma=0.9)
#####################################################################

6.PDE损失、边界条件损失、初始条件损失的设置

#####################################################################
#pde损失函数设置
def interior_loss(Mat):
    # 内点
    Mat = torch.tensor(Mat)
    x = Mat[:, 0:1].float()
    y = Mat[:, 1:2].float()
    t = Mat[:, 2:3].float()
    x.requires_grad_(True)
    y.requires_grad_(True)
    t.requires_grad_(True)
    x = x.to(device=default_device)
    y = y.to(device = default_device)
    t = t.to(device = default_device)
    u = pinn(torch.cat([t, x, y], dim=1))
    cond1 = gradients(u, t, order=2)
    cond2 = c*(gradients(u, x, order=2)+gradients(u, y, order=2))
    return loss(cond1,cond2)

def left_loss(Mat):
    # 左边界点
    Mat = torch.tensor(Mat)
    Mat[:,0:1] = torch.zeros_like(Mat[:,0:1])
    x = Mat[:, 0:1].float()
    y = Mat[:, 1:2].float()
    t = Mat[:, 2:3].float()
    zero = torch.zeros_like(t).float()
    x = x.to(device=default_device)
    y = y.to(device = default_device)
    t = t.to(device = default_device)
    zero = zero.to(device = default_device)
    u = pinn(torch.cat([t, x, y], dim=1))
    cond1 = u
    cond2 = zero
    return loss(cond1,cond2)

def right_loss(Mat):
    # 右边界点
    Mat = torch.tensor(Mat)
    Mat[:,0:1] = torch.ones_like(Mat[:,0:1])*torch.pi*4
    x = Mat[:, 0:1].float()
    y = Mat[:, 1:2].float()
    t = Mat[:, 2:3].float()
    zero = torch.zeros_like(t).float()
    x.requires_grad_(True)
    y.requires_grad_(True)
    t.requires_grad_(True)
    x = x.to(device=default_device)
    y = y.to(device = default_device)
    t = t.to(device = default_device)
    zero = zero.to(device = default_device)
    u = pinn(torch.cat([t, x, y], dim=1))
    zero = torch.zeros_like(u).requires_grad_(True)
    zero = zero.to(device = default_device)
    cond1 = u
    cond2 = zero
    return loss(cond1,cond2)

def down_loss(Mat):
    # 下边界点
    Mat = torch.tensor(Mat)
    Mat[:,1:2] = torch.zeros_like(Mat[:,1:2])
    x = Mat[:, 0:1].float()
    y = Mat[:, 1:2].float()
    t = Mat[:, 2:3].float()
    zero = torch.zeros_like(t).float()
    x.requires_grad_(True)
    y.requires_grad_(True)
    t.requires_grad_(True)
    zero.requires_grad_(True)
    x = x.to(device=default_device)
    y = y.to(device = default_device)
    t = t.to(device = default_device)
    zero = zero.to(device = default_device)
    u = pinn(torch.cat([t, x, y], dim=1))
    cond1 = u
    cond2 = zero
    return loss(cond1,cond2)

def up_loss(Mat):
    # 上边界点
    Mat = torch.tensor(Mat)
    Mat[:,1:2] = torch.ones_like(Mat[:,1:2])*torch.pi*4
    x = Mat[:, 0:1].float()
    y = Mat[:, 1:2].float()
    t = Mat[:, 2:3].float()
    zero = torch.zeros_like(t).float()
    x.requires_grad_(True)
    y.requires_grad_(True)
    t.requires_grad_(True)
    x = x.to(device=default_device)
    y = y.to(device = default_device)
    t = t.to(device = default_device)
    zero = zero.to(device = default_device)
    u = pinn(torch.cat([t, x, y], dim=1))
    cond1 = u
    cond2 = zero
    return loss(cond1,cond2)

def ini1_loss(Mat):
    # 初始条件1
    Mat = torch.tensor(Mat)
    Mat[:,2:3] = torch.zeros_like(Mat[:,2:3])
    x = Mat[:, 0:1].float()
    y = Mat[:, 1:2].float()
    t = Mat[:, 2:3].float()
    x.requires_grad_(True)
    y.requires_grad_(True)
    t.requires_grad_(True)
    x = x.to(device=default_device)
    y = y.to(device = default_device)
    t = t.to(device = default_device)
    u = pinn(torch.cat([t, x, y], dim=1))
    cond1 = u
    cond2 = 2*torch.sin(x/4)*torch.sin(y/4)
    return loss(cond1,cond2)

def ini2_loss(Mat):
    # 初始条件2
    Mat = torch.tensor(Mat)
    Mat[:,2:3] = torch.zeros_like(Mat[:,2:3])
    x = Mat[:, 0:1].float()
    y = Mat[:, 1:2].float()
    t = Mat[:, 2:3].float()
    x.requires_grad_(True)
    y.requires_grad_(True)
    t.requires_grad_(True)
    x = x.to(device=default_device)
    y = y.to(device = default_device)
    t = t.to(device = default_device)
    u = pinn(torch.cat([t, x, y], dim=1))
    cond1 = gradients(u, t, order=1)
    cond2 = 2*lamda*torch.sin(x/4)*torch.sin(y/4)
    return loss(cond1,cond2)
#####################################################################

7.计算预测波场与绘图函数的设置

#####################################################################
#计算预测结果函数
class Output_Fun():
    def __init__(self, t_b,t_e,t_num,x_b,x_e,x_num,y_b,y_e,y_num,lamda):
        self.t_b = t_b
        self.t_e = t_e
        self.t_num = t_num
        self.x_b = x_b
        self.x_e = x_e
        self.x_num = x_num
        self.y_b = y_b
        self.y_e = y_e
        self.y_num = y_num
        self.lamda = lamda
        
    def U_Pred(self,model_path):
        model = torch.load(model_path,map_location=torch.device('cpu'))
        tc = torch.linspace(self.t_b, self.t_e, self.t_num)
        xc = torch.linspace(self.x_b, self.x_e, self.x_num)
        yc = torch.linspace(self.y_b, self.y_e, self.y_num)
        tm, xm, ym = torch.meshgrid(tc, xc,yc)
        xx = xm.reshape(-1, 1)
        yy = ym.reshape(-1, 1)
        tt = tm.reshape(-1, 1)
        xy = torch.cat([tt,xx,yy], dim=1)
        u_pred = model(xy)
        return u_pred
    
    def U_True(self):
        tc = torch.linspace(self.t_b, self.t_e, self.t_num)
        xc = torch.linspace(self.x_b, self.x_e, self.x_num)
        yc = torch.linspace(self.y_b, self.y_e, self.y_num)
        tm, xm, ym = torch.meshgrid(tc, xc,yc)
        xx = xm.reshape(-1, 1)
        yy = ym.reshape(-1, 1)
        tt = tm.reshape(-1, 1)
        xy = torch.cat([tt,xx,yy], dim=1)
        u_true = torch.sin(torch.pi*xx/self.x_e)*torch.sin(torch.pi*yy/self.y_e)*2*(torch.cos(self.lamda*tt)+torch.sin(self.lamda*tt))##########)
        return u_true

    def U_Draw(self,true_mat,pred_mat):
        plt.rcParams['font.sans-serif']=['FangSong'] # 用来正常显示中文标签
        plt.rcParams['axes.unicode_minus']=False # 用来正常显示
        my_font1 = {"family":"Microsoft JhengHei", "size":35, "style":"italic"}
        my_font2 = {"family":"Microsoft JhengHei", "size":20, "style":"italic"}
        my_font3 = {"family":"Microsoft JhengHei", "size":30, "style":"italic"}
        x = torch.linspace(self.x_b,self.x_e,self.x_num)
        y = torch.linspace(self.y_b,self.y_e,self.y_num)
        X,Y = torch.meshgrid(x,y)
        dt = (self.t_e-self.t_b)/int(self.t_num-1)
        true_mat = true_mat.detach().numpy()
        pred_mat = pred_mat.detach().numpy()
        res_mat = true_mat - pred_mat
        lenspace = int(self.x_num**2)
        for i in range(0,int(self.t_num),1):
            u_i = pred_mat[i*lenspace:(i+1)*lenspace,:].reshape((self.x_num,self.y_num))
            fig = plt.figure(figsize=(15,15))
            ax = fig.add_subplot(111)
            surf = ax.pcolor(X, Y, u_i)
            cb=fig.colorbar(surf,shrink=0.6,aspect=20)
            cb.ax.tick_params(labelsize=20)  #设置色标刻度字体大小。
            cb.set_label('波场值',fontdict=my_font3) #设置colorbar的标签字体及其大小
            ax.set_xlabel('X',fontdict=my_font3)
            ax.set_ylabel('Y',fontdict=my_font3)
            t = round(dt*i,3)
            ax.set_title("第"+str(t)+"秒波场预测图像",fontdict=my_font1)
            ax.tick_params(labelsize=20)
            plt.savefig("D:/trl/图像波场/二维预测波场/第"+str(t)+"秒波场图像.jpg")
            
        for i in range(0,int(self.t_num),1):
            u_i = true_mat[i*lenspace:(i+1)*lenspace,:].reshape((self.x_num,self.y_num))
            fig = plt.figure(figsize=(15,15))
            ax = fig.add_subplot(111)
            surf = ax.pcolor(X, Y, u_i)
            cb=fig.colorbar(surf,shrink=0.6,aspect=20)
            cb.ax.tick_params(labelsize=20)  #设置色标刻度字体大小。
            cb.set_label('波场值',fontdict=my_font3) #设置colorbar的标签字体及其大小
            ax.set_xlabel('X',fontdict=my_font3)
            ax.set_ylabel('Y',fontdict=my_font3)
            t = round(dt*i,3)
            ax.set_title("第"+str(t)+"秒波场真实图像",fontdict=my_font1)
            ax.tick_params(labelsize=20)
            plt.savefig("D:/trl/图像波场/二维真实波场/第"+str(t)+"秒波场图像.jpg")
            
        for i in range(0,int(self.t_num),1):
            u_i = res_mat[i*lenspace:(i+1)*lenspace,:].reshape((self.x_num,self.y_num))
            fig = plt.figure(figsize=(15,15))
            ax = fig.add_subplot(111)
            surf = ax.pcolor(X, Y, u_i)
            cb=fig.colorbar(surf,shrink=0.6,aspect=20)
            cb.ax.tick_params(labelsize=20)  #设置色标刻度字体大小。
            cb.set_label('波场值',fontdict=my_font3) #设置colorbar的标签字体及其大小
            ax.set_xlabel('X',fontdict=my_font3)
            ax.set_ylabel('Y',fontdict=my_font3)
            t = round(dt*i,3)
            ax.set_title("第"+str(t)+"秒波场误差图像",fontdict=my_font1)
            ax.tick_params(labelsize=20)
            plt.savefig("D:/trl/图像波场/二维误差波场/第"+str(t)+"秒波场图像.jpg")
#####################################################################

8.参数设置

######################################################################参数设置
c,len_x,len_y,len_t = 2, 4*torch.pi, 4*torch.pi, torch.pi/(1/8)**(1/2)
t_b,t_e,t_num = 0.0,torch.pi/(1/8)**(1/2),50000
x_b,x_e,x_num = 0.0,4*torch.pi,50000
y_b,y_e,y_num = 0.0,4*torch.pi,50000
c = 2
lamda = c*(1/8)**(1/2)
#范围设置
t_range = [t_b, t_e]
x_range = [x_b, x_e]
y_range = [y_b, y_e]
lb = np.asarray([x_range[0], y_range[0], t_range[0]])
ub = np.asarray([x_range[1], y_range[1], t_range[1]])
#内点，边界，初始条件采样
xyt_in = LHS_Sampling(10,lb,ub)
xyt_left = LHS_Sampling(10,lb,ub)
xyt_right = LHS_Sampling(10,lb,ub)
xyt_down = LHS_Sampling(10,lb,ub)
xyt_up = LHS_Sampling(10,lb,ub)
xyt_ini1 = LHS_Sampling(10,lb,ub)
xyt_ini2 = LHS_Sampling(10,lb,ub)
#设置训练次数，设置接收损失值的列表
epochs = 10000
#接收各个边界的损失值，用于绘图
loss_i,loss_le,loss_ri,loss_d,loss_u,loss_i1,loss_i2,loss_all = [],[],[],[],[],[],[],[]
#####################################################################

9.网络的训练

#####################################################################
#训练
for i in range(epochs):
    opt.zero_grad()
    a1= interior_loss(xyt_in)
    a2 = left_loss(xyt_left)
    a3 = right_loss(xyt_right)
    a4 = down_loss(xyt_down)
    a5 = up_loss(xyt_up)
    a6 = ini1_loss(xyt_ini1)
    a7 = ini2_loss(xyt_ini2)
    l = a1+a2+a3+a4+a5+a6+a6
    loss_i.append(a1)
    loss_le.append(a2)
    loss_ri.append(a3)
    loss_d.append(a4)
    loss_u.append(a5)
    loss_i1.append(a6)
    loss_i2.append(a7)
    loss_all.append(l)
    l = a1+a2+a3+a4+a5+a6+a7
    l.backward()
    opt.step()
    if i % 100 == 0:
        print(i)
#####################################################################

10.损失函数图像的绘制

c1 = Tran_fun(loss_i)
c2 = Tran_fun(loss_le)
c3 = Tran_fun(loss_ri)
c4 = Tran_fun(loss_d)
c5 = Tran_fun(loss_u)
c6 = Tran_fun(loss_i1)
c7 = Tran_fun(loss_i2)
c8 = Tran_fun(loss_all)
#####################################################################
plt.rcParams['font.sans-serif']=['Microsoft YaHei']
plt.rcParams['axes.unicode_minus']=False # 用来正常显示
#####################################################################
#损失函数绘制
my_font1 = {"family":"Microsoft JhengHei", "size":20, "style":"italic"}
my_font2 = {"family":"Microsoft JhengHei", "size":35, "style":"italic"}
my_font3 = {"family":"Times New Roman", "size":40, "style":"italic"}
#单张
fig, ax = plt.subplots(figsize = (20,10),dpi=200)#设置像素
x = np.linspace(1,epochs,epochs)
#设置线的形和点
#ax.semilogy(x, c1,linestyle = "solid",lw=3,label="内部PDE损失")#mfc设
#ax.semilogy(x, c2,linestyle = "solid",lw=3,label="左边界损失")#mfc设
#ax.semilogy(x, c3,linestyle = "solid",lw=3,label="右边边界总损失")#mfc设
#ax.semilogy(x, c4,linestyle = "solid",lw=3,label="上边界总损失")#mfc设
#ax.semilogy(x, c5,linestyle = "solid",lw=3,label="下边界总损失")#mfc设
#ax.semilogy(x, c6,linestyle = "solid",lw=3,label="初始条件1总损失")#mfc设
#ax.semilogy(x, c7,linestyle = "solid",lw=3,label="初始条件2总损失")#mfc设
ax.semilogy(x, c8,linestyle = "solid",lw=3,label="总损失")#mfc设

ax.set_xlabel('训练次数',fontdict=my_font2)
ax.set_ylabel('Loss',fontdict=my_font3)
#x主副刻度
x_major_locator=plt.MultipleLocator(100)
x_minor_locator=plt.MultipleLocator(50)
ax.xaxis.set_minor_locator(x_minor_locator)
ax.xaxis.set_major_locator(x_major_locator)
#y主副刻度
#y_major_locator=plt.MultipleLocator(0.01)
#y_minor_locator=plt.MultipleLocator(0.001)
#ax.yaxis.set_minor_locator(y_minor_locator)
#ax.yaxis.set_major_locator(y_major_locator)

ax.tick_params(labelsize=30)
ax.set_title(label="PINN Loss",fontdict=my_font3)
#显示图例
ax.legend(prop=my_font1)
plt.savefig("D:/trl/损失函数/损失函数图像.jpg")
#####################################################################

分批绘制不同边界的损失函数图像

11.保存模型

#####################################################################
#模型保存
model = pinn
torch.save(pinn,"D:/trl/波场结构/no_resoure_wave_model2d.pth")
#####################################################################

12.求解方程并成图

##########################################################
#求解方程并绘制图像
sol = Output_Fun(t_b,t_e,t_num,x_b,x_e,x_num,y_b,y_e,y_num,lamda)
mat1 = sol.U_Pred("D:/trl/波场结构/no_resoure_wave_model2d.pth")
mat2 = sol.U_True()
sol.U_Draw(mat1,mat2)
##########################################################