End-to-end chunking in `uxarray/core`

[cmdupuis3]

Master issue: end-to-end chunked core/ with compiled kernels underneath

Tracking issue for converting uxarray/core/ to a lazy, chunked high level (xr.apply_ufunc / map_blocks) over compiled low-level kernels (numba njit / guvectorize), optimizing peak memory and vectorization together.

Findings anchored to main @ e01d71e3. Line numbers may drift; treat them as anchors, not addresses.

Why this is propagation, not a rewrite

The architecture already exists in-tree, in four places, and was never carried into core/:

Template	Location	Provides
CSR + @guvectorize segmented reduction	grid/neighbors.py:1195-1338 — _csr_neighbors() → (flat, starts, counts), layout "(n),(k),(m),(m),()->(m)", target="parallel", _rechunk_grid_dim	the kernel pattern
Sparse operator + blockwise apply	remap/weights.py:79 RemapWeights (scipy CSR) + remap/apply_weights.py:99-109	the operator pattern
Shared-kernel eager/lazy dispatch	core/aggregation.py:84-200	the dispatch pattern
njit kernel under apply_ufunc	grid/connectivity.py:136-165	the minimal pattern

xr.apply_ufunc appears at only 6 real compute call sites package-wide. map_blocks and persist: zero occurrences. That is the gap.

grid/neighbors.py:1195-1208 already states the governing rule: topology dims are core dims that dask must refuse to split, because connectivity gathers can reach any index — parallelism comes from time/lev instead.

The chunkability rule: "is it an index?", not "is it topology?"

A grid variable must be concrete iff it is used as an index. Dask indexing by a dask index forces the indexed dim into a single chunk anyway, so keeping an index array lazy buys nothing and costs a graph execution per access.

Tier	Members	Policy
Eager / numpy	face_node_connectivity, edge_node_connectivity, face_edge_connectivity, *_face_connectivity, n_nodes_per_face, all *_indices	Materialize once at first use, cache in _ds, never re-chunk. edge_node_connectivity construction is provably non-decomposable (global lexicographic edge numbering) — connectivity.py:192 already calls .compute() honestly
Blockwise / dask-able	node_x/y/z, node_lon/lat, face_lon/lat, edge_lon/lat, face_areas, bounds, face_bounds_lon/lat, edge_*_distances, max_face_radius	Follow the grid's chunk policy; build with apply_ufunc(dask="parallelized"). This is where the peak-memory win lives
Data	everything in UxDataset / UxDataArray	Already correct — chunk/chunks/persist/compute/arithmetic all preserve uxgrid

Per-dim: n_node / n_edge never chunkable. n_face not chunkable for reductions (zonal touches faces scattered across the whole index range; gradient has a 2-ring stencil), but is chunkable for element-local geometry (areas, bounds). latitudes / n_bands / radius and time / lev are freely chunkable and are the real parallel axes — currently unexploited in zonal.py.

Three constraints:

• The dense padded representation is public API. All 8 Grid.*_connectivity properties (docs/api.rst:142-151) and Grid.chunk (:193) are documented, and docs/user-guide/representation.rst documents the (n_face, n_max_face_nodes) fill-value-padded shape as the user-facing contract. A CSR form must be added alongside, not swapped in — a cached derived view on the Grid.
• A hard-eager topology policy caps grid size at RAM. A 1 km global grid is ~10⁹ faces; face_node_connectivity alone is ~24 GB, and that is the regime this targets. Wants an escape hatch and an honest error (grid.nbytes_topology), not a silent swap to dask.
• Lazy loading is format-dependent. Lazy: UGRID, MPAS, ICON, ESMF, HEALPix, EXODUS, FESOM (netCDF). Not lazy: SCRIP, GEOS, Structured, Points, FESOM (ascii). A chunk-policy API must degrade honestly on the second group.

Already landed

