def compute_loss(model, X_ic, X_bc, X_pde):
# 初始条件损失(t=0时的解)
t_ic, x_ic, h_real_ic, h_imag_ic = X_ic
with tf.GradientTape() as tape:
u_pred_ic, v_pred_ic = model(tf.stack([t_ic, x_ic], axis=1))
loss_ic = tf.reduce_mean(tf.square(u_pred_ic - h_real_ic) +
tf.reduce_mean(tf.square(v_pred_ic - h_imag_ic))
# 周期性边界条件损失(x=-5和x=5处的解和导数相等)
t_bc, x_bc_left, x_bc_right = X_bc
# 左边界预测
u_left, v_left = model(tf.stack([t_bc, x_bc_left], axis=1))
u_x_left = tape.gradient(u_left, x_bc_left)
v_x_left = tape.gradient(v_left, x_bc_left)
# 右边界预测
u_right, v_right = model(tf.stack([t_bc, x_bc_right], axis=1))
u_x_right = tape.gradient(u_right, x_bc_right)
v_x_right = tape.gradient(v_right, x_bc_right)
# 边界损失
loss_bc = (tf.reduce_mean(tf.square(u_left - u_right)) +
tf.reduce_mean(tf.square(v_left - v_right)) +
tf.reduce_mean(tf.square(u_x_left - u_x_right)) +
tf.reduce_mean(tf.square(v_x_left - v_x_right)))
# PDE残差损失(方程本身的误差)
t_pde, x_pde = X_pde
with tf.GradientTape(persistent=True) as tape:
tape.watch([t_pde, x_pde])
u_pde, v_pde = model(tf.stack([t_pde, x_pde], axis=1))
# 计算二阶导数
u_x = tape.gradient(u_pde, x_pde)
u_xx = tape.gradient(u_x, x_pde)
v_x = tape.gradient(v_pde, x_pde)
v_xx = tape.gradient(v_x, x_pde)
# 计算时间导数
u_t = tape.gradient(u_pde, t_pde)
v_t = tape.gradient(v_pde, t_pde)
# PDE残差
f_real = -v_t + 0.5 * u_xx + (u_pde**2 + v_pde**2) * u_pde
f_imag = u_t + 0.5 * v_xx + (u_pde**2 + v_pde**2) * v_pde
loss_pde = tf.reduce_mean(tf.square(f_real) + tf.square(f_imag))
# 总损失(加权求和)
total_loss = 1.0 * loss_ic + 1.0 * loss_bc + 1.0 * loss_pde
return total_loss