Avoid tiny numpy arrays in numba: optimize intersections.py (and refactor: "un-scalarize") - #1727
Open
Sevans711 wants to merge 9 commits into
Open
Avoid tiny numpy arrays in numba: optimize intersections.py (and refactor: "un-scalarize")#1727Sevans711 wants to merge 9 commits into
Sevans711 wants to merge 9 commits into
Conversation
undo many changes from #1688 because tuples are easier to read than scalars, and maintain all performance improvements. (There was some intent there to refactor back to vectors, but it got merged before that happened, with the intent to refactor being pushed as something to be handled later.) Also optimizes to avoid creating numpy array of shape (2,3) for every gca_gca_intersection() call. And, updates gca_gca_intersection() output to always be a tuple of 2 length-3 tuples, filling with NaNs, as promised in the docstring, instead of a numpy array with shape depending on number of intersections.
and other relevant intersections methods. Removes scalarized versions of related methods: _accux_constlat_scalar, _snap_const_lat_endpoint_xy, _on_minor_arc_xyz, (if you really want to provide a bunch of scalars, just convert to tuples at call site; using tuples has the same performance as a scalarized implementation but it means there is no longer any need to maintain two versions of each method.)
Sevans711
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
Sevans711
marked this pull request as ready for review
September 2, 2026 18:37
Collaborator
Author
|
Marking as ready for review because tests passed and ASV benchmarks look good. Pinging @cmdupuis3 regarding the |
Collaborator
|
@Sevans711 We should be able to compare it to main commits prior to #1688, the latest of which is 7eff90d. I can do that and report the results here if you'd like. |
Collaborator
Author
|
@cmdupuis3 Ah, yes if you are able to do that comparison and share results here that would be great! |
dylannelson
reviewed
Sep 3, 2026
dylannelson
requested changes
Sep 3, 2026
dylannelson
left a comment
dylannelson
left a comment
Member
There was a problem hiding this comment.
Looking around here, one small comment
dylannelson
approved these changes
Sep 4, 2026
dylannelson
left a comment
dylannelson
left a comment
Member
There was a problem hiding this comment.
Everything I ran came through clean. One minor line was fixed. I didn't see any blockers. Hoping tests done with Chris turn out well. Looks good on my end! Thanks!
Collaborator
|
The timings are pretty small, so some wide variability is expected, but this is what I see:
|
cmdupuis3
requested changes
Sep 4, 2026
cmdupuis3
left a comment
cmdupuis3
left a comment
Collaborator
There was a problem hiding this comment.
All tests pass, and the benchmarks show improvements. I had Claude scan for accuracy divergences and none were found, although LLMs aren't great for that. Considering that this is a sort-of reversion of #1688, that should be expected.
A few benchmarks still need to be migrated though:
-
geometry_samebody.py:270— import renamed _accux_constlat_scalar → _accux_constlat, but the px, py, nxo, nyo = ... 4-way unpack was kept; the new kernel returns two 3-tuples. TypingError, which takes down the asv class SameBodyConstLat. -
geometry_samebody.py:292— res[0, 0] 2D-indexes the now-tuple return of gca_const_lat_intersection. -
geometry_samebody.py:345— same 2D indexing in main()'s accuracy cross-check.
Collaborator
Author
@cmdupuis3 Thank you for reviewing and pointing these out! I now see that they were originally reported as "failed" in the ASV benchmarking suite results comment above. I believe I have fixed all of these with the latest commit. |
(these benchmarks need to be updated whenever intersections.py is updated!)
cmdupuis3
requested changes
Sep 4, 2026
cmdupuis3
left a comment
cmdupuis3
left a comment
Collaborator
There was a problem hiding this comment.
Two issues:
-
geometry_samebody_gcagca.py:274 uses np.max; the const-lat sibling correctly uses np.nanmax at geometry_samebody.py:296. Since a 1-intersection result is now (point, (nan,nan,nan)), np.max returns NaN — and max(0.0, nan) returns 0.0, because NaN loses every comparison.
-
geometry_samebody_gcagca.py:61 _fp64_try_gca_gca_intersection is plain @njit(cache=True), but production _try_gca_gca_intersection is @njit(cache=True, inline="always", ...). This means the FP64 version will show different behavior in benchmarks, so the FP64 version should have
inline="always"as well.
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
3 participants
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 #1726 (sub-issue of #1648)
Overview
Optimizes numba routines in
uxarray/grid/intersections.pyto avoid constructing many tiny numpy arrays inside numba routines, as discussed in #1648. ASV benchmarks seem to show roughly 2x speedup of intersections algorithms, and no performance degradations. Hard to know for sure how gca_gca intersection performance has been affected because those benchmarks are failing onmainright now (see discussion on #1688 for details), but the results here at least seem to be the right order of magnitude (roughly 1 μs; latest successes of these benchmarks on main were roughly 1 to 1.5μs).Additionally, refactors the intersections.py routines to restore prior behavior (from before #1688 merged) of using 3-vector inputs and outputs where possible. To maintain performance improvements, just use tuples instead of allocating tiny numpy arrays. The discussion in PR 1688 clarified the intent was to restore to the pre-scalarized function interfaces where possible without sacrificing performance.
A few changes are within scope of the original issue but the reasoning is not necessarily obvious directly from inspecting the code diff; clarifying here:
gca_gca_intersectionoutput (len(result) == number of intersections) did not match its docstring, which promised shape (2,3) result but filling unused rows with NaNs. This PR updates the implementation to fill unused rows with NaNs, and always return a tuple of 2 length-3 tuples, like ((x1,y1,z1),(x2,y2,z2)), filling with NaNs to represent nonexistent intersection points. (The consistent shape is necessary for the performance optimization here. Numba complains that it "Can't unify return type" if the output length depends on the number of intersections.) The test suite has been updated accordingly, to check for nans instead of len(result).get_number_of_intersectionsto clarify that it could now be applied to results fromgca_gca_intersectiontoo._accux_constlat_scalar(use_accux_constlatinstead), and_snap_const_lat_endpoint_xy(use_snap_const_lat_endpointinstead). By using tuples instead of tiny numpy arrays, the "non-scalar" versions also avoid allocation costs, so there is no performance-related need to continue maintaining the scalarized versions._on_minor_arc_xyzby moving its logic intoon_minor_arc, since the scalarized version wasn't being used anywhere anymore (except in benchmarking suite, which has been updated appropriately to useon_minor_arc). Updatedon_minor_arcdocstring to clarify inputs don't need to be numpy arrays; tuples of length 3 also work just fine.EDIT: discussions below revealed the need to incorporate the following changes, too:
_fp64_try_gca_gca_intersectionwas using "and" instead of multiplying together 0s and 1s, so fp64 vs accux wasn't actually the only difference between it and intersections.py's_try_gca_gca_intersection.PR Checklist
General
Testing & Benchmarking
Documentation and Examples
docs/api.rst; internal (private) function names start with an underscore (_)AI Disclosure
AI Usage: GitHub Copilot's inline code suggestions, plus some Claude for numba debugging questions