• Release the GIL in Python-entry-point Numba kernels #1755 released the GIL in Python-entry-point numba kernels (2 → 22 of 140), with measured threaded scaling 1.01× 1 thread → 6.53× 8 threads. Deliberately not a blanket sweep: nogil only acts at the Python→native boundary, so it is dead weight on the ~106 kernels reachable only from jitted code, including every inline="always" primitive in utils/computing.py. This was the hard prerequisite for everything below.
• Fix isel() bugs: missing indexers' coords; UxDataset.isel() silently skips slicing / fails if grid dim coords exist #1759 fixed the isel coord bugs (crash on a dataset with an n_face coordinate; silent overwrite of a user coord named face_lon).

---

Tier 0 — prerequisites

#	PR	Fix	Site	What's wrong	What to do	What we gain
0.1	#1792	Lazy _set_desired_longitude_range	coordinates.py:727, called from grid.py:257	if da.max() > 180: — bool() on a lazy reduction, ×3 (node_lon, edge_lon, face_lon). Laziness does not survive the constructor: open_grid(..., chunks=-1) executes 4 graphs before the user touches anything, and it re-fires on every Grid construction, so every isel and every copy() pays it again	xr.where(da > 180, (da+180) % 360 - 180, da) unconditionally — elementwise, chunk-parallel, no reduction. Or defer to first node_lon access behind a _lon_range_checked flag	open_grid becomes O(metadata); removes one compute from every isel and copy()
0.2		parallel=True + dask threads = oversubscription	the 15 parallel=True kernels	numba's threadpool nested under dask's → N×M thread explosion	Plan to drop parallel=True once dask owns the chunk loop	Peak memory set by chunk size rather than by n_face
0.3		ux.map_blocks + Grid.chunks	absent	xr.map_blocks(f, uxda) returns a plain xr.DataArray, and per-block objects handed to f are plain xr.DataArray — uxgrid is lost inside and outside. hasattr(grid, "chunks") == False. Grid.chunk() exists (grid.py:1728) but mutates in place, returns None, and has no readable counterpart	ux.map_blocks that re-attaches the grid and, for blocks chunked along a grid dim, hands each block the correctly isel'd sub-grid. Add Grid.chunks; make Grid.chunk() return a new Grid	The missing keystone — unlocks the whole map_blocks half
0.4		__getattribute__ tax	dataset.py:661, dataarray.py:2462	A module import + dict membership test on every attribute access. Measured 0.90 µs vs 0.05 µs for plain xarray = 18× on uxds.attrs	Module-level frozenset checked before the import; better, define the 8 accessor methods explicitly and delete the override	Graph construction touches attributes thousands of times per op — a direct tax on a lazy-graph-heavy design
0.5		Test harness for gradient	test_vector_calculus.py	core/gradient.py has zero dask/chunk references in the source and all 36 of its tests are eager numpy. It is simultaneously Tier 1's flagship target and the only module with no numpy-vs-dask equivalence net. Contrast aggregation, integrate, connectivity, zonal, remap, subset, cross_sections, which all have real equivalence tests	Write the equivalence harness before touching 1.4	Makes the flagship change reviewable
0.6		peakmem benchmarks for the two worst sites	benchmarks/	No peakmem_/track_peakmem_ benchmark exists for zonal_anomaly (1.6) or _populate_face_bounds (1.11) — the two largest peak-memory claims here. Also: the aggregation dask path has zero ASV coverage, and mpas_dyamond.py runs at true DYAMOND scale but is time-only	Add them. mpas_ocean.NeighborhoodDask (parameterized chunking ∈ {numpy, time_chunks, grid_chunks} + track_peakmem_mean) is the template — it is the only benchmark that exercises chunked execution and memory together	Makes the headline claims measurable

---

Tier 1 — high-level lazy / chunked sites

Ranked by payoff ÷ difficulty.

#	Site	Anchor	What's wrong	What to do	Chunkable	Gain
1.1	azimuthal_mean	dataarray.py:984	self.isel(n_face=faces_in_bin) without ignore_grid=True → builds a full subgrid per ring (connectivity re-derivation), then weighted_mean() re-derives face_areas on that subgrid. 11 radii = 11 subgrid constructions to produce 11 scalars; .data forces a compute per ring	ignore_grid=True + gather areas from the parent grid. Then rings are disjoint → a (n_radius, n_face) sparse operator → one apply_ufunc	n_face whole; radius output core dim	Cheapest large win. ignore_grid=True is already used in 7 other hot loops (all 3 remap modules, 4 sites in zonal.py) — this is the lone holdout
1.2	weighted_mean	dataarray.py:1089	(self * weights).sum(axis=-1) — full (…, n_face) temporary before the reduction, and axis=-1 assumes the grid dim is last (the code comment says so)	xr.dot(self, weights_da, dim=grid_dim) / total — fuses multiply+reduce, no temporary, dim-position agnostic, lazy. Exactly what integrate already does	n_face/n_edge whole	Peak-mem 2× + a correctness fix
1.3	difference → _calculate_edge_face_difference	dataarray.py:1919-1945, gradient.py:10-35	.values (no dask path at all) + np.zeros + 2 fancy gathers + diff temp + np.abs copy = ~4 edge-sized arrays live at once	@guvectorize(["void(f8[:],i8[:,:],f8[:])","void(f4[:],i8[:,:],f4[:])"], "(f),(e,two)->(e)", target="parallel") computing abs(d[c[i,0]]-d[c[i,1]]) fused with the INT_FILL_VALUE branch inline; wrap in apply_ufunc. _node_to_edge_kernel is the template	n_face/n_edge whole; leading dims free	4 temps → 0, and .difference() becomes lazy
1.4	gradient / curl / divergence / scalardotgradient	gradient.py:98, dataarray.py:1544	raise DimensionError if ndim > 1. Users loop over time in Python → n_time full grid traversals, n_time numba dispatches, geometry (cross products, arc lengths, dual-cell area) recomputed n_time times, zero dask	Reshape to (n_lead, n_face); keep prange over n_face outer, inner loop over n_lead — geometry computed once per face and reused across all leading elements. Wrap in apply_ufunc with grid arrays passed via kwargs= (identical every block), not as core-dim inputs	n_face whole (2-ring stencil); time/lev free	Flagship. ~n_lead× geometry savings + laziness. Drops four ndim>1 raises at once (gradient.py:98, dataarray.py:1668, :1770, :1861)
1.5	zonal_mean non-conservative	zonal.py:56-98	Python for lat; weights depend only on (grid, lat) — zero data dependence — yet the data is touched n_lat times. :98 is one chained dask __setitem__ layer per latitude over the full output (19 by default, 181-1801 in practice). (time, lev) is never a parallel axis	Split geometry from reduction. Build CSR (indptr, indices, w) once, then one apply_ufunc, exactly as remap/apply_weights.py:99. Weights from integrate.py:636 are already row-normalized, so the re-normalize at :86/:96 is redundant	n_face whole; latitudes free output axis; time/lev free	O(1) graph layers instead of O(n_lat)
1.6	zonal_anomaly conservative	zonal.py:464,471,482	np.zeros(uxda.shape) + np.full(uxda.shape) = 2× the full array dense float64 in RAM even when the data is dask, plus band_mean.compute() inside the band loop = one full graph execution per band, each re-reading all chunks. The docstring at :394 claims laziness	The scatter is a sparse matmul: two row-stochastic operators M (n_bands, n_face) and B (n_face, n_bands); one apply_ufunc whose kernel does blk - ((M@blk.T).T @ B.T) per block. Never .compute() in a loop	n_face whole; time/lev free	Worst peak-mem site in core/: 3×full → 1 chunk
1.7	_node_to_face_kernel	aggregation.py:112-131	(a) get_face_node_partitions (an argsort(n_face) + np.unique) runs inside the kernel → once per dask block; 100 time-chunks = 100 argsorts, each allocating an n_face int64 transient (80 MB at n_face=1e7). (b) data[..., face_nodes_par] gather materializes (…, n_faces_in_partition, e) ≈ 4-8× the output	CSR segment reduction reusing neighbors.py:1209-1214 verbatim: n=n_node, k=Σ n_nodes_per_face, m=n_face. If UGRID rows are left-packed then starts = cumsum(n_nodes_per_face) — no argsort, no get_face_node_partitions at all. median/std/var still work via neighbors.py:1237's per-thread buffer	n_node whole (already correct); n_face output core dim; leading free	Kills both problems in one move
1.8	_compute_zonal_anomaly (centroid)	zonal.py:499-512	Python for bi in range(nb): np.nonzero(band_indices == bi) then isel(...).mean() → O(n_bands·n_face) mask evals + n_bands separate dask subgraphs each touching every chunk. This is groupby(band).mean() spelled the slow way	band_indices (:494) is already the label array. argsort once → starts/counts, then reuse neighbors.py:1213's kernel unchanged	n_face whole; time/lev free	n_bands× fewer passes
1.9	get_faces_at_constant_latitude call count	zonal.py:57 → intersections.py:137	Full O(n_face) boolean scan per latitude → O(n_lat·n_face), plus an n_face bool temp allocated/freed n_lat times (10 MB each at n_face=1e7)	One @njit(parallel=True) pass over faces: searchsorted the sorted latitudes against each face's [lat_min, lat_max], counting-sort into CSR. O(n_face + nnz). Produces exactly the CSR 1.5 needs. Do not njit the screener itself — the comment at intersections.py:29-31 is right that NumPy wins there; the problem is the call count	n_face whole	Both
1.10	_populate_node_xyz / _populate_node_latlon	coordinates.py:180-211	.values → numpy → assign. Purely elementwise, no reason to be eager. Materializes 2 in + 3 out float64 n_node arrays	apply_ufunc(dask="parallelized"), or plain xarray arithmetic	n_node chunks freely	5 full arrays → chunk-sized peak; free vectorization
1.11	_populate_face_bounds	bounds.py:86-96	7 × .values into one @njit(parallel=True) kernel. Largest peak-memory event in the grid lifecycle; the kernel is prange over n_face and only ever reads nodes of that face. Compounding it, bounds/face_bounds_lon/face_bounds_lat are excluded from Grid.chunk() (absent from DESCRIPTOR_NAMES, conventions/descriptors.py:1-9), so they stay eager numpy even when a chunked grid is explicitly requested	apply_ufunc(dask="parallelized") chunked on n_face with n_node core-dim single-chunked (the aggregation.py:158-190 pattern), and add the three names to DESCRIPTOR_NAMES	n_face chunks freely (element-local)	Bounds on grids that don't fit in RAM. The bounds_array output itself (bounds.py:151) is already memory-good — the cost is the 7 whole-grid inputs
1.12	_compute_face_areas_and_jacobian	grid.py:2312-2325	.values × 5; also an arr[0] scalar probe that would compute on dask	Same apply_ufunc treatment — areas are the canonical embarrassingly-parallel-over-n_face kernel	n_face chunks freely	Chunked face_areas, lower peak
1.13	sel label→index	dataarray.py:2179, core/utils.py:204	xr.DataArray(np.arange(coord_array.sizes[dim])) — a full int64 arange(n_face) (800 MB at n_face=1e8) purely to recover positional indices	coord_array.indexes[dim].get_indexer(labels) / get_slice_bound	n/a (index space)	Peak-mem on huge grids
1.14	Whole-grid edge-array builders	grid/utils.py:148, :323	_get_cartesian_face_edge_nodes_array materializes (n_face, n_max, 2, 3) float64 = 48 bytes × n_face × n_max. But it has zero callers — already superseded by the per-face _get_cartesian_face_edge_nodes (:387) and the _array_subset variant (:270, used by zonal.py:68,316 for candidate faces only)	Delete both. No replacement needed	n/a	Removes the worst latent peak-memory trap in grid/

