PINN代码 -- 一维连续时间模型：Schrodinger方程（CPU/GPU+TensorFlow2.10）

12!

已于 2024-08-26 10:04:51 修改

阅读量1.9k

点赞数 33

文章标签： tensorflow 人工智能 python 深度学习机器学习

于 2024-08-08 12:18:56 首次发布

本文链接：https://blog.csdn.net/m0_46611848/article/details/141022035

版权

《Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations》
作者：Maziar Raissi等人
作者原代码链接：https://github.com/maziarraissi/PINNs
环境：TensorFlow1.x版本。

Schrodinger方程：

TensorFlow2.x版本与1.x不兼容，无法直接运行作者的原代码。本文修改的代码是连续时间模型中的一维薛定谔方程，即
continuous_time_inference (Schrodinger)。

文件位置详见下图。

PINN

一.环境说明

GPU环境：

Windows
Anaconda + GPU + PyCharm
Python3.9
CUDA 11.2 + cuDNN 8.1 + TensorFlow 2.10

已知GPU适配的 TensorFlow 的最高版本为 2.10，如果版本高于2.10，则需要重新配置环境。配置教程可以参考本人之前写的：《Anaconda同时安装PyTorch和TensorFlow版本教程（Windows + GPU/CUDA）》，因教程不定时修改更新，链接可能会失效，可在本人主页查询。

经本人试验，修改后的代码也能在CPU上运行，代价是计算成本高、计算时间长。

CPU环境：

Windows
CPU + PyCharm
TensorFlow 2.11 + Python3.9

在CPU环境中，TensorFlow2.x版本与Python版本适配即可。

二. 分块代码

修改的代码可以直接替换原来的continuous_time_inference (Schrodinger)/Schrodinger.py内容，也可以继续在该文件目录continuous_time_inference (Schrodinger)下配置新.py文件，比如try_Schrodinger_2.11.py。如果自己新建项目，记得拷贝数据集，修改代码中的读取路径，具体操作内容不再赘述。

说明：代码中的 model.train(10000)表示训练回合数为1000，可自行更改。

1. 库

import sys
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import scipy.io
from scipy.interpolate import griddata
from pyDOE import lhs
from mpl_toolkits.mplot3d import Axes3D
import time
import matplotlib.gridspec as gridspec
from mpl_toolkits.axes_grid1 import make_axes_locatable
import os

与原代码的区别在于绘图的方式不同。模型的训练结果通常通过绘制输出内容查看。原代码应该是TensowFlow1.x版本，利用from plotting import newfig, savefig绘图。TensorFlow2.x版本是没有newfig, savefig的，因此我考虑使用matplotlib.pyplot绘图。

np.random.seed(1234)
tf.set_random_seed(1234)

设置随机数种子，以便复现数值实验。

matplotlib.use('Agg')

# 禁用Eager Execution
tf.compat.v1.disable_eager_execution()

# 运行初期会输出当前计算任务在分配设备时的详细情况日志，不想了解可以设置日志级别为ERROR,只输出错误信息
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

设置matplotlib.use('Agg')的目的是保证在GPU上绘图时不报错，它的代价时无法显示绘图内容（即plot.show()失效），只能到文件夹里去查看图片。如果当前环境为CPU，那么可以注释这一条命令。

在模型训练初期会输出当前计算任务在分配设备时的详细情况日志，不是报错信息的话，不用理会。目前我也没找到不显示这些日志的方法。

2. PINN模型

俗话说，授人以鱼不如授人以渔。这里分享一下的的修改思路，主要有以下两点：

（1）为了保证兼容，将所有被提示不兼容的函数格式改成tf.compat.v1.(函数名)，例如作者原代码中使用的是tf.Session，我们只需要改成tf.compat.v1.Session即可。然而，有个别函数不能如此简单粗暴地修改，所以列了一个表格供参考。

原代码	我的
`tf.Session`	`tf.compat.v1.Session`
`tf.placeholder`	`tf.compat.v1.placeholder`
`tf.AdamOptimizer()`	`tf.compat.v1.train.AdamOptimizer()`
`tf.global_variables_initializer()`	`tf.compat.v1.global_variables_initializer()`
`tf.truncated_normal`	`tf.random.truncated_normal`

其实TensorFlow 2.x版本中有更简洁的函数，但由于初期刚接触TensorFlow，为了避免不必要的麻烦，主打一个先看懂代码再去改代码的心态。

（2）TensorFlow 2.x版本中移除掉的库/函数，直接删掉。比如作者设置并使用了 L-BFGS-B优化器，其原代码为：

self.optimizer = tf.contrib.opt.ScipyOptimizerInterface(self.loss, method = 'L-BFGS-B', options = {'maxiter': 50000, 'maxfun': 50000, 'maxcor': 50, 'maxls': 50, 'ftol' : 1.0 * np.finfo(float).eps})

