Improve compute_average with multi-axis aggregation

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,106 +9,170 @@
 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(
         ACTION_NAME, help="Compute the average"
     )
     average_parser.add_argument(
-        "--axis",
+        "--axes",
         type=str,
-        choices=sorted(VALID_AXIS),
+        nargs="+",
+        choices=["depth", "latitude", "longitude", "time"],
         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.",
     )
 
 
-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", axis) -> "xr.Dataset":
-    """Compute the average of all variables along a specified axis.
+class Aggregation:
+    def __init__(self, data: xr.Dataset):
+        self.ds = data
+        self.weights = get_area(data)
 
-    Three axes are supported:
+    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"``: 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.
+        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
 
-    Args:
-        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
-            ``"depth"``, ``"space"``, or ``"time"``.
+        result = average_between_layers(data, depth_min, depth_max)
 
-    Returns:
-        xr.Dataset: Dataset with the chosen dimension collapsed, containing
-        the averaged values for each variable.
+        return result
 
-    Raises:
-        ValueError: If *axis* is not one of the valid choices.
-    """
-    if axis not in VALID_AXIS.keys():
-        raise ValueError(
-            f"Axis '{axis}' is not valid. Choose from {list(VALID_AXIS.keys())}"
-        )
+    def reduce_lat_lon(self, data: xr.Dataset):
+        weights = self.weights.broadcast_like(data)
 
-    LOGGER.info(f"Computing average over axis '{axis}'")
+        averaged_dataset = xr.Dataset()
+
+        for var in data.data_vars:
+            da = data[var]
+
+            mask = da.notnull()
+            w = weights * mask
 
-    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")
+            weighted_sum = (da.fillna(0) * w).sum(
+                dim=("latitude", "longitude")
+            )
 
-    elif axis == "space":
-        weights = get_volume(data)
-        weights = weights.expand_dims({"time": data.time})
+            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"))
-            total_weights = weights.sum(dim=("latitude", "longitude"))
+            weighted_sum = (da * weights).sum(dim="latitude")
+
+            total_weights = weights.sum(dim="latitude")
+
+            averaged_dataset[var] = weighted_sum / total_weights
+
+        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")
+
             averaged_dataset[var] = weighted_sum / total_weights
 
-        LOGGER.info("Space average computed successfully")
+        result = averaged_dataset
+        return result
+
+    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 axis == "time":
-        averaged_dataset = data.mean(dim="time")
+        elif "latitude" in axes:
+            result = self.reduce_lat(result)
 
-        LOGGER.info("Time average computed successfully")
+        elif "longitude" in axes:
+            result = self.reduce_lon(result)
 
-    return averaged_dataset
+        if "time" in axes:
+            result = self.reduce_time(result)
+
+        return result
+
+
+def compute_average(data, axes, depth_min, depth_max):
+    aggregator = Aggregation(data)
+
+    return aggregator.averaging(
+        axes=axes,
+        depth_min=depth_min,
+        depth_max=depth_max,
+    )

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."