Merge pull request #1345 from UXARRAY/rajeeja/zonal-mean-analytic-band

conservative zonal means

Author: rajeeja

docs/user-guide/zonal-average.ipynb

@@ -92,3 +92,114 @@ def test_mismatched_dims():
 
     assert za.shape == (10, 19, 5)
     assert za.dims[1] == "latitudes"
+
+
+class TestConservativeZonalMean:
+    """Test conservative zonal mean functionality."""
+
+    def test_conservative_zonal_mean_basic(self, gridpath, datasetpath):
+        """Test basic conservative zonal mean with bands."""
+        grid_path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+        data_path = datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc")
+        uxds = ux.open_dataset(grid_path, data_path)
+
+        # Test with explicit bands
+        bands = np.array([-90, -30, 0, 30, 90])
+        result = uxds["psi"].zonal_mean(lat=bands, conservative=True)
+
+        # Should have one less value than bands (4 bands from 5 edges)
+        assert result.shape == (len(bands) - 1,)
+        assert np.all(np.isfinite(result.values))
+
+    def test_conservative_float_step_size(self, gridpath, datasetpath):
+        """Test conservative zonal mean with float step sizes."""
+        grid_path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+        data_path = datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc")
+        uxds = ux.open_dataset(grid_path, data_path)
+
+        # Test with float step size (5.5 degrees)
+        result = uxds["psi"].zonal_mean(lat=(-90, 90, 0.05), conservative=True)
+
+        # Should get valid results
+        assert len(result) > 0
+        assert np.all(np.isfinite(result.values))
+
+        # Test with reasonable float step size (no warning)
+        result = uxds["psi"].zonal_mean(lat=(-90, 90, 5.5), conservative=True)
+        expected_n_bands = int(np.ceil(180 / 5.5))
+        assert result.shape[0] == expected_n_bands
+        assert np.all(np.isfinite(result.values))
+
+    def test_conservative_near_pole(self, gridpath, datasetpath):
+        """Test conservative zonal mean with bands near the poles."""
+        grid_path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+        data_path = datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc")
+        uxds = ux.open_dataset(grid_path, data_path)
+
+        # Test near north pole with float step
+        bands_north = np.array([85.0, 87.5, 90.0])
+        result_north = uxds["psi"].zonal_mean(lat=bands_north, conservative=True)
+        assert result_north.shape == (2,)
+        assert np.all(np.isfinite(result_north.values))
+
+        # Test near south pole with float step
+        bands_south = np.array([-90.0, -87.5, -85.0])
+        result_south = uxds["psi"].zonal_mean(lat=bands_south, conservative=True)
+        assert result_south.shape == (2,)
+        assert np.all(np.isfinite(result_south.values))
+
+        # Test spanning pole with non-integer step
+        bands_span = np.array([88.5, 89.25, 90.0])
+        result_span = uxds["psi"].zonal_mean(lat=bands_span, conservative=True)
+        assert result_span.shape == (2,)
+        assert np.all(np.isfinite(result_span.values))
+
+    def test_conservative_step_size_validation(self, gridpath, datasetpath):
+        """Test that step size validation works correctly."""
+        grid_path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+        data_path = datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc")
+        uxds = ux.open_dataset(grid_path, data_path)
+
+        # Test negative step size
+        with pytest.raises(ValueError, match="Step size must be positive"):
+            uxds["psi"].zonal_mean(lat=(-90, 90, -10), conservative=True)
+
+        # Test zero step size
+        with pytest.raises(ValueError, match="Step size must be positive"):
+            uxds["psi"].zonal_mean(lat=(-90, 90, 0), conservative=True)
+
+    def test_conservative_full_sphere_conservation(self, gridpath, datasetpath):
+        """Test that single band covering entire sphere conserves global mean."""
+        grid_path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+        data_path = datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc")
+        uxds = ux.open_dataset(grid_path, data_path)
+
+        # Single band covering entire sphere
+        bands = np.array([-90, 90])
+        result = uxds["psi"].zonal_mean(lat=bands, conservative=True)
+
+        # Compare with global mean
+        global_mean = uxds["psi"].mean()
+
+        assert result.shape == (1,)
+        assert result.values[0] == pytest.approx(global_mean.values, rel=0.01)
+
+    def test_conservative_vs_nonconservative_comparison(self, gridpath, datasetpath):
+        """Compare conservative and non-conservative methods."""
+        grid_path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+        data_path = datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc")
+        uxds = ux.open_dataset(grid_path, data_path)
+
+        # Non-conservative at band centers
+        lat_centers = np.array([-60, 0, 60])
+        non_conservative = uxds["psi"].zonal_mean(lat=lat_centers)
+
+        # Conservative with bands
+        bands = np.array([-90, -30, 30, 90])
+        conservative = uxds["psi"].zonal_mean(lat=bands, conservative=True)
+
+        # Results should be similar but not identical
+        assert non_conservative.shape == conservative.shape
+        # Check they are in the same ballpark
+        assert np.all(np.abs(conservative.values - non_conservative.values) <
+                     np.abs(non_conservative.values) * 0.5)