self.optimizer.minimize(self.sess,  feed_dict = tf_dict, fetches = [self.loss, self.lambda_1, self.lambda_2],loss_callback = self.callback)

在TensorFlow1.x版本中，上述代码没问题。但问题在于，TensorFlow 2.x版本移除了 tf.contrib 。Adam优化器已经够用了。本懒人改代码的思路是“多一事不如少一事”，遂直接删除L-BFGS-B优化器相关内容。

以下是PINN模型结构设计：


class PhysicsInformedNN(tf.Module):
    # Initialize the class
    def __init__(self, x0, u0, v0, tb, X_f, layers, lb, ub):

        X0 = np.concatenate((x0, 0 * x0), 1)  # (x0, 0)
        X_lb = np.concatenate((0 * tb + lb[0], tb), 1)  # (lb[0], tb)
        X_ub = np.concatenate((0 * tb + ub[0], tb), 1)  # (ub[0], tb)

        self.lb = lb
        self.ub = ub

        self.x0 = X0[:, 0:1]
        self.t0 = X0[:, 1:2]

        self.x_lb = X_lb[:, 0:1]
        self.t_lb = X_lb[:, 1:2]

        self.x_ub = X_ub[:, 0:1]
        self.t_ub = X_ub[:, 1:2]

        self.x_f = X_f[:, 0:1]
        self.t_f = X_f[:, 1:2]

        self.u0 = u0
        self.v0 = v0

        # Initialize NNs
        self.layers = layers
        self.weights, self.biases = self.initialize_NN(layers)

        # tf Placeholders
        self.x0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x0.shape[1]])
        self.t0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t0.shape[1]])

        self.u0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.u0.shape[1]])
        self.v0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.v0.shape[1]])

        self.x_lb_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x_lb.shape[1]])
        self.t_lb_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t_lb.shape[1]])

        self.x_ub_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x_ub.shape[1]])
        self.t_ub_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t_ub.shape[1]])

        self.x_f_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x_f.shape[1]])
        self.t_f_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t_f.shape[1]])

        # tf Graphs
        self.u0_pred, self.v0_pred, _, _ = self.net_uv(self.x0_tf, self.t0_tf)
        self.u_lb_pred, self.v_lb_pred, self.u_x_lb_pred, self.v_x_lb_pred = self.net_uv(self.x_lb_tf, self.t_lb_tf)
        self.u_ub_pred, self.v_ub_pred, self.u_x_ub_pred, self.v_x_ub_pred = self.net_uv(self.x_ub_tf, self.t_ub_tf)
        self.f_u_pred, self.f_v_pred = self.net_f_uv(self.x_f_tf, self.t_f_tf)

        # Loss
        self.loss = tf.reduce_mean(tf.square(self.u0_tf - self.u0_pred)) + \
                    tf.reduce_mean(tf.square(self.v0_tf - self.v0_pred)) + \
                    tf.reduce_mean(tf.square(self.u_lb_pred - self.u_ub_pred)) + \
                    tf.reduce_mean(tf.square(self.v_lb_pred - self.v_ub_pred)) + \
                    tf.reduce_mean(tf.square(self.u_x_lb_pred - self.u_x_ub_pred)) + \
                    tf.reduce_mean(tf.square(self.v_x_lb_pred - self.v_x_ub_pred)) + \
                    tf.reduce_mean(tf.square(self.f_u_pred)) + \
                    tf.reduce_mean(tf.square(self.f_v_pred))

        # Optimizers
        self.optimizer_Adam = tf.compat.v1.train.AdamOptimizer()
        self.train_op_Adam = self.optimizer_Adam.minimize(self.loss)

        # tf session
        self.sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(allow_soft_placement=True,
                                                                         log_device_placement=True))

        init = tf.compat.v1.global_variables_initializer()
        self.sess.run(init)

    def initialize_NN(self, layers):
        weights = []
        biases = []
        num_layers = len(layers)
        for l in range(0, num_layers - 1):
            W = self.xavier_init(size=[layers[l], layers[l + 1]])
            b = tf.Variable(tf.zeros([1, layers[l + 1]], dtype=tf.float32), dtype=tf.float32)
            weights.append(W)
            biases.append(b)
        return weights, biases

    def xavier_init(self, size):
        in_dim = size[0]
        out_dim = size[1]
        xavier_stddev = np.sqrt(2 / (in_dim + out_dim))
        return tf.Variable(tf.random.truncated_normal([in_dim, out_dim], stddev=xavier_stddev), dtype=tf.float32)

    def neural_net(self, X, weights, biases):
        num_layers = len(weights) + 1

        H = 2.0 * (X - self.lb) / (self.ub - self.lb) - 1.0
        for l in range(0, num_layers - 2):
            W = weights[l]
            b = biases[l]
            H = tf.tanh(tf.add(tf.matmul(H, W), b))
        W = weights[-1]
        b = biases[-1]
        Y = tf.add(tf.matmul(H, W), b)
        return Y

    def net_uv(self, x, t):
        X = tf.concat([x, t], 1)

        uv = self.neural_net(X, self.weights, self.biases)
        u = uv[:, 0:1]
        v = uv[:, 1:2]

        u_x = tf.gradients(u, x)[0]
        v_x = tf.gradients(v, x)[0]

        return u, v, u_x, v_x

    def net_f_uv(self, x, t):
        u, v, u_x, v_x = self.net_uv(x, t)

        u_t = tf.gradients(u, t)[0]
        u_xx = tf.gradients(u_x, x)[0]

        v_t = tf.gradients(v, t)[0]
        v_xx = tf.gradients(v_x, x)[0]

        f_u = u_t + 0.5 * v_xx + (u ** 2 + v ** 2) * v
        f_v = v_t - 0.5 * u_xx - (u ** 2 + v ** 2) * u

        return f_u, f_v

    def callback(self, loss):
        print('Loss:', loss)

    def train(self, nIter):

        tf_dict = {self.x0_tf: self.x0, self.t0_tf: self.t0,
                   self.u0_tf: self.u0, self.v0_tf: self.v0,
                   self.x_lb_tf: self.x_lb, self.t_lb_tf: self.t_lb,
                   self.x_ub_tf: self.x_ub, self.t_ub_tf: self.t_ub,
                   self.x_f_tf: self.x_f, self.t_f_tf: self.t_f}

        start_time = time.time()
        for it in range(nIter):
            self.sess.run(self.train_op_Adam, tf_dict)

            # Print
            if it % 10 == 0:
                elapsed = time.time() - start_time
                loss_value = self.sess.run(self.loss, tf_dict)
                print('It: %d, Loss: %.3e, Time: %.2f' %
                      (it, loss_value, elapsed))
                start_time = time.time()

    def predict(self, X_star):

        tf_dict = {self.x0_tf: X_star[:, 0:1], self.t0_tf: X_star[:, 1:2]}

        u_star = self.sess.run(self.u0_pred, tf_dict)
        v_star = self.sess.run(self.v0_pred, tf_dict)

        tf_dict = {self.x_f_tf: X_star[:, 0:1], self.t_f_tf: X_star[:, 1:2]}

        f_u_star = self.sess.run(self.f_u_pred, tf_dict)
        f_v_star = self.sess.run(self.f_v_pred, tf_dict)

        return u_star, v_star, f_u_star, f_v_star

