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

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

模型问题

$$\begin{array}{rlr}-\Delta u& =f& \text{in }\Omega ,\\ u& =g& \text{on }\Gamma :=\partial \Omega .\end{array}$$
其中 $\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_rangey0, y1 = self.y_rangeXmin = 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 uclass Train:def __init__(self, train_dict):self.train_dict = train_dict  # 训练数据self.step = 0  # 训练步数# 打印训练损失def callback(self, loss_value):self.step += 1if 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 = 0max_steps = 1000loss_threshold = 4.0e-4current_loss = 1.0while n < max_steps and current_loss > loss_threshold:n += 1u_, 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_funcreturn residualdef 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 = 101num_y_points = 101# 残差点和边界点数量num_residual_points = 8000num_boundary_x_points = 100num_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 = 101num_test_y_points = 101data_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 方程，到二维，三维以及非线性方程，发展方程，积分方程等等，所有文章包含全部可运行代码。请持续关注！