Enable broadcasting in vertical interpolation (#133)

MeteoSwiss-APN · Apr 17, 2024 · 6be8648 · 6be8648
1 parent cac54a8
commit 6be8648
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 12 deletions.
diff --git a/src/idpi/operators/vertical_interpolation.py b/src/idpi/operators/vertical_interpolation.py
@@ -101,18 +101,13 @@ def interpolate_k2p(
             "interpolate_k2p: pressure field must be defined on model levels"
         )
     # Check that dimensions are the same for field and p_field
-    if (
-        field.origin_z != p_field.origin_z
-        or field.dims != p_field.dims
-        or field.size != p_field.size
-    ):
+    if field.origin_z != p_field.origin_z:
         raise RuntimeError(
-            "interpolate_k2p: field and p_field must have equal "
-            "origin, dimensions and size"
+            "interpolate_k2p: field and p_field must have equal vertical staggering"
         )
 
     # Prepare output field field_on_tc on target coordinates
-    field_on_tc = init_field_with_vcoord(field, tc, np.nan)
+    field_on_tc = init_field_with_vcoord(field.broadcast_like(p_field), tc, np.nan)
 
     # Interpolate
     # ... prepare interpolation
@@ -268,7 +263,7 @@ def interpolate_k2theta(
         )
 
     # Prepare output field field_on_tc on target coordinates
-    field_on_tc = init_field_with_vcoord(field, tc, np.nan)
+    field_on_tc = init_field_with_vcoord(field.broadcast_like(th_field), tc, np.nan)
 
     # Interpolate
     # ... prepare interpolation
@@ -375,7 +370,7 @@ def interpolate_k2any(
     )
 
     # Prepare output field field_on_tc on target coordinates
-    field_on_tc = init_field_with_vcoord(field, tc, np.nan)
+    field_on_tc = init_field_with_vcoord(field.broadcast_like(tc_field), tc, np.nan)
 
     # Interpolate
     # ... prepare interpolation

diff --git a/tests/test_idpi/test_intpl_k2any.py b/tests/test_idpi/test_intpl_k2any.py
@@ -19,7 +19,7 @@ def test_intpl_k2theta(data_dir, fieldextra):
     reader = GribReader.from_files([cdatafile, datafile])
     ds = reader.load_fieldnames(["DBZ", "HHL"])
 
-    hfl = destagger(ds["HHL"], "z")
+    hfl = destagger(ds["HHL"].squeeze(drop=True), "z")
 
     # call interpolation operator
     echo_top = interpolate_k2any(hfl, "high_fold", ds["DBZ"], tc_values, hfl)

diff --git a/tests/test_idpi/test_intpl_k2theta.py b/tests/test_idpi/test_intpl_k2theta.py
@@ -27,7 +27,7 @@ def test_intpl_k2theta(mode, data_dir, fieldextra):
     ds = reader.load_fieldnames(["P", "T", "HHL"])
 
     theta = compute_theta(ds["P"], ds["T"])
-    hfl = destagger(ds["HHL"], "z")
+    hfl = destagger(ds["HHL"].squeeze(drop=True), "z")
 
     # call interpolation operator
     t = interpolate_k2theta(ds["T"], mode, theta, tc_values, tc_units, hfl)