add numba_math.py; optimize coordinates.py and _small_angle_of_2_vectors - #1674
Merged
Sevans711 merged 15 commits intoAug 18, 2026
Merged
Conversation
E.g. locally seeing grid.construct_face_centers('welzl') time decrease from 14 s to ~150 ms, for a grid with 157k faces.
Sevans711
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
13 tasks
Sevans711
marked this pull request as ready for review
August 14, 2026 15:25
Sevans711
added
run-benchmark
cmdupuis3
requested changes
Aug 17, 2026
cmdupuis3
added
run-benchmark
Collaborator
|
Also, I think if we're going to go this route with |
Collaborator
|
Ope, wait, there's an issue with operator overloads. If you disable the numba jit, the decorator skips and Python goes back to concatenating with The joys of Python... |
Collaborator
Author
|
Thank you @cmdupuis3 for reviewing!
Yeahhh, while Python objects would be happy to support a convenient interface with overloads like that, it seems to me like for performance-optimized numba it might be better to just keep the clunkier interface. Numba compilations make things harder to reason about, so keeping it simple seems to be the way to go here. (E.g. I also originally tried To me it doesn't feel like a change to the whole repo philosophy, only a change to the philosophy of numba-compiled functions (wherein this directly causes huge speedup). Outside of njit-decorated functions, I believe uxarray should certainly still use numpy/dask/xarray arrays for dealing with vectors! |
erogluorhan
approved these changes
Aug 17, 2026
erogluorhan
left a comment
erogluorhan
left a comment
Member
There was a problem hiding this comment.
This looks good to me, so I approve it; just a single comment below though.
cmdupuis3
approved these changes
Aug 18, 2026
…ttps://github.com/UXARRAY/uxarray into sevans/add-numba_math.py--optimize-coordinates.py
Sevans711
deleted the
sevans/add-numba_math.py--optimize-coordinates.py
branch
August 18, 2026 20:38
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 #1680 and closes #1681 (both are sub-issues of #1648).
Overview
(For #1680) 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.(For #1681) 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.