map doesn't behave with to_gcm_domain

henryaddison · Dec 1, 2023 · c57cfc9 · c57cfc9
1 parent dd9d3f7
commit c57cfc9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/ml_downscaling_emulator/bin/postprocess.py b/src/ml_downscaling_emulator/bin/postprocess.py
@@ -179,7 +179,9 @@ def gcmify(
     ensemble_member: str = typer.Option(...),
 ):
     def process_samples(ds):
-        return ds.groupby("ensemble_member").map(to_gcm_domain)
+        ds = to_gcm_domain(ds.sel(ensemble_member=ensemble_member).map(to_gcm_domain))
+        ds["pred_pr"] = ds["pred_pr"].expand_dims({"ensemble_member": ensemble_member})
+        return ds
 
     process_each_sample(
         workdir,

diff --git a/src/ml_downscaling_emulator/postprocess.py b/src/ml_downscaling_emulator/postprocess.py
@@ -83,4 +83,5 @@ def to_gcm_domain(ds: xr.Dataset):
     ds = Remapcon(target_grid_filepath).run(ds)
     ds = ShiftLonBreak().run(ds)
     ds = SelectGCMDomain(subdomain="birmingham", size=9).run(ds)
+    ds = ds.drop_vars(["rotated_latitude_longitude"], errors="ignore")
     return ds