个人笔记-离散模型AMAW代码框架（改前）

import需要用到的包

import os
import time
import numpy as np
import torch
from matplotlib import pyplot as plt, gridspec
from torch import nn
from torch.autograd import Variable
from tqdm import tqdm, trange
from pyDOE import lhs
# import h5py
# import scipy.io
import argparse

随机种子和指定GPU

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
seed = 1234

torch.set_default_dtype(torch.float)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)

use_gpu = torch.cuda.is_available()
print('GPU:', use_gpu)

参数设置

parser = argparse.ArgumentParser()

parser.add_argument('--hid_layers', help='number of hidden layers', type=int, default=6) 
parser.add_argument('--hid_neurons', help='number of neurons per layer', type=int, default=128)  
.
.
.
parser.add_argument('--adam_lr', help='learning rate of adam', type=float, default=0.001)
parser.add_argument('--lbfgs_lr', help='learning rate of lbfgs', type=float, default=0.5)
.
.
.

一些用到的函数

生成神经网络的结构

def generate_layers  # 如 [ 2 128 128 128 128 128 q+1]

生成随机内部点

def random_fun  # 随机生成可移动内部点，比如离散就随机生成维度为（1000，2）的内部点，而连续就是（1000，3）

数据格式转换

def is_cuda  # 将张量（Tensor）从CPU内存移动到GPU内存
def np_tensor  # 将传入的 data_var（预期是一个NumPy数组）转换为一个PyTorch张量（Tensor），并确保这个张量的数据类型是浮点数
def np_to_tensor_to_cuda   # 先把参数从numpy变成tensor再将Tensor从CPU内存移动到GPU内存

将两个 NumPy 数组（var1 和 var2）水平堆叠（hstack）在一起。这里的 order=‘F’ 表示以 Fortran 风格（即列优先）进行展平

def hstack_data(var1, var2):   
    var_hstack = np.hstack((var1.flatten(order='F')[:, None], 
                            var2.flatten(order='F')[:, None]))
    return var_hstack

在离散模型，这里说是（固定）内部点，其实是（固定）初值条件的点，直接生成coordinates的值就行，不用先生成index再转换到coordinates

def get_col_data(lb, ub, num_x, num_y)

生成固定的边界点

def get_bc_data(num_x, num_y, x, y)

生成测试用的数据集

def get_test_data(x, y)

生成训练/测试用的数据集

def data_generate(x, y, t, lb, ub, num_x, num_y):

    x_test = get_test_data(x, y)

    (xyt_b_bot_meshgrid_train, xyt_b_top_meshgrid_train, xyt_b_left_meshgrid_train,
     xyt_b_right_meshgrid_train) = get_bc_data(num_x, num_y, x, y)

    x_f_N = get_col_data(lb, ub, num_x, num_y)

    return (x_test, xyt_b_bot_meshgrid_train, xyt_b_top_meshgrid_train, xyt_b_left_meshgrid_train,
            xyt_b_right_meshgrid_train, x_f_N)

net的类

class Net(nn.Module):
    def __init__(self, layers):
    def forward(self, x):     

Model的类

初始化函数

def __init__

返回预测的所有结果（即离散模型就是q+1个u的值，而连续模型就是1个u的值）

    def train_U(self, x):
        return self.net(x)

预测的u，因为是离散模型，这里就只取h的值 （后期画图用）

    def predict_U(self, x):
        return self.train_U(x)[:, self.q:(self.q+1)]  # h

likelihood_loss

即take the negative logarithm of the likelihood function to get the negative log-likelihood as the desired loss function
L = e − s f ϕ L f ϕ + e − s f μ L f μ + e − s b ϕ L b ϕ + e − s b μ L b μ + s f ϕ + s f μ + s b ϕ + s b μ

L=​e−sfϕ​​Lfϕ​​+e−sfμ​​Lfμ​​+e−sbϕ​​Lbϕ​​+e−sbμ​​Lbμ​​+sfϕ​​+sfμ​​+sbϕ​​+sbμ​​​

    def likelihood_loss(self, loss_f_h, loss_b_h, loss_f_mu, loss_b_mu):
        loss = (torch.exp(-self.x_f_s_h) * loss_f_h.detach() + self.x_f_s_h
                + torch.exp(-self.x_b_s_h) * loss_b_h.detach() + self.x_b_s_h
                + torch.exp(-self.x_f_s_mu) * loss_f_mu.detach() + self.x_f_s_mu
                + torch.exp(-self.x_b_s_mu) * loss_b_mu.detach() + self.x_b_s_mu)

        return loss