该模型代码包含了设置网络架构、设置优化器、自动微分、计算损失、训练模型、输出预测值，非常全面。

（1）init函数

__init__函数的作用是初始化。

①参数

def __init__(self, x0, u0, v0, tb, X_f, layers, lb, ub):

参数介绍：
x0, u0, v0: 初始条件的输入数据。
tb: 边界条件的时间数据。
X_f: 物理域内的点数据。
layers: 神经网络的层数和每层的神经元数量。
lb, ub: 输入数据的下界和上界，用于数据归一化

②拼接

X0 = np.concatenate((x0, 0 * x0), 1)
X_lb = np.concatenate((0 * tb + lb[0], tb), 1)  # (lb[0], tb)
X_ub = np.concatenate((0 * tb + ub[0], tb), 1)  # (ub[0], tb)

np.concatenate 是 NumPy 中用于拼接数组的函数，axis=1 表示沿着列（横向）拼接数组，即将两个数组在列方向上连接起来。这种操作常用于在一些算法中构造扩展矩阵，或将数据与某些附加信息（如零矩阵、其他特征）合并到一起。

③计算图和占位符

self.sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(allow_soft_placement=True, log_device_placement=False))

Session是用户使用TensorFlow时的交互式接口，提供了run()方法来执行计算图。在TensorFlow1.x版本中，用tf.Session()创建会话实例，然后用run()运行一个会话，最后通过close()函数关闭会话、释放资源。

allow_soft_placement 也是布尔型参数，若设置为True，则当运算无法在GPU上执行时，会被自动转移至CPU上执行。log_device_placement 也是布尔型参数，当它为True时，网络运行日志便会记录每个节点使用的计算设备并打印，为减少工作量，默认为False。

self.x0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x0.shape[1]])
self.t0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t0.shape[1]])

该代码placeholder是定义一个占位符，占位符是一种特殊的 TensorFlow 张量，通常用于在 TensorFlow 图（计算图）中表示输入数据，不包含实际数据。

使用 tf.compat.v1 是为了在 TensorFlow 2.x 中兼容 TensorFlow 1.x 的代码，之后的一些函数中也可以看到。

④优化器

self.optimizer_Adam = tf.compat.v1.train.AdamOptimizer()
self.train_op_Adam = self.optimizer_Adam.minimize(self.loss)

这里修改了原作者的优化器，不再使用L-BFGS-B，而是只考虑Adam。怎么不算是一种偷懒呢……

（2）xavier_init函数

