Skip to content

Avoid tiny numpy arrays in numba: optimize area.py - #1701

Open
Sevans711 wants to merge 4 commits into
mainfrom
sevans/tinynumpynumba-optimize-area.py
Open

Avoid tiny numpy arrays in numba: optimize area.py#1701
Sevans711 wants to merge 4 commits into
mainfrom
sevans/tinynumpynumba-optimize-area.py

Conversation

@Sevans711

@Sevans711 Sevans711 commented Aug 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Closes #1699 (sub-issue of #1648)

Overview

Optimizes numba routines in uxarray/grid/area.py to avoid constructing many tiny numpy arrays inside numba routines, as discussed in #1648. Locally I saw roughly 4x speedup in a few face areas computations; hopefully ASV benchmarks will show similar improvements.

Minor note: intentionally left quadrature_dg and dw as numpy arrays because those are only constructed once.

Small expansions of PR scope:

  • Now raises ValueError in face areas computations when using an invalid order, instead of silently returning zeros. Added corresponding regression test in test_areas.py.
  • Improves docstrings for helper functions in area.py, e.g. adds notes about shapes of dG and dW arrays.
  • Hoists if gaussian out of for p in range(n_weights) in _face_area_from_quadrature. This surely has negligible impact compared to the other changes, but maybe matters a tiny amount?

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)
  • [N/A] 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)

AI Disclosure

AI Usage: Claude discussions, GitHub Copilot's inline code suggestions.

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

@Sevans711 Sevans711 added scalability Related to scalability & performance efforts run-benchmark Run ASV benchmark workflow labels Aug 21, 2026
@github-actions

github-actions Bot commented Aug 21, 2026 •
edited
Loading

Copy link
Copy Markdown

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change Before [9d5bcdb] After [d0f920d] Ratio Benchmark (Parameter)
- 57.9±0.09ms 13.7±0.09ms 0.24 mpas_ocean.FaceAreas.time_face_areas('120km')
- 838k 744k 0.89 mpas_ocean.FaceAreas.track_peakmem_face_areas('480km')
- 166±0.8ms 49.6±0.4ms 0.3 mpas_ocean.Gradient.time_gradient('120km')
- 11.6±0.07ms 4.25±0.02ms 0.37 mpas_ocean.Gradient.time_gradient('480km')

Benchmarks that have stayed the same:

