Skip to content

Better chunking strategies for constlat intersections and zonal routines. - #1624

Open
cmdupuis3 wants to merge 72 commits into
mainfrom
cmd/accusphere3
Open

Better chunking strategies for constlat intersections and zonal routines.#1624
cmdupuis3 wants to merge 72 commits into
mainfrom
cmd/accusphere3

Conversation

@cmdupuis3

@cmdupuis3 cmdupuis3 commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

This PR contains two post-accusphere optimizations, eliminating low-level hard materializations by using a vector-based masking strategy rather than individual conditionals, and reduced zonal_mean peakmem by building only the candidate faces instead of the whole grid.

Partly addresses #1587

Overview

Phase A eliminated the largest peak-memory amplifier in the zonal path: the port materialized the whole-grid (n_face, n_max, 2, 3) face-edge array up front (a 5.8× blow-up over the node coordinates, with a ~23 MB build peak on a 28k-face grid) even though each latitude only touches ~1% of faces. I added an @njit(parallel=True) subset builder, _get_cartesian_face_edge_nodes_array_subset, that constructs edges only for the candidate faces of a given latitude/band — bit-identical to indexing the full array — and rewired _compute_non_conservative_zonal_mean and _compute_face_band_weights to build per-candidate subsets instead of the whole grid. Peak memory for a 180-latitude zonal_mean dropped from 23.66 MB to 0.27 MB (≈88×) and it ran ~10% faster (no full build, no per-latitude fancy-index copies), verified lossless via a git-stash A/B (the only diff was a pre-existing 4.4e-16 parallel-reduction nondeterminism) with the full zonal/cross-section suite passing.

