make fixing idempotent

henryaddison · Dec 14, 2023 · b0cb375 · b0cb375
1 parent a08c581
commit b0cb375
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/bin/check_preds.py b/bin/check_preds.py
@@ -24,7 +24,10 @@
 def fix(pred_ds):
     pred_ds[PRED_VAR] = pred_ds[PRED_VAR].assign_attrs(PRED_PR_ATTRS)
     for var in META_VARS:
-        pred_ds[var] = pred_ds[var].isel(ensemble_member=0, time=0)
+        if "ensemble_member" in pred_ds[var].dims:
+            pred_ds[var] = pred_ds[var].isel(ensemble_member=0)
+        if "time" in pred_ds[var].dims:
+            pred_ds[var] = pred_ds[var].isel(time=0)
 
     return pred_ds