Skip to content

add numba_math.py; optimize coordinates.py and _small_angle_of_2_vectors - #1674

Open
Sevans711 wants to merge 4 commits into
mainfrom
sevans/add-numba_math.py--optimize-coordinates.py
Open

add numba_math.py; optimize coordinates.py and _small_angle_of_2_vectors#1674
Sevans711 wants to merge 4 commits into
mainfrom
sevans/add-numba_math.py--optimize-coordinates.py

Conversation

@Sevans711

@Sevans711 Sevans711 commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Addresses part of #1648 but does not fully close it.

Overview

Adds uxarray/utils/numba_math.py to 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_vectors function (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 just def _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_xyz and _xyz_to_lonlat_rad.

PR Checklist

General

  • An issue is created and linked
  • Added appropriate labels (if your uxarray repo permissions allow it)
  • Filled out Overview and Expected Usage (if applicable) sections

Testing & Benchmarking

  • Adequate tests are created if there is new functionality
  • Tests are not too basic (such as simply calling a function and nothing else)
  • Tests cover all major paths in your new functions
  • If this PR could affect performance, ran ASV benchmarks and confirmed they show expected behavior (add a new benchmark if necessary)

Documentation

  • Docstrings have been added to all new functions
  • Docstrings have been updated with any function changes
  • [N/A] User (public) functions have been added to docs/api.rst
  • Internal (private) function names start with an underscore (_)

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.

  • I take responsibility for all AI-generated content in my PR.
  • I have tested all AI-generated content in my PR.

E.g. locally seeing grid.construct_face_centers('welzl') time decrease from 14 s to ~150 ms, for a grid with 157k faces.
@Sevans711 Sevans711 added scalability Related to scalability & performance efforts run-benchmark Run ASV benchmark workflow labels Aug 13, 2026
@github-actions

github-actions Bot commented Aug 13, 2026

Copy link
Copy Markdown

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change Before [7adbd4c] After [349f8c3] Ratio Benchmark (Parameter)
- 518M 337M 0.65 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
- 636M 337M 0.53 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
- 3.49±0.02s 99.0±0.3ms 0.03 mpas_ocean.ConstructFaceLatLon.time_welzl('120km')
- 225±2ms 10.7±0.3ms 0.05 mpas_ocean.ConstructFaceLatLon.time_welzl('480km')
- 443M 332M 0.75 mpas_ocean.FaceAreas.peakmem_compute_face_areas('480km')
- 468M 331M 0.71 mpas_ocean.Gradient.peakmem_gradient('480km')

Benchmarks that have stayed the same:

Change Before [7adbd4c] After [349f8c3] Ratio Benchmark (Parameter)
216±9ms 210±7ms 0.97 bench_connectivity.Connectivity.time_edge_face('120km')
12.5±0.1ms 12.5±0.1ms 1 bench_connectivity.Connectivity.time_edge_face('480km')
202±2ms 202±1ms 1 bench_connectivity.Connectivity.time_edge_node('120km')
11.3±0.07ms 11.3±0.1ms 1 bench_connectivity.Connectivity.time_edge_node('480km')
202±2ms 201±0.5ms 0.99 bench_connectivity.Connectivity.time_face_edge('120km')
11.7±0.2ms 11.6±0.07ms 0.99 bench_connectivity.Connectivity.time_face_edge('480km')
905±6ms 905±10ms 1 bench_connectivity.Connectivity.time_face_face('120km')
58.5±0.3ms 60.0±0.3ms 1.03 bench_connectivity.Connectivity.time_face_face('480km')
72.1±1μs 74.9±3μs 1.04 bench_connectivity.Connectivity.time_face_node('120km')
70.8±2μs 71.4±2μs 1.01 bench_connectivity.Connectivity.time_face_node('480km')
420±10μs 420±10μs 1 bench_connectivity.Connectivity.time_n_nodes_per_face('120km')
358±4μs 375±20μs 1.05 bench_connectivity.Connectivity.time_n_nodes_per_face('480km')
201±0.7ms 202±0.7ms 1.01 bench_connectivity.Connectivity.time_node_edge('120km')
11.6±0.07ms 11.8±0.1ms 1.01 bench_connectivity.Connectivity.time_node_edge('480km')
89.5±0.6ms 90.7±0.9ms 1.01 bench_connectivity.Connectivity.time_node_face('120km')
5.34±0.03ms 5.59±0.2ms 1.05 bench_connectivity.Connectivity.time_node_face('480km')
336M 335M 1 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
367M 367M 1 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
8.65±0.07ms 8.79±0.1ms 1.02 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
2.76±0.06ms 2.72±0.04ms 0.99 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
10.3±0.08ms 10.2±0.08ms 0.99 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
2.15±0.01ms 2.15±0ms 1 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
1.27±0.02μs 1.29±0.06μs 1.02 geometry_kernels.AccucrossKernels.time_accucross
2.78±0.03μs 2.79±0.04μs 1 geometry_kernels.AccucrossKernels.time_accucross_pair
466±20ns 471±20ns 1.01 geometry_kernels.EFTPrimitives.time_acc_sqrt_re
455±20ns 456±50ns 1 geometry_kernels.EFTPrimitives.time_diff_of_products
380±8ns 416±30ns 1.09 geometry_kernels.EFTPrimitives.time_two_prod
371±10ns 396±10ns 1.07 geometry_kernels.EFTPrimitives.time_two_sum
1.55±0.04μs 1.50±0.02μs 0.97 geometry_kernels.GCAConstLatIntersection.time_accux_constlat_kernel
1.16±0.02μs 1.19±0.01μs 1.03 geometry_kernels.GCAConstLatIntersection.time_gca_const_lat_intersection
1.92±0.06μs 1.93±0.03μs 1.01 geometry_kernels.GCAConstLatIntersection.time_try_gca_const_lat_intersection
1.74±0.03μs 1.67±0.03μs 0.96 geometry_kernels.GCAGCAIntersection.time_accux_gca_kernel
1.45±0.03μs 1.39±0.04μs 0.96 geometry_kernels.GCAGCAIntersection.time_gca_gca_intersection
2.23±0.04μs 2.17±0.06μs 0.97 geometry_kernels.GCAGCAIntersection.time_try_gca_gca_intersection
54.1±1μs 52.9±0.6μs 0.98 geometry_kernels.OrientPredicates.time_on_minor_arc
1.15±0.06μs 1.10±0.03μs 0.96 geometry_kernels.OrientPredicates.time_orient3d_on_sphere
2.73±0.1ms 2.61±0.01ms 0.96 geometry_samebody.SameBodyConstLat.time_accux_dispatch
1.17±0ms 1.18±0.01ms 1.01 geometry_samebody.SameBodyConstLat.time_accux_kernel
1.73±0.02ms 1.72±0ms 0.99 geometry_samebody.SameBodyConstLat.time_fp64_dispatch
149±0.6μs 149±0.9μs 1 geometry_samebody.SameBodyConstLat.time_fp64_kernel
32.8±0.06ms 32.2±0.04ms 0.98 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_dispatch
10.3±0.01ms 10.3±0.05ms 1 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_kernel
26.5±0.02ms 26.9±0.2ms 1.01 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_dispatch
4.91±0.05ms 4.85±0.02ms 0.99 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_kernel
859±6ms 866±1ms 1.01 import.Imports.timeraw_import_uxarray
2.81±0.02ms 2.81±0.01ms 1 mpas_ocean.CheckNorm.time_check_norm('120km')
2.31±0.01ms 2.39±0.03ms 1.03 mpas_ocean.CheckNorm.time_check_norm('480km')
852±7ms 854±2ms 1 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
54.8±0.5ms 56.2±0.6ms 1.02 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
689±10μs 670±10μs 0.97 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
602±10μs 622±6μs 1.03 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
5.58±0.01ms 5.61±0.05ms 1.01 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('120km')
3.94±0.04ms 4.05±0.03ms 1.03 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('480km')
18.2±0.04ms 18.3±0.03ms 1 mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
1.08±0.01ms 1.09±0.02ms 1.01 mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
10.7±0.04ms 10.7±0.03ms 1 mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')
756±10μs 767±5μs 1.01 mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
728±8ms 746±8ms 1.02 mpas_ocean.CrossSections.time_const_lat('120km', 1)
371±2ms 375±3ms 1.01 mpas_ocean.CrossSections.time_const_lat('120km', 2)
192±1ms 192±3ms 1 mpas_ocean.CrossSections.time_const_lat('120km', 4)
571±3ms 575±2ms 1.01 mpas_ocean.CrossSections.time_const_lat('480km', 1)
286±1ms 288±3ms 1.01 mpas_ocean.CrossSections.time_const_lat('480km', 2)
152±2ms 152±1ms 1 mpas_ocean.CrossSections.time_const_lat('480km', 4)
25.2±0.2ms 26.1±0.4ms 1.04 mpas_ocean.DualMesh.time_dual_mesh_construction('120km')
3.36±0.1ms 3.40±0.03ms 1.01 mpas_ocean.DualMesh.time_dual_mesh_construction('480km')
350M 350M 1 mpas_ocean.FaceAreas.peakmem_compute_face_areas('120km')
61.1±0.3ms 61.1±0.09ms 1 mpas_ocean.FaceAreas.time_compute_face_areas('120km')
7.25±0.01ms 7.28±0.09ms 1 mpas_ocean.FaceAreas.time_compute_face_areas('480km')
952±5ms 960±8ms 1.01 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
56.1±1ms 56.9±0.7ms 1.01 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
85.7±0.3ms 85.4±0.4ms 1 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
5.93±0.09ms 6.14±0.2ms 1.04 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
351M 352M 1 mpas_ocean.Gradient.peakmem_gradient('120km')
176±0.2ms 175±0.3ms 1 mpas_ocean.Gradient.time_gradient('120km')
12.7±0.1ms 12.7±0.06ms 1 mpas_ocean.Gradient.time_gradient('480km')
382±20μs 393±9μs 1.03 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
218±4μs 217±10μs 0.99 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
351M 351M 1 mpas_ocean.Integrate.peakmem_integrate('120km')
329M 330M 1 mpas_ocean.Integrate.peakmem_integrate('480km')
551±10μs 574±10μs 1.04 mpas_ocean.Integrate.time_integrate('120km')
516±10μs 525±20μs 1.02 mpas_ocean.Integrate.time_integrate('480km')
189±5ms 184±1ms 0.97 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
185±1ms 185±2ms 1 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
185±1ms 183±2ms 0.99 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
14.1±0.2ms 14.1±0.2ms 1 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
14.0±0.1ms 13.9±0.2ms 1 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
13.9±0.03ms 13.8±0.1ms 0.99 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
426±10μs 426±9μs 1 mpas_ocean.PointInPolygon.time_face_search_lonlat('120km')
419±10μs 420±9μs 1 mpas_ocean.PointInPolygon.time_face_search_lonlat('480km')
400±10μs 408±9μs 1.02 mpas_ocean.PointInPolygon.time_face_search_xyz('120km')
397±20μs 381±9μs 0.96 mpas_ocean.PointInPolygon.time_face_search_xyz('480km')
248±0.3ms 241±0.8ms 0.97 mpas_ocean.RemapDownsample.time_bilinear_remapping
304±3ms 291±2ms 0.96 mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
16.0±0.08ms 16.5±0.08ms 1.03 mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
1.46±0s 1.41±0.01s 0.97 mpas_ocean.RemapUpsample.time_bilinear_remapping
38.0±0.4ms 38.2±0.2ms 1 mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
12.9±0.3ms 12.9±0.3ms 1 mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
27.5±0.1ms 26.9±0.1ms 0.98 mpas_ocean.ZonalAverage.time_zonal_average('120km')
6.16±0.07ms 6.23±0.2ms 1.01 mpas_ocean.ZonalAverage.time_zonal_average('480km')
326M 329M 1.01 quad_hexagon.QuadHexagon.peakmem_open_dataset
325M 325M 1 quad_hexagon.QuadHexagon.peakmem_open_grid
7.37±0.2ms 7.77±0.4ms 1.05 quad_hexagon.QuadHexagon.time_open_dataset
6.24±0.1ms 6.42±0.1ms 1.03 quad_hexagon.QuadHexagon.time_open_grid

@Sevans711
Sevans711 marked this pull request as ready for review August 14, 2026 15:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-benchmark Run ASV benchmark workflow scalability Related to scalability & performance efforts

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant