def call(self, psi_hr, psi_gen):
vel_gen = ops.get_velocity(psi_gen, self.dx, self.dy)
vel_hr = ops.get_velocity(psi_hr, self.dx, self.dy)
vel_grad_gen = ops.get_velocity_grad(vel_gen, self.dx, self.dy)
mse_loss = get_mse_loss(psi_hr, psi_gen)
ens_loss = get_ens_loss(psi_hr, psi_gen, self.dx, self.dy)
continuity_loss = ops.get_continuity_residual(vel_grad_gen)
content_loss = (1-self.lambda_ens) * mse_loss + \
self.lambda_ens * ens_loss
gen_loss = (1-self.lambda_phys) * content_loss + \
self.lambda_phys * continuity_loss
return gen_loss
def get_ens_loss(input_hr, input_gen, dx, dy):
input_size = tf.cast(tf.size(input_hr), tf.float32)
vel_grad_gen = ops.get_velocity_grad(ops.get_velocity(input_gen, dx, dy),
dx, dy)
vel_grad_hr = ops.get_velocity_grad(ops.get_velocity(input_hr, dx, dy), dx,
dy)
vorticity_gen = ops.get_vorticity(vel_grad_gen)
vorticity_hr = ops.get_vorticity(vel_grad_hr)
ens_gen = ops.get_enstrophy(vorticity_gen)
ens_hr = ops.get_enstrophy(vorticity_hr)
# ens_hr_mean2 = tf.reduce_mean(tf.reduce_mean(
# tf.square(ens_hr), axis=1, keepdims=True), axis=2, keepdims=True)
#ens_loss = tf.reduce_mean(tf.square(ens_gen-ens_hr)/ens_hr_mean2)
ens_loss = tf.reduce_mean(tf.square(ens_gen - ens_hr) / input_size)
return ens_loss
ph_LR_ps = ops.PeriodicPadding2D(padding=24)(ph_LR_single)
fig.clear()
# 可視化
fig = plt.figure()
ax1 = fig.add_subplot(121)
#im1 = ax1.imshow(test_tf_sin[0, ..., 0], cmap='jet')
im1 = ax1.imshow(ph_LR_single[0, ..., 0], cmap='jet')
fig.colorbar(im1)
ax2 = fig.add_subplot(122)
#im2 = ax2.imshow(test_tf_sin_ps[0, ..., 0], cmap='jet')
im2 = ax2.imshow(ph_LR_ps[0, ..., 0], cmap='jet')
fig.colorbar(im2)
fig.savefig('ps_TF_test.png')
# 速度場も出してみる
vx, vy = ops.get_velocity(ph_LR_single, dx, dy)
# 可視化
fig.clear()
fig = plt.figure(figsize=(12, 6))
ax1 = fig.add_subplot(131)
#im1 = ax1.imshow(test_tf_sin[0, ..., 0], cmap='jet')
im1 = ax1.imshow(ph_LR_single[0, ..., 0], cmap='jet')
fig.colorbar(im1)
ax2 = fig.add_subplot(132)
#im2 = ax2.imshow(test_tf_sin_ps[0, ..., 0], cmap='jet')
im2 = ax2.imshow(vx[0, ..., 0], cmap='jet')
fig.colorbar(im2)
ax3 = fig.add_subplot(133)
im3 = ax3.imshow(vy[0, ..., 0], cmap='jet')
fig.colorbar(im3)
fig.savefig('vel_TF_test.png')