Avoid tiny numpy arrays in numba: optimize intersections.py (and refactor: "un-scalarize") - #1727
Open
Sevans711 wants to merge 4 commits into
Open
Avoid tiny numpy arrays in numba: optimize intersections.py (and refactor: "un-scalarize")#1727Sevans711 wants to merge 4 commits into
Sevans711 wants to merge 4 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 |
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
1 participant
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.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