return tf.Variable(tf.random.truncated_normal([in_dim, out_dim], stddev=xavier_stddev), dtype=tf.float32)

tf.random.truncated_normal表示生成一个服从截断正态分布的随机张量。

（2）neural_net函数

其核心在于：

H = tf.tanh(tf.add(tf.matmul(H, W), b))

即tanh(H*W+b)的操作，包含了神经网络层和激活函数。与作者原代码基本一致。

（3）自动微分

值得注意的是，作者本人在论文里多次强调的自动微分，在代码里其实很简单，即tf.gradients(u, x)[0]。具体来说，tf.gradients(u, x)[0]是自动计算求u关于x的偏导数，即∂u/∂x，其他部分可类推。传统的自动微分一般被包含/应用在损失的反向传播部分，这一过程我们是看不见的。

对自动微分代码的解释需要结合论文中的薛定谔方程示例。由于h分为实部u和虚部v，即h=u+iv，相应地，在训练过程中也需要分为实部和虚部两个分支。

3. 训练PINN模型

if __name__ == "__main__":

    # 列出所有物理设备
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        print(f"Available GPU devices: {gpus}")
    else:
        print("No GPU available.")

    noise = 0.0

    # 定义域
    lb = np.array([-5.0, 0.0])  # 空间域
    ub = np.array([5.0, np.pi / 2])  # 时间域

    # 训练点个数
    N0 = 50
    N_b = 50
    N_f = 20000

    layers = [2, 100, 100, 100, 100, 2]

    # 空间离散：谱傅里叶法； 时间离散：四阶Runge-Kutta法
    data = scipy.io.loadmat('../Data/NLS.mat')  # 从文件导入离散数据点

    t = data['tt'].flatten()[:, None]
    x = data['x'].flatten()[:, None]
    # Exact = np.real(data['uu']).T
    Exact = data['uu']
    Exact_u = np.real(Exact)  # 实部
    Exact_v = np.imag(Exact)  # 虚部
    Exact_h = np.sqrt(Exact_u ** 2 + Exact_v ** 2)  # h的值

    X, T = np.meshgrid(x, t)

    X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
    u_star = Exact_u.T.flatten()[:, None]
    v_star = Exact_v.T.flatten()[:, None]
    h_star = Exact_h.T.flatten()[:, None]

    idx_x = np.random.choice(x.shape[0], N0, replace=False)
    x0 = x[idx_x, :]
    u0 = Exact_u[idx_x, 0:1]
    v0 = Exact_v[idx_x, 0:1]

    idx_t = np.random.choice(t.shape[0], N_b, replace=False)
    tb = t[idx_t, :]

    X_f = lb + (ub - lb) * lhs(2, N_f)

    model = PhysicsInformedNN(x0, u0, v0, tb, X_f, layers, lb, ub)

    start_time = time.time()
    model.train(30000)

    elapsed = time.time() - start_time
    print('Training time: %.4f' % (elapsed))

    # -------------------------------- 打印模型预测值

    u_pred, v_pred, f_u_pred, f_v_pred = model.predict(X_star)
    h_pred = np.sqrt(u_pred ** 2 + v_pred ** 2)

    # 分别打印实部、虚部、总体的误差
    error_u = np.linalg.norm(u_star - u_pred, 2) / np.linalg.norm(u_star, 2)
    error_v = np.linalg.norm(v_star - v_pred, 2) / np.linalg.norm(v_star, 2)
    error_h = np.linalg.norm(h_star - h_pred, 2) / np.linalg.norm(h_star, 2)
    print('Error u: %e' % (error_u))
    print('Error v: %e' % (error_v))
    print('Error h: %e' % (error_h))

    U_pred = griddata(X_star, u_pred.flatten(), (X, T), method='cubic')
    V_pred = griddata(X_star, v_pred.flatten(), (X, T), method='cubic')
    H_pred = griddata(X_star, h_pred.flatten(), (X, T), method='cubic')

    FU_pred = griddata(X_star, f_u_pred.flatten(), (X, T), method='cubic')
    FV_pred = griddata(X_star, f_v_pred.flatten(), (X, T), method='cubic')

    # -------------------------------- Plotting ---------------------------------

    X0 = np.concatenate((x0, 0 * x0), 1)  # (x0, 0)
    X_lb = np.concatenate((0 * tb + lb[0], tb), 1)  # (lb[0], tb)
    X_ub = np.concatenate((0 * tb + ub[0], tb), 1)  # (ub[0], tb)
    X_u_train = np.vstack([X0, X_lb, X_ub])

    fig, ax = plt.subplots(1, 2, figsize=(16, 8))
    for axis in ax:
        axis.axis('off')

    # ----------- Row 0: h(t,x)
    gs0 = gridspec.GridSpec(1, 2)
    gs0.update(top=1 - 0.06, bottom=1 - 1 / 3, left=0.15, right=0.85, wspace=0)
    ax0 = plt.subplot(gs0[:, :])

    h = ax0.imshow(H_pred.T, interpolation='nearest', cmap='YlGnBu',
                   extent=[lb[1], ub[1], lb[0], ub[0]],
                   origin='lower', aspect='auto')
    divider = make_axes_locatable(ax0)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    fig.colorbar(h, cax=cax)

    ax0.plot(X_u_train[:, 1], X_u_train[:, 0], 'kx', label='Data (%d points)' % (X_u_train.shape[0]), markersize=4,
             clip_on=False)

    line = np.linspace(x.min(), x.max(), 2)[:, None]
    ax0.plot(t[75] * np.ones((2, 1)), line, 'k--', linewidth=1)
    ax0.plot(t[100] * np.ones((2, 1)), line, 'k--', linewidth=1)
    ax0.plot(t[125] * np.ones((2, 1)), line, 'k--', linewidth=1)

    ax0.set_xlabel('$t$')
    ax0.set_ylabel('$x$')
    leg = ax0.legend(frameon=False, loc='best')
    ax0.set_title('$|h(t,x)|$', fontsize=10)

    # ----------- Row 1: h(t,x) slices
    gs1 = gridspec.GridSpec(1, 3)
    gs1.update(top=1 - 1 / 3, bottom=0, left=0.1, right=0.9, wspace=0.5)

    time_labels = [75, 100, 125]  # 定义要绘制的时间点

    ax1 = plt.subplot(gs1[0, 0])
    ax1.plot(x, Exact_h[:, 75], 'b-', linewidth=2, label='Exact')
    ax1.plot(x, H_pred[75, :], 'r--', linewidth=2, label='Prediction')
    ax1.set_xlabel('$x$')
    ax1.set_ylabel('$|h(t,x)|$')
    ax1.set_title(f'$t = {t[75].item():.2f}$', fontsize=10)
    ax1.axis('square')
    ax1.set_xlim([-5.1, 5.1])
    ax1.set_ylim([-0.1, 5.1])

    ax2 = plt.subplot(gs1[0, 1])
    ax2.plot(x, Exact_h[:, 100], 'b-', linewidth=2, label='Exact')
    ax2.plot(x, H_pred[100, :], 'r--', linewidth=2, label='Prediction')
    ax2.set_xlabel('$x$')
    ax2.set_ylabel('$|h(t,x)|$')
    ax2.axis('square')
    ax2.set_xlim([-5.1, 5.1])
    ax2.set_ylim([-0.1, 5.1])
    ax2.set_title(f'$t = {t[100].item():.2f}$', fontsize=10)
    ax2.legend(loc='upper center', bbox_to_anchor=(0.5, -0.2), ncol=5, frameon=False)

    ax3 = plt.subplot(gs1[0, 2])
    ax3.plot(x, Exact_h[:, 125], 'b-', linewidth=2, label='Exact')
    ax3.plot(x, H_pred[125, :], 'r--', linewidth=2, label='Prediction')
    ax3.set_xlabel('$x$')
    ax3.set_ylabel('$|h(t,x)|$')
    ax3.axis('square')
    ax3.set_xlim([-5.1, 5.1])
    ax3.set_ylim([-0.1, 5.1])
    ax3.set_title(f'$t = {t[125].item():.2f}$', fontsize=10)

    plt.savefig('./figures/NLS.png')  # 保存图片
    # plt.show()
    plt.close()