---

Tier 2 — low-level compiled sites

2a. Where the guvectorize seam belongs

At the arc level — never below it. _cdp8 takes 16 doubles and accucross_pair 12; on arm64 AAPCS only 8 doubles pass in v0-v7 and the rest spill. These are cheap only because inline="always" erases the ABI boundary. A gufunc creates a real call boundary and would re-introduce the spill.

Site	Current	Proposed signature
intersections.py:593 gca_const_lat_intersection	1 arc + 1 z per call; called per-edge from 3 loops (integrate.py:410 Python, integrate.py:544 njit, zonal.py:194 njit)	"(two,three),()->(two,three)"
intersections.py:372 gca_gca_intersection	1 arc-pair; geometry.py:888 loops edges against a fixed ref-edge	"(two,three),(e,two,three)->(e,two,three)"
geometry.py:1103 haversine_distance	4 scalars, no allocation	@vectorize(["float64(float64,float64,float64,float64)"], target="parallel")
arcs.py:360 compute_arc_length	2 points; integrate.py:592 loops it	"(three),(three)->()"
arcs.py:452 orient3d_on_sphere	3 vectors, no batch form	"(three),(three),(k,three)->(k)"
point_in_face.py:18 _face_contains_point	1 face; looped at :117 and geometry.py:1243	"(e,two,three),(three)->()"
arcs.py:191/281 extreme_gca_latitude / extreme_gca_z	1 edge; bounds.py:270-332 calls 2× per edge per face	Blocked: extreme_type is a str. Split into _min/_max variants or an int flag first, then "(two,three),(two,two)->()"
area.py:12 calculate_face_area	1 face; zonal.py:327 Python loop per partial face	Blocked: quadrature_rule is a str. Route through _face_area_from_quadrature (area.py:68) with hoisted dG/dW

