以𝑑𝑦/𝑑𝑡=(𝑎+𝜔cos(𝜋𝑡/2) ), 𝑡∈[0, 10], 𝑎=1,𝜔=𝜋/2为例写一段用tensorflow编写PINN求解该微分方程的代码
时间: 2024-06-01 16:12:13 浏览: 83
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
# 定义神经网络模型
class PINN(tf.keras.Model):
def __init__(self):
super(PINN, self).__init__()
self.dense1 = tf.keras.layers.Dense(32, activation=tf.nn.tanh)
self.dense2 = tf.keras.layers.Dense(32, activation=tf.nn.tanh)
self.dense3 = tf.keras.layers.Dense(1)
def call(self, inputs):
x = inputs[:, 0:1]
t = inputs[:, 1:]
u = x * 0.0
inputs = tf.concat([x, t, u], axis=1)
x = self.dense1(inputs)
x = self.dense2(x)
u = self.dense3(x)
return u
# 定义损失函数
def pinn_loss(u, dudt, x, t):
with tf.GradientTape() as g:
g.watch(x)
g.watch(t)
inputs = tf.concat([x, t], axis=1)
u_pred = model(inputs)
grads = g.gradient(u_pred, inputs)
dudx = grads[:, 0:1]
dudt_pred = grads[:, 1:]
f_pred = dudt_pred - 0.5 * tf.sin(0.5 * np.pi * t) * dudx + np.pi * tf.cos(0.5 * np.pi * t) * u_pred
loss_u = tf.reduce_mean(tf.square(u - u_pred))
loss_f = tf.reduce_mean(tf.square(f_pred - dudt))
loss = loss_u + loss_f
return loss
# 定义训练函数
def train(model, x, t, u, dudt, optimizer, epochs):
for epoch in range(epochs):
with tf.GradientTape() as tape:
loss = pinn_loss(u, dudt, x, t)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
if epoch % 100 == 0:
print('Epoch %d: Loss = %f' % (epoch, loss.numpy()))
# 定义训练数据
n = 1000
x = np.linspace(0, 1, n)[:, None]
t = np.linspace(0, 10, n)[:, None]
u = np.zeros_like(x)
dudt = np.cos(0.5 * np.pi * t)
# 创建模型和优化器
model = PINN()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
# 开始训练
train(model, x, t, u, dudt, optimizer, epochs=1000)
# 预测结果并可视化
u_pred = model(tf.concat([x, t], axis=1)).numpy()
plt.plot(x, u_pred, 'r-', label='Prediction')
plt.plot(x, u, 'b--', label='Exact')
plt.legend()
plt.show()
阅读全文