Python 机器学习求解 PDE 学习项目——PINN 求解二维 Poisson 方程

本文使用 TensorFlow 1.15 环境搭建深度神经网络（PINN）求解二维 Poisson 方程:

模型问题

$\begin{align} -\Delta u &= f \quad & \text{in } \Omega,\\ u & =g \quad & \text{on } \Gamma:=\partial \Omega. \end{align}$
其中 $\Omega = [X_a,X_b]\times[Y_a,Y_b]$ 是一个二维矩形区域， $\Delta u = u_{xx}+u_{yy}, g$ 是边界条件给定的函数，可以非零.

在这里插入图片描述

代码展现

二维PINN 与一维的整体框架是类似的，只是数据的维度升高了，完整代码及其注释如下：

# PINN 求解 2D Poisson 方程
import tensorflow as tf
print(tf.__version__)

import os
#tensorflow-intel automatically set the TF_ENABLE_ONEDNN_OPTS=1
#os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
# Here, TF_ENABLE_ONEDNN_OPTS=0 should be above import tensorflow as tf
import tensorflow as tf
import numpy as np
import time
import matplotlib.pyplot as plt
import scipy.io
import math
# 定义数据集类，用于生成训练所需的数据
class Dataset:
    def __init__(self, x_range, y_range, N_res, N_bx, N_by, Nx, Ny, xa, xb, ya, yb):
        self.x_range = x_range  # x 轴范围
        self.y_range = y_range  # y 轴范围
        self.N_res = N_res  # 方程残差点数量
        self.N_bx = N_bx  # x 方向边界条件点数量
        self.N_by = N_by  # y 方向边界条件点数量
        self.Nx = Nx  # x 方向网格数量
        self.Ny = Ny  # y 方向网格数量
        self.xa = xa  # x 方向左边界
        self.xb = xb  # x 方向右边界
        self.ya = ya  # y 方向下边界
        self.yb = yb  # y 方向上边界

    # 定义边界条件函数
    # 可以求解非齐次 Dirichlet 边界条件
    def bc(self, X_b):
        U_bc = Exact(self.xa, self.xb, self.ya, self.yb)
        u_bc = U_bc.u_exact(X_b)
        return u_bc

    # 生成数据：残差点和边界条件点
    def build_data(self):
        x0, x1 = self.x_range
        y0, y1 = self.y_range

        Xmin = np.hstack((x0, y0))
        Xmax = np.hstack((x1, y1))

        ## 如果使用均匀网格，代码如下：
        """
        x_ = np.linspace(x0, x1, self.Nx).reshape((-1, 1))
        y_ = np.linspace(y0, y1, self.Ny).reshape((-1, 1))
        x, y = np.meshgrid(x_, y_)
        x = np.reshape(x, (-1, 1))
        y = np.reshape(y, (-1, 1))
        xy = np.hstack((x, y))
        X_res_input = xy
        """

        ## 为方程生成随机残差点
        x_res = x0 + (x1 - x0) * np.random.rand(self.N_res, 1)
        y_res = y0 + (y1 - y0) * np.random.rand(self.N_res, 1)
        X_res_input = np.hstack((x_res, y_res))

        # 生成 x = xa, xb 的边界条件点
        y_b = y0 + (y1 - y0) * np.random.rand(self.N_by, 1)
        x_b0 = x0 * np.ones_like(y_b)
        x_b1 = x1 * np.ones_like(y_b)
        X_b0_input = np.hstack((x_b0, y_b))
        X_b1_input = np.hstack((x_b1, y_b))

        # 生成 y = ya, yb 的边界条件点
        x_b = x0 + (x1 - x0) * np.random.rand(self.N_bx, 1)
        y_b0 = y0 * np.ones_like(x_b)
        y_b1 = y1 * np.ones_like(x_b)
        Y_b0_input = np.hstack((x_b, y_b0))
        Y_b1_input = np.hstack((x_b, y_b1))
        
        return X_res_input, X_b0_input, X_b1_input, Y_b0_input, Y_b1_input, Xmin, Xmax

# 定义精确解类，用于计算精确解
class Exact:
    def __init__(self, xa, xb, ya, yb):
        self.xa = xa  # x 方向左边界
        self.xb = xb  # x 方向右边界
        self.ya = ya  # y 方向下边界
        self.yb = yb  # y 方向上边界

    # 精确解函数
    def u_exact(self, X):
        x = X[:, 0:1]
        y = X[:, 1:2]

        u = np.sin(2 * np.pi * x  ) * np.sin(2 * np.pi * y )
        return u

