Merge pull request #88 from jsadler2/nnse-samplewise

Nnse samplewise

Author: jsadler2

river_dl/loss_functions.py

@@ -19,18 +19,38 @@ def rmse(y_true, y_pred):
     return rmse_loss
 
 
-def nse(y_true, y_pred):
+def sample_avg_nse(y_true, y_pred):
+    """
+    calculate the sample averaged nse, i.e., it will calculate the nse across
+    each of the samples (the 1st dimension of the arrays) and then average those
+    """
     y_true = tf.cast(y_true, tf.float32)
     y_pred = tf.cast(y_pred, tf.float32)
     zero_or_error = tf.where(
         tf.math.is_nan(y_true), tf.zeros_like(y_true), y_pred - y_true
     )
 
-    numerator = tf.reduce_sum(tf.square(zero_or_error))
+    # add a small value to the deviation to prevent instability
+    deviation = dev_masked(y_true) + 0.1
+
+    numerator_samplewise = tf.reduce_sum(tf.square(zero_or_error), axis=1)
+    denomin_samplewise = tf.reduce_sum(tf.square(deviation), axis=1)
+    nse_samplewise = 1 - numerator_samplewise/denomin_samplewise
+    nse_samplewise_avg = tf.reduce_sum(nse_samplewise)/tf.cast(tf.shape(y_true)[0], tf.float32)
+    return nse_samplewise_avg
+
+
+def nse(y_true, y_pred):
+    y_true = tf.cast(y_true, tf.float32)
+    y_pred = tf.cast(y_pred, tf.float32)
+    zero_or_error = tf.where(
+        tf.math.is_nan(y_true), tf.zeros_like(y_true), y_pred - y_true
+    )
 
     deviation = dev_masked(y_true)
+    numerator = tf.reduce_sum(tf.square(zero_or_error))
     denominator = tf.reduce_sum(tf.square(deviation))
-    return 1 - numerator / denominator  
+    return 1 - numerator / denominator
 
 
 def nnse(y_true, y_pred):
@@ -41,11 +61,20 @@ def nnse_loss(y_true, y_pred):
     return 1 - nnse(y_true, y_pred)
 
 
+def samplewise_nnse_loss(y_true, y_pred):
+    nnse_val = 1 / (2 - sample_avg_nse(y_true, y_pred))
+    return 1 - nnse_val
+
+
 @tf.function
 def nnse_masked_one_var(data, y_pred, var_idx):
     y_true, y_pred, weights = y_data_components(data, y_pred, var_idx)
     return nnse_loss(y_true, y_pred)
 
+@tf.function
+def nnse_one_var_samplewise(data, y_pred, var_idx):
+    y_true, y_pred, weights = y_data_components(data, y_pred, var_idx)
+    return samplewise_nnse_loss(y_true, y_pred)
 
 @tf.function
 def y_data_components(data, y_pred, var_idx):

river_dl/postproc_utils.py

@@ -35,7 +35,7 @@ def prepped_array_to_df(data_array, dates, ids, col_names):
     return df
 
 
-def take_first_half(df):
+def take_half(df, first_half=True):
     """
     filter out the second half of the dates in the predictions. this is to
     retain a "test" set of the i/o data for evaluation
@@ -47,9 +47,12 @@ def take_first_half(df):
     df.sort_index(inplace=True)
     unique_dates = df.index.unique()
     halfway_date = unique_dates[int(len(unique_dates) / 2)]
-    df_first_half = df.loc[:halfway_date]
-    df_first_half.reset_index(inplace=True)
-    return df_first_half
+    if first_half:
+        df_half = df.loc[:halfway_date]
+    else:
+        df_half = df.loc[halfway_date:]
+    df_half.reset_index(inplace=True)
+    return df_half
 
 
 def unscale_output(y_scl, y_std, y_mean, data_cols, logged_q=False):
@@ -197,11 +200,16 @@ def predict(model, io_data, partition, outfile, logged_q=False, half_tst=False):
     """
     io_data = get_data_if_file(io_data)
 
-    # evaluate training
-    if partition == "trn" or partition == "tst":
+    if partition in ["trn", "tst", "ver"]:
         pass
     else:
-        raise ValueError('partition arg needs to be "trn" or "tst"')
+        raise ValueError('partition arg needs to be "trn" or "tst" or "ver"')
+
+    if partition == "ver":
+        partition = "tst"
+        tst_partition = "ver"
+    elif partition == "tst":
+        tst_partition = "tst"
 
     num_segs = len(np.unique(io_data["ids_trn"]))
     y_pred = model.predict(io_data[f"x_{partition}"], batch_size=num_segs)
@@ -220,8 +228,12 @@ def predict(model, io_data, partition, outfile, logged_q=False, half_tst=False):
         logged_q,
     )
 
-    if half_tst and partition == "tst":
-        y_pred_pp = take_first_half(y_pred_pp)
+    if partition == "tst":
+        if half_tst and tst_partition == "tst":
+            y_pred_pp = take_half(y_pred_pp, first_half=True)
+
+        if half_tst and tst_partition == "ver":
+            y_pred_pp = take_half(y_pred_pp, first_half=False)
 
     y_pred_pp.to_feather(outfile)
     return y_pred_pp
@@ -372,13 +384,14 @@ def overall_metrics(
 
 
 def combined_metrics(
-    pred_trn, pred_tst, obs_temp, obs_flow, grp=None, outfile=None
+    pred_trn, pred_tst,  obs_temp, obs_flow, pred_ver=None, grp=None, outfile=None
 ):
     """
     calculate the metrics for flow and temp and training and test sets for a
     given grouping
     :param pred_trn: [str] path to training prediction feather file
     :param pred_tst: [str] path to testing prediction feather file
+    :param pred_tst: [str] path to verification prediction feather file
     :param obs_temp: [str] path to observations temperature zarr file
     :param obs_flow: [str] path to observations flow zarr file
     :param group: [str or list] which group the metrics should be computed for.
@@ -393,6 +406,10 @@ def combined_metrics(
     tst_temp = overall_metrics(pred_tst, obs_temp, "temp", "tst", grp)
     tst_flow = overall_metrics(pred_tst, obs_flow, "flow", "tst", grp)
     df_all = [trn_temp, tst_temp, trn_flow, tst_flow]
+    if pred_ver:
+        ver_temp = overall_metrics(pred_ver, obs_temp, "temp", "ver", grp)
+        ver_flow = overall_metrics(pred_ver, obs_flow, "flow", "ver", grp)
+        df_all.extend([ver_temp, ver_flow])
     df_all = pd.concat(df_all, axis=0)
     if outfile:
         df_all.to_csv(outfile, index=False)

river_dl/rnns.py

@@ -2,7 +2,7 @@
 from __future__ import print_function, division
 import tensorflow as tf
 from tensorflow.keras import layers
-from river_dl.loss_functions import nnse_masked_one_var
+from river_dl.loss_functions import nnse_masked_one_var, nnse_one_var_samplewise
 
 
 class LSTMModel(tf.keras.Model):
@@ -47,8 +47,8 @@ def train_step(self, data):
         with tf.GradientTape(persistent=True) as tape:
             y_pred = self(x, training=True)  # forward pass
 
-            loss_main = nnse_masked_one_var(y, y_pred, 0)
-            loss_aux = nnse_masked_one_var(y, y_pred, 1)
+            loss_main = nnse_one_var_samplewise(y, y_pred, 0)
+            loss_aux = nnse_one_var_samplewise(y, y_pred, 1)
 
         trainable_vars = self.trainable_variables