add test apply_dimension, quantile (#292)

* add test apply_dimension, quantile * run pre-commit
Open-EO · Nov 13, 2024 · afa06cc · afa06cc
1 parent 461372e
commit afa06cc
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 35 deletions.
diff --git a/openeo_processes_dask/process_implementations/cubes/apply.py b/openeo_processes_dask/process_implementations/cubes/apply.py
@@ -45,7 +45,10 @@ def apply_dimension(
             f"Provided dimension ({dimension}) not found in data.dims: {data.dims}"
         )
 
+    keepdims = False
     is_new_dim_added = target_dimension is not None
+    if is_new_dim_added:
+        keepdims = True
 
     if target_dimension is None:
         target_dimension = dimension
@@ -67,7 +70,7 @@ def apply_dimension(
             "positional_parameters": positional_parameters,
             "named_parameters": named_parameters,
             "axis": reordered_data.get_axis_num(dimension),
-            "keepdims": True,
+            "keepdims": keepdims,
             "source_transposed_axis": data.get_axis_num(dimension),
             "context": context,
         },

diff --git a/openeo_processes_dask/process_implementations/math.py b/openeo_processes_dask/process_implementations/math.py
@@ -324,6 +324,8 @@ def quantiles(
         result = np.quantile(
             data, q=probabilities, method="linear", axis=axis, keepdims=keepdims
         )
+    if axis:
+        result = np.moveaxis(result, 0, axis)
 
     return result
 

diff --git a/tests/test_apply.py b/tests/test_apply.py
@@ -78,40 +78,6 @@ def test_apply_dimension_case_1(
     )
 
 
-@pytest.mark.parametrize("size", [(6, 5, 4, 4)])
-@pytest.mark.parametrize("dtype", [np.float32])
-def test_apply_dimension_reduce_dimension(
-    temporal_interval, bounding_box, random_raster_data, 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["mean"].implementation,
-        data=ParameterReference(from_parameter="data"),
-    )
-
-    # Target dimension is null and therefore defaults to the source dimension
-    output_cube_reduced = apply_dimension(
-        data=input_cube, process=_process, dimension="x"
-    )
-
-    expected_output = (input_cube.mean(dim="x")).expand_dims("x")
-
-    general_output_checks(
-        input_cube=input_cube,
-        output_cube=output_cube_reduced,
-        verify_attrs=True,
-        verify_crs=False,
-        expected_results=expected_output,
-    )
-
-
 @pytest.mark.parametrize("size", [(6, 5, 4, 4)])
 @pytest.mark.parametrize("dtype", [np.float32])
 def test_apply_dimension_target_dimension(
@@ -238,6 +204,34 @@ def test_apply_dimension_ordering_processes(
     )
 
 
+@pytest.mark.parametrize("size", [(6, 5, 30, 4)])
+@pytest.mark.parametrize("dtype", [np.float32])
+def test_apply_dimension_quantile_processes(
+    temporal_interval, bounding_box, random_raster_data, 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",
+    )
+    probability = 4
+
+    _process_quantile = partial(
+        process_registry["quantiles"].implementation,
+        data=ParameterReference(from_parameter="data"),
+        probabilities=probability,
+    )
+
+    output_cube_quantile = apply_dimension(
+        data=input_cube,
+        process=_process_quantile,
+        dimension="t",
+    )
+    assert output_cube_quantile.shape == (6, 5, probability - 1, 4)
+
+
 @pytest.mark.parametrize("size", [(6, 5, 4, 4)])
 @pytest.mark.parametrize("dtype", [np.float32])
 def test_apply_kernel(temporal_interval, bounding_box, random_raster_data):