class Train:
    def __init__(self, train_dict):
        self.train_dict = train_dict  # 训练数据
        self.step = 0  # 训练步数

    # 打印训练损失
    def callback(self, loss_value):
        self.step += 1
        if self.step % 200 == 0:
            print(f'Loss: {loss_value:.4e}')

    # 使用 Adam 和 L-BFGS 优化器进行训练
    def nntrain(self, sess, u_pred, loss, train_adam, train_lbfgs):
        n = 0
        max_steps = 1000
        loss_threshold = 4.0e-4
        current_loss = 1.0
        while n < max_steps and current_loss > loss_threshold:
            n += 1
            u_, current_loss, _ = sess.run([u_pred, loss, train_adam], feed_dict=self.train_dict)
            # 每2^n步打印一次损失并绘制结果  
            if math.isclose(math.fmod(math.log2(n), 1), 0, abs_tol=1e-9): 
                print(f'Steps: {n}, loss: {current_loss:.4e}')

        train_lbfgs.minimize(sess, feed_dict=self.train_dict, fetches=[loss], loss_callback=self.callback)

class DNN:
    def __init__(self, layer_sizes, Xmin, Xmax):
        self.layer_sizes = layer_sizes  # 每层的节点数
        self.Xmin = Xmin  # 输入范围最小值
        self.Xmax = Xmax  # 输入范围最大值
    
    # 初始化神经网络的权重和偏置
    def hyper_initial(self):
        num_layers = len(self.layer_sizes)
        weights = []
        biases = []
        for l in range(1, num_layers):
            in_dim = self.layer_sizes[l-1]
            out_dim = self.layer_sizes[l]
            std = np.sqrt(2 / (in_dim + out_dim))
            weight = tf.Variable(tf.random_normal(shape=[in_dim, out_dim], stddev=std))
            bias = tf.Variable(tf.zeros(shape=[1, out_dim]))
            weights.append(weight)
            biases.append(bias)
        return weights, biases

    # 构建前馈神经网络
    def fnn(self, X, weights, biases):
        A = 2.0 * (X - self.Xmin) / (self.Xmax - self.Xmin) - 1.0  # 归一化输入
        num_layers = len(weights)
        for i in range(num_layers - 1):
            A = tf.tanh(tf.add(tf.matmul(A, weights[i]), biases[i]))  # 隐藏层激活函数
        Y = tf.add(tf.matmul(A, weights[-1]), biases[-1])  # 输出层
        return Y
    
    # 构建用于求解 Poisson 方程的神经网络
    def pdenn(self, x, y, weights, biases):
        u = self.fnn(tf.concat([x, y], 1), weights, biases)  # 前馈网络输出
        u_x = tf.gradients(u, x)[0]  # u 对 x 的一阶导数
        u_xx = tf.gradients(u_x, x)[0]  # u 对 x 的二阶导数
        u_y = tf.gradients(u, y)[0]  # u 对 y 的一阶导数
        u_yy = tf.gradients(u_y, y)[0]  # u 对 y 的二阶导数
        
        # 源项函数
        rhs_func = 8 * np.pi**2 * tf.sin(2 * np.pi * x  ) * tf.sin(2 * np.pi * y )

        # 残差项
        residual = -(u_xx + u_yy) - rhs_func

        return residual

def compute_errors(u_pred, u_exact):
    """
    计算数值解与精确解之间的 L2 误差和最大模误差
    :param u_pred: 数值解
    :param u_exact: 精确解
    :return: L2 误差和最大模误差
    """
    # 计算 L2 误差
    L2_error = np.sqrt(np.mean((u_pred - u_exact) ** 2))
    
    # 计算最大模误差
    max_error = np.max(np.abs(u_pred - u_exact))
    
    return L2_error, max_error

# 检查保存路径是否存在，如果不存在则创建
save_path = './Output'
if not os.path.exists(save_path):
    os.makedirs(save_path)

