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

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

Reviewers

@rajeeja rajeeja Awaiting requested review from rajeeja
@erogluorhan erogluorhan Awaiting requested review from erogluorhan
@cmdupuis3 cmdupuis3 Awaiting requested review from cmdupuis3

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

Assignees

No one assigned

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

1 participant

@Sevans711