Repository navigation
Conversation
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 <noreply@anthropic.com>
cmdupuis3
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
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 833840d) -> 1.8 ms (833840d) -> 0.5 ms 120km 169 ms -> 10.1 ms -> 1.6 ms Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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 <noreply@anthropic.com>
cmdupuis3
added
run-benchmark
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 <noreply@anthropic.com>
cmdupuis3
marked this pull request as ready for review
October 26, 2026 23:21
Sevans711
requested changes
Oct 27, 2026
Collaborator
There was a problem hiding this comment.
The changes here look good overall, and a 100x speedup is awesome, I'm excited to get this merged into main!
I have two small concerns: (1) with the tests, and (2) in the possibility for tiny differences between these results and the results from main.
(1) I can tell that some of the new tests in test_azimuthal.py look like they should match the old tests (e.g. test_monotone_profile should cover the old test_gaussian and test_inverse_gaussian I think) but I'm not really sure why the old tests needed to be deleted. It would be easier to approve this if the old tests were maintained as-is, and new tests were added, instead of completely deleting/refactoring the old tests. Refactoring the tests at the same time as the code means there's nothing left to clearly demonstrate that the old test suite's guarantees are still 100% guaranteed to hold true after the changes too. Even though it looks like the new tests cover the same things, maybe there's some subtlety I'm missing.
(2) The results after this change are not bit-for-bit identical to the results before it. For a pure performance change I would have probably expected bit-for-bit identical results, and maybe a test that demonstrates as much, as well. Any guesses as to why the values are different? This PR would be very easy to approve if the results were bit-for-bit identical, but it's more difficult to approve when there are tiny differences.
Minimal reproducer for (2):
psi = ux.tutorial.open_dataset("outCSne30-vortex")['psi']
result = psi.azimuthal_mean((10, 20), 40, 2)Running this on main and on this branch, and then subtracting the two, I see some values are exactly the same (diff=0) while some have diff=1.11e-16, 2.22e-16, or 4.44e-16.
Collaborator
Author
|
(1) They were refactored mostly to condense the code, but they are also more flexible. The old tests can be seen with these parameter combos:
I think the git diff is being stupid though, you'll have to compare them side-by-side, or maybe find a good diff tool (git's default tool is terrible even by 2010 standards). Claude could also be helpful. (2) That is true, this code doesn't reproduce exactly. The new code should be more accurate since it's doing arithmetic n float64 rather than float32. The old error should be around 10^-8, while the error now should be around 10^-16. |
cmdupuis3
added
run-benchmark
Collaborator
|
For (1), thank you for clarifying the parameter combos to get from old tests to new tests. I was also sort of hoping to just get the old tests restored since that makes review easier, but I do agree the new tests look like they should reproduce the old ones and look like they should be reasonable. If you don't want to restore the old tests, I'll just make sure to run them on this branch locally to ensure they still pass, before I approve. For (2), can you provide some sort of evidence to demonstrate that the new code is indeed always (or at least almost always) accurate (or, more accurate)? It looks like there's 1 correctness test being introduced here, as If precision is really the only reason there is any difference between here and |
Collaborator
Author
|
(1) Yeah, I'd rather not restore them. The more horizontal alternative would be to do a massive test consolidation sweep, so I think it's reasonable to clean up bits as we go. (2) Casting Try this: from fractions import Fraction
import numpy as np
import uxarray as ux
from uxarray.core.azimuthal import _ring_membership
from uxarray.grid.coordinates import _lonlat_rad_to_xyz
base = "test/meshfiles/mpas/dyamond-30km/"
uxds = ux.open_dataset(base + "gradient_grid_subset.nc", base + "gradient_data_subset.nc")
center, outer_radius, radius_step = (45, 0), 2, 0.25
radii = np.arange(0.0, outer_radius + radius_step, radius_step)
chords = 2.0 * np.sin(np.deg2rad(radii) / 2.0)
center_xyz = np.array(_lonlat_rad_to_xyz(*np.deg2rad(center)))
faces, ring_ids = _ring_membership(uxds.uxgrid._get_scipy_kd_tree(), center_xyz, chords)
areas = uxds.uxgrid.face_areas.values
def exact_mean(x, a):
"""sum(x * a) / sum(a) in exact rational arithmetic."""
return sum(Fraction(float(xi)) * Fraction(float(ai)) for xi, ai in zip(x, a)) / sum(
Fraction(float(ai)) for ai in a
)
def rel_err(value, exact):
return float(abs(Fraction(float(value)) - exact) / abs(exact))
def report(name, data):
print(f"\n{name} (data {data.dtype}, areas {areas.dtype})")
new = data.azimuthal_mean(
center_coord=center, outer_radius=outer_radius, radius_step=radius_step
).values
print(f"{'ring':>4} {'faces':>6} {'old (float32) err':>18} {'new (float64) err':>18}")
for i in range(1, radii.size):
f = faces[ring_ids == i]
x, a = data.values[f], areas[f]
exact = exact_mean(x, a)
old = (x * a).sum() / a.sum() # what weighted_mean() computed: all float32
print(f"{i:>4} {f.size:>6} {rel_err(old, exact):>18.1e} {rel_err(new[i], exact):>18.1e}")
# 1. The real test field
report("gaussian", uxds["gaussian"])
# 2. A constant field: the exact mean is the constant itself
const = uxds["gaussian"].copy(data=np.full(uxds["gaussian"].shape, 0.1, dtype=np.float32))
report("constant 0.1f", const) |
Collaborator
Fair enough!
I think there was a misunderstanding: I wasn't expecting "main casted to float64" to match "main now (which is float32)". But, your previous claim:
made me expect "main casted to float64" should match "this PR". Is that the correct understanding / can you show whether that is actually true? If it's true then I would happily agree that the only reason there is a difference is the float64 versus float32, and approve the PR on the basis of "float64 is more accurate than float32"! |
Collaborator
Author
|
Basically yes, there is no other reason why this code is more accurate than that we're using float64 instead of float32. It's not exactly the same as doing that on main, but it's close, the difference being that you'd need to have the data in float64 to have a fair comparison, but if it is, it would be nearly the same. The change in math on this branch necessarily results in some differences of a few ulps. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1848
Overview
Simple speedup for azimuthal means, and an incidental accuracy fix.
The old tests can be seen with these parameter combos:
test_gaussiantest_monotone_profile[gaussian--1]:name="gaussian",sign=-1test_inverse_gaussiantest_monotone_profile[inverse_gaussian-1]:name="inverse_gaussian",sign=1test_non_zero_hit_countstest_matches_reference[center0-None]and[center0-chunks1]:center=(45, 0),chunksisNoneor{"n_face": 50, "time": 1}test_zero_hit_countstest_matches_reference[center1-None]and[center1-chunks1]:center=(-45, 0)(outside the subset), same twochunksvaluesPR Checklist
General
Testing & Benchmarking
Documentation and Examples
docs/api.rst; internal (private) function names start with an underscore (_)AI Disclosure
AI Usage: Claude Opus 5.5