Optimize azimuthal_mean with one sparse ring operator

benchmarks/mpas_ocean.py

@@ -326,6 +326,21 @@ def time_zonal_average(self, resolution):
         self.uxds['bottomDepth'].zonal_mean(lat=(-45, 45, lat_step))
 
 
+class AzimuthalAverage(DatasetBenchmark):
+    def setup(self, resolution, *args, **kwargs):
+        super().setup(resolution, *args, **kwargs)
+        # Grid caches, so only the per-ring reduction is timed
+        _ = self.uxds.uxgrid.face_areas
+        _ = self.uxds.uxgrid._get_scipy_kd_tree()
+
+    def time_azimuthal_average(self, resolution):
+        # 61 rings around an equatorial Atlantic point, ocean on every oQU mesh.
+        # Wide enough that the ring faces, not per-call overhead, set the cost.
+        self.uxds[data_var].azimuthal_mean(
+            center_coord=(-30, 0), outer_radius=60, radius_step=1
+        )
+
+
 class ZonalAveragePeakMem:
     """Peak memory of a cold-start non-conservative zonal-mean sweep.
     Uses a fresh interpreter per sample."""

test/core/test_azimuthal.py

@@ -1,68 +1,123 @@
-import pytest
-import uxarray as ux
+import dask.array as da
 import numpy as np
+import pytest
+import xarray as xr
 
