Skip to content

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

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

Better chunking strategies for constlat intersections and zonal routines.#1624
cmdupuis3 wants to merge 73 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

github-actions Bot commented Jul 29, 2026

Copy link
Copy Markdown

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change Before [852a0c0] After [8de811e] Ratio Benchmark (Parameter)
- 523M 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'))
- 451M 329M 0.73 mpas_ocean.Gradient.peakmem_gradient('480km')
- 25.6±0.3ms 7.56±0.2ms 0.29 mpas_ocean.ZonalAverage.time_zonal_average('120km')
- 5.77±0.09ms 4.77±0.02ms 0.83 mpas_ocean.ZonalAverage.time_zonal_average('480km')

Benchmarks that have stayed the same:

Change Before [852a0c0] After [8de811e] Ratio Benchmark (Parameter)
10.3±0.05μs 10.6±0.08μs 1.03 bench_connectivity.Connectivity.time_edge_face('120km')
10.9±0.2μs 11.1±0.1μs 1.02 bench_connectivity.Connectivity.time_edge_face('480km')
10.7±0.3μs 10.4±0.06μs 0.97 bench_connectivity.Connectivity.time_edge_node('120km')
10.9±0.1μs 10.8±0.3μs 0.99 bench_connectivity.Connectivity.time_edge_node('480km')
10.3±0.2μs 10.6±0.09μs 1.03 bench_connectivity.Connectivity.time_face_edge('120km')
10.9±0.2μs 11.0±0.2μs 1.01 bench_connectivity.Connectivity.time_face_edge('480km')
10.4±0.1μs 10.6±0.1μs 1.02 bench_connectivity.Connectivity.time_face_face('120km')
10.8±0.04μs 10.9±0.2μs 1.01 bench_connectivity.Connectivity.time_face_face('480km')
21.1±0.2μs 21.2±0.04μs 1 bench_connectivity.Connectivity.time_face_node('120km')
21.9±0.3μs 22.6±0.2μs 1.04 bench_connectivity.Connectivity.time_face_node('480km')
10.4±0.2μs 10.6±0.07μs 1.02 bench_connectivity.Connectivity.time_node_edge('120km')
10.8±0.06μs 10.9±0.2μs 1.01 bench_connectivity.Connectivity.time_node_edge('480km')
10.4±0.1μs 10.5±0.2μs 1 bench_connectivity.Connectivity.time_node_face('120km')
10.9±0.1μs 10.9±0.3μs 0.99 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'))
365M 363M 0.99 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
22.7±0.08μs 23.1±0.2μs 1.02 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
10.0±0.02μs 10.1±0.07μs 1 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
10.2±0.05ms 10.1±0.04ms 0.99 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
2.14±0.02ms 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'))
953±9ns 952±4ns 1 geometry_kernels.AccucrossKernels.time_accucross
2.44±0.01μs 2.45±0.02μs 1 geometry_kernels.AccucrossKernels.time_accucross_pair
292±1ns 306±10ns 1.05 geometry_kernels.EFTPrimitives.time_acc_sqrt_re
281±2ns 279±3ns 1 geometry_kernels.EFTPrimitives.time_diff_of_products
234±0.4ns 235±2ns 1 geometry_kernels.EFTPrimitives.time_two_prod
237±0.5ns 235±1ns 0.99 geometry_kernels.EFTPrimitives.time_two_sum
1.21±0μs 1.21±0.02μs 1 geometry_kernels.GCAConstLatIntersection.time_accux_constlat_kernel
878±6ns 892±2ns 1.02 geometry_kernels.GCAConstLatIntersection.time_gca_const_lat_intersection
1.60±0.01μs 1.59±0μs 1 geometry_kernels.GCAConstLatIntersection.time_try_gca_const_lat_intersection
1.34±0.02μs 1.34±0.03μs 1 geometry_kernels.GCAGCAIntersection.time_accux_gca_kernel
1.11±0.01μs 1.11±0.01μs 1 geometry_kernels.GCAGCAIntersection.time_gca_gca_intersection
1.92±0.05μs 1.83±0μs 0.95 geometry_kernels.GCAGCAIntersection.time_try_gca_gca_intersection
49.9±0.8μs 49.1±0.6μs 0.98 geometry_kernels.OrientPredicates.time_on_minor_arc
520±10ns 506±1ns 0.97 geometry_kernels.OrientPredicates.time_orient3d_on_sphere
2.71±0.1ms 2.58±0ms 0.95 geometry_samebody.SameBodyConstLat.time_accux_dispatch
1.17±0ms 1.16±0ms 1 geometry_samebody.SameBodyConstLat.time_accux_kernel
1.72±0.01ms 1.71±0.01ms 1 geometry_samebody.SameBodyConstLat.time_fp64_dispatch
145±0.3μs 145±0.3μs 1 geometry_samebody.SameBodyConstLat.time_fp64_kernel
32.5±0.2ms 32.2±0.07ms 0.99 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_dispatch
10.2±0.03ms 10.1±0.02ms 0.99 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_kernel
26.6±0.2ms 26.4±0.03ms 0.99 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_dispatch
4.85±0.02ms 4.86±0.04ms 1 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_kernel
825±20ms 808±6ms 0.98 import.Imports.timeraw_import_uxarray
932±30ns 937±10ns 1.01 mpas_ocean.CheckNorm.time_check_norm('120km')
881±7ns 890±7ns 1.01 mpas_ocean.CheckNorm.time_check_norm('480km')
824±20ms 821±2ms 1 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
54.5±0.6ms 52.6±0.3ms 0.97 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
14.3±0.03μs 14.5±0.05μs 1.01 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
14.5±0.06μs 14.6±0.1μs 1 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
4.83±0.03ms 4.82±0.02ms 1 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('120km')
3.37±0.05ms 3.35±0.01ms 1 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('480km')
3.45±0.04s 3.45±0.01s 1 mpas_ocean.ConstructFaceLatLon.time_welzl('120km')
227±2ms 221±0.7ms 0.97 mpas_ocean.ConstructFaceLatLon.time_welzl('480km')
18.1±0.01ms 18.2±0.04ms 1 mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
1.03±0.02ms 1.02±0.02ms 0.99 mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
10.6±0.03ms 10.5±0.02ms 0.99 mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')
708±10μs 692±20μs 0.98 mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
695±10ms 691±2ms 0.99 mpas_ocean.CrossSections.time_const_lat('120km', 1)
350±3ms 351±2ms 1 mpas_ocean.CrossSections.time_const_lat('120km', 2)
183±4ms 180±0.4ms 0.99 mpas_ocean.CrossSections.time_const_lat('120km', 4)
540±0.5ms 543±2ms 1.01 mpas_ocean.CrossSections.time_const_lat('480km', 1)
274±0.9ms 276±0.8ms 1.01 mpas_ocean.CrossSections.time_const_lat('480km', 2)
140±0.6ms 139±0.3ms 0.99 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)
338M 337M 1 mpas_ocean.CrossSectionsPeakMem.peakmem_const_lat('480km', 4)
24.4±0.3ms 24.4±0.2ms 1 mpas_ocean.DualMesh.time_dual_mesh_construction('120km')
3.11±0.1ms 3.15±0.09ms 1.01 mpas_ocean.DualMesh.time_dual_mesh_construction('480km')
945±6ms 941±5ms 1 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
53.7±2ms 51.7±0.3ms 0.96 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
83.4±2ms 83.1±0.4ms 1 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
5.71±0.2ms 5.33±0.1ms 0.93 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
350M 350M 1 mpas_ocean.Gradient.peakmem_gradient('120km')
174±2ms 173±1ms 0.99 mpas_ocean.Gradient.time_gradient('120km')
12.4±0.1ms 12.2±0.1ms 0.99 mpas_ocean.Gradient.time_gradient('480km')
225±0.8μs 225±1μs 1 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
131±1μs 132±0.7μs 1 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
350M 349M 1 mpas_ocean.Integrate.peakmem_integrate('120km')
329M 329M 1 mpas_ocean.Integrate.peakmem_integrate('480km')
214±2μs 217±0.6μs 1.01 mpas_ocean.Integrate.time_integrate('120km')
197±3μs 200±3μs 1.02 mpas_ocean.Integrate.time_integrate('480km')
181±2ms 181±0.9ms 1 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
183±1ms 182±0.8ms 0.99 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
180±1ms 181±0.7ms 1.01 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
13.4±0.2ms 13.1±0.06ms 0.98 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
13.5±0.4ms 13.2±0.1ms 0.97 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
13.3±0.2ms 13.2±0.2ms 1 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
352±1μs 352±2μs 1 mpas_ocean.PointInPolygon.time_face_search_lonlat('120km')
354±6μs 358±2μs 1.01 mpas_ocean.PointInPolygon.time_face_search_lonlat('480km')
337±3μs 334±1μs 0.99 mpas_ocean.PointInPolygon.time_face_search_xyz('120km')
333±2μs 336±2μs 1.01 mpas_ocean.PointInPolygon.time_face_search_xyz('480km')
240±1ms 242±0.9ms 1.01 mpas_ocean.RemapDownsample.time_bilinear_remapping
285±2ms 289±1ms 1.02 mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
4.34±0.09ms 4.29±0.01ms 0.99 mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
1.44±0.01s 1.43±0s 0.99 mpas_ocean.RemapUpsample.time_bilinear_remapping
36.3±0.7ms 36.3±0.7ms 1 mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
9.31±0.1ms 9.31±0.07ms 1 mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
378M 356M 0.94 mpas_ocean.ZonalAveragePeakMem.peakmem_zonal_average('120km')
339M 339M 1 mpas_ocean.ZonalAveragePeakMem.peakmem_zonal_average('480km')
326M 325M 1 quad_hexagon.QuadHexagon.peakmem_open_dataset
324M 324M 1 quad_hexagon.QuadHexagon.peakmem_open_grid
6.93±0.08ms 7.02±0.1ms 1.01 quad_hexagon.QuadHexagon.time_open_dataset
5.89±0.06ms 5.93±0.02ms 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

@rajeeja rajeeja left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Verified independently: subset builder is bit-identical to whole-grid builder (tested HEALPix + mixed-polygon MPAS), vectorized screeners match brute-force reference, both dask/AttributeError bugs reproduce on main and are fixed here. Two small comments below, non-blocking on correctness but worth addressing before merge.

Comment thread uxarray/core/zonal.py
nedge = n_nodes_per_face[f]
weights[pos] = _compute_band_overlap_area(
faces_edge_nodes_xyz[f, :nedge], zmin, zmax
if partial.size:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

New branch (zero partial faces). Add a test for a band with only fully-contained faces.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Not sure what you mean here... The weights are already computed for all faces overlapping the band, this branch only reweights partial faces, so a strict fully-contained set of faces is already accounted for by the time this branch is reached.

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