Added predict_diff

Author: tnipen

climetlab_maelstrom_yr/yr.py

@@ -49,6 +49,7 @@ def __init__(
         limit_predictors=None,
         probabilistic_target=False,
         normalize=False,
+        predict_diff=False,
         verbose=False,
     ):
         """
@@ -60,6 +61,8 @@ def __init__(
             pattern (str): Pattern for filenames
             probabilistic_target (bool): If true, include target std as the second target parameter
             normalize (bool): If true, normalize the data
+            predict_diff (bool): If true, change the target to be the difference between the target
+                and raw forecast
             verbose (bool): Show debug statements if True
         """
         if size not in ["5GB", "5TB"]:
@@ -70,8 +73,9 @@ def __init__(
 
         self.size = size
         self.parameter = parameter
-        self.do_normalize = normalize
         self.probabilistic_target = probabilistic_target
+        self.do_normalize = normalize
+        self.do_predict_diff = predict_diff
         self.verbose = verbose
 
         is_url = location.find("://") >= 0
@@ -244,6 +248,10 @@ def preprocess(self, ds):
         else:
             data_vars["targets"] = (("leadtime", "y", "x", "target"), np.expand_dims(ds.variables["target_mean"], -1))
 
+        if self.do_predict_diff:
+            I = np.where(coords["predictor"] == "air_temperature_2m")[0][0]
+            data_vars["targets"][1][..., 0] -= data_vars["predictors"][1][..., I]
+
         if self.do_normalize:
             for i, name in enumerate(coords["predictor"]):
                 Yr.normalize(data_vars["predictors"][1][..., i], name)

tests/test_yr.py

@@ -19,8 +19,11 @@ def test_read():
     )
     ds = cmlds.to_xarray()
     I = np.where(ds["predictor"] == "air_temperature_2m")[0][0]
-    value0 = ds["predictors"][0, 0, 0, 0, I].values
-    np.testing.assert_almost_equal(value0, 0.33352637)
+    value = ds["predictors"][0, 0, 0, 0, I].values
+    np.testing.assert_almost_equal(value, 0.33352637)
+
+    value = ds["targets"][0, 0, 0, 0, 0].values
+    np.testing.assert_almost_equal(value, 0.3591344)
 
 def test_normalize_dataset():
     cmlds = cml.load_dataset(
@@ -29,12 +32,17 @@ def test_normalize_dataset():
         dates=["2020-03-01", "2020-03-02"],
         # location=f"{dir}/",
         probabilistic_target=False,
+        predict_diff=True,
         normalize=True,
     )
     ds = cmlds.to_xarray()
     I = np.where(ds["predictor"] == "air_temperature_2m")[0][0]
-    value0 = ds["predictors"][0, 0, 0, 0, I].values
-    np.testing.assert_almost_equal(value0, (0.33352637 - 5.388952960570653)/7.4335476655246735)
+    value = ds["predictors"][0, 0, 0, 0, I].values
+    np.testing.assert_almost_equal(value, (0.33352637 - 5.388952960570653)/7.4335476655246735)
+
+    # Check that the target has been normalized by air_temperature_2m
+    value = ds["targets"][0, 0, 0, 0, 0].values
+    np.testing.assert_almost_equal(value, 0.3591344 - 0.33352637)
 
 def test_normalize_functions():
     ar = np.array([0, 1, 2], np.float32)
@@ -43,7 +51,6 @@ def test_normalize_functions():
     np.testing.assert_array_almost_equal(ar, [0, 1, 2])
 
 
-
 if __name__ == "__main__":
     from climetlab.testing import main