add numba_math.py; optimize coordinates.py and _small_angle_of_2_vectors - #1674
Open
Sevans711 wants to merge 4 commits into
Open
add numba_math.py; optimize coordinates.py and _small_angle_of_2_vectors#1674Sevans711 wants to merge 4 commits into
_small_angle_of_2_vectors#1674Sevans711 wants to merge 4 commits into
Conversation
E.g. locally seeing grid.construct_face_centers('welzl') time decrease from 14 s to ~150 ms, for a grid with 157k faces.
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
13 tasks
Sevans711
marked this pull request as ready for review
August 14, 2026 15:25
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.
Addresses part of #1648 but does not fully close it.
Overview
Adds
uxarray/utils/numba_math.pyto provide numba routines primarily for 3-vectors. Applying these directly to tuples to avoid constructing lots of tiny numpy arrays inside numba routines can cause huge speedup in some cases.Optimizes numba routines in
uxarray/grid/coordinates.py, and the_small_angle_of_2_vectorsfunction (which is called directly by some routines in coordinates.py). As per discussion in #1648, this PR does not yet attempt to make these optimizations more broadly, but instead just focuses on improving a single area of the code.After these changes, locally I am seeing grid.construct_face_centers('welzl') time decrease from ~14 s to ~150 ms, for a grid with 157k faces. ASV benchmarks similarly show huge improvement (>=20x speedup) there.
Sidenote: attempted to make a cleverer / more convenient interface, such as
def _numba_add3(*vectors)which could add multiple vectors together, instead of justdef _numba_add3(u, v)but that leads to extremely confusing numba errors in case of any type mismatch between values in tuples. Seems like a better idea to keep this simple (yet maybe slightly less convenient) interface to avoid those confusing errors.Tiny expansion of PR scope: also fixes type-hinting for
_lonlat_rad_to_xyzand_xyz_to_lonlat_rad.PR Checklist
General
Testing & Benchmarking
Documentation
docs/api.rst_)AI Disclosure
AI Usage: GitHub Copilot's inline code suggestions; discussions with Claude Sonnet and ChatGPT to help understand and debug numba routines and confusing numba errors.