roll back resample cube spatial (#314)

* update align_arrays for resample cube spatial * run pre-commit * roll back resample cube spatial update * roll back tests
Open-EO · Feb 19, 2025 · a5dadcb · a5dadcb
1 parent acf74eb
commit a5dadcb
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 121 deletions.
diff --git a/openeo_processes_dask/process_implementations/cubes/resample.py b/openeo_processes_dask/process_implementations/cubes/resample.py
@@ -114,91 +114,67 @@ def resample_spatial(
 def resample_cube_spatial(
     data: RasterCube, target: RasterCube, method="near", options=None
 ) -> RasterCube:
-    if target.openeo.y_dim is None or target.openeo.x_dim is None:
+    methods_list = [
+        "near",
+        "bilinear",
+        "cubic",
+        "cubicspline",
+        "lanczos",
+        "average",
+        "mode",
+        "max",
+        "min",
+        "med",
+        "q1",
+        "q3",
+    ]
+
+    if (
+        data.openeo.y_dim is None
+        or data.openeo.x_dim is None
+        or target.openeo.y_dim is None
+        or target.openeo.x_dim is None
+    ):
         raise DimensionMissing(
-            f"Spatial dimension missing for target dataset: {target} "
+            f"Spatial dimension missing from data or target. Available dimensions for data: {data.dims} for target: {target.dims}"
         )
 
-    target_resolution, target_crs = None, None
-    if hasattr(target, "rio"):
-        if hasattr(target.rio, "resolution"):
-            if type(target.rio.resolution()) in [tuple, list]:
-                target_resolution = target.rio.resolution()[0]
-            else:
-                target_resolution = target.rio.resolution()
-        if hasattr(target.rio, "crs"):
-            target_crs = target.rio.crs
-    if not target_crs:
-        raise OpenEOException(f"Projection not found in target dataset: {target} ")
-    if not target_resolution:
-        raise OpenEOException(f"Resolution not found in target dataset: {target} ")
-
-    resampled_data = resample_spatial(
-        data=data, projection=target_crs, resolution=target_resolution, method=method
-    )
+    # ODC reproject requires y to be before x
+    required_dim_order = (..., data.openeo.y_dim, data.openeo.x_dim)
 
-    def align_arrays(a, b, res):
-        desc = 1
-        if len(a) == 1 and len(b) == 1:
-            return a
-        elif len(a) == 1 and len(b) > 1:
-            if b[0] > b[1]:
-                desc = -1
-        elif len(a) > len(b):
-            b0 = np.absolute(b[0] - a).argmin()
-            blen = len(b)
-            close = np.isclose(b, a[b0 : b0 + blen], atol=res * 0.99)
-            if close.all():
-                return a[b0 : b0 + blen]
-            else:
-                raise Exception("Coordinates could not be aligned! ")
-        elif len(a) == len(b):
-            if b[0] > b[1]:
-                if a[0] < a[1]:
-                    a = np.flip(a)
-            else:
-                if a[0] > a[1]:
-                    a = np.flip(a)
-            close = np.isclose(b, a, atol=res * 0.99)
-            if close.all():
-                return a
-            else:
-                raise Exception("Coordinates could not be aligned! ")
-        else:
-            if b[0] > b[1]:
-                desc = -1
-                if a[0] < a[1]:
-                    a = np.flip(a)
-            else:
-                if a[0] > a[1]:
-                    a = np.flip(a)
-        a0 = np.absolute(b - a[0]).argmin()
-        alen = len(a)
-        close = np.isclose(a, b[a0 : a0 + alen], atol=res * 0.99)
-        if not close.all():
-            raise Exception("Coordinates could not be aligned! ")
-
-        new_b = np.arange(b[0], a[0], res * desc)
-        new_b = np.append(new_b, a)
-        new_b = np.append(new_b, np.flip(np.arange(b[-1], a[-1], -res * desc)))
-        if len(b) != len(new_b):
-            raise Exception("Coordinates could not be aligned! ")
-        return new_b
-
-    x_data = resampled_data[resampled_data.openeo.x_dim[0]].values
-    x_target = target[target.openeo.x_dim[0]].values
-    if not (len(x_data) == len(x_target) and (x_data == x_target).all()):
-        resampled_data[resampled_data.openeo.x_dim[0]] = align_arrays(
-            x_target, x_data, target_resolution
+    data_reordered = data.transpose(*required_dim_order, missing_dims="ignore")
+    target_reordered = target.transpose(*required_dim_order, missing_dims="ignore")
+
+    if method == "near":
+        method = "nearest"
+
+    elif method not in methods_list:
+        raise Exception(
+            f'Selected resampling method "{method}" is not available! Please select one of '
+            f"[{', '.join(methods_list)}]"
         )
 
-    y_data = resampled_data[resampled_data.openeo.y_dim[0]].values
-    y_target = target[target.openeo.y_dim[0]].values
-    if not (len(y_data) == len(y_target) and (y_data == y_target).all()):
-        resampled_data[resampled_data.openeo.y_dim[0]] = align_arrays(
-            y_target, y_data, target_resolution
+    resampled_data = data_reordered.odc.reproject(
+        target_reordered.odc.geobox, resampling=method
+    )
+
+    resampled_data.rio.write_crs(target_reordered.rio.crs, inplace=True)
+
+    try:
+        # odc.reproject renames the coordinates according to the geobox, this undoes that.
+        resampled_data = resampled_data.rename(
+            {"longitude": data.openeo.x_dim, "latitude": data.openeo.y_dim}
         )
+    except ValueError:
+        pass
+
+    # Order axes back to how they were before
+    resampled_data = resampled_data.transpose(*data.dims)
 
+    # Ensure that attrs except crs are copied over
+    for k, v in data.attrs.items():
+        if k.lower() != "crs":
+            resampled_data.attrs[k] = v
     return resampled_data
 
 

diff --git a/tests/test_resample.py b/tests/test_resample.py
@@ -141,50 +141,6 @@ def test_resample_cube_spatial(
     assert output_cube.odc.spatial_dims == ("y", "x")
 
 
-@pytest.mark.parametrize(
-    "output_crs",
-    [
-        3587,
-        "32633",
-        "+proj=aeqd +lat_0=53 +lon_0=24 +x_0=5837287.81977 +y_0=2121415.69617 +datum=WGS84 +units=m +no_defs",
-    ],
-)
-@pytest.mark.parametrize("output_res", [5, 30, 60])
-@pytest.mark.parametrize("size", [(30, 30, 20, 4)])
-@pytest.mark.parametrize("dtype", [np.float32])
-def test_resample_cube_spatial_small(
-    output_crs, output_res, temporal_interval, bounding_box, random_raster_data
-):
-    """Test to ensure resolution gets changed correctly."""
-    input_cube = create_fake_rastercube(
-        data=random_raster_data,
-        spatial_extent=bounding_box,
-        temporal_extent=temporal_interval,
-        bands=["B02", "B03", "B04", "B08"],
-        backend="dask",
-    )
-
-    resampled_cube = resample_spatial(
-        data=input_cube, projection=output_crs, resolution=output_res
-    )
-
-    output_cube = resample_cube_spatial(
-        data=input_cube, target=resampled_cube[10:60, 20:150, :, :], method="average"
-    )
-
-    general_output_checks(
-        input_cube=input_cube,
-        output_cube=output_cube,
-        expected_dims=input_cube.dims,
-        verify_attrs=False,
-        verify_crs=False,
-    )
-
-    assert list(output_cube.shape) == list(resampled_cube.shape)
-    assert (output_cube["x"].values == resampled_cube["x"].values).all()
-    assert (output_cube["y"].values == resampled_cube["y"].values).all()
-
-
 @pytest.mark.parametrize("size", [(6, 5, 30, 4)])
 @pytest.mark.parametrize("dtype", [np.float64])
 @pytest.mark.parametrize(