Change Before [9d5bcdb] After [d0f920d] Ratio Benchmark (Parameter)
232±0.3ms 240±7ms 1.03 bench_connectivity.Connectivity.time_edge_face('120km')
13.7±0.05ms 13.5±0.09ms 0.99 bench_connectivity.Connectivity.time_edge_face('480km')
232±0.4ms 232±0.6ms 1.00 bench_connectivity.Connectivity.time_edge_node('120km')
12.8±0.2ms 12.8±0.3ms 1.00 bench_connectivity.Connectivity.time_edge_node('480km')
243±6ms 232±0.2ms 0.95 bench_connectivity.Connectivity.time_face_edge('120km')
14.0±0.7ms 13.4±0.04ms 0.96 bench_connectivity.Connectivity.time_face_edge('480km')
834±2ms 830±2ms 0.99 bench_connectivity.Connectivity.time_face_face('120km')
51.5±0.6ms 51.6±0.5ms 1.00 bench_connectivity.Connectivity.time_face_face('480km')
66.2±4μs 63.6±2μs 0.96 bench_connectivity.Connectivity.time_face_node('120km')
62.3±1μs 62.0±2μs 1.00 bench_connectivity.Connectivity.time_face_node('480km')
403±8μs 389±5μs 0.97 bench_connectivity.Connectivity.time_n_nodes_per_face('120km')
319±4μs 308±5μs 0.97 bench_connectivity.Connectivity.time_n_nodes_per_face('480km')
241±10ms 233±0.6ms 0.96 bench_connectivity.Connectivity.time_node_edge('120km')
13.0±0.06ms 13.4±0.3ms 1.03 bench_connectivity.Connectivity.time_node_edge('480km')
84.2±3ms 82.6±2ms 0.98 bench_connectivity.Connectivity.time_node_face('120km')
4.92±0.05ms 4.90±0.06ms 1.00 bench_connectivity.Connectivity.time_node_face('480km')
7.55±0.03ms 7.61±0.03ms 1.01 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
2.52±0.1ms 2.49±0.09ms 0.99 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
10.4±10s 9.41±10ms ~0.00 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
1.99±0.02ms 1.97±0.01ms 0.99 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
57.3k 57.3k 1.00 face_bounds.FaceBounds.track_nbytes_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
12.3k 12.3k 1.00 face_bounds.FaceBounds.track_nbytes_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
123k 123k 1.00 face_bounds.FaceBounds.track_nbytes_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
128 128 1.00 face_bounds.FaceBounds.track_nbytes_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
1.27M 1.27M 1.00 face_bounds.FaceBounds.track_nbytes_grid_with_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
50.1k 50.1k 1.00 face_bounds.FaceBounds.track_nbytes_grid_with_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
1.48M 1.48M 1.00 face_bounds.FaceBounds.track_nbytes_grid_with_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
712 712 1.00 face_bounds.FaceBounds.track_nbytes_grid_with_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
1.98M 1.98M 1.00 face_bounds.FaceBounds.track_peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
1.98M 1.98M 1.00 face_bounds.FaceBounds.track_peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
2.15M 2.15M 1.00 face_bounds.FaceBounds.track_peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
38.3k 38.3k 1.00 face_bounds.FaceBounds.track_peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
336M 336M 1.00 face_bounds.FaceBoundsColdStartRss.peakmem_open_and_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
367M 365M 0.99 face_bounds.FaceBoundsColdStartRss.peakmem_open_and_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
337M 337M 1.00 face_bounds.FaceBoundsColdStartRss.peakmem_open_and_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
336M 337M 1.00 face_bounds.FaceBoundsColdStartRss.peakmem_open_and_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
1.16±0.06μs 1.20±0.04μs 1.03 geometry_kernels.AccucrossKernels.time_accucross
2.61±0.04μs 2.64±0.04μs 1.01 geometry_kernels.AccucrossKernels.time_accucross_pair
406±9ns 401±20ns 0.99 geometry_kernels.EFTPrimitives.time_acc_sqrt_re
396±10ns 395±10ns 1.00 geometry_kernels.EFTPrimitives.time_diff_of_products
340±9ns 370±30ns 1.09 geometry_kernels.EFTPrimitives.time_two_prod
375±40ns 351±8ns 0.94 geometry_kernels.EFTPrimitives.time_two_sum
1.49±0.05μs 1.42±0.03μs 0.95 geometry_kernels.GCAConstLatIntersection.time_accux_constlat_kernel
1.01±0.03μs 1.03±0.03μs 1.01 geometry_kernels.GCAConstLatIntersection.time_gca_const_lat_intersection
1.82±0.09μs 1.77±0.02μs 0.98 geometry_kernels.GCAConstLatIntersection.time_try_gca_const_lat_intersection
1.53±0.03μs 1.56±0.02μs 1.02 geometry_kernels.GCAGCAIntersection.time_accux_gca_kernel
1.27±0.03μs 1.27±0.01μs 1.00 geometry_kernels.GCAGCAIntersection.time_gca_gca_intersection
2.00±0.03μs 2.02±0.02μs 1.01 geometry_kernels.GCAGCAIntersection.time_try_gca_gca_intersection
38.6±0.09μs 37.6±0.4μs 0.97 geometry_kernels.OrientPredicates.time_on_minor_arc
691±30ns 706±30ns 1.02 geometry_kernels.OrientPredicates.time_orient3d_on_sphere
2.92±0.1ms 2.80±0.01ms 0.96 geometry_samebody.SameBodyConstLat.time_accux_dispatch
1.33±0ms 1.33±0ms 1.00 geometry_samebody.SameBodyConstLat.time_accux_kernel
1.80±0.01ms 1.80±0.01ms 1.00 geometry_samebody.SameBodyConstLat.time_fp64_dispatch
155±5μs 153±0.8μs 0.99 geometry_samebody.SameBodyConstLat.time_fp64_kernel
32.6±0.06ms 32.5±0.02ms 1.00 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_dispatch
15.6±2ms 10.8±0.08ms ~0.69 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_kernel
27.2±0ms 28.3±0.6ms 1.04 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_dispatch
4.36±0ms 4.61±0.09ms 1.06 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_kernel
805±8ms 800±5ms 0.99 import.Imports.timeraw_import_uxarray
292M 293M 1.00 import.Imports.track_peakmem_import_uxarray
2.34±0.03ms 2.30±0.03ms 0.98 mpas_ocean.CheckNorm.time_check_norm('120km')
1.87±0.05ms 1.81±0.02ms 0.97 mpas_ocean.CheckNorm.time_check_norm('480km')
733±5ms 720±3ms 0.98 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
45.3±0.3ms 45.6±0.2ms 1.01 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
631±8μs 635±7μs 1.01 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
559±5μs 550±7μs 0.98 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
5.13±0.08ms 5.10±0.02ms 0.99 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('120km')
3.60±0.07ms 3.60±0.02ms 1.00 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('480km')
102±0.3ms 102±0.3ms 1.01 mpas_ocean.ConstructFaceLatLon.time_welzl('120km')
10.4±0.05ms 10.2±0.3ms 0.98 mpas_ocean.ConstructFaceLatLon.time_welzl('480km')
19.7±0.02ms 19.7±0.02ms 1.00 mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
1.14±0.01ms 1.14±0.01ms 0.99 mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
10.6±0.01ms 10.7±0.03ms 1.00 mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')
772±5μs 743±10μs 0.96 mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
525±3ms 523±1ms 1.00 mpas_ocean.CrossSections.time_const_lat('120km', 1)
266±0.4ms 267±0.4ms 1.01 mpas_ocean.CrossSections.time_const_lat('120km', 2)
138±0.9ms 137±0.9ms 0.99 mpas_ocean.CrossSections.time_const_lat('120km', 4)
463±0.7ms 464±2ms 1.00 mpas_ocean.CrossSections.time_const_lat('480km', 1)
233±2ms 233±0.8ms 1.00 mpas_ocean.CrossSections.time_const_lat('480km', 2)
121±0.7ms 121±0.6ms 1.00 mpas_ocean.CrossSections.time_const_lat('480km', 4)
356M 356M 1.00 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('120km', 1)
356M 356M 1.00 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('120km', 2)
355M 356M 1.00 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('120km', 4)
339M 339M 1.00 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('480km', 1)
339M 339M 1.00 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('480km', 2)
339M 339M 1.00 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('480km', 4)
22.7±0.3ms 22.2±0.07ms 0.98 mpas_ocean.DualMesh.time_dual_mesh_construction('120km')
2.58±0.02ms 2.59±0.02ms 1.00 mpas_ocean.DualMesh.time_dual_mesh_construction('480km')
4.82±5s 3.79±0.4ms ~0.00 mpas_ocean.FaceAreas.time_face_areas('480km')
229k 229k 1.00 mpas_ocean.FaceAreas.track_nbytes_face_areas('120km')
14.3k 14.3k 1.00 mpas_ocean.FaceAreas.track_nbytes_face_areas('480km')
2.12M 2.12M 1.00 mpas_ocean.FaceAreas.track_peakmem_face_areas('120km')
844±8ms 835±3ms 0.99 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
51.5±0.8ms 50.2±1ms 0.97 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
73.6±0.1ms 73.4±0.3ms 1.00 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
5.75±0.1ms 5.65±0.08ms 0.98 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
457k 457k 1.00 mpas_ocean.Gradient.track_nbytes_gradient('120km')
28.7k 28.7k 1.00 mpas_ocean.Gradient.track_nbytes_gradient('480km')
5.08M 5.08M 1.00 mpas_ocean.Gradient.track_peakmem_gradient('120km')
328k 328k 1.00 mpas_ocean.Gradient.track_peakmem_gradient('480km')
352M 352M 1.00 mpas_ocean.GradientColdStartRss.peakmem_gradient('120km')
331M 331M 1.00 mpas_ocean.GradientColdStartRss.peakmem_gradient('480km')
374±5μs 360±6μs 0.96 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
210±8μs 206±9μs 0.98 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
604±6μs 584±8μs 0.97 mpas_ocean.Integrate.time_integrate('120km')
512±6μs 518±20μs 1.01 mpas_ocean.Integrate.time_integrate('480km')
18.4M 18.4M 1.00 mpas_ocean.Integrate.track_nbytes_integrate('120km')
1.2M 1.2M 1.00 mpas_ocean.Integrate.track_nbytes_integrate('480km')
187±0.7ms 185±0.8ms 0.99 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
186±0.7ms 187±1ms 1.00 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
186±1ms 187±0.9ms 1.01 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
14.4±0.1ms 13.9±0.1ms 0.97 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
14.3±0.3ms 13.9±0.05ms 0.97 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
14.0±0.06ms 14.1±0.1ms 1.01 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
364±5μs 371±20μs 1.02 mpas_ocean.PointInPolygon.time_face_search_lonlat('120km')
344±7μs 349±6μs 1.02 mpas_ocean.PointInPolygon.time_face_search_lonlat('480km')
339±7μs 333±8μs 0.98 mpas_ocean.PointInPolygon.time_face_search_xyz('120km')
332±9μs 340±20μs 1.02 mpas_ocean.PointInPolygon.time_face_search_xyz('480km')
206±1ms 203±2ms 0.98 mpas_ocean.RemapDownsample.time_bilinear_remapping
228±2ms 226±1ms 0.99 mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
15.8±0.2ms 15.6±0.09ms 0.98 mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
1.13±0.01s 1.13±0s 1.00 mpas_ocean.RemapUpsample.time_bilinear_remapping
36.5±0.2ms 36.3±0.4ms 0.99 mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
11.7±0.3ms 11.3±0.1ms 0.97 mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
8.93±0.04ms 8.86±0.07ms 0.99 mpas_ocean.ZonalAverage.time_zonal_average('120km')
4.35±0.03ms 4.31±0.02ms 0.99 mpas_ocean.ZonalAverage.time_zonal_average('480km')
357M 360M 1.01 mpas_ocean.ZonalAveragePeakMem.peakmem_zonal_average('120km')
341M 341M 1.00 mpas_ocean.ZonalAveragePeakMem.peakmem_zonal_average('480km')
6.44±0.1ms 6.33±0.01ms 0.98 quad_hexagon.QuadHexagon.time_open_dataset
5.42±0.1ms 5.34±0.03ms 0.99 quad_hexagon.QuadHexagon.time_open_grid
408 408 1.00 quad_hexagon.QuadHexagon.track_nbytes_open_dataset
392 392 1.00 quad_hexagon.QuadHexagon.track_nbytes_open_grid
73.5k 73.5k 1.00 quad_hexagon.QuadHexagon.track_peakmem_open_dataset
72.7k 72.8k 1.00 quad_hexagon.QuadHexagon.track_peakmem_open_grid

@Sevans711 Sevans711 self-assigned this Aug 24, 2026
cmdupuis3
cmdupuis3 approved these changes Sep 1, 2026
View reviewed changes

@cmdupuis3 cmdupuis3 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good, no issues in testing or benchmarking, and the approach here is fine.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cmdupuis3 cmdupuis3 cmdupuis3 approved these changes
@rajeeja rajeeja Awaiting requested review from rajeeja
@erogluorhan erogluorhan Awaiting requested review from erogluorhan

At least 2 approving reviews are required to merge this pull request.

Assignees

@Sevans711 Sevans711

Labels

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

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Avoid tiny numpy arrays in numba: optimize area.py

2 participants

@Sevans711 @cmdupuis3