该部分代码包含两部分，一部分为模型训练，另一部分为绘图。我认为原作者的代码已经够简洁明了了，所以没有再加什么注释。

数据处理的注意事项是需要保证数据类型一致。没有小数点的数据，TensorFlow会默认其为int型，有小数点的则为float型。在计算过程中，数据类型不同会导致计算无法正确执行，所以在导入数据后，应该确认数据类型，或者手动设置所有数据的类型。

让人很难过的是，由于我对TensorFlow并不熟悉，在前期我尝试了很多种保存模型的方式，无一例外都失败了。从我使用TensorFlow和PyTorch这两种深度学习框架的感受来说，在保存模型及调用上，PyTorch对新手更友好。目前用PINN模型，只能用一次就重新跑一次代码。如果有什么更好的解决方法，欢迎在评论区留言或者私信我，不胜感激！

什么，你问我为什么不用PyTorch写PINN？我当然有尝试啦！目前我写的PyTorch代码能够运行，但PINN模型拟合效果很差，在正则化的情况下也很难避免梯度消失，我还不知道该怎么改进。如果模型跑成功了，我再写一篇分享吧。

三. 完整代码

修改的代码可以直接替换原来的continuous_time_inference (Schrodinger)/Schrodinger.py，也可以继续在该文件目录continuous_time_inference (Schrodinger)下配置新文件.py文件，例如try_Schrodinger_2.11.py。如果自己新建项目，记得拷贝数据集，修改代码中的读取路径，具体操作内容不再赘述。

训练前说明：代码中的 model.train(10000)表示训练回合数为1000，可自行更改。

