From 833840d13d2da9be1792f72322133b31ec1565fa Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 23 Oct 2026 16:37:59 -0500 Subject: [PATCH 1/4] Compute azimuthal_mean with one sparse ring operator azimuthal_mean built a full subgrid for every ring (isel without ignore_grid), recomputed face_areas on it, and forced a compute per ring. The rings are disjoint, so the per-ring area-weighted means now form one sparse (n_radius, n_ring_faces) operator over the ring faces, gathered once with ignore_grid=True and weighted by the parent grid's areas. It is applied through a single apply_ufunc, so dask-backed input stays lazy. Ring membership, hit counts, output dims/coords/attrs and NaN handling (NaN faces add nothing to the sum but still count toward the ring area) are unchanged. Accumulation is now float64: on float32 grids the old float32 sums were off by up to ~3e-7 relative against a float64 reference; the new result matches it to ~3e-16. Adds AzimuthalAverage to the oQU benchmarks (21 rings, out to 20 degrees): 480km 43.9 ms -> 0.7 ms, 120km 50.0 ms -> 1.2 ms. Co-Authored-By: Claude Opus 5.5 --- benchmarks/mpas_ocean.py | 14 +++++++ test/core/test_azimuthal.py | 80 ++++++++++++++++++++++++++++++++++++- uxarray/core/azimuthal.py | 74 ++++++++++++++++++++++++++++++++++ uxarray/core/dataarray.py | 35 +++------------- 4 files changed, 172 insertions(+), 31 deletions(-) create mode 100644 uxarray/core/azimuthal.py diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 8d8687e64..62144e62c 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -326,6 +326,20 @@ 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): + # 21 rings around an equatorial Atlantic point, ocean on every oQU mesh + self.uxds[data_var].azimuthal_mean( + center_coord=(-30, 0), outer_radius=20, radius_step=1 + ) + + class ZonalAveragePeakMem: """Peak memory of a cold-start non-conservative zonal-mean sweep. Uses a fresh interpreter per sample.""" diff --git a/test/core/test_azimuthal.py b/test/core/test_azimuthal.py index 1ecb26d04..05e34a24f 100644 --- a/test/core/test_azimuthal.py +++ b/test/core/test_azimuthal.py @@ -1,7 +1,8 @@ -import pytest -import uxarray as ux +import dask.array as da import numpy as np +import xarray as xr +import uxarray as ux def test_gaussian(gridpath, datasetpath): @@ -66,3 +67,78 @@ def test_zero_hit_counts(gridpath, datasetpath): assert np.all(hit_counts == 0) print(res) + + +def _reference_azimuthal_mean(uxda, center_coord, outer_radius, radius_step): + """Per-ring float64 oracle: nansum(x * area) / sum(area) over each ring.""" + from uxarray.grid.coordinates import _lonlat_rad_to_xyz + + kdtree = uxda.uxgrid._get_scipy_kd_tree() + lon, lat = np.deg2rad(center_coord) + center = np.array(_lonlat_rad_to_xyz(lon, lat)) + radii = np.arange(0.0, outer_radius + radius_step, radius_step) + chords = 2.0 * np.sin(np.deg2rad(radii) / 2.0) + + areas = uxda.uxgrid.face_areas.values.astype(np.float64) + values = uxda.transpose(..., "n_face").values.astype(np.float64) + out = np.full(values.shape[:-1] + (radii.size,), np.nan) + previous = set() + for i, chord in enumerate(chords): + within = set(kdtree.query_ball_point(center, chord)) + ring = np.array(sorted(within - previous), dtype=int) + previous = within + if ring.size: + out[..., i] = ( + np.nansum(values[..., ring] * areas[ring], axis=-1) / areas[ring].sum() + ) + return out + + +def _multi_dim(uxds): + """(n_face, time) data with the grid dim first and a few NaNs.""" + stacked = ( + xr.concat( + [uxds["gaussian"].to_xarray(), uxds["inverse_gaussian"].to_xarray()], + dim="time", + ) + .transpose("n_face", "time") + .copy() + ) + stacked[[3, 50, 120], 1] = np.nan + return ux.UxDataArray(stacked.assign_coords(time=[10, 20]), uxgrid=uxds.uxgrid) + + +def test_matches_reference(gridpath, datasetpath): + uxds = ux.open_dataset( + gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"), + datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc"), + ) + uxda = _multi_dim(uxds) + + res = uxda.azimuthal_mean(center_coord=(45, 0), outer_radius=2, radius_step=0.25) + + assert res.dims == ("radius", "time") + np.testing.assert_array_equal(res["time"], [10, 20]) + expected = _reference_azimuthal_mean(uxda, (45, 0), 2, 0.25) + np.testing.assert_allclose( + res.transpose("time", "radius").values, expected, rtol=1e-12 + ) + + +def test_dask_lazy_and_matches_numpy(gridpath, datasetpath): + uxds = ux.open_dataset( + gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"), + datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc"), + ) + uxda = _multi_dim(uxds) + kwargs = dict( + center_coord=(45, 0), outer_radius=2, radius_step=0.25, return_hit_counts=True + ) + + eager, eager_hits = uxda.azimuthal_mean(**kwargs) + lazy, lazy_hits = uxda.chunk(n_face=50, time=1).azimuthal_mean(**kwargs) + + assert isinstance(lazy.data, da.Array) + assert lazy.chunks[lazy.get_axis_num("time")] == (1, 1) + np.testing.assert_array_equal(lazy_hits, eager_hits) + np.testing.assert_allclose(lazy.values, eager.values, rtol=1e-14, equal_nan=True) diff --git a/uxarray/core/azimuthal.py b/uxarray/core/azimuthal.py new file mode 100644 index 000000000..5ee555680 --- /dev/null +++ b/uxarray/core/azimuthal.py @@ -0,0 +1,74 @@ +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. + """ + kdtree = uxda.uxgrid._get_scipy_kd_tree() + n_radius = chord_radii.size + + # Each ring holds the faces newly inside its bounding circle: r_(i-1) < d <= r_i + ring_faces = [] + faces_processed = np.array([], dtype=np.int_) + for r_chord in chord_radii: + within = np.array(kdtree.query_ball_point(center_xyz, r_chord), dtype=np.int_) + if within.size: + within.sort() + + ring = np.setdiff1d(within, faces_processed, assume_unique=True) + ring_faces.append(ring) + + if ring.size: + faces_processed = within # cumulative set for next iteration + + hit_count = np.array([ring.size for ring in ring_faces], dtype=np.int_) + + # 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. + faces = np.concatenate(ring_faces) + ring_ids = np.repeat(np.arange(n_radius), hit_count) + 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), + ) + + ring_data = uxda.isel(n_face=faces, ignore_grid=True).to_xarray() + if ring_data.chunks is not None: + # the sparse multiply spans every ring face, so keep them in a single chunk + ring_data = ring_data.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, "empty": hit_count == 0}, + output_dtypes=[np.result_type(weights.dtype, ring_data.dtype)], + dask_gufunc_kwargs={"output_sizes": {"radius": n_radius}}, + ) + + dims = ["radius" if dim == "n_face" else dim for dim in uxda.dims] + return means.transpose(*dims).data, hit_count + + +def _apply_ring_weights(values, weights, empty): + """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[:, empty] = np.nan + return means.reshape(values.shape[:-1] + (weights.shape[0],)) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index fde1a968a..e2f055808 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -14,6 +14,7 @@ from uxarray.constants import GRID_DIMS 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, @@ -994,6 +995,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 @@ -1009,8 +1013,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)) @@ -1020,37 +1022,12 @@ 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 + means, hit_count = _compute_azimuthal_mean(self, center_xyz, chord_radii) - # swap the leading 'radius' axis into the former n_face position + # 'radius' takes 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} From a5858784f79e4666ca0ba042f9b9a37a07a83055 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 23 Oct 2026 17:06:07 -0500 Subject: [PATCH 2/4] Find azimuthal rings with one kd-tree query The ring loop queried the kd-tree once per radius for the whole disc, then sorted it and ran setdiff1d against the previous disc, so its cost grew with rings x faces inside. At 60 degrees on oQU120 it was 93% of the runtime. One query at the outermost circle now finds every candidate face, and each is binned by its squared chord distance, summed in the kd-tree's own order. The tree accepts or rejects whole nodes on distance bounds that carry their own rounding, so within an ulp or so of a circle it can disagree with d**2 <= r**2, in either direction. Faces within 1e-12 of a ring's edge are settled by querying the tree at that radius, as the per-radius loop did, which keeps ring membership identical. Checked against the old loop on oQU480, oQU120 and the dyamond subset over 4050 cases, 15,000 of them exact ties on a ring's edge; results are bit-identical to the previous commit. The AzimuthalAverage benchmark now runs out to 60 degrees (61 rings), where the ring faces rather than per-call overhead set the cost: 480km 141 ms (before 833840d1) -> 1.8 ms (833840d1) -> 0.5 ms 120km 169 ms -> 10.1 ms -> 1.6 ms Co-Authored-By: Claude Opus 5.5 --- benchmarks/mpas_ocean.py | 5 ++- test/core/test_azimuthal.py | 55 +++++++++++++++++++++++++++ uxarray/core/azimuthal.py | 75 +++++++++++++++++++++++++++---------- 3 files changed, 114 insertions(+), 21 deletions(-) diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 62144e62c..ebf9bace0 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -334,9 +334,10 @@ def setup(self, resolution, *args, **kwargs): _ = self.uxds.uxgrid._get_scipy_kd_tree() def time_azimuthal_average(self, resolution): - # 21 rings around an equatorial Atlantic point, ocean on every oQU mesh + # 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=20, radius_step=1 + center_coord=(-30, 0), outer_radius=60, radius_step=1 ) diff --git a/test/core/test_azimuthal.py b/test/core/test_azimuthal.py index 05e34a24f..7780ff61b 100644 --- a/test/core/test_azimuthal.py +++ b/test/core/test_azimuthal.py @@ -1,8 +1,10 @@ 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 def test_gaussian(gridpath, datasetpath): @@ -142,3 +144,56 @@ def test_dask_lazy_and_matches_numpy(gridpath, datasetpath): assert lazy.chunks[lazy.get_axis_num("time")] == (1, 1) np.testing.assert_array_equal(lazy_hits, eager_hits) np.testing.assert_allclose(lazy.values, eager.values, rtol=1e-14, equal_nan=True) + + +def _per_radius_rings(kdtree, center_xyz, chord_radii): + """The original ring assignment: one kd-tree query per radius.""" + faces, ring_ids = [], [] + processed = np.array([], dtype=np.int_) + for i, r_chord in enumerate(chord_radii): + within = np.sort( + np.array(kdtree.query_ball_point(center_xyz, r_chord), dtype=np.int_) + ) + ring = np.setdiff1d(within, processed, assume_unique=True) + if ring.size: + processed = within + faces.append(ring) + ring_ids.append(np.full(ring.size, i)) + return np.concatenate(faces), np.concatenate(ring_ids) + + +@pytest.mark.parametrize( + "outer_radius, radius_step", [(2, 0.1), (60, 1), (250, 10), (360, 3)] +) +def test_ring_membership_matches_per_radius_queries( + gridpath, outer_radius, radius_step +): + from uxarray.grid.coordinates import _lonlat_rad_to_xyz + + uxgrid = ux.open_grid(gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc")) + kdtree = uxgrid._get_scipy_kd_tree() + radii = np.arange(0.0, outer_radius + radius_step, radius_step) + chord_radii = 2.0 * np.sin(np.deg2rad(radii) / 2.0) + + for lon, lat in [(45, 0), (44.7, 0.6), (-45, 0)]: + center = np.array(_lonlat_rad_to_xyz(np.deg2rad(lon), np.deg2rad(lat))) + expected = _per_radius_rings(kdtree, center, chord_radii) + actual = _ring_membership(kdtree, center, chord_radii) + np.testing.assert_array_equal(actual[0], expected[0]) + np.testing.assert_array_equal(actual[1], expected[1]) + + +def test_ring_membership_ties_on_ring_edges(gridpath): + """Radii exactly at face distances, where the kd-tree's pruning rounding decides.""" + uxgrid = ux.open_grid(gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc")) + kdtree = uxgrid._get_scipy_kd_tree() + rng = np.random.default_rng(0) + + for _ in range(200): + center = kdtree.data[rng.integers(uxgrid.n_face)].copy() + offsets = kdtree.data[rng.integers(uxgrid.n_face, size=5)] - center + chord_radii = np.sort(np.append(np.sqrt((offsets * offsets).sum(axis=1)), 0.0)) + expected = _per_radius_rings(kdtree, center, chord_radii) + actual = _ring_membership(kdtree, center, chord_radii) + np.testing.assert_array_equal(actual[0], expected[0]) + np.testing.assert_array_equal(actual[1], expected[1]) diff --git a/uxarray/core/azimuthal.py b/uxarray/core/azimuthal.py index 5ee555680..168995d1e 100644 --- a/uxarray/core/azimuthal.py +++ b/uxarray/core/azimuthal.py @@ -4,6 +4,10 @@ import xarray as xr from scipy import sparse +# Squared chord distances are at most 4, so this is ~1000 ulps: well past the +# kd-tree's rounding, yet faces this close to a ring's edge are rare +_TIE_TOLERANCE = 1e-12 + def _compute_azimuthal_mean(uxda, center_xyz, chord_radii): """Computes area-weighted means over the rings between successive bounding circles. @@ -11,30 +15,15 @@ def _compute_azimuthal_mean(uxda, center_xyz, chord_radii): 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. """ - kdtree = uxda.uxgrid._get_scipy_kd_tree() n_radius = chord_radii.size - - # Each ring holds the faces newly inside its bounding circle: r_(i-1) < d <= r_i - ring_faces = [] - faces_processed = np.array([], dtype=np.int_) - for r_chord in chord_radii: - within = np.array(kdtree.query_ball_point(center_xyz, r_chord), dtype=np.int_) - if within.size: - within.sort() - - ring = np.setdiff1d(within, faces_processed, assume_unique=True) - ring_faces.append(ring) - - if ring.size: - faces_processed = within # cumulative set for next iteration - - hit_count = np.array([ring.size for ring in ring_faces], dtype=np.int_) + 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. - faces = np.concatenate(ring_faces) - ring_ids = np.repeat(np.arange(n_radius), hit_count) areas = uxda.uxgrid.face_areas.values[faces] ring_area = np.bincount(ring_ids, weights=areas, minlength=n_radius) weights = sparse.csr_matrix( @@ -62,6 +51,54 @@ def _compute_azimuthal_mean(uxda, center_xyz, chord_radii): return means.transpose(*dims).data, 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 and their ring indices, ordered by ring and then face. + """ + # The running max keeps the rings past a peak in chord length (radii + # beyond 180 degrees) empty, as their circles add no faces + chord_bounds = np.maximum.accumulate(chord_radii) + chord_bounds = chord_bounds * chord_bounds + + # One query at the outermost circle, then bin by squared chord distance, + # summed in the kd-tree's own order + faces = np.array( + kdtree.query_ball_point(center_xyz, chord_radii.max()), dtype=np.int_ + ) + diff = kdtree.data[faces] - center_xyz + dist2 = diff[:, 0] * diff[:, 0] + diff[:, 1] * diff[:, 1] + diff[:, 2] * diff[:, 2] + ring_ids = np.searchsorted(chord_bounds, dist2, side="left") + # the outer query has the final say on whether a face is in the outermost ring + outermost = np.searchsorted(chord_bounds, chord_bounds[-1], side="left") + ring_ids = np.minimum(ring_ids, outermost) + + # The tree prunes whole nodes on distances that carry rounding of their + # own, so within an ulp or so of a circle it can disagree with + # dist2 <= r**2. Settle faces that close by asking the tree at that circle. + near_circle = np.full(faces.size, -1) + for circle in (ring_ids - 1, ring_ids): # the circles just inside and outside + circle = np.clip(circle, 0, outermost) + # first circle of a plateau, where a face on that bound first enters + circle = np.searchsorted(chord_bounds, chord_bounds[circle], side="left") + near = np.abs(dist2 - chord_bounds[circle]) <= _TIE_TOLERANCE + near_circle[near] = circle[near] + # the outer query already settled the outermost circle + tied_circles = np.unique( + near_circle[(near_circle >= 0) & (near_circle < outermost)] + ) + for circle in tied_circles: + tied = near_circle == circle + inside = np.isin( + faces[tied], kdtree.query_ball_point(center_xyz, chord_radii[circle]) + ) + next_circle = np.searchsorted(chord_bounds, chord_bounds[circle], side="right") + ring_ids[tied] = np.where(inside, circle, next_circle) + + order = np.lexsort((faces, ring_ids)) + return faces[order], ring_ids[order] + + def _apply_ring_weights(values, weights, empty): """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 From e97abdac396593922a221edc9633ad7768787970 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 26 Oct 2026 17:20:57 -0500 Subject: [PATCH 3/4] Simplify azimuthal ring finding and read lazy data contiguously Fancy-indexing a lazily backed (netCDF) variable reads face by face, so gathering the ring faces took 87 ms on oQU120 against 1.6 ms for the whole call on loaded data; the benchmark fixture loads data, hiding it. The span of ring faces is now read contiguously and indexed in memory. Dask input with no face inside outer_radius raised ZeroDivisionError in dask's gufunc rechunking; n_face is already one chunk, so allow_rechunk is now off. Ring membership is now d <= r on the computed squared distances, as documented, rather than reproducing the kd-tree's pruning rounding at ring edges. The single query is padded so the tree can't drop faces on the outer circle. Only faces exactly on a ring edge can change ring. oQU, 61 rings loaded lazy 480km 0.40 -> 0.33 ms 5.9 -> 0.40 ms 120km 1.57 -> 0.92 ms 86.7 -> 1.0 ms Tests share one brute-force ring oracle; the eager, dask and hit-count checks run as one parametrized reference test, inside and outside the grid subset. Co-Authored-By: Claude Opus 5.5 --- test/core/test_azimuthal.py | 250 +++++++++++++----------------------- uxarray/core/azimuthal.py | 77 ++++------- uxarray/core/dataarray.py | 11 +- 3 files changed, 118 insertions(+), 220 deletions(-) diff --git a/test/core/test_azimuthal.py b/test/core/test_azimuthal.py index 7780ff61b..aa6cb3735 100644 --- a/test/core/test_azimuthal.py +++ b/test/core/test_azimuthal.py @@ -5,195 +5,119 @@ 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") - ) - - 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 + datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc"), ) +def _center_xyz(lon, lat): + return np.array(_lonlat_rad_to_xyz(np.deg2rad(lon), np.deg2rad(lat))) -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") - ) +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) - 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) +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 - 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("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 ) - # 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'] + # gaussian decreases away from the center, inverse_gaussian increases + diffs = sign * res[1:].diff("radius").values + np.testing.assert_array_less(-1e-12, diffs[np.isfinite(diffs)]) - # No hits - assert np.all(hit_counts == 0) - print(res) - - -def _reference_azimuthal_mean(uxda, center_coord, outer_radius, radius_step): - """Per-ring float64 oracle: nansum(x * area) / sum(area) over each ring.""" - from uxarray.grid.coordinates import _lonlat_rad_to_xyz - - kdtree = uxda.uxgrid._get_scipy_kd_tree() - lon, lat = np.deg2rad(center_coord) - center = np.array(_lonlat_rad_to_xyz(lon, lat)) - radii = np.arange(0.0, outer_radius + radius_step, radius_step) - chords = 2.0 * np.sin(np.deg2rad(radii) / 2.0) - - areas = uxda.uxgrid.face_areas.values.astype(np.float64) - values = uxda.transpose(..., "n_face").values.astype(np.float64) - out = np.full(values.shape[:-1] + (radii.size,), np.nan) - previous = set() - for i, chord in enumerate(chords): - within = set(kdtree.query_ball_point(center, chord)) - ring = np.array(sorted(within - previous), dtype=int) - previous = within - if ring.size: - out[..., i] = ( - np.nansum(values[..., ring] * areas[ring], axis=-1) / areas[ring].sum() - ) - return out - - -def _multi_dim(uxds): - """(n_face, time) data with the grid dim first and a few NaNs.""" +@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"].to_xarray(), uxds["inverse_gaussian"].to_xarray()], - dim="time", - ) + xr.concat([uxds["gaussian"], uxds["inverse_gaussian"]], dim="time") .transpose("n_face", "time") .copy() ) stacked[[3, 50, 120], 1] = np.nan - return ux.UxDataArray(stacked.assign_coords(time=[10, 20]), uxgrid=uxds.uxgrid) - + uxda = ux.UxDataArray(stacked.assign_coords(time=[10, 20]), uxgrid=uxds.uxgrid) + if chunks: + uxda = uxda.chunk(chunks) -def test_matches_reference(gridpath, datasetpath): - uxds = ux.open_dataset( - gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"), - datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc"), + res, hits = uxda.azimuthal_mean( + center_coord=center, outer_radius=2, radius_step=0.25, return_hit_counts=True ) - uxda = _multi_dim(uxds) - - res = uxda.azimuthal_mean(center_coord=(45, 0), outer_radius=2, radius_step=0.25) 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]) - expected = _reference_azimuthal_mean(uxda, (45, 0), 2, 0.25) - np.testing.assert_allclose( - res.transpose("time", "radius").values, expected, rtol=1e-12 - ) + if chunks: + assert isinstance(res.data, da.Array) + assert res.chunks[1] == (1, 1) - -def test_dask_lazy_and_matches_numpy(gridpath, datasetpath): - uxds = ux.open_dataset( - gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"), - datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc"), - ) - uxda = _multi_dim(uxds) - kwargs = dict( - center_coord=(45, 0), outer_radius=2, radius_step=0.25, return_hit_counts=True + # 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) ) - - eager, eager_hits = uxda.azimuthal_mean(**kwargs) - lazy, lazy_hits = uxda.chunk(n_face=50, time=1).azimuthal_mean(**kwargs) - - assert isinstance(lazy.data, da.Array) - assert lazy.chunks[lazy.get_axis_num("time")] == (1, 1) - np.testing.assert_array_equal(lazy_hits, eager_hits) - np.testing.assert_allclose(lazy.values, eager.values, rtol=1e-14, equal_nan=True) - - -def _per_radius_rings(kdtree, center_xyz, chord_radii): - """The original ring assignment: one kd-tree query per radius.""" - faces, ring_ids = [], [] - processed = np.array([], dtype=np.int_) - for i, r_chord in enumerate(chord_radii): - within = np.sort( - np.array(kdtree.query_ball_point(center_xyz, r_chord), dtype=np.int_) - ) - ring = np.setdiff1d(within, processed, assume_unique=True) - if ring.size: - processed = within - faces.append(ring) - ring_ids.append(np.full(ring.size, i)) - return np.concatenate(faces), np.concatenate(ring_ids) - - -@pytest.mark.parametrize( - "outer_radius, radius_step", [(2, 0.1), (60, 1), (250, 10), (360, 3)] -) -def test_ring_membership_matches_per_radius_queries( - gridpath, outer_radius, radius_step -): - from uxarray.grid.coordinates import _lonlat_rad_to_xyz - - uxgrid = ux.open_grid(gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc")) - kdtree = uxgrid._get_scipy_kd_tree() - radii = np.arange(0.0, outer_radius + radius_step, radius_step) - chord_radii = 2.0 * np.sin(np.deg2rad(radii) / 2.0) - - for lon, lat in [(45, 0), (44.7, 0.6), (-45, 0)]: - center = np.array(_lonlat_rad_to_xyz(np.deg2rad(lon), np.deg2rad(lat))) - expected = _per_radius_rings(kdtree, center, chord_radii) - actual = _ring_membership(kdtree, center, chord_radii) - np.testing.assert_array_equal(actual[0], expected[0]) - np.testing.assert_array_equal(actual[1], expected[1]) - - -def test_ring_membership_ties_on_ring_edges(gridpath): - """Radii exactly at face distances, where the kd-tree's pruning rounding decides.""" - uxgrid = ux.open_grid(gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc")) - kdtree = uxgrid._get_scipy_kd_tree() + 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(200): - center = kdtree.data[rng.integers(uxgrid.n_face)].copy() - offsets = kdtree.data[rng.integers(uxgrid.n_face, size=5)] - center - chord_radii = np.sort(np.append(np.sqrt((offsets * offsets).sum(axis=1)), 0.0)) - expected = _per_radius_rings(kdtree, center, chord_radii) - actual = _ring_membership(kdtree, center, chord_radii) - np.testing.assert_array_equal(actual[0], expected[0]) - np.testing.assert_array_equal(actual[1], expected[1]) + 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]) diff --git a/uxarray/core/azimuthal.py b/uxarray/core/azimuthal.py index 168995d1e..a529a90a5 100644 --- a/uxarray/core/azimuthal.py +++ b/uxarray/core/azimuthal.py @@ -4,10 +4,6 @@ import xarray as xr from scipy import sparse -# Squared chord distances are at most 4, so this is ~1000 ulps: well past the -# kd-tree's rounding, yet faces this close to a ring's edge are rare -_TIE_TOLERANCE = 1e-12 - def _compute_azimuthal_mean(uxda, center_xyz, chord_radii): """Computes area-weighted means over the rings between successive bounding circles. @@ -31,10 +27,15 @@ def _compute_azimuthal_mean(uxda, center_xyz, chord_radii): shape=(n_radius, faces.size), ) - ring_data = uxda.isel(n_face=faces, ignore_grid=True).to_xarray() - if ring_data.chunks is not None: + # 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 = ring_data.chunk({"n_face": -1}) + ring_data = span.isel(n_face=faces - start).chunk({"n_face": -1}) means = xr.apply_ufunc( _apply_ring_weights, @@ -42,64 +43,42 @@ def _compute_azimuthal_mean(uxda, center_xyz, chord_radii): input_core_dims=[["n_face"]], output_core_dims=[["radius"]], dask="parallelized", - kwargs={"weights": weights, "empty": hit_count == 0}, + kwargs={"weights": weights}, output_dtypes=[np.result_type(weights.dtype, ring_data.dtype)], - dask_gufunc_kwargs={"output_sizes": {"radius": n_radius}}, + # 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).data, hit_count + 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 and their ring indices, ordered by ring and then face. + 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 - chord_bounds = np.maximum.accumulate(chord_radii) - chord_bounds = chord_bounds * chord_bounds + bounds = np.maximum.accumulate(chord_radii) ** 2 - # One query at the outermost circle, then bin by squared chord distance, - # summed in the kd-tree's own order + # 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()), dtype=np.int_ - ) - diff = kdtree.data[faces] - center_xyz - dist2 = diff[:, 0] * diff[:, 0] + diff[:, 1] * diff[:, 1] + diff[:, 2] * diff[:, 2] - ring_ids = np.searchsorted(chord_bounds, dist2, side="left") - # the outer query has the final say on whether a face is in the outermost ring - outermost = np.searchsorted(chord_bounds, chord_bounds[-1], side="left") - ring_ids = np.minimum(ring_ids, outermost) - - # The tree prunes whole nodes on distances that carry rounding of their - # own, so within an ulp or so of a circle it can disagree with - # dist2 <= r**2. Settle faces that close by asking the tree at that circle. - near_circle = np.full(faces.size, -1) - for circle in (ring_ids - 1, ring_ids): # the circles just inside and outside - circle = np.clip(circle, 0, outermost) - # first circle of a plateau, where a face on that bound first enters - circle = np.searchsorted(chord_bounds, chord_bounds[circle], side="left") - near = np.abs(dist2 - chord_bounds[circle]) <= _TIE_TOLERANCE - near_circle[near] = circle[near] - # the outer query already settled the outermost circle - tied_circles = np.unique( - near_circle[(near_circle >= 0) & (near_circle < outermost)] + kdtree.query_ball_point( + center_xyz, chord_radii.max() * (1 + 1e-9), return_sorted=True + ), + dtype=np.int_, ) - for circle in tied_circles: - tied = near_circle == circle - inside = np.isin( - faces[tied], kdtree.query_ball_point(center_xyz, chord_radii[circle]) - ) - next_circle = np.searchsorted(chord_bounds, chord_bounds[circle], side="right") - ring_ids[tied] = np.where(inside, circle, next_circle) - - order = np.lexsort((faces, ring_ids)) - return faces[order], ring_ids[order] + 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, empty): +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() @@ -107,5 +86,5 @@ def _apply_ring_weights(values, weights, empty): # 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[:, empty] = np.nan + means[:, np.diff(weights.indptr) == 0] = np.nan # empty rings return means.reshape(values.shape[:-1] + (weights.shape[0],)) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index e2f055808..eba1a3cd8 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1024,23 +1024,18 @@ def azimuthal_mean( means, hit_count = _compute_azimuthal_mean(self, center_xyz, chord_radii) - # 'radius' takes the former n_face position - face_axis = self.dims.index("n_face") - dims = list(self.dims) - dims[face_axis] = "radius" - 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 From c794f5666a83ee6583d97b56e56d09fcb6e27492 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 26 Oct 2026 17:50:36 -0500 Subject: [PATCH 4/4] Fail the azimuthal profile tests on any NaN past radius 0 The consolidated monotone test dropped non-finite differences, so a NaN in a non-empty ring passed silently. Only the empty ring at radius 0 is NaN, and res[1:] already excludes it, so neither profile needs the filter; the old test_gaussian had none. Co-Authored-By: Claude Opus 5.5 --- test/core/test_azimuthal.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/core/test_azimuthal.py b/test/core/test_azimuthal.py index aa6cb3735..34c548d7c 100644 --- a/test/core/test_azimuthal.py +++ b/test/core/test_azimuthal.py @@ -41,9 +41,9 @@ def test_monotone_profile(uxds, name, sign): center_coord=(45, 0), outer_radius=2, radius_step=0.5 ) - # gaussian decreases away from the center, inverse_gaussian increases - diffs = sign * res[1:].diff("radius").values - np.testing.assert_array_less(-1e-12, diffs[np.isfinite(diffs)]) + # 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) @pytest.mark.parametrize("chunks", [None, {"n_face": 50, "time": 1}])