PINN求解二维泊松方程:u''(x)+u''(y)=-2pipisin(pix)sin(piy);边界上有u(x,y)=0,真解为u=sin(pix)sin(pi*y)的代码
时间: 2024-05-22 21:15:58 浏览: 179
# 导入必要的库
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense, Lambda
# 定义模型参数
x_min, x_max = 0, 1
y_min, y_max = 0, 1
Nx, Ny = 20, 20
dx, dy = (x_max - x_min) / (Nx - 1), (y_max - y_min) / (Ny - 1)
# 生成网格点
x = np.linspace(x_min, x_max, Nx)
y = np.linspace(y_min, y_max, Ny)
X, Y = np.meshgrid(x, y)
# 定义边界条件
u_b = np.zeros((Ny, Nx))
u_b[0, :] = np.sin(np.pi * x)
u_b[-1, :] = np.sin(np.pi * x) * np.exp(-np.pi)
u_b[:, 0] = 0
u_b[:, -1] = 0
# 定义输入和输出
x = Input(shape=(2,))
u = Dense(20, activation='tanh')(x)
u = Dense(20, activation='tanh')(u)
u = Dense(20, activation='tanh')(u)
u = Dense(1)(u)
# 定义边界损失
def boundary_loss(model, X, Y, u_b):
u_pred = model(tf.concat([tf.reshape(X, [-1, 1]), tf.reshape(Y, [-1, 1])], axis=1))
u_pred = tf.reshape(u_pred, [Ny, Nx])
return tf.reduce_mean(tf.square(u_pred[0, :] - u_b[0, :])) + \
tf.reduce_mean(tf.square(u_pred[-1, :] - u_b[-1, :])) + \
tf.reduce_mean(tf.square(u_pred[:, 0] - u_b[:, 0])) + \
tf.reduce_mean(tf.square(u_pred[:, -1] - u_b[:, -1]))
# 定义PDE损失
def pde_loss(model, X, Y):
u_pred = model(tf.concat([tf.reshape(X, [-1, 1]), tf.reshape(Y, [-1, 1])], axis=1))
u_x = tf.gradients(u_pred, X)[0]
u_xx = tf.gradients(u_x, X)[0]
u_y = tf.gradients(u_pred, Y)[0]
u_yy = tf.gradients(u_y, Y)[0]
f = -2*np.pi*np.sin(np.pi*X)*np.sin(np.pi*Y)
return tf.reduce_mean(tf.square(u_xx + u_yy - f))
# 定义总损失
def loss(model, X, Y, u_b):
return boundary_loss(model, X, Y, u_b) + pde_loss(model, X, Y)
# 定义优化器和训练步数、批次大小
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
n_epochs, batch_size = 10000, 100
# 定义模型
model = Model(inputs=x, outputs=u)
# 训练模型
for epoch in range(n_epochs):
idx = np.random.choice(Nx * Ny, batch_size, replace=False)
X_batch = np.hstack([X.reshape(-1, 1)[idx], Y.reshape(-1, 1)[idx]])
with tf.GradientTape() as tape:
loss_value = loss(model, X_batch[:, 0], X_batch[:, 1], u_b)
grads = tape.gradient(loss_value, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
if epoch % 100 == 0:
print("Epoch {}: Loss = {}".format(epoch, loss_value.numpy()))
# 计算预测结果
u_pred = model.predict(np.hstack([X.reshape(-1, 1), Y.reshape(-1, 1)]))
u_pred = np.reshape(u_pred, (Ny, Nx))
# 绘制结果
fig, ax = plt.subplots()
cs = ax.contourf(X, Y, u_pred, cmap='coolwarm')
cbar = fig.colorbar(cs)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title("PINN Solution to 2D Poisson Equation")
plt.show()
阅读全文