End-to-end chunking in `uxarray/core`
I maintainer di solito rispondono entro 1 giorno
@cmdupuis3 ci sta già lavorando.
Dal 25/9/2026.
Valutazione
Questa issue non è ancora stata valutata.
Descrizione
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.*_connectivityproperties (docs/api.rst:142-151) andGrid.chunk(:193) are documented, anddocs/user-guide/representation.rstdocuments 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_connectivityalone 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
- #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:
nogilonly acts at the Python→native boundary, so it is dead weight on the ~106 kernels reachable only from jitted code, including everyinline="always"primitive inutils/computing.py. This was the hard prerequisite for everything below. - #1759 fixed the
iselcoord bugs (crash on a dataset with ann_facecoordinate; silent overwrite of a user coord namedface_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 - (([email protected]).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 |
- Lingua principale
- Python
- Stelle
- 228
- Fork
- 55
- Merge medio
- 5g 22h
- PR unite (30g)
- 15
Preparare l'ambiente
- Nessun Dockerfile né file Docker Compose
- Ha un modello di pull request
- Leggi la guida per i contributori
Come iniziare
- Leggi tutta la issue e poi la guida ai contributi del progetto.
- Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
- Fai un fork del repository e lavora su un branch.
- Apri una pull request che faccia riferimento al numero della issue.
Altre issue di UXARRAY/uxarray
-
developer experience
Difficoltà 1/5 Meno di un'ora Idoneità per principianti 90/100
I maintainer di solito rispondono entro 1 giorno
-
zonal_mean() and weighted_mean() should return the same type of objectForse già presa @Sevans711 l’ha presa 7 giorni fa. Apertabug
Difficoltà 2/5 1-3 ore Idoneità per principianti 72/100
UXARRAY/uxarray#1733 · 1 assegnatario ·
I maintainer di solito rispondono entro 1 giorno
-
documentation
Difficoltà 2/5 1-3 ore Idoneità per principianti 70/100
I maintainer di solito rispondono entro 1 giorno
-
`_set_desired_longitude_range` forces materialization in `open_grid`Forse già presa @cmdupuis3 l’ha presa 3 giorni fa. Apertascalability
UXARRAY/uxarray#1792 · 1 assegnatario ·
I maintainer di solito rispondono entro 1 giorno
-
Create a front-facing UXarray websiteForse già presa @dylannelson l’ha presa 4 giorni fa. Apertadocumentation
UXARRAY/uxarray#1787 · 1 assegnatario ·
I maintainer di solito rispondono entro 1 giorno
Tutte le issue di UXARRAY/uxarray
Issue simili
-
Difficoltà 2/5 1-3 ore Idoneità per principianti 88/100
I maintainer di solito rispondono entro 1 giorno
-
Difficoltà 2/5 1-3 ore Idoneità per principianti 84/100
spec-kitty/spec-kitty#5319 ·
I maintainer di solito rispondono entro 1 giorno
-
backend::vllm diffusion multimodal
Difficoltà 2/5 1-3 ore Idoneità per principianti 72/100
I maintainer di solito rispondono entro 1 giorno
-
Difficoltà 2/5 1-3 ore Idoneità per principianti 84/100
openai/openai-agents-python#5229 ·
I maintainer di solito rispondono entro 1 giorno
-
Difficoltà 2/5 1-3 ore Idoneità per principianti 84/100
I maintainer di solito rispondono entro 1 giorno