+import uxarray as ux
+from uxarray.core.azimuthal import _ring_membership
+from uxarray.grid.coordinates import _lonlat_rad_to_xyz
 
 
-def test_gaussian(gridpath, datasetpath):
-    uxds = ux.open_dataset(
+@pytest.fixture
+def uxds(gridpath, datasetpath):
+    return ux.open_dataset(
         gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
-        datasetpath("mpas", "dyamond-30km", "gradient_data_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 _center_xyz(lon, lat):
+    return np.array(_lonlat_rad_to_xyz(np.deg2rad(lon), np.deg2rad(lat)))
 
 
+def _chords(outer_radius, radius_step):
+    radii = np.arange(0.0, outer_radius + radius_step, radius_step)
+    return 2.0 * np.sin(np.deg2rad(radii) / 2.0)
 
-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")
-    )
+def _reference_rings(points, center_xyz, chord_radii):
+    """Brute-force ring of each point: the first i with d <= r_i, or -1."""
+    dist2 = ((points - center_xyz) ** 2).sum(axis=1)
+    ring_of = np.full(len(points), -1)
+    # squared as an array, as a numpy scalar's ** can round differently
+    for i, bound in reversed(list(enumerate(chord_radii**2))):
+        ring_of[dist2 <= bound] = i
+    return ring_of
 
-    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")
+@pytest.mark.parametrize("name, sign", [("gaussian", -1), ("inverse_gaussian", 1)])
+def test_monotone_profile(uxds, name, sign):
+    res = uxds[name].azimuthal_mean(
+        center_coord=(45, 0), outer_radius=2, radius_step=0.5
     )
 
-    res, hit_counts = uxds['inverse_gaussian'].azimuthal_mean(center_coord=(45, 0), outer_radius=2, radius_step=0.5, return_hit_counts=True)
+    # gaussian decreases away from the center, inverse_gaussian increases. Only
+    # the empty ring at radius 0 is NaN, so any other NaN fails the comparison.
+    np.testing.assert_array_less(-1e-12, sign * res[1:].diff("radius").values)
 
-    # 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")
+@pytest.mark.parametrize("chunks", [None, {"n_face": 50, "time": 1}])
+@pytest.mark.parametrize("center", [(45, 0), (-45, 0)])  # in and outside the subset
+def test_matches_reference(uxds, center, chunks):
+    # (n_face, time) data with the grid dim first and a few NaNs
+    stacked = (
+        xr.concat([uxds["gaussian"], uxds["inverse_gaussian"]], dim="time")
+        .transpose("n_face", "time")
+        .copy()
     )
+    stacked[[3, 50, 120], 1] = np.nan
+    uxda = ux.UxDataArray(stacked.assign_coords(time=[10, 20]), uxgrid=uxds.uxgrid)
+    if chunks:
+        uxda = uxda.chunk(chunks)
 
-    # 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']
+    res, hits = uxda.azimuthal_mean(
+        center_coord=center, outer_radius=2, radius_step=0.25, return_hit_counts=True
+    )
 
-    # No hits
-    assert np.all(hit_counts == 0)
+    assert res.dims == ("radius", "time")
+    assert hits.dims == ("radius",) and hits.sizes["radius"] == res.sizes["radius"]
+    np.testing.assert_array_equal(res["time"], [10, 20])
+    if chunks:
+        assert isinstance(res.data, da.Array)
+        assert res.chunks[1] == (1, 1)
 
-    print(res)
+    # Per-ring float64 oracle: nansum(x * area) / sum(area)
+    ring_of = _reference_rings(
+        uxda.uxgrid._get_scipy_kd_tree().data, _center_xyz(*center), _chords(2, 0.25)
+    )
+    areas = uxda.uxgrid.face_areas.values.astype(np.float64)
+    values = uxda.values.astype(np.float64)
+    expected = np.full(res.shape, np.nan)
+    expected_hits = np.zeros(res.sizes["radius"], dtype=int)
+    for i in range(res.sizes["radius"]):
+        ring = ring_of == i
+        expected_hits[i] = ring.sum()
+        if ring.any():
+            weighted = np.nansum(values[ring] * areas[ring, None], axis=0)
+            expected[i] = weighted / areas[ring].sum()
+
+    np.testing.assert_array_equal(hits, expected_hits)
+    np.testing.assert_allclose(res.values, expected, rtol=1e-12, equal_nan=True)
+    if center == (45, 0):
+        assert np.all(hits[1:] > 1)
+    else:
+        assert np.all(hits == 0)
+
+
+def _ring_edge_cases(kdtree, n_cases=200):
+    """Radii exactly at face distances, so faces sit on ring edges, where the
+    kd-tree's own rounding can disagree with d <= r."""
+    rng = np.random.default_rng(0)
+    for _ in range(n_cases):
+        center = kdtree.data[rng.integers(kdtree.n)] + rng.normal(scale=1e-3, size=3)
+        offsets = kdtree.data[rng.integers(kdtree.n, size=5)] - center
+        yield center, np.sort(np.append(np.sqrt((offsets**2).sum(axis=1)), 0.0))
+
+
+# (outer_radius, radius_step), past 180 degrees too, or None for ring-edge ties
+@pytest.mark.parametrize("radii", [(2, 0.1), (60, 1), (250, 10), (360, 3), None])
+def test_ring_membership(gridpath, radii):
+    kdtree = ux.open_grid(
+        gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc")
+    )._get_scipy_kd_tree()
+    if radii is None:
+        cases = _ring_edge_cases(kdtree)
+    else:
+        chords = _chords(*radii)
+        cases = [(_center_xyz(*c), chords) for c in [(45, 0), (44.7, 0.6), (-45, 0)]]
+
+    for center, chords in cases:
+        ring_of = _reference_rings(kdtree.data, center, chords)
+        faces, ring_ids = _ring_membership(kdtree, center, chords)
+        np.testing.assert_array_equal(faces, np.flatnonzero(ring_of >= 0))
+        np.testing.assert_array_equal(ring_ids, ring_of[faces])

uxarray/core/azimuthal.py

@@ -0,0 +1,90 @@
+import math
+
+import numpy as np
+import xarray as xr
+from scipy import sparse
+
+
+def _compute_azimuthal_mean(uxda, center_xyz, chord_radii):
+    """Computes area-weighted means over the rings between successive bounding circles.
+
+    Returns the means, with ``n_face`` replaced by ``radius`` in place (lazy
+    when ``uxda`` is dask-backed), and the number of faces in each ring.
+    """
+    n_radius = chord_radii.size
+    faces, ring_ids = _ring_membership(
+        uxda.uxgrid._get_scipy_kd_tree(), center_xyz, chord_radii
+    )
+    hit_count = np.bincount(ring_ids, minlength=n_radius)
+
+    # Rings are disjoint, so all per-ring means form one sparse
+    # (n_radius, n_ring_faces) operator. Areas come from the parent grid, and
+    # only the ring faces are gathered, so no per-ring subgrid is ever built.
+    areas = uxda.uxgrid.face_areas.values[faces]
+    ring_area = np.bincount(ring_ids, weights=areas, minlength=n_radius)
+    weights = sparse.csr_matrix(
+        (areas / ring_area[ring_ids], (ring_ids, np.arange(faces.size))),
+        shape=(n_radius, faces.size),
+    )
+
+    # Read the span of ring faces contiguously, then pick them out of it: a
+    # lazily backed (e.g. netCDF) variable reads fancy indices face by face
+    start, stop = (faces[0], faces[-1] + 1) if faces.size else (0, 0)
+    span = uxda.to_xarray().isel(n_face=slice(start, stop))
+    if span.chunks is None:
+        ring_data = span.load().isel(n_face=faces - start)
+    else:
+        # the sparse multiply spans every ring face, so keep them in a single chunk
+        ring_data = span.isel(n_face=faces - start).chunk({"n_face": -1})
+
+    means = xr.apply_ufunc(
+        _apply_ring_weights,
+        ring_data,
+        input_core_dims=[["n_face"]],
+        output_core_dims=[["radius"]],
+        dask="parallelized",
+        kwargs={"weights": weights},
+        output_dtypes=[np.result_type(weights.dtype, ring_data.dtype)],
+        # n_face is already one chunk; dask's rechunk divides by zero when it's empty
+        dask_gufunc_kwargs={
+            "output_sizes": {"radius": n_radius},
+            "allow_rechunk": False,
+        },
+    )
+
+    dims = ["radius" if dim == "n_face" else dim for dim in uxda.dims]
+    return means.transpose(*dims), hit_count
+
+
+def _ring_membership(kdtree, center_xyz, chord_radii):
+    """Assigns each face to the first ring whose bounding circle holds it: r_(i-1) < d <= r_i.
+
+    Returns the faces, in ascending order, and their ring indices.
+    """
+    # The running max keeps the rings past a peak in chord length (radii
+    # beyond 180 degrees) empty, as their circles add no faces
+    bounds = np.maximum.accumulate(chord_radii) ** 2
+
+    # One padded query finds every candidate, so the tree's rounding can't drop
+    # faces on the outer circle; the squared distances alone then bin them
+    faces = np.array(
+        kdtree.query_ball_point(
+            center_xyz, chord_radii.max() * (1 + 1e-9), return_sorted=True
+        ),
+        dtype=np.int_,
+    )
+    dist2 = ((kdtree.data[faces] - center_xyz) ** 2).sum(axis=1)
+    inside = dist2 <= bounds[-1]
+    return faces[inside], np.searchsorted(bounds, dist2[inside])
+
+
+def _apply_ring_weights(values, weights):
+    """Applies the sparse ring operator along the trailing (ring face) dimension."""
+    # NaN faces add nothing to the weighted sum but still count toward the ring
+    # area, matching the skipna sum of the former per-ring weighted_mean()
+    values = np.where(np.isnan(values), 0, values)
+    # explicit leading size: -1 is ambiguous when no face falls in any ring
+    flat = values.reshape(math.prod(values.shape[:-1]), values.shape[-1])
+    means = (weights @ flat.T).T
+    means[:, np.diff(weights.indptr) == 0] = np.nan  # empty rings
+    return means.reshape(values.shape[:-1] + (weights.shape[0],))

uxarray/core/dataarray.py

@@ -15,6 +15,7 @@
 from uxarray.conventions.attrs import _drop_value_describing_attrs
 from uxarray.core.aggregation import _uxda_grid_aggregate
 from uxarray.core.arithmetic import UxSupportsArithmetic
+from uxarray.core.azimuthal import _compute_azimuthal_mean
 from uxarray.core.gradient import (
     _calculate_edge_face_difference,
     _calculate_edge_node_difference,
@@ -1020,6 +1021,9 @@ def azimuthal_mean(
         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.
+
+        Dask-backed data is reduced lazily, with the faces inside ``outer_radius``
+        gathered into a single ``n_face`` chunk.
         """
         from uxarray.grid.coordinates import _lonlat_rad_to_xyz
 
@@ -1035,8 +1039,6 @@ def azimuthal_mean(
                 f"outer_radius must be a positive scalar during azimuthal_mean(); got outer_radius={outer_radius}"
             )
 
-        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))
@@ -1046,50 +1048,20 @@ def azimuthal_mean(
         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)
+        means, hit_count = _compute_azimuthal_mean(self, center_xyz, chord_radii)
 
         hit_count = xr.DataArray(
             data=hit_count, dims="radius", coords={"radius": radii_deg}
         )
 
-        # Assign coords from `self` to the result except one that corresponds to `dims[face_axis]`
+        # Assign coords from `self` to the result except those on n_face
         new_coords = _preserve_valid_coords(self, "n_face")
         # Add radii_deg to the resulting coords
         new_coords["radius"] = radii_deg
 
         uxda = xr.DataArray(
-            means,
-            dims=dims,
+            means.data,
+            dims=means.dims,
             coords=new_coords,
             name=self.name + "_azimuthal_mean"
             if self.name is not None