diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 62dc2935b..0a105f390 100644 --- a/benchmarks/mpas_ocean.py +++ b/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.""" diff --git a/test/core/test_azimuthal.py b/test/core/test_azimuthal.py index 1ecb26d04..34c548d7c 100644 --- a/test/core/test_azimuthal.py +++ b/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]) diff --git a/uxarray/core/azimuthal.py b/uxarray/core/azimuthal.py new file mode 100644 index 000000000..a529a90a5 --- /dev/null +++ b/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],)) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index a10f4ab1c..898f7ad79 100644 --- a/uxarray/core/dataarray.py +++ b/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