From 614f0c72cccb7ff0bf1b7c20bc4a5a29520de2a8 Mon Sep 17 00:00:00 2001
From: Sam Evans <47793072+Sevans711@users.noreply.github.com>
Date: Thu, 1 Oct 2026 11:22:12 -0400
Subject: [PATCH 1/6] add uxds means, fix uxda mean docs & weighted_mean

adds UxDataset.zonal_mean, zonal_average, zonal_anomaly, azimuthal_mean, azimuthal_average, weighted_mean (see #1800).

improves UxDataArray.zonal_mean docstring.

fixes UxDataArray.weighted_mean bugs (see #1800) and improves its docstring.

Adds private methods: _grid_dims and _grid_dim, and UxDataset._apply_across_data_vars
---
 uxarray/core/dataarray.py | 138 +++++++++++++---------
 uxarray/core/dataset.py   | 238 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 320 insertions(+), 56 deletions(-)

diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py
index 05757b4d1..eefdb8f29 100644
--- a/uxarray/core/dataarray.py
+++ b/uxarray/core/dataarray.py
@@ -186,6 +186,30 @@ def uxgrid(self, ugrid_obj: Grid):
             )
         self._uxgrid = ugrid_obj
 
+    @property
+    def _grid_dims(self) -> set[str]:
+        """set of all grid dimensions associated with self.
+        This is a (possibly-empty) subset of {"n_face", "n_edge", "n_node"}.
+        """
+        return set(d for d in GRID_DIMS if d in self.dims)
+
+    @property
+    def _grid_dim(self) -> str:
+        """name of the single grid dimension associated with self.
+        This is "n_face", "n_edge", or "n_node" if exactly one is present in self.dims,
+        else raises DataCenteringError.
+        """
+        grid_dims = self._grid_dims
+        if len(grid_dims) == 1:
+            return grid_dims.pop()
+        else:
+            if len(grid_dims) == 0:
+                grid_dims = "none"
+            raise DataCenteringError(
+                f"Expected {type(self).__name__} with exactly 1 grid dimension, but got {grid_dims}, "
+                f"in self.dims={self.dims}. Known grid dimensions are: {GRID_DIMS}."
+            )
+
     @property
     def data_mapping(self):
         """Returns which grid element a data variable is mapped to.
@@ -659,29 +683,35 @@ def integrate(
         return uxda
 
     def zonal_mean(self, lat=(-90, 90, 10), conservative: bool = False, **kwargs):
-        """Compute non-conservative or conservative averages of a face-centered variable along lines of constant latitude or latitude bands.
+        """Returns averages of face-centered variables along lines or bands of constant latitude.
 
-        A zonal mean in UXarray operates differently depending on the ``conservative`` flag:
+        The weighting method and the output size depend on the ``conservative`` flag:
 
-        - **Non-conservative**: Calculates the mean by sampling face values at specific latitude lines and weighting each contribution by the length of the line where each face intersects that latitude.
-        - **Conservative**: Preserves integral quantities by calculating the mean by sampling face values within latitude bands and weighting contributions by their area overlap with latitude bands.
+        - ``conservative=False``: weight contributions by each face's overlap with a given
+          line of constant latitude. ``lat`` indicates which lines to use.
+          (E.g., the default lat=(-90, 90, 10) produces a result with 19 latitudes,
+          -90, -80, ..., 90, corresponding to means at -90, -80, ..., 90 degrees.)
+        - ``conservative=True``: weight contributions by each face's overlap with a given
+          band of latitude. ``lat`` indicates the edges of the bands to use.
+          (E.g., the default lat=(-90, 90, 10) produces a result with 18 latitudes,
+          -85, -75, ..., 85, corresponding to means over the bands from
+          -90 to -80, -80 to -70, ..., and 80 to 90 degrees.)
+          Using ``conservative=True`` preserves integral quantities.
 
         Parameters
         ----------
-        lat : tuple, float, or array-like, default=(-90, 90, 10)
-            Latitude specification:
-                - tuple (start, end, step): For non-conservative, computes means at intervals of `step`.
-                For conservative, creates band edges via np.arange(start, end+step, step).
-                - float: Single latitude for non-conservative averaging
-                - array-like: For non-conservative, latitudes to sample. For conservative, band edges.
+        lat : int, float, tuple of length 3, list, or np.ndarray, default=(-90, 90, 10)
+            Latitudes at which to compute means, either as latitude lines (conservative=False)
+            or latitude band edges (conservative=True). If int or float, treated as [lat] instead.
+            If tuple of (start, end, step), treated as np.arange(start, end+step, step) instead.
         conservative : bool, default=False
             If True, performs conservative (area-weighted) zonal averaging over latitude bands.
             If False, performs non-conservative (intersection-weighted) averaging at latitude lines.
 
         Returns
         -------
-        UxDataArray
-            Contains zonal means with a new 'latitudes' dimension and corresponding coordinates.
+        xr.DataArray
+            Contains zonal means values, with a new 'latitudes' dimension and coordinates.
             Name will be original_name + '_zonal_mean' or 'zonal_mean' if unnamed.
 
         Examples
@@ -1022,53 +1052,45 @@ def azimuthal_mean(
     azimuthal_average = azimuthal_mean
 
     def weighted_mean(self, weights=None):
-        """Computes a weighted mean.
+        """Returns a weighted mean. If weights are not provided:
 
-        This function calculates the weighted mean of a variable,
-        using the specified `weights`. If no weights are provided, it will automatically select
-        appropriate weights based on whether the variable is face-centered or edge-centered. If
-        the variable is neither face nor edge-centered a warning is raised, and an unweighted mean is computed instead.
+        - For face-centered data, use face areas as weights (i.e., area-weighted mean).
+        - For edge-centered data, use edge lengths as weights.
+        - For node-centered data, crash with DataCenteringError.
 
-        Parameters
-        ----------
-        weights : np.ndarray or None, optional
-            The weights to use for the weighted mean calculation. If `None`, the function will
-            determine weights based on the variable's association:
+        If weights are provided, they are treated as weights along the grid dimension,
+        ("n_face", "n_edge", or "n_node") and must be 1D with appropriate length.
 
-            - For face-centered variables: uses `self.uxgrid.face_areas.data`
-            - For edge-centered variables: uses `self.uxgrid.edge_node_distances.data`
+        Mathematically equivalent to sum(self * weights) / sum(weights),
+        where the first sum is taken along the grid dimension.
 
-            If the variable is neither face-centered nor edge-centered, a warning is raised, and
-            an unweighted mean is computed instead. User-defined weights should match the shape
-            of the data variable's last dimension.
+        Parameters
+        ----------
+        weights : array-like or None, optional
+            1D array of weights to apply along the data's grid dimension.
+            If None, use face areas, edge lengths, or crash, as described above.
+            If xr.DataArray, the dimension must match the data's grid dimension.
 
         Returns
         -------
         UxDataArray
             A new `UxDataArray` object representing the weighted mean of the input variable. The
-            result is attached to the same `uxgrid` attribute as the original variable.
+            result's `uxgrid` matches the input's uxgrid, even though the result does not have a
+            grid dimension anymore and the result's data does not lie on a grid. This can cause
+            issues with functionality like result.plot(); consider calling result.to_xarray() as
+            a workaround. (Might be updated in a future release to return xr.DataArray instead.)
 
         Example
         -------
         >>> weighted_mean = uxds["t2m"].weighted_mean()
 
-
         Raises
         ------
-        AssertionError
-            If user-defined `weights` are provided and the shape of `weights` does not match
-            the shape of the data variable's last dimension.
-
-        Warnings
-        --------
-        UserWarning
-            Raised when attempting to compute a weighted mean on a variable without associated
-            weights. An unweighted mean will be computed in this case.
-
-        Notes
-        -----
-        - The weighted mean is computed along the last dimension of the data variable, which is
-          assumed to be the geometry dimension (e.g., faces, edges, or nodes).
+        DataCenteringError (subclass of ValueError)
+            If the data is node-centered and no weights are provided.
+        DimensionError (subclass of ValueError)
+            If provided weights which are not 1D, or incompatible with the data due to
+            having wrong size or (if xr.DataArray) wrong dimension name.
         """
         if weights is None:
             if self._face_centered():
@@ -1076,23 +1098,29 @@ def weighted_mean(self, weights=None):
             elif self._edge_centered():
                 weights = self.uxgrid.edge_node_distances.data
             else:
-                warnings.warn(
-                    "Attempting to perform a weighted mean calculation on a variable that does not have"
-                    "associated weights. Weighted mean is only supported for face or edge centered "
-                    "variables. Performing an unweighted mean."
+                raise DataCenteringError(
+                    "weighted_mean() cannot automatically infer weights for node-centered data. "
+                    "Consider providing weights or re-centering the data first."
                 )
         else:
             # user-defined weights
-            assert weights.shape[-1] == self.shape[-1]
+            if not isinstance(weights, xr.DataArray):
+                weights = np.asanyarray(weights)
+            if weights.ndim != 1:
+                raise DimensionError(f"Expected 1D weights, got ndim={weights.ndim}.")
+            if isinstance(weights, xr.DataArray):
+                if weights.dims[0] != self._grid_dim:
+                    raise DimensionError(
+                        f"Expected xr.DataArray weights dimension to match the data's grid dimension "
+                        f"({self._grid_dim!r}), but got weights with dimension {weights.dims[0]!r}."
+                    )
+        if not isinstance(weights, xr.DataArray):
+            # convert to xr.DataArray to ensure operations below align dims properly
+            weights = xr.DataArray(weights, dims=(self._grid_dim,))
 
-        # compute the total weight
         total_weight = weights.sum()
-
-        # compute the weighted mean, with an assumption on the index of dimension (last one is geometry)
-        weighted_mean = (self * weights).sum(axis=-1) / total_weight
-
-        # create a UxDataArray and return it
-        return UxDataArray(weighted_mean, uxgrid=self.uxgrid)
+        weighted_mean = (self * weights).sum(self._grid_dim) / total_weight
+        return type(self)(weighted_mean, uxgrid=self.uxgrid)
 
     def topological_mean(
         self,
diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py
index 9de3be7bd..d26b583bd 100644
--- a/uxarray/core/dataset.py
+++ b/uxarray/core/dataset.py
@@ -19,7 +19,7 @@
     _resolve_coordinate_labels_to_indices,
     _validate_indexers,
 )
-from uxarray.errors import DimensionError, GridInvalidError
+from uxarray.errors import DataCenteringError, DimensionError, GridInvalidError
 from uxarray.formatting_html import dataset_repr
 from uxarray.grid import Grid
 from uxarray.grid.dual import construct_dual
@@ -202,6 +202,30 @@ def uxgrid(self, ugrid_obj: Grid):
             )
         self._uxgrid = ugrid_obj
 
+    @property
+    def _grid_dims(self) -> set[str]:
+        """set of all grid dimensions associated with self.
+        This is a (possibly-empty) subset of {"n_face", "n_edge", "n_node"}.
+        """
+        return set(d for d in GRID_DIMS if d in self.dims)
+
+    @property
+    def _grid_dim(self) -> str:
+        """name of the single grid dimension associated with self.
+        This is "n_face", "n_edge", or "n_node" if exactly one is present in self.dims,
+        else raises DataCenteringError.
+        """
+        grid_dims = self._grid_dims
+        if len(grid_dims) == 1:
+            return grid_dims.pop()
+        else:
+            if len(grid_dims) == 0:
+                grid_dims = "none"
+            raise DataCenteringError(
+                f"Expected {type(self).__name__} with exactly 1 grid dimension, but got {grid_dims}, "
+                f"in self.dims={self.dims}. Known grid dimensions are: {GRID_DIMS}."
+            )
+
     def _calculate_binary_op(self, *args, **kwargs):
         """Override to make the result a complete instance of
         ``uxarray.UxDataset``."""
@@ -979,3 +1003,215 @@ def fillna(self, value: Any):
         return UxDataset(super().fillna(value), uxgrid=self._uxgrid)
 
     fillna.__doc__ = xr.Dataset.fillna.__doc__
+
+    # --- methods which just apply iteratively across data_vars --- #
+
+    def _apply_across_data_vars(
+        self, method, grid_dims=GRID_DIMS, *method_args, **method_kwargs
+    ):
+        """Apply method to each relevant data_var in self with any grid_dims,
+        returning a Dataset formed by the results, keeping all other data_vars unchanged.
+
+        If any data_var has a grid_dim (one of GRID_DIMS) not in `grid_dims`, crash with
+        DataCenteringError. (E.g. "n_node" but method doesn't support node-centered data.)
+        Also crash with DataCenteringError if any coordinate has a grid_dim not in `grid_dims`.
+        Also crash with DataCenteringError if none of the supported grid_dims appear in self.
+
+        Parameters
+        ----------
+        method : str or callable
+            The method to apply to each relevant data_var.
+            str --> attribute of data_var, e.g. 'zonal_mean' --> data_var.zonal_mean()
+            callable --> called with data_var as first argument.
+        grid_dims : list of str
+            The grid dimensions supported by `method`.
+            The method is applied to all data_vars containing at least one of these dims.
+            If any data_var has a grid_dim not in `grid_dims`, crash with DataCenteringError.
+        remaining args and kwargs get passed to `method`.
+
+        Returns
+        -------
+        UxDataset or xr.Dataset
+            Results of applying `method` where relevant, keeping other data_vars unchanged.
+            Type is UxDataset if `method` returns UxDataArray for any data_var, else xr.Dataset.
+        """
+        assert not isinstance(
+            grid_dims, str
+        )  # made a typo if grid_dims is a single string!
+        _grid_dims_as_a_set = set(grid_dims)
+
+        # eager crash if self not centered properly for method
+        if not any(dim in self.dims for dim in _grid_dims_as_a_set):
+            raise DataCenteringError(
+                f"Expected {type(self).__name__}.dims to contain least 1 grid dimension "
+                f"supported by {method!r}, i.e. one of {grid_dims}, but got dims={self.dims}."
+            )
+
+        # eager crash if any data_var not centered properly
+        for name, da in self.data_vars.items():
+            da_grid_dims = set(da.dims).intersection(GRID_DIMS)
+            if len(da_grid_dims) > 0 and not (da_grid_dims <= _grid_dims_as_a_set):
+                raise DataCenteringError(
+                    f"Data variable {name!r} has grid dimension(s) {da_grid_dims} which are "
+                    f"not supported by method {method!r}. Supported grid dimensions: {grid_dims}."
+                )
+
+        # separately track coords not attached to any data_var
+        bonus_coords = {}
+        for coord in self.coords:
+            for data_var in self.data_vars.values():
+                if coord in data_var.coords:
+                    break
+            else:  # didn't break
+                bonus_coords[coord] = self.coords[coord]
+
+        # eager crash if any bonus coord not centered properly
+        for name, da in bonus_coords.items():
+            da_grid_dims = set(da.dims).intersection(GRID_DIMS)
+            if len(da_grid_dims) > 0 and not (da_grid_dims <= _grid_dims_as_a_set):
+                raise DataCenteringError(
+                    f"Coordinate {name!r} has grid dimension(s) {da_grid_dims} which are "
+                    f"not supported by method {method!r}. Supported grid dimensions: {grid_dims}"
+                )
+
+        # actually compute the results (for data_vars)
+        results = {}
+        for name, da in self.data_vars.items():
+            if any(dim in da.dims for dim in GRID_DIMS):
+                if isinstance(method, str):
+                    results[name] = getattr(da, method)(*method_args, **method_kwargs)
+                else:
+                    results[name] = method(da, *method_args, **method_kwargs)
+            else:
+                results[name] = da
+
+        # return UxDataset if any result is a UxDataArray, else xr.Dataset
+        _ux_kws = {}
+        if any(isinstance(res, UxDataArray) for res in results.values()):
+            cls = type(self)
+            _ux_kws["uxgrid"] = self._uxgrid
+        else:
+            cls = xr.Dataset
+        return cls(results, coords=bonus_coords, attrs=self.attrs, **_ux_kws)
+
+    def zonal_mean(
+        self, lat=(-90, 90, 10), *, conservative: bool = False, **kw_uxda_zonal_mean
+    ):
+        """Returns averages of face-centered data_vars along lines or bands of constant latitude,
+        as an xr.Dataset with new "latitudes" dimension.
+        Only affects data_vars with grid dimensions; everything else is passed through unchanged.
+
+        The weighting method and the output size depend on the ``conservative`` flag:
+
+        - ``conservative=False``: weight contributions by each face's overlap with a given
+          line of constant latitude. ``lat`` indicates which lines to use.
+          (E.g., the default lat=(-90, 90, 10) produces a result with 19 latitudes,
+          -90, -80, ..., 90, corresponding to means at -90, -80, ..., 90 degrees.)
+        - ``conservative=True``: weight contributions by each face's overlap with a given
+          band of latitude. ``lat`` indicates the edges of the bands to use.
+          (E.g., the default lat=(-90, 90, 10) produces a result with 18 latitudes,
+          -85, -75, ..., 85, corresponding to means over the bands from
+          -90 to -80, -80 to -70, ..., and 80 to 90 degrees.)
+          Using ``conservative=True`` preserves integral quantities.
+
+        For more details about parameters, see :meth:`UxDataArray.zonal_mean`.
+        """
+        kw = dict(lat=lat, conservative=conservative, **kw_uxda_zonal_mean)
+        return self._apply_across_data_vars("zonal_mean", ("n_face",), **kw)
+
+    def zonal_average(
+        self, lat=(-90, 90, 10), *, conservative: bool = False, **kw_uxda_zonal_mean
+    ):
+        """Alias of zonal_mean. For full docstring, see :meth:`UxDataset.zonal_mean`."""
+        return self.zonal_mean(lat=lat, conservative=conservative, **kw_uxda_zonal_mean)
+
+    def zonal_anomaly(self, lat=(-90, 90, 10), *, conservative: bool = False):
+        """Return UxDataset with zonal anomaly of face-centered data_vars (values minus zonal means),
+        Only affects data_vars with grid dimensions; everything else is passed through unchanged.
+
+        For more details about parameters, see :meth:`UxDataArray.zonal_anomaly`.
+        """
+        kw = dict(lat=lat, conservative=conservative)
+        return self._apply_across_data_vars("zonal_anomaly", ("n_face",), **kw)
+
+    def azimuthal_mean(
+        self,
+        center_coord,
+        outer_radius: int | float,
+        radius_step: int | float,
+        *,
+        return_hit_counts: bool = False,
+    ) -> xr.Dataset:
+        """Return averages of face-centered data_vars along circles of constant great-circle distance
+        from ``center_coord`` (lon, lat), as an xr.Dataset with new "radius" dimension.
+        Only affects data_vars with grid dimensions; everything else is passed through unchanged.
+
+        (Does not yet support ``return_hit_counts`` option.)
+
+        For more details about parameters, see :meth:`UxDataArray.azimuthal_mean`.
+        """
+        if return_hit_counts:
+            raise NotImplementedError(
+                f"{type(self).__name__}.azimuthal_mean(..., return_hit_counts=True)"
+            )
+        kw = dict(
+            center_coord=center_coord,
+            outer_radius=outer_radius,
+            radius_step=radius_step,
+        )
+        return self._apply_across_data_vars("azimuthal_mean", ("n_face",), **kw)
+
+    azimuthal_average = azimuthal_mean  # alias
+
+    def weighted_mean(self, weights=None):
+        """Return UxDataset of weighted means of data_vars. If weights are not provided:
+
+        - For face-centered data, use face areas as weights (i.e., area-weighted mean).
+        - For edge-centered data, use edge lengths as weights.
+        - For node-centered data, crash with DataCenteringError.
+
+        If weights are provided, they are treated as weights along the grid dimension,
+        ("n_face", "n_edge", or "n_node") and must be 1D with appropriate length.
+        If weights are provided, and multiple grid dimensions appear throughout data_vars,
+        raise DataCenteringError, or NotImplementedError if weights is a 1D DataArray
+        with a grid dimension that appears in this dataset's dims.
+
+        Mathematically equivalent to sum(self * weights) / sum(weights),
+        where the first sum is taken along the grid dimension.
+
+        For more details and full docstring, see :meth:`UxDataArray.weighted_mean`.
+        """
+        if weights is None:
+            return self._apply_across_data_vars(
+                "weighted_mean", ("n_face", "n_edge"), weights=None
+            )
+        # else:
+        grid_dims = self._grid_dims
+        if len(grid_dims) > 1:
+            if (
+                isinstance(weights, xr.DataArray)
+                and len(weights.dims) == 1
+                and weights.dims[0] in grid_dims
+            ):
+                raise NotImplementedError(
+                    "uxds.weighted_mean(..., weights=1D DataArray with dim in uxds's grid dims) "
+                    f"when uxds (a {type(self).__name__}) has multiple grid dims."
+                )
+                # This case has an intuitive, unambiguous way it could be implemented:
+                # simply apply weighted_mean() only along data_vars with the grid dim from weights,
+                # skipping data_vars with other grid dims.
+                # Doing this would require updating _apply_across_data_vars accordingly first, because
+                # currently it crashes if any data_var has a grid dim not in its input ``grid_dims``.
+            else:
+                raise DataCenteringError(
+                    f"uxds.weighted_mean(..., weights=...) when uxds (a {type(self).__name__}) "
+                    f"has multiple grid dims ({grid_dims}). Provide weights=None, or try again "
+                    f"with a subset of data_vars that all share the same grid dim."
+                )
+                # Note: after filling in the implementation of the 1D xr.DataArray with grid_dim
+                # case above, update this error message too, to also suggest providing weights as
+                # an xr.DataArray with the desired grid dim.
+        else:
+            return self._apply_across_data_vars(
+                "weighted_mean", grid_dims, weights=weights
+            )

From 86e8afb23167e2eff4a0247049689ebb4f1259a1 Mon Sep 17 00:00:00 2001
From: Sam Evans <47793072+Sevans711@users.noreply.github.com>
Date: Thu, 1 Oct 2026 12:08:21 -0400
Subject: [PATCH 2/6] tests for weighted_mean fixes (see #1800)

---
 test/core/test_weighted_mean.py | 72 +++++++++++++++++++++++++++++++++
 1 file changed, 72 insertions(+)

diff --git a/test/core/test_weighted_mean.py b/test/core/test_weighted_mean.py
index 44ac9ffb1..fa912b501 100644
--- a/test/core/test_weighted_mean.py
+++ b/test/core/test_weighted_mean.py
@@ -6,6 +6,7 @@
 import numpy.testing as nt
 
 import uxarray as ux
+import xarray as xr
 
 def test_quad_hex_face_centered(gridpath, datasetpath):
     """Compares the weighted average computation for the quad hexagon grid
@@ -94,6 +95,19 @@ def test_quad_hex_edge_centered_dask(gridpath, test_data_dir):
     # ensure values are within 3 decimal points of each other
     nt.assert_almost_equal(computed_result.values, expected_weighted_mean, decimal=3)
 
+def test_weighted_mean_crash_if_node_centered_and_no_weights():
+    """Ensure weighted_mean crashes if the data is node-centered and no weights are provided.
+    Regression test for bug 2 of issue #1800.
+    """
+    ds = ux.tutorial.open_dataset('quad-hexagon-random-node')
+    arr = ds['random_data_node']
+    ERRMSG_ARR = r"weighted_mean\(\) cannot automatically infer weights for node-centered data."
+    with pytest.raises(ux.errors.DataCenteringError, match=ERRMSG_ARR):
+        arr.weighted_mean()
+    ERRMSG_DS = r"Expected UxDataset\.dims to contain least 1 grid dimension supported by 'weighted_mean'"
+    with pytest.raises(ux.errors.DataCenteringError, match=ERRMSG_DS):
+        ds.weighted_mean()
+
 def test_csne30_equal_area(gridpath, datasetpath):
     """Compute the weighted average with a grid that has equal-area faces and
     compare the result to the regular mean."""
@@ -139,3 +153,61 @@ def test_csne30_equal_area_dask(gridpath, datasetpath, chunk_size):
 
     # ensure values are within 3 decimal points of each other
     nt.assert_almost_equal(computed_result.values, expected_weighted_mean, decimal=3)
+
+def test_weighted_mean_if_provided_weights():
+    """Ensure can provide weights to weighted_mean and get correct results."""
+    ds_face = ux.tutorial.open_dataset('quad-hexagon-random-face')
+    ds_node = ux.tutorial.open_dataset('quad-hexagon-random-node')
+    ds_edge = ux.tutorial.open_dataset('quad-hexagon-random-edge')
+    for ds in [ds_face, ds_node, ds_edge]:
+        assert len(ds.dims) == 1 and len(ds.data_vars) == 1
+        arr = ds.to_array('variable').isel(variable=0)
+        weights_values = np.arange(arr.size)
+        weights_nparr = weights_values
+        weights_list = list(weights_nparr)
+        weights_da = xr.DataArray(weights_nparr, dims=arr.dims)
+        expected_weighted_mean = (arr.values * weights_values).sum() / weights_values.sum()
+        # ^just .sum() is fine because arr is 1D; don't need to worry about other dims.
+        for weights in [weights_nparr, weights_list, weights_da]:
+            weighted_mean = arr.weighted_mean(weights=weights)
+            assert weighted_mean.ndim == 0   # arr is 1D; now no dims remain.
+            nt.assert_equal(weighted_mean.item(), expected_weighted_mean)
+        # repeat checks but for UxDataset:
+        for weights in [weights_nparr, weights_list, weights_da]:
+            weighted_mean_ds = ds.weighted_mean(weights=weights)
+            weighted_mean = weighted_mean_ds.to_array('variable').isel(variable=0)
+            assert weighted_mean.ndim == 0
+            nt.assert_equal(weighted_mean.item(), expected_weighted_mean)
+        # special cases for xr.DataArray weights: wrong dim, or has a scalar coord
+        # wrong dim --> need to crash.
+        weights_da_wrongdim = xr.DataArray(weights_nparr, dims=['wrongdim'])
+        with pytest.raises(ux.errors.DimensionError):
+            arr.weighted_mean(weights=weights_da_wrongdim)
+        with pytest.raises(ux.errors.DimensionError):
+            ds.weighted_mean(weights=weights_da_wrongdim)
+        # has a scalar coord --> result should have that coord, too.
+        weights_da_scalarcoord = xr.DataArray(weights_nparr, dims=arr.dims, coords={'scalarcoord': 7})
+        weighted_mean = arr.weighted_mean(weights=weights_da_scalarcoord)
+        assert weighted_mean.coords['scalarcoord'] == 7
+        weighted_mean_ds = ds.weighted_mean(weights=weights_da_scalarcoord)
+        assert weighted_mean_ds.coords['scalarcoord'] == 7
+
+def test_weighted_mean_doesnt_care_about_dim_order():
+    """Ensure weighted mean does not care about dimension order.
+    E.g., grid dim doesn't need to be the last dim, in order to get correct results.
+    Regression test for bug 1 of issue #1800.
+    """
+    arr = ux.tutorial.open_dataset("outCSne30-timeseries")['psi']
+    assert arr.dims == ('time', 'n_face')   # (assert original has 'n_face' dim last.)
+    order0_result = arr.weighted_mean()
+    order1_result = arr.transpose('n_face', 'time').weighted_mean()
+    assert order0_result.dims == ('time',) == order1_result.dims
+    np.allclose(order0_result, order1_result, atol=0, rtol=1e-13)
+    # ensure correctness even if two dims have same size:
+    n_time = arr.sizes['time']
+    sliced = arr.isel(n_face=slice(n_time))
+    assert sliced.sizes == {'time': n_time, 'n_face': n_time}
+    order0_result = sliced.weighted_mean()
+    order1_result = sliced.transpose('n_face', 'time').weighted_mean()
+    assert order0_result.dims == ('time',) == order1_result.dims
+    np.allclose(order0_result, order1_result, atol=0, rtol=1e-13)

From 9d8796b2283a377d87a50fac390ab39c36f069b1 Mon Sep 17 00:00:00 2001
From: Sam Evans <47793072+Sevans711@users.noreply.github.com>
Date: Thu, 1 Oct 2026 13:12:03 -0400
Subject: [PATCH 3/6] tests for uxds.zonal_mean, anomaly, weighted_mean

(sidenote: some UxDataset.weighted_mean tests are already covered by test_weighted_mean.py)
---
 test/core/test_dataset.py       | 286 ++++++++++++++++++++++++++++++++
 test/core/test_weighted_mean.py |  19 ++-
 2 files changed, 303 insertions(+), 2 deletions(-)

diff --git a/test/core/test_dataset.py b/test/core/test_dataset.py
index f58bcbccc..1b28a1bcd 100644
--- a/test/core/test_dataset.py
+++ b/test/core/test_dataset.py
@@ -321,3 +321,289 @@ def test_uxgrid_None_is_invalid_in_uxdataset():
     # it also applies (for non-None non-Grid objects) during __init__:
     with pytest.raises(TypeError):
         ux.UxDataset({'arr1': xr.DataArray([4,5], dims=['n_face'])}, uxgrid=[1,2])
+
+
+class TestUxDatasetMimicsUxDataArrayMethods:
+    """Testing behavior of UxDataset methods which simply apply UxDataArray methods iteratively,
+    such as UxDataset.zonal_mean().
+    """
+    def _uxds_face_with_just_psi(self):
+        uxds = ux.tutorial.open_dataset("outCSne30-timeseries")
+        assert set(uxds.data_vars) == {'psi'}
+        assert 'n_face' in uxds.dims
+        return uxds
+
+    def _uxds_face_with_more_vars_and_coords(self):
+        uxds = ux.tutorial.open_dataset("outCSne30-timeseries")
+        assert set(uxds.data_vars) == {'psi'}
+        uxds = uxds.assign(psi2 = uxds['psi']*2)
+        uxds = uxds.assign(unrelated_time_var = xr.DataArray(10*np.arange(uxds.sizes['time']), dims=['time']))
+        uxds = uxds.assign(unrelated_scalar_var = 7)
+        uxds = uxds.assign_coords(unrelated_scalar_coord = 70)
+        # (will also want to check what happens to coords which are not used by any data_vars!)
+        uxds = uxds.assign_coords(unused_dim_coord = xr.DataArray([1,2,3], dims='unused_dim'))
+        assert set(uxds.data_vars) == {'psi', 'psi2', 'unrelated_time_var', 'unrelated_scalar_var'}
+        return uxds
+
+    def _uxds_hex_face(self):
+        uxds = ux.tutorial.open_dataset('quad-hexagon-random-face')
+        assert set(uxds.data_vars) == {'random_data_face'}
+        assert 'n_face' in uxds.dims
+        return uxds
+
+    def _uxds_hex_edge(self):
+        uxds = ux.tutorial.open_dataset('quad-hexagon-random-edge')
+        assert set(uxds.data_vars) == {'random_data_edge'}
+        assert 'n_edge' in uxds.dims
+        return uxds
+
+    def _uxds_hex_node(self):
+        uxds = ux.tutorial.open_dataset('quad-hexagon-random-node')
+        assert set(uxds.data_vars) == {'random_data_node'}
+        assert 'n_node' in uxds.dims
+        return uxds
+
+    def _uxds_hex_face_and_node(self):
+        arr_face = self._uxds_hex_face()['random_data_face']
+        arr_node = self._uxds_hex_node()['random_data_node']
+        uxds = ux.UxDataset({'face_data': arr_face, 'node_data': arr_node}, uxgrid=arr_face.uxgrid)
+        return uxds
+
+    def _uxds_hex_face_and_edge(self):
+        arr_face = self._uxds_hex_face()['random_data_face']
+        arr_edge = self._uxds_hex_edge()['random_data_edge']
+        uxds = ux.UxDataset({'face_data': arr_face, 'edge_data': arr_edge}, uxgrid=arr_face.uxgrid)
+        return uxds
+
+    def _uxds_hex_node_and_edge(self):
+        arr_node = self._uxds_hex_node()['random_data_node']
+        arr_edge = self._uxds_hex_edge()['random_data_edge']
+        uxds = ux.UxDataset({'node_data': arr_node, 'edge_data': arr_edge}, uxgrid=arr_node.uxgrid)
+        return uxds
+
+    def _uxds_hex_face_and_node_and_edge(self):
+        arr_face = self._uxds_hex_face()['random_data_face']
+        arr_node = self._uxds_hex_node()['random_data_node']
+        arr_edge = self._uxds_hex_edge()['random_data_edge']
+        uxds = ux.UxDataset({'face_data': arr_face, 'node_data': arr_node, 'edge_data': arr_edge}, uxgrid=arr_face.uxgrid)
+        return uxds
+
+    def test_uxds_mimics_uxda_zonal_mean(self):
+        """Ensure UxDataset.zonal_mean() mimics UxDataArray.zonal_mean() for each variable."""
+        ds = self._uxds_face_with_just_psi()
+        psi_result = ds['psi'].zonal_mean()
+        assert ds.zonal_mean()['psi'].equals(psi_result)
+        assert isinstance(psi_result, xr.DataArray)
+        assert isinstance(ds.zonal_mean(), xr.Dataset)
+
+        # quick sanity check: zonal_mean and zonal_average are aliases.
+        # (doing this here instead of making a separate test for zonal_average...)
+        assert ds['psi'].zonal_average().equals(ds['psi'].zonal_mean())
+        assert ds.zonal_average().equals(ds.zonal_mean())
+
+        ds = self._uxds_face_with_more_vars_and_coords()
+        psi_result = ds['psi'].zonal_mean()
+        psi2_result = ds['psi2'].zonal_mean()
+        ds_result = ds.zonal_mean()
+        assert ds_result['psi'].equals(psi_result)
+        assert ds_result['psi2'].equals(psi2_result)
+        assert 'unrelated_scalar_var' not in ds['psi'] and 'unrelated_scalar_var' not in psi_result.coords
+        assert ds_result['unrelated_scalar_var'].equals(ds['unrelated_scalar_var'])
+        assert 'unrelated_time_var' not in ds['psi'] and 'unrelated_time_var' not in psi_result.coords
+        assert ds_result['unrelated_time_var'].equals(ds['unrelated_time_var'])
+        assert ds_result.coords['unrelated_scalar_coord'].equals(ds.coords['unrelated_scalar_coord'])
+        assert ds_result.coords['unused_dim_coord'].equals(ds.coords['unused_dim_coord'])
+
+        # zonal mean doesn't support edge-centered data
+        ds = self._uxds_hex_edge()
+        arr = ds['random_data_edge']
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            arr.zonal_mean()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.zonal_mean()
+
+        # zonal mean doesn't support node-centered data
+        ds = self._uxds_hex_node()
+        arr = ds['random_data_node']
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            arr.zonal_mean()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.zonal_mean()
+
+        # also should crash if any data_vars contain unsupported centering,
+        # even if some data_vars are centered at supported location (faces).
+        ds = self._uxds_hex_face_and_node()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.zonal_mean()
+        ds = self._uxds_hex_face_and_edge()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.zonal_mean()
+
+    def test_uxds_mimics_uxda_zonal_anomaly(self):
+        """Ensure UxDataset.zonal_anomaly() mimics UxDataArray.zonal_anomaly() for each variable."""
+        ds = self._uxds_face_with_just_psi()
+        psi_result = ds['psi'].zonal_anomaly()
+        assert ds.zonal_anomaly()['psi'].equals(psi_result)
+        assert isinstance(psi_result, xr.DataArray)
+        assert isinstance(ds.zonal_anomaly(), xr.Dataset)
+
+        ds = self._uxds_face_with_more_vars_and_coords()
+        psi_result = ds['psi'].zonal_anomaly()
+        psi2_result = ds['psi2'].zonal_anomaly()
+        ds_result = ds.zonal_anomaly()
+        assert ds_result['psi'].equals(psi_result)
+        assert ds_result['psi2'].equals(psi2_result)
+        assert 'unrelated_scalar_var' not in ds['psi'] and 'unrelated_scalar_var' not in psi_result.coords
+        assert ds_result['unrelated_scalar_var'].equals(ds['unrelated_scalar_var'])
+        assert 'unrelated_time_var' not in ds['psi'] and 'unrelated_time_var' not in psi_result.coords
+        assert ds_result['unrelated_time_var'].equals(ds['unrelated_time_var'])
+        assert ds_result.coords['unrelated_scalar_coord'].equals(ds.coords['unrelated_scalar_coord'])
+        assert ds_result.coords['unused_dim_coord'].equals(ds.coords['unused_dim_coord'])
+
+        # zonal anomaly doesn't support edge-centered data
+        ds = self._uxds_hex_edge()
+        arr = ds['random_data_edge']
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            arr.zonal_anomaly()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.zonal_anomaly()
+
+        # zonal anomaly doesn't support node-centered data
+        ds = self._uxds_hex_node()
+        arr = ds['random_data_node']
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            arr.zonal_anomaly()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.zonal_anomaly()
+
+        # also should crash if any data_vars contain unsupported centering,
+        # even if some data_vars are centered at supported location (faces).
+        ds = self._uxds_hex_face_and_node()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.zonal_anomaly()
+        ds = self._uxds_hex_face_and_edge()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.zonal_anomaly()
+
+    def test_uxds_mimics_uxda_azimuthal_mean(self):
+        """Ensure UxDataset.azimuthal_mean() mimics UxDataArray.azimuthal_mean() for each variable."""
+        kw_psi = dict(center_coord=(45, 0), outer_radius=50, radius_step=10)
+        kw_hex = dict(center_coord=(0, 0), outer_radius=0.3, radius_step=0.1)
+
+        ds = self._uxds_face_with_just_psi()
+        psi_result = ds['psi'].azimuthal_mean(**kw_psi)
+        assert ds.azimuthal_mean(**kw_psi)['psi'].equals(psi_result)
+        assert isinstance(psi_result, xr.DataArray)
+        assert isinstance(ds.azimuthal_mean(**kw_psi), xr.Dataset)
+
+        # quick sanity check: azimuthal_mean and azimuthal_average are aliases.
+        # (doing this here instead of making a separate test for azimuthal_average...)
+        assert ds['psi'].azimuthal_average(**kw_psi).equals(ds['psi'].azimuthal_mean(**kw_psi))
+        assert ds.azimuthal_average(**kw_psi).equals(ds.azimuthal_mean(**kw_psi))
+
+        ds = self._uxds_face_with_more_vars_and_coords()
+        psi_result = ds['psi'].azimuthal_mean(**kw_psi)
+        psi2_result = ds['psi2'].azimuthal_mean(**kw_psi)
+        ds_result = ds.azimuthal_mean(**kw_psi)
+        assert ds_result['psi'].equals(psi_result)
+        assert ds_result['psi2'].equals(psi2_result)
+        assert 'unrelated_scalar_var' not in ds['psi'] and 'unrelated_scalar_var' not in psi_result.coords
+        assert ds_result['unrelated_scalar_var'].equals(ds['unrelated_scalar_var'])
+        assert 'unrelated_time_var' not in ds['psi'] and 'unrelated_time_var' not in psi_result.coords
+        assert ds_result['unrelated_time_var'].equals(ds['unrelated_time_var'])
+        assert ds_result.coords['unrelated_scalar_coord'].equals(ds.coords['unrelated_scalar_coord'])
+        assert ds_result.coords['unused_dim_coord'].equals(ds.coords['unused_dim_coord'])
+
+        # azimuthal mean doesn't support edge-centered data
+        ds = self._uxds_hex_edge()
+        arr = ds['random_data_edge']
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            arr.azimuthal_mean(**kw_hex)
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.azimuthal_mean(**kw_hex)
+
+        # azimuthal mean doesn't support node-centered data
+        ds = self._uxds_hex_node()
+        arr = ds['random_data_node']
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            arr.azimuthal_mean(**kw_hex)
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.azimuthal_mean(**kw_hex)
+
+        # also should crash if any data_vars contain unsupported centering,
+        # even if some data_vars are centered at supported location (faces).
+        ds = self._uxds_hex_face_and_node()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.azimuthal_mean(**kw_hex)
+        ds = self._uxds_hex_face_and_edge()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.azimuthal_mean(**kw_hex)
+
+    def test_uxds_mimics_uxda_weighted_mean(self):
+        """Ensure UxDataset.weighted_mean() mimics UxDataArray.weighted_mean() for each variable."""
+        ds = self._uxds_face_with_just_psi()
+        psi_result = ds['psi'].weighted_mean()
+        assert ds.weighted_mean()['psi'].equals(psi_result)
+        assert isinstance(psi_result, ux.UxDataArray)
+        assert isinstance(ds.weighted_mean(), ux.UxDataset)
+
+        ds = self._uxds_face_with_more_vars_and_coords()
+        psi_result = ds['psi'].weighted_mean()
+        psi2_result = ds['psi2'].weighted_mean()
+        ds_result = ds.weighted_mean()
+        assert ds_result['psi'].equals(psi_result)
+        assert ds_result['psi2'].equals(psi2_result)
+        assert 'unrelated_scalar_var' not in ds['psi'] and 'unrelated_scalar_var' not in psi_result.coords
+        assert ds_result['unrelated_scalar_var'].equals(ds['unrelated_scalar_var'])
+        assert 'unrelated_time_var' not in ds['psi'] and 'unrelated_time_var' not in psi_result.coords
+        assert ds_result['unrelated_time_var'].equals(ds['unrelated_time_var'])
+        assert ds_result.coords['unrelated_scalar_coord'].equals(ds.coords['unrelated_scalar_coord'])
+        assert ds_result.coords['unused_dim_coord'].equals(ds.coords['unused_dim_coord'])
+
+        # weighted_mean does support edge-centered data
+        ds = self._uxds_hex_edge()
+        arr = ds['random_data_edge']
+        arr_result = arr.weighted_mean()
+        ds_result = ds.weighted_mean()
+        assert ds_result['random_data_edge'].equals(arr_result)
+
+        # weighted_mean also supports ds with both edge & face data,
+        # but only if weights not provided
+        ds = self._uxds_hex_face_and_edge()
+        arr_face_result = ds['face_data'].weighted_mean()
+        arr_edge_result = ds['edge_data'].weighted_mean()
+        ds_result = ds.weighted_mean()
+        assert ds_result['face_data'].equals(arr_face_result)
+        assert ds_result['edge_data'].equals(arr_edge_result)
+        # when weights provided, usually make DataCenteringError:
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.weighted_mean(weights=np.arange(ds.sizes['n_face']))
+        # but, if weights are a 1D xr.DataArrray with dim='n_face' or 'n_edge',
+        # raise NotImplementedError instead (a clear unambiguous implementation is
+        # possible in these case, but just isn't implemented yet).
+        weights_face = xr.DataArray(np.arange(ds.sizes['n_face']), dims=['n_face'])
+        weights_edge = xr.DataArray(np.arange(ds.sizes['n_edge']), dims=['n_edge'])
+        with pytest.raises(NotImplementedError):
+            ds.weighted_mean(weights=weights_face)
+        with pytest.raises(NotImplementedError):
+            ds.weighted_mean(weights=weights_edge)
+
+        # weighted_mean doesn't support node-centered data
+        ds = self._uxds_hex_node()
+        arr = ds['random_data_node']
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            arr.weighted_mean()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.weighted_mean()
+
+        # also should crash if any data_vars contain unsupported centering,
+        # even if some data_vars are centered at supported location (faces).
+        ds = self._uxds_hex_face_and_node()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.weighted_mean()
+        ds = self._uxds_hex_node_and_edge()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.weighted_mean()
+        ds = self._uxds_hex_face_and_node_and_edge()
+        with pytest.raises(uxarray.errors.DataCenteringError):
+            ds.weighted_mean()
diff --git a/test/core/test_weighted_mean.py b/test/core/test_weighted_mean.py
index fa912b501..013ed5ec8 100644
--- a/test/core/test_weighted_mean.py
+++ b/test/core/test_weighted_mean.py
@@ -23,6 +23,10 @@ def test_quad_hex_face_centered(gridpath, datasetpath):
     # ensure values are within 3 decimal points of each other
     nt.assert_almost_equal(result.values, expected_weighted_mean, decimal=3)
 
+    # check can compute weighted mean from UxDataset too:
+    result_ds = uxds.weighted_mean()
+    assert result_ds['t2m'].equals(result)
+
 def test_quad_hex_face_centered_dask(gridpath, datasetpath):
     """Compares the weighted average computation for the quad hexagon grid
     using a face centered data variable on a dask-backed UxDataset & Grid to the expected value computed by
@@ -50,6 +54,10 @@ def test_quad_hex_face_centered_dask(gridpath, datasetpath):
     # ensure values are within 3 decimal points of each other
     nt.assert_almost_equal(computed_result.values, expected_weighted_mean, decimal=3)
 
+    # check can compute weighted mean from UxDataset too:
+    result_ds = uxds.weighted_mean().compute()
+    assert result_ds['t2m'].equals(computed_result)
+
 def test_quad_hex_edge_centered(gridpath, test_data_dir):
     """Compares the weighted average computation for the quad hexagon grid
     using an edge centered data variable to the expected value computed by
@@ -65,6 +73,10 @@ def test_quad_hex_edge_centered(gridpath, test_data_dir):
 
     nt.assert_equal(result, expected_weighted_mean)
 
+    # check can compute weighted mean from UxDataset too:
+    result_ds = uxds.weighted_mean()
+    assert result_ds['random_data_edge'].equals(result)
+
 def test_quad_hex_edge_centered_dask(gridpath, test_data_dir):
     """Compares the weighted average computation for the quad hexagon grid
     using an edge centered data variable on a dask-backed UxDataset & Grid to the expected value computed by
@@ -89,12 +101,15 @@ def test_quad_hex_edge_centered_dask(gridpath, test_data_dir):
     assert isinstance(computed_result.data, np.ndarray)
 
     # expected weighted average computed by hand
-    expected_weighted_mean = (uxds[
-                                  'random_data_edge'].values * uxds.uxgrid.edge_node_distances).sum() / uxds.uxgrid.edge_node_distances.sum()
+    expected_weighted_mean = (uxds['random_data_edge'].values * uxds.uxgrid.edge_node_distances).sum() / uxds.uxgrid.edge_node_distances.sum()
 
     # ensure values are within 3 decimal points of each other
     nt.assert_almost_equal(computed_result.values, expected_weighted_mean, decimal=3)
 
+    # check can compute weighted mean from UxDataset too:
+    result_ds = uxds.weighted_mean().compute()
+    assert result_ds['random_data_edge'].equals(computed_result)
+
 def test_weighted_mean_crash_if_node_centered_and_no_weights():
     """Ensure weighted_mean crashes if the data is node-centered and no weights are provided.
     Regression test for bug 2 of issue #1800.

From af07f74ab3bd9df34ed65e47d4c40317f9842f0e Mon Sep 17 00:00:00 2001
From: Sam Evans <47793072+Sevans711@users.noreply.github.com>
Date: Thu, 1 Oct 2026 15:47:24 -0400
Subject: [PATCH 4/6] add user guide examples for UxDataset methods

---
 docs/user-guide/azimuthal-average.ipynb | 20 ++++++++
 docs/user-guide/weighted_mean.ipynb     | 61 +++++++++++++++++++++++++
 docs/user-guide/zonal-average.ipynb     | 23 +++++++++-
 3 files changed, 103 insertions(+), 1 deletion(-)

diff --git a/docs/user-guide/azimuthal-average.ipynb b/docs/user-guide/azimuthal-average.ipynb
index 2f439adae..b113b1f92 100644
--- a/docs/user-guide/azimuthal-average.ipynb
+++ b/docs/user-guide/azimuthal-average.ipynb
@@ -138,6 +138,26 @@
     "azim_mean_psi"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "1f2b787a",
+   "metadata": {},
+   "source": [
+    "`azimuthal_mean()` can also be called directly on a UxDataset to compute the zonal mean for every face-centered variable (the example dataset here has only 1 variable, but this can be especially convenient for a dataset with many variables), although it doesn't yet support the `return_hit_counts` parameter:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24048119",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "azim_mean_uxds = uxds.azimuthal_mean()\n",
+    "print('Equivalent to UxDataArray.azimuthal_mean:', azim_mean_uxds[\"psi\"].equals(azim_mean_psi)\n",
+    "azim_mean_uxds"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "9ecf6cb2-9575-4e13-8dea-d8ec50149a7a",
diff --git a/docs/user-guide/weighted_mean.ipynb b/docs/user-guide/weighted_mean.ipynb
index b15fdb5b2..49d9cff35 100644
--- a/docs/user-guide/weighted_mean.ipynb
+++ b/docs/user-guide/weighted_mean.ipynb
@@ -20,6 +20,7 @@
     "import warnings\n",
     "\n",
     "import cartopy.crs as ccrs\n",
+    "import numpy as np\n",
     "import xarray as xr\n",
     "\n",
     "import uxarray as ux\n",
@@ -179,6 +180,66 @@
     "unweighted_result = uxds_edge[\"random_data_edge\"].mean()\n",
     "unweighted_result.values"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e32ecee9",
+   "metadata": {},
+   "source": [
+    "## Weighted Mean of a UxDataset\n",
+    "\n",
+    "It is also simple to take the weighted mean all relevant data_vars in a UxDataset (skipping anything without a grid dimension).  \n",
+    "Non-grid dimensions are also handled in the intuitive way."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "272a836a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# construct an example dataset, with nontrivial extra dimensions, vars, and coords:\n",
+    "edge_data = uxds_edge[\"random_data_edge\"].compute()  # so they are visible below.\n",
+    "face_data = uxds_face[\"random_data_face\"].compute()  # compute() to load values,\n",
+    "times = xr.DataArray([0,1,2], dims='time', coords={'time': 10*np.arange(3)})\n",
+    "uxds = ux.UxDataset({\"edge_data\": edge_data,\n",
+    "                     \"face_data\": face_data,\n",
+    "                     \"face_time_data\": face_data * times,\n",
+    "                     \"times\": times,\n",
+    "                     \"time_plus_5\": times + 5,\n",
+    "                     \"scalar_var\": 7},\n",
+    "                     coords={\"scalar_coord\": 9},\n",
+    "                     uxgrid = uxds_face.uxgrid)\n",
+    "uxds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2a371bdc",
+   "metadata": {},
+   "source": [
+    "Taking the weighted mean aggregates values along the grid dimensions (\"n_face\" and \"n_edge\") appropriately but keeps all other dimensions, variables, and coordinates unchanged."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3b5dd3f2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result = uxds.weighted_mean()\n",
+    "result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "78915269",
+   "metadata": {},
+   "source": [
+    "Note in particular that the relevant results here agree with corresponding results above. The result's \"face_data\" and \"edge_data\" agree with the face and edge examples above, and its \"face_time_data\" values are equal to `[0,1,2] * result[\"face_data\"]` because `uxds[\"face_time_data\"]` was constructed by multiplying the face data by `times` (==`[0,1,2]`)."
+   ]
   }
  ],
  "metadata": {
diff --git a/docs/user-guide/zonal-average.ipynb b/docs/user-guide/zonal-average.ipynb
index 88471d8f2..cab9fdce1 100644
--- a/docs/user-guide/zonal-average.ipynb
+++ b/docs/user-guide/zonal-average.ipynb
@@ -22,7 +22,8 @@
     "- [3. Conservative Zonal Averaging](#3-conservative-zonal-averaging) — area-weighted bands, conservation checks, and comparisons.\n",
     "- [4. Combined Plots](#4-combined-plots) — pair global maps with their zonal means for context.\n",
     "- [5. HEALPix Zonal Averaging (Conservative vs Non-Conservative)](#5-healpix-zonal-averaging-conservative-vs-non-conservative) — run the same workflow on a different grid.\n",
-    "- [6. 2D Zonal Means on NE30 (RELHUM)](#6-2d-zonal-means-on-ne30-relhum) — build latitude–height slices and inspect the differences.\n"
+    "- [6. 2D Zonal Means on NE30 (RELHUM)](#6-2d-zonal-means-on-ne30-relhum) — build latitude–height slices and inspect the differences.\n",
+    "- [7. Zonal Anomalies](#7-zonal-anomalies) - values on each face minus zonal mean within the corresponding latitude band"
    ]
   },
   {
@@ -120,6 +121,26 @@
     "zonal_mean_psi"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "5223d24c",
+   "metadata": {},
+   "source": [
+    "`zonal_mean()` can also be called directly on a UxDataset to compute the zonal mean for every face-centered variable (the example dataset here has only 1 variable, but this can be especially convenient for a dataset with many variables):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "231d67b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zonal_mean_uxds = uxds.zonal_mean()\n",
+    "print('Equivalent to UxDataArray.zonal_mean:', zonal_mean_uxds[\"psi\"].equals(zonal_mean_psi)\n",
+    "zonal_mean_uxds"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "65194a38c76e8a62",

From 69aa5233fe76a25b04e0c7d6dd9b77869ee063ba Mon Sep 17 00:00:00 2001
From: Sam Evans <47793072+Sevans711@users.noreply.github.com>
Date: Thu, 1 Oct 2026 15:51:04 -0400
Subject: [PATCH 5/6] pre-commit ruff formatting

---
 docs/user-guide/azimuthal-average.ipynb |  5 +-
 docs/user-guide/weighted_mean.ipynb     | 22 +++---
 docs/user-guide/zonal-average.ipynb     | 91 +++++++++++++++++++------
 3 files changed, 87 insertions(+), 31 deletions(-)

diff --git a/docs/user-guide/azimuthal-average.ipynb b/docs/user-guide/azimuthal-average.ipynb
index b113b1f92..36660b092 100644
--- a/docs/user-guide/azimuthal-average.ipynb
+++ b/docs/user-guide/azimuthal-average.ipynb
@@ -154,7 +154,10 @@
    "outputs": [],
    "source": [
     "azim_mean_uxds = uxds.azimuthal_mean()\n",
-    "print('Equivalent to UxDataArray.azimuthal_mean:', azim_mean_uxds[\"psi\"].equals(azim_mean_psi)\n",
+    "print(\n",
+    "    \"Equivalent to UxDataArray.azimuthal_mean:\",\n",
+    "    azim_mean_uxds[\"psi\"].equals(azim_mean_psi),\n",
+    ")\n",
     "azim_mean_uxds"
    ]
   },
diff --git a/docs/user-guide/weighted_mean.ipynb b/docs/user-guide/weighted_mean.ipynb
index 49d9cff35..d332b6bc9 100644
--- a/docs/user-guide/weighted_mean.ipynb
+++ b/docs/user-guide/weighted_mean.ipynb
@@ -202,15 +202,19 @@
     "# construct an example dataset, with nontrivial extra dimensions, vars, and coords:\n",
     "edge_data = uxds_edge[\"random_data_edge\"].compute()  # so they are visible below.\n",
     "face_data = uxds_face[\"random_data_face\"].compute()  # compute() to load values,\n",
-    "times = xr.DataArray([0,1,2], dims='time', coords={'time': 10*np.arange(3)})\n",
-    "uxds = ux.UxDataset({\"edge_data\": edge_data,\n",
-    "                     \"face_data\": face_data,\n",
-    "                     \"face_time_data\": face_data * times,\n",
-    "                     \"times\": times,\n",
-    "                     \"time_plus_5\": times + 5,\n",
-    "                     \"scalar_var\": 7},\n",
-    "                     coords={\"scalar_coord\": 9},\n",
-    "                     uxgrid = uxds_face.uxgrid)\n",
+    "times = xr.DataArray([0, 1, 2], dims=\"time\", coords={\"time\": 10 * np.arange(3)})\n",
+    "uxds = ux.UxDataset(\n",
+    "    {\n",
+    "        \"edge_data\": edge_data,\n",
+    "        \"face_data\": face_data,\n",
+    "        \"face_time_data\": face_data * times,\n",
+    "        \"times\": times,\n",
+    "        \"time_plus_5\": times + 5,\n",
+    "        \"scalar_var\": 7,\n",
+    "    },\n",
+    "    coords={\"scalar_coord\": 9},\n",
+    "    uxgrid=uxds_face.uxgrid,\n",
+    ")\n",
     "uxds"
    ]
   },
diff --git a/docs/user-guide/zonal-average.ipynb b/docs/user-guide/zonal-average.ipynb
index cab9fdce1..6145b4ecc 100644
--- a/docs/user-guide/zonal-average.ipynb
+++ b/docs/user-guide/zonal-average.ipynb
@@ -137,7 +137,10 @@
    "outputs": [],
    "source": [
     "zonal_mean_uxds = uxds.zonal_mean()\n",
-    "print('Equivalent to UxDataArray.zonal_mean:', zonal_mean_uxds[\"psi\"].equals(zonal_mean_psi)\n",
+    "print(\n",
+    "    \"Equivalent to UxDataArray.zonal_mean:\",\n",
+    "    zonal_mean_uxds[\"psi\"].equals(zonal_mean_psi),\n",
+    ")\n",
     "zonal_mean_uxds"
    ]
   },
@@ -770,58 +773,104 @@
   {
    "cell_type": "markdown",
    "id": "1af6beaf",
-   "source": "## 7. Zonal Anomalies\n\nA zonal anomaly is the per-face departure from the mean of its latitude band. `zonal_anomaly` returns a `UxDataArray` with the same dims and dtype as the input (integer dtypes are promoted to float so empty bands can hold `NaN`).\n\n- **Centroid mode** (`conservative=False`, default): each face is assigned to one band by its centroid latitude (`np.digitize`). The unweighted per-band mean is exactly zero.\n- **Conservative mode** (`conservative=True`): faces straddling band edges contribute to multiple bands by area overlap (reusing the `zonal_mean` weight kernel), so per-band means are small but not exactly zero.\n\n### Step 7.1: Compute the centroid-mode anomaly",
-   "metadata": {}
+   "metadata": {},
+   "source": "## 7. Zonal Anomalies\n\nA zonal anomaly is the per-face departure from the mean of its latitude band. `zonal_anomaly` returns a `UxDataArray` with the same dims and dtype as the input (integer dtypes are promoted to float so empty bands can hold `NaN`).\n\n- **Centroid mode** (`conservative=False`, default): each face is assigned to one band by its centroid latitude (`np.digitize`). The unweighted per-band mean is exactly zero.\n- **Conservative mode** (`conservative=True`): faces straddling band edges contribute to multiple bands by area overlap (reusing the `zonal_mean` weight kernel), so per-band means are small but not exactly zero.\n\n### Step 7.1: Compute the centroid-mode anomaly"
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "c0347122",
-   "source": "anomaly = uxds[\"psi\"].zonal_anomaly(lat=(-90, 90, 10))\nanomaly",
    "metadata": {},
-   "execution_count": null,
-   "outputs": []
+   "outputs": [],
+   "source": [
+    "anomaly = uxds[\"psi\"].zonal_anomaly(lat=(-90, 90, 10))\n",
+    "anomaly"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "1783125a",
-   "source": "### Step 7.2: Verify the per-band sum-to-zero property\n\nIn centroid mode every populated band has an unweighted mean of exactly zero. We can confirm by binning faces the same way `zonal_anomaly` does (`np.digitize`) and reducing each band.",
-   "metadata": {}
+   "metadata": {},
+   "source": "### Step 7.2: Verify the per-band sum-to-zero property\n\nIn centroid mode every populated band has an unweighted mean of exactly zero. We can confirm by binning faces the same way `zonal_anomaly` does (`np.digitize`) and reducing each band."
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "ed0f731f",
-   "source": "bands = np.arange(-90, 91, 10)\nface_lat = uxds.uxgrid.face_lat.values\nband_idx = np.clip(np.digitize(face_lat, bands) - 1, 0, len(bands) - 2)\n\nfor bi in range(len(bands) - 1):\n    mask = band_idx == bi\n    if not mask.any():\n        continue\n    band_mean = float(anomaly.values[mask].mean())\n    print(\n        f\"band [{bands[bi]:+4d}, {bands[bi + 1]:+4d})  \"\n        f\"n_faces={int(mask.sum()):4d}  mean={band_mean:+.2e}\"\n    )",
    "metadata": {},
-   "execution_count": null,
-   "outputs": []
+   "outputs": [],
+   "source": [
+    "bands = np.arange(-90, 91, 10)\n",
+    "face_lat = uxds.uxgrid.face_lat.values\n",
+    "band_idx = np.clip(np.digitize(face_lat, bands) - 1, 0, len(bands) - 2)\n",
+    "\n",
+    "for bi in range(len(bands) - 1):\n",
+    "    mask = band_idx == bi\n",
+    "    if not mask.any():\n",
+    "        continue\n",
+    "    band_mean = float(anomaly.values[mask].mean())\n",
+    "    print(\n",
+    "        f\"band [{bands[bi]:+4d}, {bands[bi + 1]:+4d})  \"\n",
+    "        f\"n_faces={int(mask.sum()):4d}  mean={band_mean:+.2e}\"\n",
+    "    )"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "41dbb45d",
-   "source": "### Step 7.3: Visualize the anomaly field\n\nPlot the anomaly map and the zonal mean it was subtracted from side by side using a diverging colormap centered at zero.",
-   "metadata": {}
+   "metadata": {},
+   "source": "### Step 7.3: Visualize the anomaly field\n\nPlot the anomaly map and the zonal mean it was subtracted from side by side using a diverging colormap centered at zero."
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "84cd278e",
-   "source": "vmax = float(np.nanmax(np.abs(anomaly.values)))\nanomaly_map = anomaly.plot(\n    cmap=\"RdBu_r\",\n    periodic_elements=\"split\",\n    clim=(-vmax, vmax),\n    title=\"Zonal Anomaly (psi - zonal mean)\",\n).opts(width=525, height=400, colorbar=True)\n\nzm = uxds[\"psi\"].zonal_mean(lat=(-90, 90, 10))\nzm_df = zm.to_dataframe(name=\"zonal_mean\").reset_index()\nzm_panel = zm_df.hvplot.line(\n    x=\"zonal_mean\",\n    y=\"latitudes\",\n    line_width=2,\n    title=\"Zonal Mean (subtracted)\",\n    ylim=(-90, 90),\n    width=400,\n    height=400,\n).opts(show_grid=True)\n\n(anomaly_map + zm_panel).cols(2)",
    "metadata": {},
-   "execution_count": null,
-   "outputs": []
+   "outputs": [],
+   "source": [
+    "vmax = float(np.nanmax(np.abs(anomaly.values)))\n",
+    "anomaly_map = anomaly.plot(\n",
+    "    cmap=\"RdBu_r\",\n",
+    "    periodic_elements=\"split\",\n",
+    "    clim=(-vmax, vmax),\n",
+    "    title=\"Zonal Anomaly (psi - zonal mean)\",\n",
+    ").opts(width=525, height=400, colorbar=True)\n",
+    "\n",
+    "zm = uxds[\"psi\"].zonal_mean(lat=(-90, 90, 10))\n",
+    "zm_df = zm.to_dataframe(name=\"zonal_mean\").reset_index()\n",
+    "zm_panel = zm_df.hvplot.line(\n",
+    "    x=\"zonal_mean\",\n",
+    "    y=\"latitudes\",\n",
+    "    line_width=2,\n",
+    "    title=\"Zonal Mean (subtracted)\",\n",
+    "    ylim=(-90, 90),\n",
+    "    width=400,\n",
+    "    height=400,\n",
+    ").opts(show_grid=True)\n",
+    "\n",
+    "(anomaly_map + zm_panel).cols(2)"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "15c921d4",
-   "source": "### Step 7.4: Compare centroid vs conservative anomaly\n\nConservative mode blends straddling faces across band edges, so its anomaly differs slightly from the centroid version. The difference shows where face geometry crosses band boundaries.",
-   "metadata": {}
+   "metadata": {},
+   "source": "### Step 7.4: Compare centroid vs conservative anomaly\n\nConservative mode blends straddling faces across band edges, so its anomaly differs slightly from the centroid version. The difference shows where face geometry crosses band boundaries."
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "36a9096c",
-   "source": "anomaly_cons = uxds[\"psi\"].zonal_anomaly(lat=(-90, 90, 10), conservative=True)\ndiff = anomaly_cons - anomaly\n\nprint(f\"max |centroid|     = {float(np.nanmax(np.abs(anomaly.values))):.4f}\")\nprint(f\"max |conservative| = {float(np.nanmax(np.abs(anomaly_cons.values))):.4f}\")\nprint(f\"max |difference|   = {float(np.nanmax(np.abs(diff.values))):.4f}\")",
    "metadata": {},
-   "execution_count": null,
-   "outputs": []
+   "outputs": [],
+   "source": [
+    "anomaly_cons = uxds[\"psi\"].zonal_anomaly(lat=(-90, 90, 10), conservative=True)\n",
+    "diff = anomaly_cons - anomaly\n",
+    "\n",
+    "print(f\"max |centroid|     = {float(np.nanmax(np.abs(anomaly.values))):.4f}\")\n",
+    "print(f\"max |conservative| = {float(np.nanmax(np.abs(anomaly_cons.values))):.4f}\")\n",
+    "print(f\"max |difference|   = {float(np.nanmax(np.abs(diff.values))):.4f}\")"
+   ]
   }
  ],
  "metadata": {

From f244f20312b5cdebd7b86cd32ff4dbec222532c7 Mon Sep 17 00:00:00 2001
From: Sam Evans <47793072+Sevans711@users.noreply.github.com>
Date: Thu, 1 Oct 2026 17:58:56 -0400
Subject: [PATCH 6/6] add UxDataset zonal and means to api.rst

---
 docs/api.rst | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/docs/api.rst b/docs/api.rst
index 3e8f59b82..b46ab1cef 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -568,6 +568,7 @@ Azimuthal aggregations apply an aggregation (i.e. averaging) along circles of co
    :toctree: generated/
 
    UxDataArray.azimuthal_mean
+   UxDataset.azimuthal_mean
 
 
 Neighborhood
@@ -630,6 +631,9 @@ Zonal Average
    UxDataArray.zonal_average
    UxDataArray.zonal_mean
    UxDataArray.zonal_anomaly
+   UxDataset.zonal_average
+   UxDataset.zonal_mean
+   UxDataset.zonal_anomaly
 
 
 Weighted
@@ -638,6 +642,7 @@ Weighted
    :toctree: generated/
 
    UxDataArray.weighted_mean
+   UxDataset.weighted_mean
 
 
 Spherical Geometry