代码如下：

import sys
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import scipy.io
from scipy.interpolate import griddata
from pyDOE import lhs
from mpl_toolkits.mplot3d import Axes3D
import time
import matplotlib.gridspec as gridspec
from mpl_toolkits.axes_grid1 import make_axes_locatable
import os

np.random.seed(1234)
tf.random.set_seed(1234)

# 禁用Eager Execution
tf.compat.v1.disable_eager_execution()

# 运行初期会输出当前计算任务在分配设备时的详细情况日志，不想了解可以设置日志级别为ERROR,只输出错误信息
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

class PhysicsInformedNN(tf.Module):
    # Initialize the class
    def __init__(self, x0, u0, v0, tb, X_f, layers, lb, ub):

        X0 = np.concatenate((x0, 0 * x0), 1)  # (x0, 0)
        X_lb = np.concatenate((0 * tb + lb[0], tb), 1)  # (lb[0], tb)
        X_ub = np.concatenate((0 * tb + ub[0], tb), 1)  # (ub[0], tb)

        self.lb = lb
        self.ub = ub

        self.x0 = X0[:, 0:1]
        self.t0 = X0[:, 1:2]

        self.x_lb = X_lb[:, 0:1]
        self.t_lb = X_lb[:, 1:2]

        self.x_ub = X_ub[:, 0:1]
        self.t_ub = X_ub[:, 1:2]

        self.x_f = X_f[:, 0:1]
        self.t_f = X_f[:, 1:2]

        self.u0 = u0
        self.v0 = v0

        # Initialize NNs
        self.layers = layers
        self.weights, self.biases = self.initialize_NN(layers)

        # tf Placeholders
        self.x0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x0.shape[1]])
        self.t0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t0.shape[1]])

        self.u0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.u0.shape[1]])
        self.v0_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.v0.shape[1]])

        self.x_lb_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x_lb.shape[1]])
        self.t_lb_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t_lb.shape[1]])

        self.x_ub_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x_ub.shape[1]])
        self.t_ub_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t_ub.shape[1]])

        self.x_f_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x_f.shape[1]])
        self.t_f_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t_f.shape[1]])

        # tf Graphs
        self.u0_pred, self.v0_pred, _, _ = self.net_uv(self.x0_tf, self.t0_tf)
        self.u_lb_pred, self.v_lb_pred, self.u_x_lb_pred, self.v_x_lb_pred = self.net_uv(self.x_lb_tf, self.t_lb_tf)
        self.u_ub_pred, self.v_ub_pred, self.u_x_ub_pred, self.v_x_ub_pred = self.net_uv(self.x_ub_tf, self.t_ub_tf)
        self.f_u_pred, self.f_v_pred = self.net_f_uv(self.x_f_tf, self.t_f_tf)

        # Loss
        self.loss = tf.reduce_mean(tf.square(self.u0_tf - self.u0_pred)) + \
                    tf.reduce_mean(tf.square(self.v0_tf - self.v0_pred)) + \
                    tf.reduce_mean(tf.square(self.u_lb_pred - self.u_ub_pred)) + \
                    tf.reduce_mean(tf.square(self.v_lb_pred - self.v_ub_pred)) + \
                    tf.reduce_mean(tf.square(self.u_x_lb_pred - self.u_x_ub_pred)) + \
                    tf.reduce_mean(tf.square(self.v_x_lb_pred - self.v_x_ub_pred)) + \
                    tf.reduce_mean(tf.square(self.f_u_pred)) + \
                    tf.reduce_mean(tf.square(self.f_v_pred))

        # Optimizers
        self.optimizer_Adam = tf.compat.v1.train.AdamOptimizer()
        self.train_op_Adam = self.optimizer_Adam.minimize(self.loss)

        # tf session
        self.sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(allow_soft_placement=True,
                                                                         log_device_placement=True))

        init = tf.compat.v1.global_variables_initializer()
        self.sess.run(init)

    def initialize_NN(self, layers):
        weights = []
        biases = []
        num_layers = len(layers)
        for l in range(0, num_layers - 1):
            W = self.xavier_init(size=[layers[l], layers[l + 1]])
            b = tf.Variable(tf.zeros([1, layers[l + 1]], dtype=tf.float32), dtype=tf.float32)
            weights.append(W)
            biases.append(b)
        return weights, biases

    def xavier_init(self, size):
        in_dim = size[0]
        out_dim = size[1]
        xavier_stddev = np.sqrt(2 / (in_dim + out_dim))
        return tf.Variable(tf.random.truncated_normal([in_dim, out_dim], stddev=xavier_stddev), dtype=tf.float32)

    def neural_net(self, X, weights, biases):
        num_layers = len(weights) + 1

        H = 2.0 * (X - self.lb) / (self.ub - self.lb) - 1.0
        for l in range(0, num_layers - 2):
            W = weights[l]
            b = biases[l]
            H = tf.tanh(tf.add(tf.matmul(H, W), b))
        W = weights[-1]
        b = biases[-1]
        Y = tf.add(tf.matmul(H, W), b)
        return Y

    def net_uv(self, x, t):
        X = tf.concat([x, t], 1)

        uv = self.neural_net(X, self.weights, self.biases)
        u = uv[:, 0:1]
        v = uv[:, 1:2]

        u_x = tf.gradients(u, x)[0]
        v_x = tf.gradients(v, x)[0]

        return u, v, u_x, v_x

    def net_f_uv(self, x, t):
        u, v, u_x, v_x = self.net_uv(x, t)

        u_t = tf.gradients(u, t)[0]
        u_xx = tf.gradients(u_x, x)[0]

        v_t = tf.gradients(v, t)[0]
        v_xx = tf.gradients(v_x, x)[0]

        f_u = u_t + 0.5 * v_xx + (u ** 2 + v ** 2) * v
        f_v = v_t - 0.5 * u_xx - (u ** 2 + v ** 2) * u

        return f_u, f_v

    def callback(self, loss):
        print('Loss:', loss)

    def train(self, nIter):

        tf_dict = {self.x0_tf: self.x0, self.t0_tf: self.t0,
                   self.u0_tf: self.u0, self.v0_tf: self.v0,
                   self.x_lb_tf: self.x_lb, self.t_lb_tf: self.t_lb,
                   self.x_ub_tf: self.x_ub, self.t_ub_tf: self.t_ub,
                   self.x_f_tf: self.x_f, self.t_f_tf: self.t_f}

        start_time = time.time()
        for it in range(nIter):
            self.sess.run(self.train_op_Adam, tf_dict)

            # Print
            if it % 10 == 0:
                elapsed = time.time() - start_time
                loss_value = self.sess.run(self.loss, tf_dict)
                print('It: %d, Loss: %.3e, Time: %.2f' %
                      (it, loss_value, elapsed))
                start_time = time.time()

    def predict(self, X_star):

        tf_dict = {self.x0_tf: X_star[:, 0:1], self.t0_tf: X_star[:, 1:2]}

        u_star = self.sess.run(self.u0_pred, tf_dict)
        v_star = self.sess.run(self.v0_pred, tf_dict)

        tf_dict = {self.x_f_tf: X_star[:, 0:1], self.t_f_tf: X_star[:, 1:2]}

        f_u_star = self.sess.run(self.f_u_pred, tf_dict)
        f_v_star = self.sess.run(self.f_v_pred, tf_dict)

        return u_star, v_star, f_u_star, f_v_star