# 定义保存和绘图类
class SavePlot:
    def __init__(self, session, x_range, y_range, num_x_points, num_y_points, xa, xb, ya, yb):
        self.x_range = x_range  # x 轴范围
        self.y_range = y_range  # y 轴范围
        self.num_x_points = num_x_points  # x 方向上的测试点数量
        self.num_y_points = num_y_points  # y 方向上的测试点数量
        self.session = session  # TensorFlow 会话
        self.xa = xa  # x 方向左边界
        self.xb = xb  # x 方向右边界
        self.ya = ya  # y 方向下边界
        self.yb = yb  # y 方向上边界

    # 保存并绘制预测和精确解
    def save_and_plot(self, u_pred, x_res_train, y_res_train):
        # 生成测试点
        x_test = np.linspace(self.x_range[0], self.x_range[1], self.num_x_points).reshape((-1, 1))
        y_test = np.linspace(self.y_range[0], self.y_range[1], self.num_y_points).reshape((-1, 1))
        x_test_grid, y_test_grid = np.meshgrid(x_test, y_test)
        x_test_grid = np.reshape(x_test_grid, (-1, 1))
        y_test_grid = np.reshape(y_test_grid, (-1, 1))
        
        # 创建测试字典
        test_feed_dict = {x_res_train: x_test_grid, y_res_train: y_test_grid}
        
        # 在测试网格上进行预测
        u_test = self.session.run(u_pred, feed_dict=test_feed_dict)
        u_test = np.reshape(u_test, (y_test.shape[0], x_test.shape[0]))
        u_test = np.transpose(u_test)

        # 保存预测结果到文件
        np.savetxt(os.path.join(save_path, 'u_pred.txt'), u_test, fmt='%e')

        # 绘制预测结果并保存图片
        plt.imshow(u_test, cmap='rainbow', aspect='auto')
        plt.colorbar()
        plt.title('Numerical Solution')
        plt.xlabel('X-axis')
        plt.ylabel('Y-axis')
        plt.savefig(os.path.join(save_path, 'u_pred.png'))
        plt.show()
        plt.close()

        # 计算并保存精确解
        exact_solution = Exact(self.xa, self.xb, self.ya, self.yb)
        u_exact = exact_solution.u_exact(np.hstack((x_test_grid, y_test_grid)))
        u_exact = np.reshape(u_exact, (y_test.shape[0], x_test.shape[0]))
        u_exact = np.transpose(u_exact)
        np.savetxt(os.path.join(save_path, 'u_exact.txt'), u_exact, fmt='%e')

        # 绘制精确解并保存图片
        plt.imshow(u_exact, cmap='rainbow', aspect='auto')
        plt.colorbar()
        plt.title('Exact Solution')
        plt.xlabel('X-axis')
        plt.ylabel('Y-axis')
        plt.savefig(os.path.join(save_path, 'u_exact.png'))
        plt.show()
        plt.close()

下面是主程序：

import os
import tensorflow as tf
import numpy as np
import time
import matplotlib.pyplot as plt

# 设置随机种子以确保可重复性
np.random.seed(1234)
tf.set_random_seed(1234)