test/grid/integrate/test_zonal.py

@@ -76,17 +76,19 @@ def test_dataarray_methods(gridpath, datasetpath):
     # plot.scatter() is an xarray method
     assert hasattr(uxds.plot, 'scatter')
 
+import hvplot.xarray  # registers .hvplot accessor
+
 def test_line(gridpath):
     mesh_path = gridpath("mpas", "QU", "oQU480.231010.nc")
     uxds = ux.open_dataset(mesh_path, mesh_path)
-    _plot_line = uxds['bottomDepth'].zonal_average().plot.line()
+    _plot_line = uxds['bottomDepth'].zonal_average().hvplot.line()
     assert isinstance(_plot_line, hv.Curve)
 
 def test_scatter(gridpath):
     mesh_path = gridpath("mpas", "QU", "oQU480.231010.nc")
     uxds = ux.open_dataset(mesh_path, mesh_path)
-    _plot_line = uxds['bottomDepth'].zonal_average().plot.scatter()
-    assert isinstance(_plot_line, hv.Scatter)
+    _plot_scatter = uxds['bottomDepth'].zonal_average().hvplot.scatter()
+    assert isinstance(_plot_scatter, hv.Scatter)
 
 
 

test/test_plot.py

@@ -20,7 +20,10 @@
     _compute_gradient,
 )
 from uxarray.core.utils import _map_dims_to_ugrid
-from uxarray.core.zonal import _compute_non_conservative_zonal_mean
+from uxarray.core.zonal import (
+    _compute_conservative_zonal_mean_bands,
+    _compute_non_conservative_zonal_mean,
+)
 from uxarray.cross_sections import UxDataArrayCrossSectionAccessor
 from uxarray.formatting_html import array_repr
 from uxarray.grid import Grid