2b. Tiny-array allocation inside @njit

Same anti-pattern #1727 fixed in intersections.py. computing.py and intersections.py are clean; everything below is untouched.

Site	Allocation	Frequency	Fix
gradient.py:407-414	5 × np.empty(max_stencil) inside prange(n_face)	per face; 5e7 mallocs at n_face=1e7. max_stencil ≈ 16-48, so allocation commonly dominates a loop body this small	Blocked prange: allocate per block, inner range(lo, hi). Portable, needs no thread-id API
bounds.py:391,458,463,483,488 + :203-205,256-284,322-354	~6-8 tiny arrays per edge per face, plus np.copy(old_box) on every insert_pt_in_latlonbox	_construct_face_bounds is the single biggest remaining instance in the repo	Tuple-ize insert_pt_in_latlonbox to 4 scalars; drop n1n2_cart/n1n2_lonlat once extreme_gca_* takes tuples
geometry.py _newton_quadrilateral (:1380, loop at :1434)	np.empty(3) + np.empty((3,3)) ×2 per Newton iteration, max_iterations=150 → up to 600 allocs per quad	per quad face in bilinear remap	Hold as scalars / unrolled locals — the 3×3 is fixed size
geometry.py:722-783	10 tiny arrays per call to pole_point_inside_polygon, 4 of them compile-time constants	2× per face from bounds.py:190-191	Hoist constants to module scope; pass ref edges as tuples
geometry.py:1235-1239	5 arrays per triangle inside the n>4 fan loop	per triangle per polygon	Inline the 3-edge construction — the connectivity arrays are constants
point_in_face.py:42-71	2× np.allclose, 2 subtractions, 2× np.linalg.norm, np.cross = 7 temps per edge	per edge × per candidate face × per point	Rewrite on scalars/tuples via numba_math
arcs.py:256-263, :337-344	np.empty(3) filled from _normalize_xyz_scalar — which already returns a tuple — and only [2] is ever read	per GCA edge, ×2	Delete both: _, _, z = _normalize_xyz_scalar(...)
arcs.py:76-87, grid/utils.py:65,71	np.cross + subtractions + np.dot; a fresh np.array([0.,0.,1.]) per call	per edge	_numba_cross3 / _numba_dot3 on tuples
geometry.py:1320-1331	np.column_stack + np.linalg.det + np.linalg.inv on a 3×3 → LAPACK call per triangle	per triangle	Closed-form 3×3 inverse (already hand-written a few lines below)
area.py:476,710	get_gauss_quadrature_dg allocates the entire table (up to 33×3 + 33) inside njit	_get_all_face_area_from_coords:260 hoists it correctly; calculate_face_area:54-59 does not, so zonal.py:327 rebuilds it per partial face	Make calculate_face_area a thin wrapper over the hoisting path

