From 5f6d095c7b45c7cc769e8f276a0e117cb582b5f3 Mon Sep 17 00:00:00 2001
From: katharinastarzer21 <katharina.starzer@eodc.eu>
Date: Thu, 28 Nov 2024 16:08:25 +0100
Subject: [PATCH] Draft: apply_polygon and test_apply_polygon (#298)

* adding apply_polygon and tests

* update apply_polygon

* delete apply_polygon_large test

* updated submodule

* update pyproject.toml
---
 .../process_implementations/cubes/apply.py    | 47 ++++++++++++++++++-
 openeo_processes_dask/specs/openeo-processes  |  2 +-
 pyproject.toml                                |  2 +-
 tests/test_apply.py                           | 45 ++++++++++++++++++
 4 files changed, 93 insertions(+), 3 deletions(-)

diff --git a/openeo_processes_dask/process_implementations/cubes/apply.py b/openeo_processes_dask/process_implementations/cubes/apply.py
index ecad9cf..d39b36d 100644
--- a/openeo_processes_dask/process_implementations/cubes/apply.py
+++ b/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
diff --git a/openeo_processes_dask/specs/openeo-processes b/openeo_processes_dask/specs/openeo-processes
index 45f52ab..c7bf70b 160000
--- a/openeo_processes_dask/specs/openeo-processes
+++ b/openeo_processes_dask/specs/openeo-processes
@@ -1 +1 @@
-Subproject commit 45f52ab1cc7ce00e6afbe78a6a9827c51f086a42
+Subproject commit c7bf70b0b4d8bb4e5fcdaed2bba4d8a860fa337e
diff --git a/pyproject.toml b/pyproject.toml
index 9258967..4f87115 100644
--- a/pyproject.toml
+++ b/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 <lukas.weidenholzer@eodc.eu>", "Sean Hoyal <sean.hoyal@eodc.eu>", "Valentina Hutter <valentina.hutter@eodc.eu>"]
 maintainers = ["EODC Staff <support@eodc.eu>"]
diff --git a/tests/test_apply.py b/tests/test_apply.py
index 7a6812d..cb5d10e 100644
--- a/tests/test_apply.py
+++ b/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()