Better chunking strategies for constlat intersections and zonal routines. by cmdupuis3 · Pull Request #1624 · UXARRAY/uxarray

cmdupuis3 · 2026-07-27T16:41:35Z

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

Closes #1509

…rite intersections, add 241 baseline testsgit status! - most came from accusphere

…tion docstring

…us/mask, dispatcher

…r_sq<0

- 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)

Closes #1509

…rite intersections, add 241 baseline testsgit status! - most came from accusphere

…tion docstring

…us/mask, dispatcher

…r_sq<0

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.

…ranch-free)

…k validity per AccuSphGeom

…hGeom

…ace PR

… issue

…int-in-face

…onstLat scope

…_lat_intersection

cmdupuis3 · 2026-07-28T21:34:30Z

pre-commit.ci autofix

for more information, see https://pre-commit.ci

rajeeja added 30 commits June 8, 2026 15:47

Port AccuSphGeom EFT algorithms and add spherical geometry user guide

878e075

Closes #1509

Fix pre-commit: remove unused imports and variable

1bbf349

Fix pre-commit: split semicolons in notebook cells

1034a18

Address Hongyu's review: port AccuSphGeom compensated arithmetic, rew…

fb74e75

…rite intersections, add 241 baseline testsgit status! - most came from accusphere

Fix RTD: remove RST-invalid numbered list from gca_const_lat_intersec…

256b347

…tion docstring

Fix RTD notebook kernel metadata

e8999e6

Address AccuSphGeom review feedback

0220276

Separate intersection kernels into three layers: numerical core, stat…

0fe327c

…us/mask, dispatcher

Fix NaN/Inf propagation in intersection kernels for denom=0 and plana…

94c06fe

…r_sq<0

Trim computing.py docstring; drop inline=always from L1 kernels

f3c8618

Revert tornado pin; root cause is openssl 3.6.3 on Windows

59a8020

Port AccuSphGeom EFT algorithms and add spherical geometry user guide

35c85a8

Closes #1509

Fix pre-commit: remove unused imports and variable

6c1497e

Fix pre-commit: split semicolons in notebook cells

57373d0

Address Hongyu's review: port AccuSphGeom compensated arithmetic, rew…

3a55e5d

…rite intersections, add 241 baseline testsgit status! - most came from accusphere

Fix RTD: remove RST-invalid numbered list from gca_const_lat_intersec…

683bc9a

…tion docstring

Fix RTD notebook kernel metadata

3ae95f3

Address AccuSphGeom review feedback

7e35cee

Separate intersection kernels into three layers: numerical core, stat…

1f844e1

…us/mask, dispatcher

Fix NaN/Inf propagation in intersection kernels for denom=0 and plana…

ec81e16

…r_sq<0

Merge branch 'main' into rajeeja/accusphere

896d968

Merge branch 'main' into rajeeja/accusphere

aab2bbd

Merge branch 'main' into rajeeja/accusphere

16cb171

style: apply ruff formatting to computing.py

69c6d48

rajeeja and others added 23 commits July 16, 2026 17:56

Cap thread-scaling sweep at performance cores to avoid E-core artifact

b9f5bb3

Fix thread-scaling plot ticks to show only 1,2,4,8

f61555b

o Fix pre-commit

146672c

o doc fixes

e9b394f

o Remove benchmark image from repo

a7d8e88

Split out point-in-face and lat-lon bounds to a separate PR

a5fd603

Merge remote-tracking branch 'origin/main' into rajeeja/accusphere

d2550bc

o computing: match AccuSphGeom acc_sqrt_re exactly (residual order, b…

e87873f

…ranch-free)

o intersections: compensated norm in AccuXGCA kernel, branch-free mas…

4af6ed1

…k validity per AccuSphGeom

o arcs: branch-free on_minor_arc int mask at tol 1e-8 matching AccuSp…

ff2d6e0

…hGeom

o benchmarks: drop point-in-polygon kernel bench, moved to point-in-f…

3e3a972

…ace PR

o benchmarks: drop thread-scaling constlat, moved to separate scaling…

5c80ff7

… issue

o docs: move spherical-geometry notebook out, depends on split-out po…

66b907b

…int-in-face

o docs: drop spherical-geometry notebook from user-guide toctree

8e6c88a

o test_plot: drop to_raster_auto_extent, unrelated to AccuXGCA/AccuXC…

64ecad4

…onstLat scope

o fix on_minor_arc antipodal false positive

793f351

o add on_minor_arc degeneracy regression tests

f78f5ff

fix diff_of_products addition grouping to match AccuSphGeom eft.hpp

1d00adb

Merge branch 'main' into rajeeja/accusphere

9d88f83

o intersections: collapse pos/neg-valid branches via XOR in gca_const…

48831eb

…_lat_intersection

Parallel partial edge_nodes

79e962d

Lazy intersections

bf0abbe

zonal and cross-section peakmem benchmarks

18ee69b

cmdupuis3 self-assigned this Jul 27, 2026

cmdupuis3 added the scalability Related to scalability & performance efforts label Jul 27, 2026

cmdupuis3 added 2 commits July 28, 2026 14:51

Merge main into cmd/accusphere3

5cfc64c

forgot to save; fixes unmerged code

e8ad82d

[pre-commit.ci] auto fixes from pre-commit.com hooks

74fb517

for more information, see https://pre-commit.ci

cmdupuis3 added the benchmarking Related to benchmarks, memory usage, and/or time profiling label Jul 28, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

cmdupuis3 commented Jul 27, 2026 •

edited by rajeeja

Loading

Uh oh!

cmdupuis3 commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

cmdupuis3 commented Jul 27, 2026 • edited by rajeeja Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

PR Checklist

Uh oh!

cmdupuis3 commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cmdupuis3 commented Jul 27, 2026 •

edited by rajeeja

Loading