if __name__ == "__main__":

    # 列出所有物理设备，检查GPU是否可用
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        print(f"Available GPU devices: {gpus}")
    else:
        print("No GPU available.")

    noise = 0.0

    # 定义域
    lb = np.array([-5.0, 0.0])  # 空间域
    ub = np.array([5.0, np.pi / 2])  # 时间域

    # 训练点个数
    N0 = 50
    N_b = 50
    N_f = 20000

    layers = [2, 100, 100, 100, 100, 2]

    # 空间离散：谱傅里叶法； 时间离散：四阶Runge-Kutta法
    data = scipy.io.loadmat('../Data/NLS.mat') # 从文件导入离散数据点

    t = data['tt'].flatten()[:, None]
    x = data['x'].flatten()[:, None]
    # Exact = np.real(data['uu']).T
    Exact = data['uu']
    Exact_u = np.real(Exact)  # 实部
    Exact_v = np.imag(Exact)  # 虚部
    Exact_h = np.sqrt(Exact_u**2 + Exact_v**2)  # h的值

    X, T = np.meshgrid(x, t)

    X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
    u_star = Exact_u.T.flatten()[:, None]
    v_star = Exact_v.T.flatten()[:, None]
    h_star = Exact_h.T.flatten()[:, None]

    idx_x = np.random.choice(x.shape[0], N0, replace=False)
    x0 = x[idx_x,:]
    u0 = Exact_u[idx_x,0:1]
    v0 = Exact_v[idx_x,0:1]

    idx_t = np.random.choice(t.shape[0], N_b, replace=False)
    tb = t[idx_t, :]

    X_f = lb + (ub - lb) * lhs(2, N_f)

    model = PhysicsInformedNN(x0, u0, v0, tb, X_f, layers, lb, ub)

    start_time = time.time()
    model.train(10000)

    elapsed = time.time() - start_time
    print('Training time: %.4f' % (elapsed))

    # -------------------------------- 打印模型预测值

    u_pred, v_pred, f_u_pred, f_v_pred = model.predict(X_star)
    h_pred = np.sqrt(u_pred ** 2 + v_pred ** 2)

    # 分别打印实部、虚部、总体的误差
    error_u = np.linalg.norm(u_star - u_pred, 2) / np.linalg.norm(u_star, 2)
    error_v = np.linalg.norm(v_star - v_pred, 2) / np.linalg.norm(v_star, 2)
    error_h = np.linalg.norm(h_star - h_pred, 2) / np.linalg.norm(h_star, 2)
    print('Error u: %e' % (error_u))
    print('Error v: %e' % (error_v))
    print('Error h: %e' % (error_h))

    U_pred = griddata(X_star, u_pred.flatten(), (X, T), method='cubic')
    V_pred = griddata(X_star, v_pred.flatten(), (X, T), method='cubic')
    H_pred = griddata(X_star, h_pred.flatten(), (X, T), method='cubic')

    FU_pred = griddata(X_star, f_u_pred.flatten(), (X, T), method='cubic')
    FV_pred = griddata(X_star, f_v_pred.flatten(), (X, T), method='cubic')

    # -------------------------------- Plotting ---------------------------------

    X0 = np.concatenate((x0, 0 * x0), 1)  # (x0, 0)
    X_lb = np.concatenate((0 * tb + lb[0], tb), 1)  # (lb[0], tb)
    X_ub = np.concatenate((0 * tb + ub[0], tb), 1)  # (ub[0], tb)
    X_u_train = np.vstack([X0, X_lb, X_ub])

    fig, ax = plt.subplots(1, 2, figsize=(16, 8))
    for axis in ax:
        axis.axis('off')

    # ----------- Row 0: h(t,x)
    gs0 = gridspec.GridSpec(1, 2)
    gs0.update(top=1 - 0.06, bottom=1 - 1 / 3, left=0.15, right=0.85, wspace=0)
    ax0 = plt.subplot(gs0[:, :])

    h = ax0.imshow(H_pred.T, interpolation='nearest', cmap='YlGnBu',
                   extent=[lb[1], ub[1], lb[0], ub[0]],
                   origin='lower', aspect='auto')
    divider = make_axes_locatable(ax0)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    fig.colorbar(h, cax=cax)

    ax0.plot(X_u_train[:, 1], X_u_train[:, 0], 'kx', label='Data (%d points)' % (X_u_train.shape[0]), markersize=4,
             clip_on=False)

    line = np.linspace(x.min(), x.max(), 2)[:, None]
    ax0.plot(t[75] * np.ones((2, 1)), line, 'k--', linewidth=1)
    ax0.plot(t[100] * np.ones((2, 1)), line, 'k--', linewidth=1)
    ax0.plot(t[125] * np.ones((2, 1)), line, 'k--', linewidth=1)

    ax0.set_xlabel('$t$')
    ax0.set_ylabel('$x$')
    leg = ax0.legend(frameon=False, loc='best')
    ax0.set_title('$|h(t,x)|$', fontsize=10)

    # ----------- Row 1: h(t,x) slices
    gs1 = gridspec.GridSpec(1, 3)
    gs1.update(top=1 - 1 / 3, bottom=0, left=0.1, right=0.9, wspace=0.5)

    time_labels = [75, 100, 125]  # 定义要绘制的时间点

    ax1 = plt.subplot(gs1[0, 0])
    ax1.plot(x, Exact_h[:, 75], 'b-', linewidth=2, label='Exact')
    ax1.plot(x, H_pred[75, :], 'r--', linewidth=2, label='Prediction')
    ax1.set_xlabel('$x$')
    ax1.set_ylabel('$|h(t,x)|$')
    ax1.set_title(f'$t = {t[75].item():.2f}$', fontsize=10)
    ax1.axis('square')
    ax1.set_xlim([-5.1, 5.1])
    ax1.set_ylim([-0.1, 5.1])

    ax2 = plt.subplot(gs1[0, 1])
    ax2.plot(x, Exact_h[:, 100], 'b-', linewidth=2, label='Exact')
    ax2.plot(x, H_pred[100, :], 'r--', linewidth=2, label='Prediction')
    ax2.set_xlabel('$x$')
    ax2.set_ylabel('$|h(t,x)|$')
    ax2.axis('square')
    ax2.set_xlim([-5.1, 5.1])
    ax2.set_ylim([-0.1, 5.1])
    ax2.set_title(f'$t = {t[100].item():.2f}$', fontsize=10)
    ax2.legend(loc='upper center', bbox_to_anchor=(0.5, -0.2), ncol=5, frameon=False)

    ax3 = plt.subplot(gs1[0, 2])
    ax3.plot(x, Exact_h[:, 125], 'b-', linewidth=2, label='Exact')
    ax3.plot(x, H_pred[125, :], 'r--', linewidth=2, label='Prediction')
    ax3.set_xlabel('$x$')
    ax3.set_ylabel('$|h(t,x)|$')
    ax3.axis('square')
    ax3.set_xlim([-5.1, 5.1])
    ax3.set_ylim([-0.1, 5.1])
    ax3.set_title(f'$t = {t[125].item():.2f}$', fontsize=10)

    plt.savefig('./figures/NLS.png')  # 保存图片
    plt.show()
    plt.close()