Skip to content

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

Draft
cmdupuis3 wants to merge 68 commits into
mainfrom
cmd/accusphere3
Draft

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

Waiting on accusphere merge to main #1513 first.

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.
rajeeja and others added 28 commits July 13, 2026 13:43
Move the uxarray kernel imports from module level into each benchmark
class's setup, so a stale or mismatched environment build only errors the
affected benchmark instead of aborting collection of every benchmark in
the directory. Thanks @cmdupuis3 for catching the asv import failure.
* Honor quadrature kwargs in calculate_total_face_area

calculate_total_face_area accepted quadrature_rule, order and
latitude_adjusted_area but ignored them, always returning the cached
default-parameter face_areas. As a result the gaussian/corrected path
produced the same total as the triangular one.

Keep the cached fast path for the default parameters (which also preserves
the equal-area values used for HEALPix grids) and route any non-default
quadrature settings through _compute_face_areas so the requested rule,
order and latitude adjustment actually take effect.

* Restore matplotlib backend after HoloViews matplotlib plot (#1538)

* Restore matplotlib backend after HoloViews matplotlib plot

plot(backend='matplotlib') calls hv.extension('matplotlib'), which switches
the active matplotlib backend and clobbers the IPython inline display hook,
silently breaking subsequent native matplotlib/xarray .plot() calls. Restore
the original matplotlib backend right after the HoloViews extension switch;
HoloViews objects still display via Store.current_backend, so this is safe.

Closes #1537

* Address review: capture backend at switch, accurate docstring, effective test

* Reconfigure IPython inline display hook when restoring backend

mpl.use() restores the matplotlib backend name but does not re-register
IPython's inline display integration that hv.extension('matplotlib')
clobbers. In a Jupyter kernel without an explicit %matplotlib inline,
native matplotlib/xarray .plot() calls after a uxarray matplotlib plot
still failed to render. Re-run configure_inline_support when the restored
backend is inline so the display hook is reinstated. See #1537.

* Restore matplotlib backend via IPython shell reactivation to fix inline display

---------

Co-authored-by: Orhan Eroglu <32553057+erogluorhan@users.noreply.github.com>

* Allow SCRIP reader to respect units w/r/t radians (#1433)

* Allow SCRIP reader to respect units wrt to radians

* Add test file and unit test for SCRIP radian coordinate handling

* minor formatting changes for scrip radians fix

Moves meshfiles/scrip/scrip_radians.nc to meshfiles/scrip/scrip_radians/scrip_radians_grid.nc to match style of other meshfiles naming schemes.

Renames the new _scrip._convert_to_degrees() to _scrip._values_in_degrees(). It doesn't always convert; and it also does more than just converting, because it gives numpy array from DataArray. Clarified docstring.

ran the following, so it should now pass ruff checks: pre-commit run --all-files

---------

Co-authored-by: Sam Evans <s7evans11@gmail.com>
Co-authored-by: Sam Evans <47793072+Sevans711@users.noreply.github.com>
Co-authored-by: Rajeev Jain <rajeeja@gmail.com>

* Support Python 3.14 and upgrade YAC to v3.18 (with DNN remapping) (#1563)

* Upgrade YAC to v3.18, expose DNN remapping, drop pathlib backport

- upgrade YAC CI to v3.18.0 on Python 3.14 (numba>=0.63, py3.14 classifier,
  cython>=3.1 via conda for the bindings build)
- fix add_average: reduction_type -> weight_type (renamed in YAC v3.18)
- expose distance-nearest-neighbour (yac_method='dnn', new in YAC v3.15) + test
- remove the obsolete 'pathlib' backport dependency: it shadows the stdlib and
  breaks tools that import pathlib on Python 3.10+ (broke YAC's Cython build)

* Reference #1561 in numba version pin comment

---------

Co-authored-by: Christopher Dupuis <45972964+cmdupuis3@users.noreply.github.com>

* Bump actions/download-artifact in the actions group across 1 directory (#1572)

Bumps the actions group with 1 update in the / directory: [actions/download-artifact](https://github.com/actions/download-artifact).


Updates `actions/download-artifact` from 7 to 8
- [Release notes](https://github.com/actions/download-artifact/releases)
- [Commits](actions/download-artifact@v7...v8)

---
updated-dependencies:
- dependency-name: actions/download-artifact
  dependency-version: '8'
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* [pre-commit.ci] pre-commit autoupdate (#1565)

updates:
- [github.com/astral-sh/ruff-pre-commit: v0.15.20 → v0.15.21](astral-sh/ruff-pre-commit@v0.15.20...v0.15.21)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Accusphere: gca_gca benchmarks and better arc sampling

* Accusphere: fix sum_of_squares accuracy bug + tests

* Accusphere: numba fma refactor and fix validation tautology

* Accusphere: scalarized _counts_as_crossing

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: vakudo <127726617+vakudo@users.noreply.github.com>
Co-authored-by: Sam Evans <47793072+Sevans711@users.noreply.github.com>
Co-authored-by: Rajeev Jain <rajeeja@gmail.com>
Co-authored-by: Orhan Eroglu <32553057+erogluorhan@users.noreply.github.com>
Co-authored-by: zarzycki <colin.zarzycki@gmail.com>
Co-authored-by: Sam Evans <s7evans11@gmail.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@cmdupuis3 cmdupuis3 self-assigned this Jul 27, 2026
@cmdupuis3 cmdupuis3 added the scalability Related to scalability & performance efforts label Jul 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

scalability Related to scalability & performance efforts

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants