Count (#274)

* add count process * add count process * add count process * update tests * add more count tests
Open-EO · Sep 10, 2024 · 1c52aa9 · 1c52aa9
1 parent ebd456c
commit 1c52aa9
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 1 deletion.
diff --git a/openeo_processes_dask/process_implementations/arrays.py b/openeo_processes_dask/process_implementations/arrays.py
@@ -355,5 +355,8 @@ def count(
     if condition is True:
         return np.nansum(np.ones_like(data), axis=axis, keepdims=keepdims)
     if callable(condition):
-        count = condition(data)
+        if not context:
+            context = {}
+        context.pop("x", None)
+        count = condition(x=data, **context)
         return np.nansum(count, axis=axis, keepdims=keepdims)
diff --git a/tests/test_arrays.py b/tests/test_arrays.py
@@ -500,3 +500,42 @@ def test_count(temporal_interval, bounding_box, random_raster_data, process_regi
     )
     assert output_cube.dims == ("x", "y", "t")
     xr.testing.assert_equal(output_cube, xr.zeros_like(output_cube) + 4)
+
+    _process = partial(
+        process_registry["count"].implementation,
+        data=ParameterReference(from_parameter="data"),
+        condition=process_registry["gt"].implementation,
+    )
+    output_cube = reduce_dimension(
+        data=input_cube,
+        reducer=_process,
+        dimension="bands",
+        context={"y": -100},
+    )
+    general_output_checks(
+        input_cube=input_cube,
+        output_cube=output_cube,
+        verify_attrs=False,
+        verify_crs=True,
+    )
+    assert output_cube.dims == ("x", "y", "t")
+    xr.testing.assert_equal(output_cube, xr.zeros_like(output_cube) + 4)
+
+    _process = partial(
+        process_registry["count"].implementation,
+        data=ParameterReference(from_parameter="data"),
+        condition=process_registry["is_infinite"].implementation,
+    )
+    output_cube = reduce_dimension(
+        data=input_cube,
+        reducer=_process,
+        dimension="bands",
+    )
+    general_output_checks(
+        input_cube=input_cube,
+        output_cube=output_cube,
+        verify_attrs=False,
+        verify_crs=True,
+    )
+    assert output_cube.dims == ("x", "y", "t")
+    xr.testing.assert_equal(output_cube, xr.zeros_like(output_cube))