From 7ba6f50a2b2113d30ff2bc89b4f2d4aa99a76ef4 Mon Sep 17 00:00:00 2001 From: Nada Akkari Date: Wed, 3 Jun 2026 20:09:13 +0100 Subject: [PATCH 1/4] Add depth_min/max support to depth averaging --- src/medunda/actions/compute_average.py | 61 ++++++++++++++++++++------ 1 file changed, 47 insertions(+), 14 deletions(-) diff --git a/src/medunda/actions/compute_average.py b/src/medunda/actions/compute_average.py index 5a4ba0b..669ca9e 100644 --- a/src/medunda/actions/compute_average.py +++ b/src/medunda/actions/compute_average.py @@ -20,11 +20,25 @@ def configure_parser(subparsers): ACTION_NAME, help="Compute the average" ) average_parser.add_argument( - "--axis", + "--axes", type=str, + # nargs="+", choices=sorted(VALID_AXIS), required=True, - help="Choose the axis on which the average will be computed.", + help="Axes on which the average will be computed.", + ) + + average_parser.add_argument( + "--depth-min", + type=float, + required=False, + help="Minimum depth for the vertical average.", + ) + average_parser.add_argument( + "--depth-max", + type=float, + required=False, + help="Maximum depth for the vertical average.", ) @@ -51,7 +65,12 @@ def get_volume(data: "xr.Dataset") -> "xr.DataArray": ) -def compute_average(data: "xr.Dataset", axis) -> "xr.Dataset": +def compute_average( + data: "xr.Dataset", + axes: str, + depth_min: float = None, + depth_max: float = None, +) -> "xr.Dataset": """Compute the average of all variables along a specified axis. Three axes are supported: @@ -67,7 +86,7 @@ def compute_average(data: "xr.Dataset", axis) -> "xr.Dataset": data (xr.Dataset): Input dataset. Must include ``depth``, ``latitude``, ``longitude``, and ``time`` coordinates as required by the chosen axis. - axis (str): Axis along which to compute the average. One of + axes (str): Axes along which to compute the average. One of ``"depth"``, ``"space"``, or ``"time"``. Returns: @@ -75,22 +94,34 @@ def compute_average(data: "xr.Dataset", axis) -> "xr.Dataset": the averaged values for each variable. Raises: - ValueError: If *axis* is not one of the valid choices. + ValueError: If *axes* is not one of the valid choices. """ - if axis not in VALID_AXIS.keys(): + + if axes not in VALID_AXIS.keys(): raise ValueError( - f"Axis '{axis}' is not valid. Choose from {list(VALID_AXIS.keys())}" + f"Axes '{axes}' is not valid. Choose from {list(VALID_AXIS.keys())}" ) - LOGGER.info(f"Computing average over axis '{axis}'") + LOGGER.info(f"Computing average over axes '{axes}'") + + if axes == "depth": + if depth_min is not None: + depth_min = float(depth_min) + else: + depth_min = float(data.depth.min()) + + if depth_max is not None: + depth_max = float(depth_max) + else: + depth_max = float(data.depth.max()) - if axis == "depth": - depth_min = float(data.depth.min()) - depth_max = float(data.depth.max()) averaged_dataset = average_between_layers(data, depth_min, depth_max) - LOGGER.info("Depth average computed successfully") + LOGGER.info( + f"Depth interval: [{depth_min if depth_min is not None else float(data.depth.min())}, " + f"{depth_max if depth_max is not None else float(data.depth.max())}]" + ) - elif axis == "space": + elif axes == "space": weights = get_volume(data) weights = weights.expand_dims({"time": data.time}) @@ -100,12 +131,14 @@ def compute_average(data: "xr.Dataset", axis) -> "xr.Dataset": da = data[var] weighted_sum = (da * weights).sum(dim=("latitude", "longitude")) + total_weights = weights.sum(dim=("latitude", "longitude")) + averaged_dataset[var] = weighted_sum / total_weights LOGGER.info("Space average computed successfully") - elif axis == "time": + elif axes == "time": averaged_dataset = data.mean(dim="time") LOGGER.info("Time average computed successfully") From a9860aecdbfd22e030ed075fdff579a0e2599f50 Mon Sep 17 00:00:00 2001 From: nakkari Date: Tue, 16 Jun 2026 19:38:38 +0200 Subject: [PATCH 2/4] Fix compute_average: allow aggregation over multiple axes and improve spatial averaging --- src/medunda/actions/compute_average.py | 194 ++++++++++++++----------- 1 file changed, 113 insertions(+), 81 deletions(-) diff --git a/src/medunda/actions/compute_average.py b/src/medunda/actions/compute_average.py index 669ca9e..b617fe2 100644 --- a/src/medunda/actions/compute_average.py +++ b/src/medunda/actions/compute_average.py @@ -1,4 +1,5 @@ import logging +from collections.abc import Sequence import medunda.tools.lazy_imports.bitsea.mask as bitsea from medunda.actions.average_between_layers import average_between_layers @@ -8,12 +9,6 @@ LOGGER = logging.getLogger(__name__) ACTION_NAME = "compute_average" -VALID_AXIS = { - "depth": ["depth"], - "space": ["latitude", "longitude"], - "time": ["time"], -} - def configure_parser(subparsers): average_parser = subparsers.add_parser( @@ -22,8 +17,8 @@ def configure_parser(subparsers): average_parser.add_argument( "--axes", type=str, - # nargs="+", - choices=sorted(VALID_AXIS), + nargs="+", + choices=["depth", "latitude", "longitude", "time"], required=True, help="Axes on which the average will be computed.", ) @@ -42,105 +37,142 @@ def configure_parser(subparsers): ) -def get_volume(data: "xr.Dataset") -> "xr.DataArray": - """Compute the cell volume""" +def get_area(data: "xr.Dataset") -> "xr.DataArray": + """Compute the cell area""" + data_var = list(data.data_vars)[0] + if "time" in data[data_var].dims: reference = data[data_var].isel(time=0) else: reference = data[data_var] + tmask = np.logical_not(np.isnan(reference)) + mask = bitsea.Mask.from_xarray(dataset=xr.Dataset({"tmask": tmask})) + area = xr.DataArray(mask.area, dims=("latitude", "longitude")) - e3t = xr.DataArray(mask.e3t, dims=("depth", "latitude", "longitude")) - vol_cell = area * e3t + return xr.DataArray( - vol_cell.transpose("depth", "latitude", "longitude"), - dims=("depth", "latitude", "longitude"), + area, + dims=("latitude", "longitude"), coords={ - "depth": data.depth, "latitude": data.latitude, "longitude": data.longitude, }, ) -def compute_average( - data: "xr.Dataset", - axes: str, - depth_min: float = None, - depth_max: float = None, -) -> "xr.Dataset": - """Compute the average of all variables along a specified axis. - - Three axes are supported: - - * ``"depth"``: Computes the depth-weighted vertical average over the full - depth column using :func:`~medunda.actions.average_between_layers.average_between_layers`. - * ``"space"``: Computes a volume-weighted spatial average over all - (latitude, longitude) grid points using the cell volumes derived from - the grid mask. - * ``"time"``: Computes a simple arithmetic mean over the time dimension. - - Args: - data (xr.Dataset): Input dataset. Must include ``depth``, - ``latitude``, ``longitude``, and ``time`` coordinates as required - by the chosen axis. - axes (str): Axes along which to compute the average. One of - ``"depth"``, ``"space"``, or ``"time"``. - - Returns: - xr.Dataset: Dataset with the chosen dimension collapsed, containing - the averaged values for each variable. - - Raises: - ValueError: If *axes* is not one of the valid choices. - """ - - if axes not in VALID_AXIS.keys(): - raise ValueError( - f"Axes '{axes}' is not valid. Choose from {list(VALID_AXIS.keys())}" - ) - - LOGGER.info(f"Computing average over axes '{axes}'") - - if axes == "depth": - if depth_min is not None: - depth_min = float(depth_min) - else: - depth_min = float(data.depth.min()) - - if depth_max is not None: - depth_max = float(depth_max) - else: - depth_max = float(data.depth.max()) - - averaged_dataset = average_between_layers(data, depth_min, depth_max) - LOGGER.info( - f"Depth interval: [{depth_min if depth_min is not None else float(data.depth.min())}, " - f"{depth_max if depth_max is not None else float(data.depth.max())}]" - ) - - elif axes == "space": - weights = get_volume(data) - weights = weights.expand_dims({"time": data.time}) +class Aggregation: + def __init__(self, data: xr.Dataset): + self.ds = data + self.weights = get_area(data) + + def reduce_depth(self, data: xr.Dataset, depth_min=None, depth_max=None): + if "depth" not in data.dims: + LOGGER.warning("Depth dimensions not found, skipping depth") + + depth_min = float(data.depth.min()) if depth_min is None else depth_min + depth_max = float(data.depth.max()) if depth_max is None else depth_max + + result = average_between_layers(data, depth_min, depth_max) + + return result + + def reduce_lat_lon(self, data: xr.Dataset): + weights = self.weights.broadcast_like(data) + + averaged_dataset = xr.Dataset() + + for var in data.data_vars: + da = data[var] + + mask = da.notnull() + w = weights * mask + + weighted_sum = (da.fillna(0) * w).sum( + dim=("latitude", "longitude") + ) + + total_weights = w.sum(dim=("latitude", "longitude")) + + averaged_dataset[var] = weighted_sum / total_weights + + result = averaged_dataset + return result + + def reduce_lat(self, data: xr.Dataset): + weights = self.weights + + if "time" in data.dims: + weights = weights.expand_dims({"time": data.time}) averaged_dataset = xr.Dataset() for var in data.data_vars: da = data[var] - weighted_sum = (da * weights).sum(dim=("latitude", "longitude")) + weighted_sum = (da * weights).sum(dim="latitude") - total_weights = weights.sum(dim=("latitude", "longitude")) + total_weights = weights.sum(dim="latitude") averaged_dataset[var] = weighted_sum / total_weights - LOGGER.info("Space average computed successfully") + result = averaged_dataset + return result + + def reduce_lon(self, data: xr.Dataset): + weights = self.weights + + if "time" in data.dims: + weights = weights.expand_dims({"time": data.time}) + + averaged_dataset = xr.Dataset() + + for var in data.data_vars: + da = data[var] + + weighted_sum = (da * weights).sum(dim="longitude") + + total_weights = weights.sum(dim="longitude") - elif axes == "time": - averaged_dataset = data.mean(dim="time") + averaged_dataset[var] = weighted_sum / total_weights - LOGGER.info("Time average computed successfully") + result = averaged_dataset + return result - return averaged_dataset + def reduce_time(self, data: xr.Dataset): + result = data.mean(dim="time") + return result + + def averaging( + self, axes: Sequence[str], depth_min=None, depth_max=None + ) -> xr.Dataset: + result = self.ds + + if "depth" in axes: + result = self.reduce_depth(result, depth_min, depth_max) + + if "latitude" in axes and "longitude" in axes: + result = self.reduce_lat_lon(result) + + elif "latitude" in axes: + result = self.reduce_lat(result) + + elif "longitude" in axes: + result = self.reduce_lon(result) + + if "time" in axes: + result = self.reduce_time(result) + + return result + + +def compute_average(data, axes, depth_min, depth_max): + aggragator = Aggregation(data) + + return aggragator.averaging( + axes=axes, + depth_min=depth_min, + depth_max=depth_max, + ) From 873f812859733dc48a1d08574685d79864c7447f Mon Sep 17 00:00:00 2001 From: nakkari Date: Tue, 16 Jun 2026 19:41:19 +0200 Subject: [PATCH 3/4] Update test for compute_average --- tests/actions/test_compute_average.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/actions/test_compute_average.py b/tests/actions/test_compute_average.py index 9e37d10..c7f46aa 100644 --- a/tests/actions/test_compute_average.py +++ b/tests/actions/test_compute_average.py @@ -15,7 +15,9 @@ def test_compute_average(data4d): latitude_levels = data4d.latitude.shape[0] longitude_levels = data4d.longitude.shape[0] - ds = compute_average(data=data4d, axis="depth") + ds = compute_average( + data=data4d, axes="depth", depth_min=None, depth_max=None + ) assert "T" in ds.data_vars, "Variable 'T' not found in output dataset." assert "S" in ds.data_vars, "Variable 'S' not found in output dataset." From 9dceaf41fd4101924972f95c0aeb902ab11e7203 Mon Sep 17 00:00:00 2001 From: Nada Akkari Date: Wed, 17 Jun 2026 10:59:38 +0200 Subject: [PATCH 4/4] Fix typo --- src/medunda/actions/compute_average.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/medunda/actions/compute_average.py b/src/medunda/actions/compute_average.py index b617fe2..55decbd 100644 --- a/src/medunda/actions/compute_average.py +++ b/src/medunda/actions/compute_average.py @@ -169,9 +169,9 @@ def averaging( def compute_average(data, axes, depth_min, depth_max): - aggragator = Aggregation(data) + aggregator = Aggregation(data) - return aggragator.averaging( + return aggregator.averaging( axes=axes, depth_min=depth_min, depth_max=depth_max,