Non-Conservative Zonal Mean

uxarray/core/dataarray.py

@@ -35,6 +35,7 @@
 from uxarray.remap import UxDataArrayRemapAccessor
 from uxarray.cross_sections import UxDataArrayCrossSectionAccessor
 from uxarray.core.aggregation import _uxda_grid_aggregate
+from uxarray.core.zonal import _compute_zonal_mean
 
 import warnings
 from warnings import warn
@@ -428,6 +429,38 @@ def integrate(
 
         return uxda
 
+    def zonal_mean(self, lat=(-90, 90, 5), method="fast"):
+        if not self._face_centered() and not self._edge_centered():
+            raise NotImplementedError(
+                "Zonal mean computations are currently only supported for face-centered or edge-centered data variables."
+            )
+
+        if isinstance(lat, tuple):
+            # zonal mean over a range of latitudes
+            latitudes = np.arange(lat[0], lat[1] + lat[2], lat[2])
+
+        elif isinstance(lat, (float, int)):
+            # zonal mean over a single latitude
+            latitudes = [lat]
+
+        # Compute Zonal Mean
+        res = _compute_zonal_mean(
+            uxda=self, data_mapping=self.dims[-1], latitudes=latitudes, method=method
+        )
+
+        dims = list(self.dims[:-1]) + ["latitudes"]
+
+        # Result is stored and returned as a UxDataArray
+        uxda = UxDataArray(
+            res,
+            uxgrid=self.uxgrid,
+            dims=dims,
+            coords={"latitudes": latitudes},
+            name=self.name + "_zonal_mean" if self.name is not None else "zonal_mean",
+        )
+
+        return uxda
+
     def topological_mean(
         self,
         destination: Literal["node", "edge", "face"],

uxarray/core/zonal.py

@@ -0,0 +1,25 @@
+import numpy as np
+import dask.array as da
+
+
+def _compute_zonal_mean(uxda, data_mapping, latitudes, method="fast"):
+    # Select the appropriate function based on the data_mapping
+    if data_mapping == "n_face":
+        get_indices_func = uxda.uxgrid.get_faces_at_constant_latitude
+    else:
+        get_indices_func = uxda.uxgrid.get_edges_at_constant_latitude
+
+    # Create a result array based on whether Dask is being used
+    shape = uxda.shape[:-1] + (len(latitudes),)
+    if isinstance(uxda.data, da.Array):
+        # Create a Dask array for storing results
+        averages = da.empty(shape, dtype=uxda.dtype)
+    else:
+        # Create a NumPy array for storing results
+        averages = np.empty(shape, dtype=uxda.dtype)
+
+    for i, lat in enumerate(latitudes):
+        indices = get_indices_func(lat, method)
+        averages[..., i] = uxda[..., indices].data.mean(axis=-1)
+
+    return averages