true_loss

即比likelihood的loss少加自适应权重 $s_{f_{\varphi }}+s_{f_{\mu }}+s_{b_{\varphi }}+s_{b_{\mu }}$，loss function如下：
L = e − s f ϕ L f ϕ + e − s f μ L f μ + e − s b ϕ L b ϕ + e − s b μ L b μ

L=​e−sfϕ​​Lfϕ​​+e−sfμ​​Lfμ​​+e−sbϕ​​Lbϕ​​+e−sbμ​​Lbμ​​​

    def true_loss(self, loss_f_h, loss_b_h, loss_f_mu, loss_b_mu):
        true_loss_value = (torch.exp(-self.x_f_s_h.detach()) * loss_f_h + torch.exp(-self.x_b_s_h.detach()) * loss_b_h
                           + torch.exp(-self.x_f_s_mu.detach()) * loss_f_mu + torch.exp(-self.x_b_s_mu.detach()) * loss_b_mu)

        return true_loss_value

对预测的边界点的q+1个u求出u_x，u_y （离散比连续的这一步复杂很多）

    def ub_x_h(self, boundary_data)
    def ub_y_h(self, boundary_data)
    def ub_x_mu(self, boundary_data)
    def ub_y_mu(self, boundary_data)

得到预测的边界点的q+1个u

 def get_b_pred(self, boundary_data)

定义正方形右边的公式（含theta_s)

可移动到Model的外面

    @staticmethod
    def partial_derivative_phi(phi)

构造 ${\mathcal{L}}_{b_{\varphi }}$, ${\mathcal{L}}_{b_{\mu }}$, ${\mathcal{L}}_{f_{\varphi }}$, ${\mathcal{L}}_{f_{\mu }}$

这里只构造了边界的loss，内部点的在epoch_loss里面直接构造了

def loss_boundary_h(self)
def loss_boundary_mu(self)

画出离散型预测的u和U0的区别

这一步应该可以移到外面

def draw_loss_f(self, current_t_num, x_f, U00, U11)

生成${\mathcal{L}}_{b_{\varphi }}$, ${\mathcal{L}}_{b_{\mu }}$, ${\mathcal{L}}_{f_{\varphi }}$, ${\mathcal{L}}_{f_{\mu }}$（这一步在ADAM和LBFGS都会用到）

def epoch_loss(self)

LBFGS的loss

这里先跑一下epoch_loss得到${\mathcal{L}}_{b_{\varphi }}$, ${\mathcal{L}}_{b_{\mu }}$, ${\mathcal{L}}_{f_{\varphi }}$, ${\mathcal{L}}_{f_{\mu }}$，然后把${\mathcal{L}}_{b_{\varphi }}$, ${\mathcal{L}}_{b_{\mu }}$, ${\mathcal{L}}_{f_{\varphi }}$, ${\mathcal{L}}_{f_{\mu }}$当作true_loss的输入,得到$\mathcal{L}=e^{-s_{f_{\varphi }}}{\mathcal{L}}_{f_{\varphi }}+e^{-s_{f_{\mu }}}{\mathcal{L}}_{f_{\mu }}+e^{-s_{b_{\varphi }}}{\mathcal{L}}_{b_{\varphi }}+e^{-s_{b_{\mu }}}{\mathcal{L}}_{b_{\mu }}$

def LBFGS_epoch_loss(self)

画预测的u

这一步应该可以移到外面

def draw_evaluate(self, x_test, pred, exact)

把一些可移动内部点移动到固定内部点

这一个看看能不能移动到Model的外面

def move_some_movable_points_to_fixed_points(self)

传统的PINN， AM的PINN， AM-AW的PINN

run_AM: 决定选点
run_AM_AW1：likelihood_loss在这，决定选点

AM-AW的具体步骤：

1. 让自适应权重变成能够记录梯度属性的参数
   %%%%%%%%%%%%%%%%%%%%%%%%%