2c. math.sqrt / np.sqrt compile to a libm call; ** 0.5 compiles to native fsqrt

Disassembling the njit body (arm64, numba 0.65.1):

x ** 0.5      →  fsqrt d1, d0                      # 15 asm lines, no stack frame
np.sqrt(x)    →  adrp x8, _sqrt@GOTPAGE ; blr x8   # 18 lines, full frame, indirect call
math.sqrt(x)  →  adrp x8, _sqrt@GOTPAGE ; blr x8   # identical

The comment at numba_math.py:85-86 is empirically right but misattributes the cause — it is libm-call vs hardware fsqrt, not float32.

Site	Why it matters	Size
area.py:367,394,435,467	per quadrature point × per sub-triangle × per face — order-4 Gaussian is 16 pts × (n−2) tris. Highest-count sqrt in the repo; largest absolute win	S
computing.py:564 (acc_sqrt_re)	inline="always", called from every _accux_gca / _accux_constlat — an external call in the innermost EFT path. Caveat: must preserve the documented value<0 → nan / root==0 → nan behaviour; pow(-0.0, 0.5) = +0.0 vs sqrt(-0.0) = -0.0 could flip an inf sign. Needs a targeted test	M
geometry.py:44,47 (error_radius)	called O(k²) from _unique_points; also hoist the loop-invariant denominator	S
geometry.py:1414	compare norm2 < 1e-30 instead — no sqrt at all	S
arcs.py:379, numba_math.py:75	per arc / 3× per call	S

---

Compensated (AccuSphGeom) math

The compensated core is in good shape: fastmath is 0 occurrences package-wide, there is no allocation, the EFT structure is correct, and (hi, lo) pairs never materialize — they live as scalars and UniTuples in registers, _sum_of_squares_c is length-specialized at compile time, and inline="always" erases every call boundary. The 2× is confined to L1/registers, so chunking makes compensated math cheaper, not dearer: intermediates are per-element and die inside the kernel, never reaching a chunk boundary.

The rule to hold: keep the compensated span inside one gufunc and emit only the collapsed hi+lo result. _accux_gca:301-303 and _accux_constlat:487-490 already do this. A refactor that emits separate hi and lo arrays across a gufunc boundary would double the chunk footprint and lose inline="always".

accucross, acc_sqrt_re, _sum_of_squares_c, and diff_of_products are used in grid/intersections.py and grid/arcs.py — and nowhere in core/gradient.py, which is exactly where they are needed:

Site	Anchor	What's wrong	What to do	Size
_compute_gradients_on_faces cross product	gradient.py:469-474	Naive cross_x = f1y*f2z - f1z*f2y (×3) + sqrt(Σc²). face1/face2 are centroids of two faces sharing an edge, so θ ≈ 1e-3 (3 km) down to 1e-5 rad. ‖a×b‖ = sinθ ≈ θ, and each component is a difference of O(1) products yielding O(θ) ⇒ ~10-17 bits lost. Those bits are the normal direction, which is what gets dotted into normal_lon/normal_lat	accucross (computing.py:245) + _sum_of_squares_c (:495) + acc_sqrt_re (:530). All @njit(inline="always"), so they inline into the prange body with no call overhead. This is the case the computing.py module docstring says they exist for	S
_dual_cell_area triple product	gradient.py:337-347	Worse cancellation: fan triangles are three adjacent centroids, so the Van Oosterom–Strackee triple ≈ θ² ≈ 1e-6…1e-10 from differences of O(1) products while denom → 4. All the information is in triple. The area divides the gradient, so error propagates 1:1. The docstring already documents a 19%→1% error hunt — precision here is load-bearing	diff_of_products (computing.py:200) per component, or accucross + _normal_dot_value (arcs.py:388), written for exactly this	S

---

Hazards to guard during the work

