Speed up rhef via sort-and-group inner loop

[GillySpace27]

Motivation

On large solar maps ~sunkit_image.radial.rhef spends most of its time bucketing pixels into bins, not equalizing ranks. The per-bin loop rebuilds a (map_r >= lo & map_r < hi) boolean mask on every iteration, so the work scales O(N × nbins) — at a 4096² image with 2048 bins that is ~33 G boolean operations before any ranking happens. The kernel cost shows up before the user sees any output and grows with how fine they want their bins, which is exactly the wrong scaling for an algorithm whose whole point is "rank pixels per radial bin."

This PR removes that bottleneck. As the original author of rhef in sunkit-image, this had been on my back burner for a while; the downstream PUNCH analysis I drive now hit the wall hard enough to force the fix.

Summary

Replace the per-bin mask loop with a single sort-and-group pass:

1. One searchsorted on the bin lower edges gives every pixel its bin index in O(N log nbins).
2. One stable argsort groups same-bin pixels into contiguous slices.
3. The per-bin ranking step is then identical to the old code, just indexed into the sorted-slice view instead of via a boolean mask.

Equivalence

Output is byte-identical to the original implementation across both supported ranking methods (scipy, numpy) and across application_radius, upsilon, and overlapping-bin configurations. A new test_rhef_matches_reference_loop equivalence suite re-implements the old per-bin loop inline as the oracle and asserts both np.array_equal(out[finite], ref[finite]) AND np.array_equal(np.isfinite(out), np.isfinite(ref)) — the latter catching the single-corner-pixel edge case where r == edges_hi[-1] lands in no bin under < hi semantics.

Performance

Measured on a synthetic Map with method="scipy", scipy 1.13, numpy 2.0, macOS arm64:

N	nbins	old	new	speedup
2048	256	1.84 s	0.90 s	2.0×
2048	720	3.75 s	0.92 s	4.1×
2048	2048	9.10 s	0.95 s	9.6×
4096	256	5.46 s	2.51 s	2.2×
4096	720	11.26 s	2.53 s	4.5×
4096	2048	27.53 s	2.61 s	10.5×

The kernel cost no longer scales with bin count; wins grow with both image size and bin density.

cProfile on the 4096²/720 case shows the remaining ~2.5 s is dominated by find_pixel_radii / blackout_pixels_above_radius's WCS pixel-to-world lookup (~2.1 s of the 2.5 s — called twice). The inner RHEF kernel itself is ~150 ms. Caching that WCS step across the two calls is the subject of a follow-up PR (#325).

Notes

• No changes to the public API: signature, parameter defaults, units, return type, and plot_settings["norm"] side-effect are all preserved.
• The towncrier fragment changelog/324.feature.rst is included.
• The PR also carries a small drive-by fix (separate commit) pinning skimage.morphology.white_tophat's mode='reflect' argument in sunkit_image/stara.py, which was failing on py314-devdeps against the scikit-image 2.0 deprecation. The fix is conditional on the kwarg's availability so it stays compatible with the project's minimum scikit-image (0.20). Happy to split it into its own PR if maintainers prefer.

cc @sunpy/sunkit-image-maintainers

[GillySpace27]

Friendly ping for a review whenever someone has a cycle — CI is green (21/21) and the change is self-contained: it swaps the per-bin boolean-mask loop in rhef for a single sort-and-group pass (one searchsorted to bin, one stable argsort to group), dropping the bucketing cost from O(N·nbins) to O(N·log nbins). Output is unchanged — test_rhef_matches_reference_loop re-implements the old per-bin loop inline as an oracle and asserts byte-identical results across all ranking methods and parameter combinations.

@ayshih since you're already in this corner of the code on #325, your eyes would be very welcome here too.

[ayshih]

The actual implementation looks good to me. Can you confirm that sunkit_image.tests.test_radial.test_fig_rhef would be able to catch if this new implementation had significant errors?

My primary suggestions are about documentation. In my opinion, the comments in the function's code should not refer to the old implementation, but the comments in the test module needs to elaborate more about the motivation/history given the complexity of the tests

[GillySpace27]

Re: confirming test_fig_rhef would catch a broken kernel — yes. It runs rhef on a real AIA 171 map and renders the result through @figure_test's baseline-hash comparison, so any numerically significant change to the per-bin ranking shifts pixels in the output image and fails the hash check. It's the real-data complement to the synthetic equivalence tests, which assert bit-identical output (np.array_equal) against an inlined faithful copy of the original per-bin loop across all three ranking methods plus the edge cases (empty bins, fill propagation, application_radius, overlapping bins, upsilon). I expanded that oracle test's docstring to record this rationale.

🤖 Posted by Claude Code

[GillySpace27]

Thanks for the review! I think your comments were against 6ace76e; I've since pushed 65913cc (now the PR head) which addresses all of them:

• Overlapping bins — documented as a Notes section in the rhef docstring (highest-index bin whose lower edge the pixel lies above; ranked only if it's also below that bin's upper edge, else left at fill). The inner-loop comments now describe that design directly instead of referring to the removed per-bin loop.
• Clip comment — expanded to explain why clipping a -1 index to 0 is safe: in_bin is already False for those pixels (the flat_b >= 0 test failed), so the clipped edges_hi read can't flip a genuinely out-of-range pixel back in.
• Test narrative — the reference-oracle docstring now spells out the motivation/history (why an inlined copy of the original loop, and that it keeps proving bit-identity against future refactors), and points to test_fig_rhef as the real-data backstop.
• :487 — no reason for the asymmetry; all three equivalence tests now use np.array_equal for consistency.
• stara / white_tophat — the skimage 2.0 deprecation filter is moved out of the library and into pytest.ini's filterwarnings ignore list.
• Changelog — reworded per your suggestion.

On your test_fig_rhef question: yes. It's a hash-based @figure_test that runs rhef on real AIA 171 data (method="scipy", upsilon=None) and compares the rendered figure against a stored hash, so any significant change in the output pixels changes the hash and fails. Between that real-data backstop and test_rhef_matches_reference_loop (which asserts bit-identical output vs. the inlined original loop across all three ranking methods plus application_radius/upsilon/overlap configs), I'm confident a regression would be caught.

[GillySpace27]

The failing py313-figure / figure-tests check is unrelated to this PR: it is test_occult2_fig[array|map], which this PR does not touch (the diff is confined to radial.py, stara.py, tests/test_radial.py, pytest.ini, and the changelog).

The difference image is blank (RMS < tolerance) — a hash-only mismatch: the committed baseline sunkit_image/tests/figure_hashes_mpl_3108_ft_261_sunpy_710_astropy_720.json predates the matplotlib/freetype in the current CI image, so the OCCULT-2 figure renders pixel-identically but hashes differently. main fails this same check (this branch is 0 commits behind main), so it is pre-existing repo-wide drift rather than a regression here — likely a maintainer-side figure-hash refresh.

The RHEF kernel change is bit-identical to the previous implementation across all three ranking methods, covered by the (passing) test_radial equivalence suite; no RHEF figure test is among the failures.

🤖 Posted by Claude Code

[ayshih]

Aside from the failing figure tests (see #326), this PR looks good to me

[GillySpace27]

Aside from the failing figure tests (see #326), this PR looks good to me

How should I go about fixing this? Do we wait for #326? Or do I rebase onto that one maybe?

Thanks!