added special Sup3rCC TRH loss

sup3r/utilities/loss_metrics.py

@@ -1,6 +1,6 @@
 """Loss metrics for Sup3r"""
 
-from tensorflow.keras.losses import MeanSquaredError
+from tensorflow.keras.losses import MeanSquaredError, MeanAbsoluteError
 import tensorflow as tf
 
 
@@ -171,3 +171,39 @@ def __call__(self, x1, x2):
  x1_coarse = tf.reduce_mean(x1, axis=(1, 2))
  x2_coarse = tf.reduce_mean(x2, axis=(1, 2))
  return self.MSE_LOSS(x1_coarse, x2_coarse)
+
+
+class TRHLoss(tf.keras.losses.Loss):
+ """Loss class for Temperature and Relative Humidity Sup3rCC GAN that
+ encourages accuracy of the min/max values in the timeseries"""
+
+ MAE_LOSS = MeanAbsoluteError()
+
+ def __call__(self, x1, x2):
+ """Custom Sup3rCC content loss function for Temp + RH
+
+ Parameters
+ ----------
+ x1 : tf.tensor
+ synthetic generator output
+ (n_observations, spatial_1, spatial_2, temporal, features)
+ x2 : tf.tensor
+ high resolution data
+ (n_observations, spatial_1, spatial_2, temporal, features)
+
+ Returns
+ -------
+ tf.tensor
+ 0D tensor with loss value
+ """
+ x1_min = tf.reduce_min(x1, axis=3)
+ x2_min = tf.reduce_min(x2, axis=3)
+
+ x1_max = tf.reduce_max(x1, axis=3)
+ x2_max = tf.reduce_max(x2, axis=3)
+
+ mae = self.MAE_LOSS(x1, x2)
+ mae_min = self.MAE_LOSS(x1_min, x2_min)
+ mae_max = self.MAE_LOSS(x1_max, x2_max)
+
+ return mae + mae_min + mae_max

tests/test_custom_loss.py

@@ -3,7 +3,7 @@
 import numpy as np
 import tensorflow as tf
 
-from sup3r.utilities.loss_metrics import MmdMseLoss, CoarseMseLoss
+from sup3r.utilities.loss_metrics import MmdMseLoss, CoarseMseLoss, TRHLoss
 
 
 def test_mmd_loss():
@@ -52,3 +52,26 @@ def test_coarse_mse_loss():
  assert mse.numpy().size == 1
  assert coarse_mse.numpy().size == 1
  assert mse.numpy() > 10 * coarse_mse.numpy()
+
+
+def test_trh_loss():
+ """Test custom Sup3rCC Temp + RH loss function that looks at min/max values
+ in the timeseries."""
+ trh_loss = TRHLoss()
+
+ x = np.zeros((1, 1, 1, 72, 1))
+ y = np.zeros((1, 1, 1, 72, 1))
+
+ # loss should be dominated by special min/max values
+ x[..., 24, 0] = 20
+ y[..., 25, 0] = 25
+ loss = trh_loss(x, y)
+ assert loss.numpy() > 5
+ assert loss.numpy() < 6
+
+ # loss should be dominated by special min/max values
+ x[..., 24, 0] = -10
+ y[..., 25, 0] = -15
+ loss = trh_loss(x, y)
+ assert loss.numpy() > 5
+ assert loss.numpy() < 6