Phase B moved the six edge/face screeners (constant_lat/lon_intersections_no_extreme, constant_lat/lon_intersections_face_bounds, faces_within_lat/lon_bounds) off @njit and onto plain vectorized NumPy, drawing the boundary so the low-level Numba kernels stay for real geometry while these memory-bound elementwise predicates use NumPy — which is ~2.1× faster here and, unlike an njit kernel that forces a full .values materialization, composes with dask (a single _flatnonzero helper leans on NumPy's array_function protocol so a dask mask reduces block-wise, no explicit dask branch). Call sites pass .data instead of .values so a chunked grid stays lazy, and edge coordinates are gathered positionally (node_z.data[conn.ravel()].reshape(...)) to stay chunk-friendly. Along the way it fixed two real latent bugs: get_edges_at_constant_latitude referenced a nonexistent self.edge_node_z (it raised AttributeError on every call), and both edge paths crashed on chunked grids because xarray can't vindex with a dask indexer. Results are bit-identical to the original per-element loops across 300 randomized trials, dask==numpy, and it's committed as bf0abbe "Lazy intersections".

PR Checklist

General

  • An issue is linked created and linked
  • Add appropriate labels
  • Filled out Overview and Expected Usage (if applicable) sections

Testing

  • Adequate tests are created if there is new functionality
  • Tests cover all possible logical paths in your function
  • Tests are not too basic (such as simply calling a function and nothing else)

Documentation

  • Docstrings have been added to all new functions
  • Docstrings have updated with any function changes

rajeeja added 30 commits June 8, 2026 15:47
…rite intersections, add 241 baseline testsgit status! - most came from accusphere
- benchmarks/geometry_kernels.py: ASV micro-benchmarks for all three
  layers of the EFT intersection stack (_accux_gca, _try_gca_gca_intersection,
  gca_gca_intersection, _accux_constlat, _try_gca_const_lat_intersection,
  gca_const_lat_intersection) plus EFT primitives and point-in-polygon;
  all functions warmed before timing so results reflect steady-state cost
- test/test_plot.py: add test_to_raster_auto_extent verifying that the
  axis limits change and the raster contains finite data
…rectness fixes

Review comments addressed:
- Remove "near-double precision" / "sufficient" overclaims; say "roughly twice
  as accurate" and note the robustness tier boundary clearly
- Explain _lon_bounds_from_vertices is required for UXarray antimeridian
  encoding and cannot be removed
- Add block comment before _no_extreme functions clarifying they are
  pre-existing edge screeners unrelated to the EFT stack
- Document SoS as explicit future work in _point_in_polygon_sphere docstring
- L2 pos_fin/neg_fin: replace ternary with int(); exploit neg=-pos symmetry
- Label computation: drop dead local*0 term, use integer mask arithmetic
- Remove vertex-lat snap from bounds: _face_location_info already captures
  interior arc extrema accurately via the compensated kernel
- _ON_MINOR_ARC_TOL: document intentional 1e-10 vs C++ 1e-8 divergence

Bug fixes:
- on_minor_arc: add antipodal-endpoint guard; a x b = 0 for antipodal inputs
  so every point on the great circle passes the collinearity test (false pos)
- bounds.py: replace mask arithmetic use_ext*z_ext + (1-use_ext)*z_edge with
  plain if/else; 0*NaN = NaN propagates when norm=0, if/else does not
- _point_in_polygon_sphere: ray-nudge now restarts the loop from i=0 so all
  edges are counted with the same ray (mid-loop nudge corrupted crossing parity)

Cleanup:
- Remove _flip_sign, _SIGN_NEG, _SIGN_POS, _SIGN_ZERO dead code from
  point_in_face.py; inline literals in _counts_as_crossing
- Remove _SNAP_TOL_DEG constant and snap_tol_deg parameter throughout bounds.py
- Notebook: fix Grid.get_point_on_face -> get_faces_containing_point; remove
  incorrect geometry.py row from Section 4 table; add accucross_pair and
  acc_sqrt_re to Section 2 building-blocks table
…PI name

- ci/environment.yml: pin tornado<6.5.7 to avoid ssl.SSLError in panel 1.9.3
  on Python 3.11 Windows (conda-forge regression, 2026-06-10)
- intersections.py: remove _gca_gca_intersection_cartesian shim (dead code);
  add comment explaining _snap_const_lat_endpoint snap_sq constant
- test_intersections.py: update 4 call sites to use gca_gca_intersection directly
- spherical-geometry-accuracy.ipynb: fix stale Grid.get_point_on_face ->
  Grid.get_faces_containing_point (2 occurrences)
…rite intersections, add 241 baseline testsgit status! - most came from accusphere
Reconcile diverged accusphere branch. Resolutions:
- intersections.py: restore inline=always on L1 kernels (_accux_constlat,
  _accux_gca) for allocation scalar-replacement
- point_in_face.py: keep restart-loop ray casting (consistent parity),
  adopt named sign constants, drop unused _flip_sign
- arcs.py: keep antipodal-endpoint guard in on_minor_arc
- bounds.py: keep vertex-latitude snapping (snap_tol_deg) path
- computing.py: keep detailed docstring with SIAM/EGUsphere references
Add an LLVM fma intrinsic and route two_prod through a single fused
multiply-add for its error term on hardware that supports it, selected at
import time and validated to be bit-exact against the Veltkamp split. Falls
back to the portable Veltkamp form otherwise, so there is no hard FMA
dependency.

The FMA path is ~2x faster in the compensated geometry kernels (each
two_prod drops from ~17 flops to one FMADD) and is numerically identical:
all 241 AccuSphGeom baseline cases pass unchanged.
Add _accux_constlat_scalar, which takes the arc endpoints as six scalars and
returns the candidate coordinates as scalars instead of two np.empty(3)
arrays. _accux_constlat now wraps it so the array API is unchanged.

Returning scalars lets Numba keep the candidates in registers, so a batch
loop over many edges does no per-point heap allocation. On a 16M-point
const-lat sweep this is ~2.7x faster than the array-returning path and drops
the AccuX/FP64 cost ratio from ~19x to ~7x. Bit-identical results; all 241
AccuSphGeom baseline cases pass.
@cmdupuis3 cmdupuis3 self-assigned this Jul 27, 2026
@cmdupuis3 cmdupuis3 added the scalability Related to scalability & performance efforts label Jul 27, 2026
@cmdupuis3

Copy link
Copy Markdown
Collaborator Author

pre-commit.ci autofix

@cmdupuis3 cmdupuis3 added the benchmarking Related to benchmarks, memory usage, and/or time profiling label Jul 28, 2026
@cmdupuis3 cmdupuis3 added the run-benchmark Run ASV benchmark workflow label Jul 29, 2026
@github-actions

Copy link
Copy Markdown

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change Before [852a0c0] After [c227843] Ratio Benchmark (Parameter)
- 522M 336M 0.64 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
- 631M 335M 0.53 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
- 452M 329M 0.73 mpas_ocean.Gradient.peakmem_gradient('480km')
- 26.4±0.5ms 7.39±0.03ms 0.28 mpas_ocean.ZonalAverage.time_zonal_average('120km')
- 5.79±0.04ms 4.87±0.1ms 0.84 mpas_ocean.ZonalAverage.time_zonal_average('480km')

Benchmarks that have stayed the same:

Change Before [852a0c0] After [c227843] Ratio Benchmark (Parameter)
10.3±0.1μs 10.4±0.07μs 1.01 bench_connectivity.Connectivity.time_edge_face('120km')
10.8±0.1μs 11.3±0.09μs 1.05 bench_connectivity.Connectivity.time_edge_face('480km')
10.5±0.3μs 10.5±0.07μs 1 bench_connectivity.Connectivity.time_edge_node('120km')
10.8±0.06μs 10.9±0.2μs 1.01 bench_connectivity.Connectivity.time_edge_node('480km')
10.3±0.1μs 10.4±0.1μs 1.01 bench_connectivity.Connectivity.time_face_edge('120km')
11.0±0.1μs 11.0±0.1μs 1 bench_connectivity.Connectivity.time_face_edge('480km')
10.1±0.08μs 10.4±0.2μs 1.03 bench_connectivity.Connectivity.time_face_face('120km')
10.8±0.07μs 11.0±0.07μs 1.02 bench_connectivity.Connectivity.time_face_face('480km')
21.0±0.1μs 21.8±0.2μs 1.04 bench_connectivity.Connectivity.time_face_node('120km')
22.8±0.4μs 22.7±0.3μs 1 bench_connectivity.Connectivity.time_face_node('480km')
10.4±0.07μs 10.4±0.1μs 1 bench_connectivity.Connectivity.time_node_edge('120km')
10.9±0.2μs 11.1±0.1μs 1.02 bench_connectivity.Connectivity.time_node_edge('480km')
10.3±0.06μs 10.8±0.07μs 1.04 bench_connectivity.Connectivity.time_node_face('120km')
11.0±0.1μs 11.0±0.2μs 1.01 bench_connectivity.Connectivity.time_node_face('480km')
334M 334M 1 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
363M 365M 1.01 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
22.8±0.3μs 22.9±0.1μs 1 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
10.4±0.05μs 9.85±0.04μs 0.94 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
10.2±0.09ms 10.2±0.02ms 1 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
2.13±0.01ms 2.15±0.01ms 1.01 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
944±20ns 952±1ns 1.01 geometry_kernels.AccucrossKernels.time_accucross
2.44±0.01μs 2.48±0.02μs 1.02 geometry_kernels.AccucrossKernels.time_accucross_pair
292±1ns 294±1ns 1.01 geometry_kernels.EFTPrimitives.time_acc_sqrt_re
281±2ns 278±0.4ns 0.99 geometry_kernels.EFTPrimitives.time_diff_of_products
233±0.1ns 233±2ns 1 geometry_kernels.EFTPrimitives.time_two_prod
237±0.5ns 239±1ns 1.01 geometry_kernels.EFTPrimitives.time_two_sum
1.22±0μs 1.21±0.01μs 1 geometry_kernels.GCAConstLatIntersection.time_accux_constlat_kernel
872±10ns 880±10ns 1.01 geometry_kernels.GCAConstLatIntersection.time_gca_const_lat_intersection
1.60±0.01μs 1.64±0μs 1.02 geometry_kernels.GCAConstLatIntersection.time_try_gca_const_lat_intersection
1.34±0μs 1.33±0.01μs 0.99 geometry_kernels.GCAGCAIntersection.time_accux_gca_kernel
1.13±0.01μs 1.14±0.01μs 1.01 geometry_kernels.GCAGCAIntersection.time_gca_gca_intersection
1.88±0μs 1.85±0μs 0.99 geometry_kernels.GCAGCAIntersection.time_try_gca_gca_intersection
49.3±0.3μs 49.3±0.2μs 1 geometry_kernels.OrientPredicates.time_on_minor_arc
509±0.8ns 512±5ns 1.01 geometry_kernels.OrientPredicates.time_orient3d_on_sphere
2.70±0.1ms 2.59±0.01ms 0.96 geometry_samebody.SameBodyConstLat.time_accux_dispatch
1.17±0ms 1.17±0ms 1 geometry_samebody.SameBodyConstLat.time_accux_kernel
1.71±0ms 1.72±0ms 1 geometry_samebody.SameBodyConstLat.time_fp64_dispatch
146±0.3μs 146±0.2μs 1 geometry_samebody.SameBodyConstLat.time_fp64_kernel
32.3±0.01ms 32.2±0.1ms 1 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_dispatch
10.2±0.04ms 10.2±0.05ms 1 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_kernel
26.4±0.03ms 26.8±0.4ms 1.02 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_dispatch
4.92±0.04ms 4.91±0.03ms 1 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_kernel
821±10ms 810±3ms 0.99 import.Imports.timeraw_import_uxarray
927±3ns 942±20ns 1.02 mpas_ocean.CheckNorm.time_check_norm('120km')
877±10ns 890±10ns 1.01 mpas_ocean.CheckNorm.time_check_norm('480km')
839±10ms 822±5ms 0.98 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
53.7±0.2ms 52.6±0.6ms 0.98 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
14.2±0.08μs 14.3±0.09μs 1 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
14.7±0.09μs 14.7±0.1μs 1 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
4.86±0.01ms 4.89±0.03ms 1.01 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('120km')
3.33±0.02ms 3.40±0.04ms 1.02 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('480km')
3.49±0.03s 3.46±0.01s 0.99 mpas_ocean.ConstructFaceLatLon.time_welzl('120km')
224±0.8ms 222±2ms 0.99 mpas_ocean.ConstructFaceLatLon.time_welzl('480km')
18.2±0.02ms 18.2±0.01ms 1 mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
1.03±0.01ms 1.03±0.01ms 1 mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
10.6±0.03ms 10.6±0.01ms 1 mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')
724±10μs 730±6μs 1.01 mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
703±3ms 700±10ms 1 mpas_ocean.CrossSections.time_const_lat('120km', 1)
353±2ms 352±0.7ms 1 mpas_ocean.CrossSections.time_const_lat('120km', 2)
182±2ms 185±1ms 1.01 mpas_ocean.CrossSections.time_const_lat('120km', 4)
542±4ms 558±3ms 1.03 mpas_ocean.CrossSections.time_const_lat('480km', 1)
274±0.5ms 278±4ms 1.02 mpas_ocean.CrossSections.time_const_lat('480km', 2)
140±0.9ms 143±0.6ms 1.02 mpas_ocean.CrossSections.time_const_lat('480km', 4)
354M 354M 1 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('120km', 1)
354M 354M 1 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('120km', 2)
354M 354M 1 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('120km', 4)
338M 337M 1 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('480km', 1)
337M 338M 1 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('480km', 2)
337M 337M 1 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('480km', 4)
24.6±0.08ms 24.5±0.09ms 1 mpas_ocean.DualMesh.time_dual_mesh_construction('120km')
3.29±0.1ms 3.30±0.04ms 1.01 mpas_ocean.DualMesh.time_dual_mesh_construction('480km')
949±2ms 940±3ms 0.99 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
54.9±0.2ms 53.8±1ms 0.98 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
85.0±0.3ms 82.7±0.3ms 0.97 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
5.63±0.04ms 5.85±0.2ms 1.04 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
350M 350M 1 mpas_ocean.Gradient.peakmem_gradient('120km')
175±0.4ms 174±0.3ms 0.99 mpas_ocean.Gradient.time_gradient('120km')
12.5±0.07ms 12.5±0.1ms 1 mpas_ocean.Gradient.time_gradient('480km')
225±2μs 225±0.7μs 1 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
134±0.9μs 133±1μs 0.99 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
349M 349M 1 mpas_ocean.Integrate.peakmem_integrate('120km')
329M 329M 1 mpas_ocean.Integrate.peakmem_integrate('480km')
216±0.6μs 215±0.3μs 1 mpas_ocean.Integrate.time_integrate('120km')
196±2μs 200±0.4μs 1.02 mpas_ocean.Integrate.time_integrate('480km')
186±2ms 181±2ms 0.97 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
180±1ms 180±1ms 1 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
183±1ms 182±0.8ms 1 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
13.7±0.08ms 13.5±0.06ms 0.99 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
13.9±0.05ms 13.4±0.1ms 0.97 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
13.6±0.09ms 13.5±0.09ms 0.99 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
353±4μs 350±2μs 0.99 mpas_ocean.PointInPolygon.time_face_search_lonlat('120km')
353±2μs 354±1μs 1 mpas_ocean.PointInPolygon.time_face_search_lonlat('480km')
334±1μs 335±2μs 1 mpas_ocean.PointInPolygon.time_face_search_xyz('120km')
332±0.9μs 337±2μs 1.02 mpas_ocean.PointInPolygon.time_face_search_xyz('480km')
239±0.8ms 245±0.3ms 1.02 mpas_ocean.RemapDownsample.time_bilinear_remapping
288±0.9ms 284±1ms 0.99 mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
4.31±0.01ms 4.29±0.02ms 1 mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
1.43±0.01s 1.43±0.01s 1 mpas_ocean.RemapUpsample.time_bilinear_remapping
37.1±0.2ms 37.2±0.4ms 1 mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
9.39±0.2ms 9.35±0.08ms 1 mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
378M 356M 0.94 mpas_ocean.ZonalAveragePeakMem.peakmem_zonal_average('120km')
340M 339M 1 mpas_ocean.ZonalAveragePeakMem.peakmem_zonal_average('480km')
324M 324M 1 quad_hexagon.QuadHexagon.peakmem_open_dataset
324M 324M 1 quad_hexagon.QuadHexagon.peakmem_open_grid
7.00±0.06ms 6.95±0.1ms 0.99 quad_hexagon.QuadHexagon.time_open_dataset
5.84±0.06ms 5.89±0.06ms 1.01 quad_hexagon.QuadHexagon.time_open_grid

@cmdupuis3
cmdupuis3 requested a review from rajeeja July 29, 2026 22:38
@cmdupuis3
cmdupuis3 marked this pull request as ready for review July 29, 2026 22:38
@rajeeja

rajeeja commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Thanks for putting this together. I'll take a look at this and post my comments soon

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

Labels

benchmarking Related to benchmarks, memory usage, and/or time profiling 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.

3 participants