store raw predictions as well as transformed ones when sampling

henryaddison · Mar 8, 2024 · 768d024 · 768d024
1 parent d74e255
commit 768d024
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/bin/predict.py b/bin/predict.py
@@ -164,8 +164,12 @@ def np_samples_to_xr(np_samples, target_transform, coords, cf_data_vars):
     # add ensemble member axis to np samples
     np_samples = np_samples[np.newaxis, :]
     pred_pr_var = (pred_pr_dims, np_samples, pred_pr_attrs)
-
-    data_vars = {**cf_data_vars, "target_pr": pred_pr_var}
+    raw_pred_var = (
+        pred_pr_dims,
+        np_samples,
+        {"grid_mapping": "rotated_latitude_longitude"},
+    )
+    data_vars = {**cf_data_vars, "target_pr": pred_pr_var, "raw_pred": raw_pred_var}
 
     samples_ds = target_transform.invert(
         xr.Dataset(data_vars=data_vars, coords=coords, attrs={})