Draft: apply_polygon and test_apply_polygon

openeo_processes_dask/process_implementations/cubes/apply.py

@@ -3,8 +3,16 @@
 import numpy as np
 import scipy.ndimage
 import xarray as xr
+from shapely.geometry import MultiPolygon, Polygon, shape
+from shapely.ops import unary_union
 
-from openeo_processes_dask.process_implementations.data_model import RasterCube
+from openeo_processes_dask.process_implementations.cubes.mask_polygon import (
+    mask_polygon,
+)
+from openeo_processes_dask.process_implementations.data_model import (
+    RasterCube,
+    VectorCube,
+)
 from openeo_processes_dask.process_implementations.exceptions import (
     DimensionNotAvailable,
     KernelDimensionsUneven,
@@ -165,3 +173,40 @@ def convolve(data, kernel, mode="constant", cval=0, fill_value=0):
         cval = 0
 
     return convolve(data, kernel, mode, cval, replace_invalid) * factor
+
+
+def apply_polygon(
+    data: RasterCube,
+    polygons: Union[VectorCube, dict],
+    process: Callable,
+    mask_value: Optional[Union[int, float, str, None]] = None,
+    context: Optional[dict] = None,
+) -> RasterCube:
+    if isinstance(polygons, dict) and polygons.get("type") == "FeatureCollection":
+        polygon_geometries = [
+            shape(feature["geometry"]) for feature in polygons["features"]
+        ]
+    elif isinstance(polygons, dict) and polygons.get("type") in [
+        "Polygon",
+        "MultiPolygon",
+    ]:
+        polygon_geometries = [shape(polygons)]
+    else:
+        raise ValueError(
+            "Unsupported polygons format. Expected GeoJSON-like FeatureCollection or Polygon."
+        )
+
+    unified_polygon = unary_union(polygon_geometries)
+
+    if isinstance(unified_polygon, MultiPolygon) and len(unified_polygon.geoms) < len(
+        polygon_geometries
+    ):
+        raise Exception("GeometriesOverlap")
+
+    masked_data = mask_polygon(data, polygons, replacement=np.nan)
+
+    processed_data = apply(masked_data, process, context=context)
+
+    result = mask_polygon(processed_data, polygons, replacement=mask_value)
+
+    return result

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "openeo-processes-dask"
-version = "2024.11.3"
+version = "2024.11.4"
 description = "Python implementations of many OpenEO processes, dask-friendly by default."
 authors = ["Lukas Weidenholzer <[email protected]>", "Sean Hoyal <[email protected]>", "Valentina Hutter <[email protected]>"]
 maintainers = ["EODC Staff <[email protected]>"]

tests/test_apply.py

@@ -10,6 +10,10 @@
     apply,
     apply_dimension,
     apply_kernel,
+    apply_polygon,
+)
+from openeo_processes_dask.process_implementations.cubes.mask_polygon import (
+    mask_polygon,
 )
 from tests.general_checks import assert_numpy_equals_dask_numpy, general_output_checks
 from tests.mockdata import create_fake_rastercube
@@ -482,3 +486,44 @@ def test_apply_dimension_cummin_process(
     ).compute()
 
     assert np.isnan(output_cube_cummin_with_nan[0, 0, 16, 0].values)
+
+
+@pytest.mark.parametrize("size", [(6, 5, 4, 4)])
+@pytest.mark.parametrize("dtype", [np.float32])
+def test_apply_polygon(
+    temporal_interval,
+    bounding_box,
+    random_raster_data,
+    polygon_geometry_small,
+    process_registry,
+):
+    input_cube = create_fake_rastercube(
+        data=random_raster_data,
+        spatial_extent=bounding_box,
+        temporal_extent=temporal_interval,
+        bands=["B02", "B03", "B04", "B08"],
+        backend="dask",
+    )
+
+    _process = partial(
+        process_registry["add"].implementation,
+        y=1,
+        x=ParameterReference(from_parameter="x"),
+    )
+
+    output_cube = apply_polygon(
+        data=input_cube,
+        polygons=polygon_geometry_small,
+        process=_process,
+        mask_value=np.nan,
+    )
+
+    assert len(output_cube.dims) == len(input_cube.dims)
+
+    expected_result_mask = mask_polygon(data=input_cube, mask=polygon_geometry_small)
+
+    expected_results = expected_result_mask + 1
+
+    assert len(output_cube.dims) == len(expected_results.dims)
+
+    assert output_cube.all() == expected_results.all()