Add Azimuthal Averaging (#1255)

* created docstring for azimuthal_mean

* wrote azimuthal mean computation in public function

* draft azimuthal mean ready to test

* fixed typos, return hit count for radial bins

* made azimuthal mean more robust to axis ordering

* added central coord as attribute in output UxDataArray of azimuthal_mean

* run pre-commit

* update parameters and set default values to nan

* add tests, update to use kdtree

* update api

* update type hint

* Update docs/api.rst

Co-authored-by: Copilot <[email protected]>

* Update docs/api.rst

Co-authored-by: Copilot <[email protected]>

* Update test/core/test_azimuthal.py

Co-authored-by: Copilot <[email protected]>

* return an xr.DataArray, update hit counts

* add validation

* fix api.rst

---------

Co-authored-by: Philip Chmielowiec <[email protected]>
Co-authored-by: Philip Chmielowiec <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

docs/api.rst

@@ -453,6 +453,17 @@ on each face.
    UxDataArray.topological_all
    UxDataArray.topological_any
 
+Azimuthal
+~~~~~~~~~
+
+Azimuthal aggregations apply an aggregation (i.e. averaging) along circles of constant great-circle distance from a specified point on the sphere.
+
+
+.. autosummary::
+   :toctree: generated/
+
+   UxDataArray.azimuthal_mean
+
 Zonal Average
 ~~~~~~~~~~~~~
 .. autosummary::

test/core/test_azimuthal.py

@@ -0,0 +1,68 @@
+import pytest
+import uxarray as ux
+import numpy as np
+
+
+
+def test_gaussian(gridpath, datasetpath):
+    uxds = ux.open_dataset(
+        gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+        datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+    )
+
+    res = uxds['gaussian'].azimuthal_mean(center_coord=(45, 0), outer_radius=2, radius_step=0.5)
+
+    # Expects decreasing values from center
+    valid_vals = res[1:]
+
+    np.testing.assert_array_less(
+        valid_vals.diff("radius").values, 1e-12
+    )
+
+
+
+def test_inverse_gaussian(gridpath, datasetpath):
+
+    uxds =  ux.open_dataset(
+        gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+        datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+    )
+
+    res = uxds['inverse_gaussian'].azimuthal_mean(center_coord=(45, 0), outer_radius=2, radius_step=0.5)
+
+    # Expects increasing values from center
+    atol = 1e-12
+    diffs = res[1:].diff("radius").values
+    diffs = diffs[np.isfinite(diffs)]
+    np.testing.assert_array_less(-atol, diffs)
+
+def test_non_zero_hit_counts(gridpath, datasetpath):
+    uxds =  ux.open_dataset(
+        gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+        datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+    )
+
+    res, hit_counts = uxds['inverse_gaussian'].azimuthal_mean(center_coord=(45, 0), outer_radius=2, radius_step=0.5, return_hit_counts=True)
+
+    # At least one hit after the first circle
+    assert np.all(hit_counts[1:] > 1)
+
+    assert 'radius' in hit_counts.dims
+    assert hit_counts.sizes['radius'] == res.sizes['radius']
+
+def test_zero_hit_counts(gridpath, datasetpath):
+    uxds =  ux.open_dataset(
+        gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+        datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+    )
+
+    # Outside of grid domain
+    res, hit_counts = uxds['inverse_gaussian'].azimuthal_mean(center_coord=(-45, 0), outer_radius=2, radius_step=0.5, return_hit_counts=True)
+
+    assert 'radius' in hit_counts.dims
+    assert hit_counts.sizes['radius'] == res.sizes['radius']
+
+    # No hits
+    assert np.all(hit_counts == 0)
+
+    print(res)

uxarray/core/dataarray.py

@@ -588,6 +588,127 @@ def zonal_mean(self, lat=(-90, 90, 10), **kwargs):
     # Alias for 'zonal_mean', since this name is also commonly used.
     zonal_average = zonal_mean
 
+    def azimuthal_mean(
+        self,
+        center_coord,
+        outer_radius: int | float,
+        radius_step: int | float,
+        return_hit_counts: bool = False,
+    ):
+        """Compute averages along circles of constant great-circle distance from a point.
+
+        Parameters
+        ----------
+        center_coord: tuple, list, ndarray
+            Longitude and latitude of the center of the bounding circle
+        outer_radius: scalar, int, float
+            The maximum radius, in great-circle degrees, at which the azimuthal mean will be computed.
+        radius_step: scalar, int, float
+            Means will be computed at intervals of `radius_step` on the interval [0, outer_radius]
+        return_hit_counts: bool, false
+            Indicates whether to return the number of hits at each radius
+
+        Returns
+        -------
+        azimuthal_mean: xr.DataArray
+            Contains a variable with a dimension 'radius' corresponding to the azimuthal average.
+        hit_counts: xr.DataArray
+            The number of hits at each radius
+
+
+        Examples
+        --------
+        # Range from 0° to 5° at 0.5° intervals, around the central point lon,lat=10,50
+        >>> az = uxds["var"].azimuthal_mean(
+        ...     center_coord=(10, 50), outer_radius=5.0, radius_step=0.5
+        ... )
+        >>> az.plot(title="Azimuthal Mean")
+
+        Notes
+        -----
+        Only supported for face-centered data variables. Candidate faces are determined
+        using bounding circles - for radii = [r1, r2, r3, ...] faces whose centers lie at distance d,
+        r2 < d <= r3 are included in calculations for r3.
+        """
+        from uxarray.grid.coordinates import _lonlat_rad_to_xyz
+
+        if not self._face_centered():
+            raise ValueError(
+                "Azimuthal mean computations are currently only supported for face-centered data variables."
+            )
+
+        if outer_radius <= 0:
+            raise ValueError("Radius must be a positive scalar.")
+
+        kdtree = self.uxgrid._get_scipy_kd_tree()
+
+        lon_deg, lat_deg = map(float, np.asarray(center_coord))
+        center_xyz = np.array(
+            _lonlat_rad_to_xyz(np.deg2rad(lon_deg), np.deg2rad(lat_deg))
+        )
+
+        radii_deg = np.arange(0.0, outer_radius + radius_step, radius_step, dtype=float)
+        radii_rad = np.deg2rad(radii_deg)
+        chord_radii = 2.0 * np.sin(radii_rad / 2.0)
+
+        faces_processed = np.array([], dtype=np.int_)
+        means = np.full(
+            (radii_deg.size, *self.to_xarray().isel(drop=True, n_face=0).shape), np.nan
+        )
+        hit_count = np.zeros_like(radii_deg, dtype=np.int_)
+
+        for ii, r_chord in enumerate(chord_radii):
+            # indices of faces within the bounding circle for this radius
+            within = np.array(
+                kdtree.query_ball_point(center_xyz, r_chord), dtype=np.int_
+            )
+            if within.size:
+                within.sort()
+
+            # include only the new ring: r_(i-1) < d <= r_i
+            faces_in_bin = np.setdiff1d(within, faces_processed, assume_unique=True)
+            hit_count[ii] = faces_in_bin.size
+
+            if hit_count[ii] == 0:
+                continue
+
+            faces_processed = within  # cumulative set for next iteration
+
+            tpose = self.isel(n_face=faces_in_bin).transpose(..., "n_face")
+            means[ii, ...] = tpose.weighted_mean().data
+
+        # swap the leading 'radius' axis into the former n_face position
+        face_axis = self.dims.index("n_face")
+        dims = list(self.dims)
+        dims[face_axis] = "radius"
+        means = np.moveaxis(means, 0, face_axis)
+
+        hit_count = xr.DataArray(
+            data=hit_count, dims="radius", coords={"radius": radii_deg}
+        )
+
+        uxda = xr.DataArray(
+            means,
+            dims=dims,
+            coords={"radius": radii_deg},
+            name=self.name + "_azimuthal_mean"
+            if self.name is not None
+            else "azimuthal_mean",
+            attrs={
+                "azimuthal_mean": True,
+                "center_lon": lon_deg,
+                "center_lat": lat_deg,
+                "radius_units": "degrees",
+            },
+        )
+
+        if return_hit_counts:
+            return uxda, hit_count
+        else:
+            return uxda
+
+    azimuthal_average = azimuthal_mean
+
     def weighted_mean(self, weights=None):
         """Computes a weighted mean.