diff --git a/openeo_processes_dask/process_implementations/cubes/apply.py b/openeo_processes_dask/process_implementations/cubes/apply.py index 1c59706..2d902b1 100644 --- 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 index e8d8b31..c4a8fa9 100644 --- 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 index fb396b4..5582705 100644 --- 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):