@@ -510,16 +513,20 @@ def integrate(
 
         return uxda
 
-    def zonal_mean(self, lat=(-90, 90, 10), **kwargs):
-        """Compute averages along lines of constant latitude.
+    def zonal_mean(self, lat=(-90, 90, 10), conservative: bool = False, **kwargs):
+        """Compute non-conservative or conservative averages along lines of constant latitude or latitude bands.
 
         Parameters
         ----------
         lat : tuple, float, or array-like, default=(-90, 90, 10)
-            Latitude values in degrees. Can be specified as:
-            - tuple (start, end, step): Computes means at intervals of `step` in range [start, end]
-            - float: Computes mean for a single latitude
-            - array-like: Computes means for each specified latitude
+            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.
+        conservative : bool, default=False
+            If True, performs conservative (area-weighted) zonal averaging over latitude bands.
+            If False, performs traditional (non-conservative) averaging at latitude lines.
 
         Returns
         -------
@@ -529,62 +536,125 @@ def zonal_mean(self, lat=(-90, 90, 10), **kwargs):
 
         Examples
         --------
-        # All latitudes from -90° to 90° at 10° intervals
+        # Non-conservative averaging from -90° to 90° at 10° intervals by default
         >>> uxds["var"].zonal_mean()
 
-        # Single latitude at 30°
+        # Single latitude (non-conservative) over 30° latitude
         >>> uxds["var"].zonal_mean(lat=30.0)
 
-        # Range from -60° to 60° at 10° intervals
-        >>> uxds["var"].zonal_mean(lat=(-60, 60, 10))
+        # Conservative averaging over latitude bands
+        >>> uxds["var"].zonal_mean(lat=(-60, 60, 10), conservative=True)
+
+        # Conservative with explicit band edges
+        >>> uxds["var"].zonal_mean(lat=[-90, -30, 0, 30, 90], conservative=True)
 
         Notes
         -----
-        Only supported for face-centered data variables. Candidate faces are determined
-        using spherical bounding boxes - faces whose bounds contain the target latitude
-        are included in calculations.
+        Only supported for face-centered data variables.
+
+        Conservative averaging preserves integral quantities and is recommended for
+        physical analysis. Non-conservative averaging samples at latitude lines.
         """
         if not self._face_centered():
             raise ValueError(
                 "Zonal mean computations are currently only supported for face-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])
-            latitudes = np.clip(latitudes, -90, 90)
-        elif isinstance(lat, (float, int)):
-            # zonal mean over a single latitude
-            latitudes = [lat]
-        elif isinstance(lat, (list, np.ndarray)):
-            # zonal mean over an array of arbitrary latitudes
-            latitudes = np.asarray(lat)
-        else:
-            raise ValueError(
-                "Invalid value for 'lat' provided. Must either be a single scalar value, tuple (min_lat, max_lat, step), or array-like."
+        face_axis = self.dims.index("n_face")
+
+        if not conservative:
+            # Non-conservative (traditional) zonal averaging
+            if isinstance(lat, tuple):
+                start, end, step = lat
+                if step <= 0:
+                    raise ValueError("Step size must be positive.")
+                if step < 0.1:
+                    warnings.warn(
+                        f"Very small step size ({step}°) may lead to performance issues...",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                num_points = int(round((end - start) / step)) + 1
+                latitudes = np.linspace(start, end, num_points)
+                latitudes = np.clip(latitudes, -90, 90)
+            elif isinstance(lat, (float, int)):
+                latitudes = [lat]
+            elif isinstance(lat, (list, np.ndarray)):
+                latitudes = np.asarray(lat)
+            else:
+                raise ValueError(
+                    "Invalid value for 'lat' provided. Must be a scalar, tuple (min_lat, max_lat, step), or array-like."
+                )
+
+            res = _compute_non_conservative_zonal_mean(
+                uxda=self, latitudes=latitudes, **kwargs
             )
 
-        res = _compute_non_conservative_zonal_mean(
-            uxda=self, latitudes=latitudes, **kwargs
-        )
+            dims = list(self.dims)
+            dims[face_axis] = "latitudes"
+
+            return xr.DataArray(
+                res,
+                dims=dims,
+                coords={"latitudes": latitudes},
+                name=self.name + "_zonal_mean"
+                if self.name is not None
+                else "zonal_mean",
+                attrs={"zonal_mean": True, "conservative": False},
+            )
 
-        face_axis = self.dims.index("n_face")
-        dims = list(self.dims)
-        dims[face_axis] = "latitudes"
+        else:
+            # Conservative zonal averaging
+            if isinstance(lat, tuple):
+                start, end, step = lat
+                if step <= 0:
+                    raise ValueError(
+                        "Step size must be positive for conservative averaging."
+                    )
+                if step < 0.1:
+                    warnings.warn(
+                        f"Very small step size ({step}°) may lead to performance issues...",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                num_points = int(round((end - start) / step)) + 1
+                edges = np.linspace(start, end, num_points)
+                edges = np.clip(edges, -90, 90)
+            elif isinstance(lat, (list, np.ndarray)):
+                edges = np.asarray(lat, dtype=float)
+            else:
+                raise ValueError(
+                    "For conservative averaging, 'lat' must be a tuple (start, end, step) or array-like band edges."
+                )
 
-        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",
-            attrs={"zonal_mean": True},
-        )
+            if edges.ndim != 1 or edges.size < 2:
+                raise ValueError("Band edges must be 1D with at least two values")
 
-        return uxda
+            res = _compute_conservative_zonal_mean_bands(self, edges)
+
+            # Use band centers as coordinate values
+            centers = 0.5 * (edges[:-1] + edges[1:])
+
+            dims = list(self.dims)
+            dims[face_axis] = "latitudes"
+
+            return xr.DataArray(
+                res,
+                dims=dims,
+                coords={"latitudes": centers},
+                name=self.name + "_zonal_mean"
+                if self.name is not None
+                else "zonal_mean",
+                attrs={
+                    "zonal_mean": True,
+                    "conservative": True,
+                    "lat_band_edges": edges,
+                },
+            )
 
-    # Alias for 'zonal_mean', since this name is also commonly used.
-    zonal_average = zonal_mean
+    def zonal_average(self, lat=(-90, 90, 10), conservative: bool = False, **kwargs):
+        """Alias of zonal_mean; prefer `zonal_mean` for primary API."""
+        return self.zonal_mean(lat=lat, conservative=conservative, **kwargs)
 
     def azimuthal_mean(
         self,