2. 创建AM的for loop
   2.1 创建三个不同的优化器(L-BFGS的优化器，ADAM的优化器，自适应权重的优化器）
   %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   2.2 创建一个ADAM的for loop
   2.2.1 通过epoch_loss得到${\mathcal{L}}_{b_{\varphi }}$, ${\mathcal{L}}_{b_{\mu }}$, ${\mathcal{L}}_{f_{\varphi }}$,${\mathcal{L}}_{f_{\mu }}$
   -----------------------------------------------------------------------------
   2.2.2 optimizer_adam.zero_grad() #使用optimizer_adam优化器更新$\theta$之前，先清除之前的梯度信息
   2.2.3 通过true_loss(${\mathcal{L}}_{b_{\varphi }}$, ${\mathcal{L}}_{b_{\mu }}$, ${\mathcal{L}}_{f_{\varphi }}$, ${\mathcal{L}}_{f_{\mu }}$)得$\mathcal{L}=e^{-s_{f_{\varphi }}}{\mathcal{L}}_{f_{\varphi }}+e^{-s_{f_{\mu }}}{\mathcal{L}}_{f_{\mu }}+e^{-s_{b_{\varphi }}}{\mathcal{L}}_{b_{\varphi }}+e^{-s_{b_{\mu }}}{\mathcal{L}}_{b_{\mu }}$
   2.2.4 loss.backward() #计算2.2.3的$\mathcal{L}$关于模型参数的梯度
   2.2.5 optimizer_adam.step() #使用2.2.4计算出的梯度来更新网络模型参数$\theta$
   -------------------------------------------------------------------------------
   2.2.6 optimizer_adam_weight.zero_grad() #使用optimizer_adam_weight优化器更新$s_{f_{\varphi }},s_{f_{\mu }},s_{b_{\varphi }},s_{b_{\mu }}$之前，先清除之前的梯度信息
   2.2.7 通过likelihood_loss(${\mathcal{L}}_{b_{\varphi }}$, ${\mathcal{L}}_{b_{\mu }}$, ${\mathcal{L}}_{f_{\varphi }}$, ${\mathcal{L}}_{f_{\mu }}$)得
   $\mathcal{L}=e^{-s_{f_{\varphi }}}{\mathcal{L}}_{f_{\varphi }}+e^{-s_{f_{\mu }}}{\mathcal{L}}_{f_{\mu }}+e^{-s_{b_{\varphi }}}{\mathcal{L}}_{b_{\varphi }}+e^{-s_{b_{\mu }}}{\mathcal{L}}_{b_{\mu }}+s_{f_{\varphi }}+s_{f_{\mu }}+s_{b_{\varphi }}+s_{b_{\mu }}$
   2.2.8 loss.backward() #计算2.2.7的$\mathcal{L}$关于模型参数的梯度
   2.2.9 optimizer_adam_weight.step() #使用2.2.8计算出的梯度来更新自适应权重$s_{f_{\varphi }},s_{f_{\mu }},s_{b_{\varphi }},s_{b_{\mu }}$
   （即后面继续回到第2.2步开始新的迭代直至达到adam_iter的数）

   2.3 打印 print(‘Adam done!’)

   2.4 self.optimizer_LBFGS.step(self.LBFGS_epoch_loss)
   2.5 打印 print(‘LBFGS done!’)

   2.6 根据p(x)改变可移动内部点
   （即后面继续回到第2步开始新的一轮）

def run_baseline(self)
def run_AM(self, current_t_num)
def run_AM_AW1(self, current_t_num)

train函数

这里生成可移动的内部点，定义自适应权重$s_{f_{\varphi }}$,$s_{b_{\varphi }}$,$s_{f_{\mu }}$,$s_{b_{\mu }}$

def train(self, current_t_num)

---

---

又回到外面的一些函数

生成初值条件

def is_in_red_square(x_f)
def get_initial_h(x_f)

对内部点的u求导

def uf_x_h(h, col_pts)
def uf_y_h(h, col_pts)

IRK回溯回来的U0 （这里离散跟连续不一样）

def x_f_loss_fun_h(x_f, train_U)
def x_f_loss_fun_mu(x_f, train_U)

画图

def draw_exact_count(AM_count, current_t_num)
def draw_pred_q_count(AM_count, current_t_num)
def draw_pred_q_current(time_current)
def draw_pred_q()
def draw_current_exact(pic_num)
def draw_exact()
def draw_exact_points_count(AM_count, points, N_points=None, show_exact=False)
def draw_epoch_loss_count(AM_count, current_t_num)
def draw_exact_points_current(current_t_num, points, N_points=None, show_exact=False)
def draw_exact_points(points, N_points=None, show_exact=False)
def draw_epoch_loss_current(loss_current_num)
def draw_epoch_loss()
def draw_epoch_w()

定义不同case对应不同的上届和下界

def define_lb_ub(initial_condition_type)

main

参数，train