#	Hazard	Guard
H1	Cross-file stale numba cache via inline="always". numba validates a cache entry against the (st_mtime, st_size) of the defining file only, so intersections.*.nbi is checked against intersections.py. Editing computing.py — an EFT algorithm — leaves arcs.py / intersections.py / geometry.py cached objects valid with the old body inlined ⇒ silent numerical divergence	grid/utils.py:524-526 already documents this for the sort helpers; computing.py has no such note. Add the note and a clean target / pre-commit hook clearing uxarray/*/__pycache__/*.nb[ic]. Critical for work that touches computing.py repeatedly
H2	fastmath regression. It licenses reassociation, which collapses e = (a - (s - bp)) + (b - bp) to 0 and destroys every EFT in the file	CI grep-guard: fastmath forbidden in computing.py and every file importing from it (arcs, intersections, point_in_face). Also guard error_model="numpy" on _accux_gca / _accux_constlat — it is load-bearing for branch-free inf masking, not decoration
H3	import uxarray.utils.computing costs ~1.0 s with a fully warm numba cache (import numba 0.27 s, then computing 1.03 s; _validate_fma()'s 2000-sample loop is only 0.008 s of it). The cost is that _HAS_FMA = _validate_fma() at module scope forces eager entry into numba's compile/cache-load machinery on every import uxarray	Make it lazy (memoize behind a function, or persist the bool beside the .nbi). Drop n_samples 2000→200 as a freebie. Saves ~0.7-1.0 s off every import for users who never touch geometry
H4	Per-worker JIT compile tax under a process scheduler. A zonal test measured 12.51 s cold vs <1.33 s warm — pure @njit(cache=True) first-process compile cost	If apply_ufunc/map_blocks invoke these kernels once per dask worker process, the tax recurs per worker unless the on-disk cache is shared and writable. Test explicitly under a process-based scheduler before claiming a speedup — and don't read single-run wall-clock as steady-state cost anywhere in this work

---

Correctness bugs found in passing

Independent of the refactor; each is small and self-contained.

#	Bug	Site
C1	weighted_mean crashes on the path its own docstring promises. When data is neither face- nor edge-centered the else branch only warns — weights stays None, then total_weight = weights.sum() → AttributeError. The docstring promises "an unweighted mean is computed instead"	dataarray.py:1082-1089
C2	_map_dims_to_ugrid mutates the grid's _source_dims_dict in place, so two datasets sharing one Grid corrupt each other. match_chunks_to_ugrid and UxDataset.from_xarray both read this dict	core/utils.py:91,103
C3	topological_all/topological_any return float64 on the face destination but bool on edge. np.empty() defaults to float64 and output_dtypes=[np.float64] is hardcoded, so node→face all/any yields 0.0/1.0; the node→edge path probes dtype correctly. Also 8× memory on a boolean field, and float32 input silently becomes float64 output	aggregation.py:119,182 vs :285-289
C4	Grid.copy() is shallow despite the docstring, and drops subset state — is_subset → False, inverse_indices → None. So uxds.copy(deep=True) silently downgrades a subset grid to a non-subset one	grid.py:2041-2048
C5	Data chunks silently dropped when passing an xr.Dataset — the grid gets chunked but the data does not, exactly inverted from the request	api.py:458-459
C6	chunks= silently ignored on the xr.Dataset → open_grid branch, yet open_grid still returns chunks, so the user gets an eager grid + lazy data with no warning	api.py:125
C7	match_chunks_to_ugrid mutates the caller's dict — {'n_face':100} becomes {'n_face':100,'nCells':100} after open_grid	core/utils.py:130
C8	Accessor _preserve_methods leaks across subclasses. cls._preserve_methods.add(...) never shadows, so registering on one accessor leaks into BaseAccessor and every other accessor. Fix: cls._preserve_methods = set(cls._preserve_methods) | {name}	accessors.py:140-158
C9	Size-collision dim inference silently mis-maps. {grid._ds.sizes[n]: n for n in ("n_face","n_node","n_edge")} — on a size collision the later key wins, so any data dim whose length coincidentally equals a grid size is silently swapped to a grid dim	core/utils.py:94-103
C10	NameError on the structured branch. dim_ordered is assigned inside a for var_name in ds.data_vars loop and used after it	core/utils.py:61-73
C11	Retry-with-same-engine. except Exception → retry with kwargs.pop("engine", "netcdf4"); if the user passed engine="h5netcdf" and it failed, it retries with the same engine, and swallows chunk-spec errors into a confusing chained trace	core/utils.py:35-49
C12	Dead code. _check_face_on_boundary has zero callers; _check_node_on_boundary_and_gather_node_neighbors is referenced only from a commented-out block. Both are @njit(cache=True) ⇒ compile time and cache entries for nothing	gradient.py:52
C13	zonal_mean returns a bare xr.DataArray, losing uxgrid, while the neighbouring zonal_anomaly returns a UxDataArray. Possibly intentional (the result is off-grid) but undocumented and inconsistent	dataarray.py:753,798
C14	_flatnonzero silently materializes dask input via np.flatnonzero, defeating the "May be NumPy or dask array" docstrings on the screeners	intersections.py:34
C15	Grid.__eq__ computes both grids in full via .equals() on node_lon, node_lat, face_node_connectivity. Called per-pair in ux.concat, so a lazy concat of 50 files pays 100 full grid reads. Short-circuit on sizes, source_grid_spec, and dask .name token identity first	grid.py:751-775

---

Ragged topology — where a CSR form would pay

Connectivity is stored as dense 2-D arrays padded with INT_FILL_VALUE. Per the public-API constraint above, CSR must be added alongside as a cached derived view, not swapped in.

Persistent allocations worth a CSR companion: face_edge_connectivity (connectivity.py:336-338 and :490-492), node_face_connectivity (:554-556 — the strongest case, since it is built from a Python dict inversion at :543-547 where CSR is already implicit before densification), face_face_connectivity (:646-648), node_edge_connectivity (:698), and the dual-mesh equivalent (dual.py:42-44). Not worth it: edge_face_connectivity (:409, fixed width 2), HEALPix face_node_connectivity (always quads), bounds_array (fixed (n_face,2,2), already built without a giant intermediate).

CSR primitives already exist three times over — neighbors.py:1296-1336 _csr_neighbors(), remap/weights.py (scipy CSR), io/_structured.py:42-43 (coo_matrix) — and have never been pointed at the mesh topology itself.

Also dead or duplicated: close_face_nodes (connectivity.py:46, zero callers), the full-grid edge builders (1.14), and geometry.py:180-201/:458-486/:564-587 which re-read face_node_connectivity/n_nodes_per_face a second time inside if projection: branches instead of reusing the read above.

---

Suggested sequencing

Phase	Work	Rationale
0 — unblock	0.1 lazy _set_desired_longitude_range · 0.3 ux.map_blocks + Grid.chunks · 0.5 gradient test harness · 0.6 peakmem benchmarks · H1 cache-clean hook · H2 fastmath CI guard	Nothing else pays off until the constructor stops computing and there is a harness to review against. H1 must land before anyone edits computing.py
1 — cheap wins	1.1 ignore_grid=True · 1.2 xr.dot · 1.3 difference guvectorize · 1.8 band groupby · 1.14 + C12 + dead-code deletions · 2b gradient.py:407 blocked prange · 2c area.py sqrt · C1-C11	All S, all independent. Establishes the pattern at low risk and shrinks the surface first
2 — flagship	1.4 batch-axis gradient + apply_ufunc	Drops four ndim>1 raises at once and unblocks curl / divergence / scalardotgradient
3 — the CSR operator	One shared sparse (n_bins, n_face) operator reused by 1.5, 1.6 and 1.1, built by 1.9; 1.7 aggregation CSR; a cached CSR view of the ragged connectivity	The largest structural change. Removes the dask-__setitem__ graph explosion and both full-shape allocations. CSR stays additive — the dense properties keep their current contract
4 — grid lifecycle	1.10-1.12 coordinates / bounds / areas via apply_ufunc; the three-tier chunk policy; add bounds to DESCRIPTOR_NAMES	Where "grids that don't fit in RAM" becomes true
5 — accuracy + polish	Wire accucross / acc_sqrt_re into gradient.py; remaining 2b sites; H3 lazy _HAS_FMA; C13-C15	Accuracy work follows the structural work so it is only done once

[rajeeja]

Thanks, great details and timely!