def main():
    # 定义计算域范围
    x_range = [-0.5, 1.5]
    y_range = [-1.0, 1.0]
    
    # 网格点数量
    num_x_points = 101
    num_y_points = 101
    
    # 残差点和边界点数量
    num_residual_points = 8000
    num_boundary_x_points = 100
    num_boundary_y_points = 100
    
    # 边界范围
    xa = x_range[0]
    xb = x_range[1]
    ya = y_range[0]
    yb = y_range[1]

    # 创建数据集对象
    data = Dataset(x_range, y_range, num_residual_points, num_boundary_x_points, num_boundary_y_points, num_x_points, num_y_points, xa, xb, ya, yb)
    
    # 生成数据
    X_res, X_b0, X_b1, Y_b0, Y_b1, Xmin, Xmax = data.build_data()

    # 定义神经网络的层结构
    layers = [2] + 5 * [40] + [1]

    # 定义输入占位符
    x_res_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    y_res_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    X_x_b0_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    X_y_b0_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    X_x_b1_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    X_y_b1_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    Y_x_b0_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    Y_y_b0_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    Y_x_b1_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)
    Y_y_b1_train = tf.placeholder(shape=[None, 1], dtype=tf.float32)

    # 创建物理信息神经网络（PINN）
    pinn = DNN(layers, Xmin, Xmax)
    weights, biases = pinn.hyper_initial()

    # 预测解
    u_pred = pinn.fnn(tf.concat([x_res_train, y_res_train], 1), weights, biases)

    # 计算残差
    f_pred = pinn.pdenn(x_res_train, y_res_train, weights, biases)

    # 边界条件预测 (x = xa, xb)
    u_x_b0_pred = pinn.fnn(tf.concat([X_x_b0_train, X_y_b0_train], 1), weights, biases)
    u_x_b1_pred = pinn.fnn(tf.concat([X_x_b1_train, X_y_b1_train], 1), weights, biases)

    # 边界条件预测 (y = ya, yb)
    u_y_b0_pred = pinn.fnn(tf.concat([Y_x_b0_train, Y_y_b0_train], 1), weights, biases)
    u_y_b1_pred = pinn.fnn(tf.concat([Y_x_b1_train, Y_y_b1_train], 1), weights, biases)

    # 定义损失函数
    loss = 0.1 * tf.reduce_mean(tf.square(f_pred)) + \
           tf.reduce_mean(tf.square(u_x_b0_pred)) + \
           tf.reduce_mean(tf.square(u_x_b1_pred)) + \
           tf.reduce_mean(tf.square(u_y_b0_pred)) + \
           tf.reduce_mean(tf.square(u_y_b1_pred))

    # 定义优化器
    train_adam = tf.train.AdamOptimizer(0.0008).minimize(loss)
    train_lbfgs = tf.contrib.opt.ScipyOptimizerInterface(loss,
                                                         method="L-BFGS-B",
                                                         options={'maxiter': 10000 ,
                                                                  'ftol': 1.0 * np.finfo(float).eps})

    # 创建 TensorFlow 会话
    session = tf.Session()
    session.run(tf.global_variables_initializer())
    
    # 创建训练字典
    train_feed_dict = {x_res_train: X_res[:, 0:1], y_res_train: X_res[:, 1:2], 
                       X_x_b0_train: X_b0[:, 0:1], X_y_b0_train: X_b0[:, 1:2],
                       X_x_b1_train: X_b1[:, 0:1], X_y_b1_train: X_b1[:, 1:2], 
                       Y_x_b0_train: Y_b0[:, 0:1], Y_y_b0_train: Y_b0[:, 1:2], 
                       Y_x_b1_train: Y_b1[:, 0:1], Y_y_b1_train: Y_b1[:, 1:2]}

    # 创建训练模型
    model = Train(train_feed_dict)
    
    # 记录训练时间
    start_time = time.perf_counter()
    model.nntrain(session, u_pred, loss, train_adam, train_lbfgs)
    stop_time = time.perf_counter()
    print('训练时间为 %.3f 秒' % (stop_time - start_time))

    # 保存预测数据和图像
    num_test_x_points = 101
    num_test_y_points = 101
    data_saver = SavePlot(session, x_range, y_range, num_test_x_points, num_test_y_points, xa, xb, ya, yb)
    data_saver.save_and_plot(u_pred, x_res_train, y_res_train)

    
        # 计算误差
    x_test = np.linspace(x_range[0], x_range[1], num_test_x_points).reshape((-1, 1))
    y_test = np.linspace(y_range[0], y_range[1], num_test_y_points).reshape((-1, 1))
    x_t, y_t = np.meshgrid(x_test, y_test)
    x_t = np.reshape(x_t, (-1, 1))
    y_t = np.reshape(y_t, (-1, 1))
    test_dict = {x_res_train: x_t, y_res_train: y_t}
    u_test_pred = session.run(u_pred, feed_dict=test_dict)  # 预测在均匀网格上的解
    Exact_sln = Exact(xa, xb, ya, yb)
    u_test_exact = Exact_sln.u_exact(np.hstack((x_t, y_t)))

    # 计算误差
    L2_error, max_error = compute_errors(u_test_pred, u_test_exact)
    print('L2 Error: %.6e' % L2_error)
    print('Max Error: %.7e' % max_error)
    
    
if __name__ == '__main__':
    main()

程序中已经写好了详细的注释，关于优化器与 TF 会话(session) 的相关知识请各位移步 TensorFlow 优化器使用。另外建议读者对比阅读我之前总结的一维PINN 算法的实现，理解一维二维的本质不同，更高维的 PDE 求解也就不在话下了。

运行结果

在这里插入图片描述

效果不错！

~~-----------------------------------------------------------------------------------------------~~
本专栏目标从简单的一维 Poisson 方程，到对流扩散方程，Burges 方程，到二维，三维以及非线性方程，发展方程，积分方程等等，所有文章包含全部可运行代码。请持续关注！