High-resolution morton location channel for ragged/CSR fields

[espg]

Closes #87

Implements the high-resolution morton location channel per the approved plan (plan comment, amended by the start-with-(a) directive). All six phases are complete — phase 6 landed after mortie 0.8.5 shipped the numpy-level point-encode surface (espg/mortie#100, sign-off) — so this now Closes #87.

What / approach

A kind: ragged field can declare location: leaf_id to carry a per-centroid high-resolution location: HealpixGrid.assign now emits point-kind order-29 morton words (mortie Kind::Point, distinct from an area cell), the t-digest reducer collapses each centroid's member points to their deepest enclosing cell via mortie.common_ancestor, and the resulting (k,) uint64 location vector is stored as a CSR companion array ({field}/{shard}/locations) sharing offsets/cell_ids with the float32 digest — one logical multivariable field, two physical arrays, each channel lossless in its own dtype (the 2^53 float ceiling vs ~2^61.6 order-29 words rules out any interleaved float layout; see the layout discussion).

Dependency change: mortie floor >=0.8.3 → >=0.8.5 (phase 6; the 0.8.5 release carries both common_ancestor and the numpy-level geo2mort(..., points=True) from espg/mortie#100 — bump signed off in espg's comment).

Point and area words share path prefixes, so clip2order coarsening — and every existing dense output — is bit-identical (regression-tested). Value-only ragged fields write byte-identical CSR output (byte-level store-snapshot guard).

Per-phase map:

1. grids/healpix.py — assign emits point-kind words (initially via MortonIndexArray.from_latlon(points=True) + the morton_words adapter; phase 6 swapped this to the numpy-level surface). Coarsen-equivalence + point-kind tests.
2. stats/tdigest.py — build_tdigest(values, delta, locations=None) co-sorts locations with values and reduces each centroid's members via common_ancestor (_centroid_ancestors); merge_tdigests(..., locations1=, locations2=) folds mixed-order per-centroid locations the same way. Value-only calls return the bare digest unchanged.
3. config.py / processing/aggregate.py / processing/worker.py — location: validated (ragged-only, function-only, cell-resolution-only, HEALPix-grid-only, known column, reducer must accept a locations kwarg); get_output_signature gains a location key (output_field_signature includes it only when set, so existing shard-map signatures stay byte-identical); calculate_cell_statistics passes the cell's leaf_id slice and accepts the (payload, locations) pair; the ragged sink delivers (values_list, cell_ids, locations_list) triples for located fields, 2-tuples otherwise.
4. csr.py / processing/write.py / readers/tdigest_tensor.py — write_csr(..., locations_list=) writes {field}/locations (uint64) sharing offsets (per-cell length equality enforced); read_csr(..., locations=True) opt-in read-back; read_locations reader yields (morton_index, cell_id, locations) aligned with read_raw_values, reading only the three arrays it needs; byte-identity snapshot test proves value-only stores are unchanged key-for-key and byte-for-byte.
5. configs/atl03_tdigest_located_healpix.yaml — shipped located template (differs from the value-only template only by location: leaf_id, asserted in test) + end-to-end test: synthetic obs → point-kind assign → located digest with forced merges → CSR → read_locations → containment holds for every contributing observation; docstrings updated (stats/tdigest.py, csr.py).
6. pyproject.toml + grids/healpix.py — mortie floor >=0.8.5; assign swapped from the pandas-ExtensionArray wrapper to the numpy-level geo2mort(lats, lons, order=29, points=True) (word-identical, verified; benchmark below). The explicit order= stays as the drift self-check (point encoding is order-29-only, so mortie raises if HEALPIX_REF_ORDER ever moves).

Post-phase, the adversarial self-review's findings were folded in d93efe8 (guards for unvalidated configs, uint64 strictness, reader IO, helper dedupe, new-code mypy/codespell) — see the review summary for the full fold/standing list.

Reconciles with main: merge 2853844 (post-#70; tests/test_config.py resolved additively) and merge d7a869c (post-#152). The #152 merge intersects the located seam once: the new streaming buffered path (StreamingAggregator, issue #148) does not thread per-cell locations, so validate_streaming now rejects located fields with a clear error instead of silently dropping the channel (tested); the pooled path is untouched. Located streaming support is a natural follow-up (the located merge_tdigests law already exists) — flagged under "Questions for review".

Phases

• Phase 1 — point-kind order-29 encode in assign
• Phase 2 — located t-digest reducer
• Phase 3 — location: schema attribute + aggregation plumbing
• Phase 4 — uint64 companion CSR array sharing offsets + reader read-back + byte-identity regression
• Phase 5 — end-to-end located config + test + docs
• Adversarial self-review folded (d93efe8)
• Phase 6 — numpy-level geo2mort(..., points=True) swap + mortie >=0.8.5 floor (6fc4794)

Point-encode surface (phase 6 benchmark)

1M random points, mortie 0.8.5, same host as the phase-1 numbers:

path	ms / 1M pts	vs area
geo2mort (area)	47.1	1.00x
phase 6: geo2mort(..., points=True) (numpy-level)	53.0	1.13x
phase 1–5: from_latlon(points=True) + morton_words unwrap	59.5	1.26x
phase-1 rejected alternative: public np.asarray/to_numpy unwrap	~560	~12x

The swap removes the pandas wrapper entirely: ~11% faster than the adapter path, word-identical output (verified over 10k points), and no private-attribute reliance — the remaining 1.13x over area is the point kernel itself.

How tested

• Grid: point-kind words differ from area words; cells_of/shards_of bit-identical to clip2order on area words over 5000 random points; a lone point survives common_ancestor unchanged (tests/test_grids.py); phase-6 swap word-identical to the wrapper path.
• Reducer: digest invariance with/without locations; exact point-word round-trip in the weight-1 regime; forced-merge containment; NaN co-drop; empty/mismatch/dtype edges; located merge with mixed orders (tests/test_tdigest.py, 2 new classes).
• Plumbing: located pair through calculate_cell_statistics; ragged sink triple through process_shard; unlocated 2-tuple contract untouched; validation matrix for location: incl. reducer-signature and params-collision rejection (tests/test_processing.py, tests/test_config.py); streaming rejection of located fields (tests/test_streaming.py).
• Storage: locations round-trip; shared-offsets slicing; empty-cell alignment; length/misalignment errors; missing-channel error; byte-identity snapshot — a located write adds only locations keys and every shared key is byte-identical to the value-only write (tests/test_csr.py); reader alignment with read_raw_values (tests/test_readers.py).
• End-to-end: shipped config loads/validates and differs from the value-only template only by the location key; synthetic pipeline run with forced merges proves containment for every observation (tests/test_processing.py::TestLocatedRaggedAggregation::test_end_to_end_located_write_then_read).
• Full suite green after the phase-6 swap on the twice-reconciled tree: 1336 passed, 26 skipped (uv run pytest -v, mortie 0.8.5 resolved via uv sync).

Questions for review

• Located fields + aggregation.streaming: currently rejected at validate_streaming (clear error) because the buffered per-cell state doesn't carry locations. The located merge_tdigests law exists, so threading it through StreamingAggregator is a small follow-up if 88S-scale located products are wanted — say the word and it becomes its own issue/PR.
• Per-centroid common_ancestor cost: a located build is ~+24% per cell (Python loop, ~1 µs/centroid). A segmented rust-side reduce in mortie (common_ancestor(words, offsets)) would erase it — candidate for a follow-up mortie ask now that numpy-level point-kind geo2mort: geo2mort(..., points=True) espg/mortie#100 has landed.
• Pre-existing lint/format failures on main (not fixed here): ruff check fails N818 on src/zagg/registry.py:64 (UnknownCapability), and ruff format --check would reformat 6 files — both reproduce on a clean main checkout with the synced env (CI's ruff job is green, so it's local-vs-CI ruff version drift); flagged rather than fixed per repo conventions.
• config.py module size: already past the ~1000-line guidance before this change; this PR adds ~75 lines of validation. A split (e.g. validation into config_validate.py) should be its own refactor PR if wanted.
• _aggregate_chunk_cells return arity grew from 4 to 5 (ragged_locations added). In-repo callers and stubs updated (incl. re-verifying against 88S stress shard: pins, 900s timeout, chunk-boundary extraction, a-priori read plan, streaming tdigest (issue #148) #152's restructured worker); flagging in case anything external monkeypatches it.

[espg]

🤖 from Claude (review)

Fresh-context adversarial self-review of the full phases 1–5 diff (8 finder angles + verification). Findings folded in d93efe8; the rest are recorded here with dispositions.

Folded (commit d93efe8)

• Silent location-channel drop on unvalidated configs — a location + resolution: chunk field built without validate_config (direct PipelineConfig, Lambda dict payload) reached write_ragged_to_zarr's chunk-res branch, which discarded locations_list. Now raises loudly (write.py), with test.
• Late runtime crashes for validatable config errors — params: {locations: ...} collided with the injected kwarg (TypeError per cell), and a located field naming a reducer without a locations parameter (e.g. function: mean) crashed per populated cell. Both now rejected at validate_config via param-key check + inspect.signature introspection, with tests.
• Silent garbage morton words — build_tdigest/merge_tdigests/the aggregate coercion all cast locations to uint64 silently, so a mis-declared float column (or a buggy reducer returning negative int64) would truncate/wrap into plausible-looking words. All three sites now require uint64 outright, with tests.
• False rejection of default-grid configs — a config with no output.grid block defaults to healpix everywhere else (grids/__init__.py factory), but the located-field check read the absent type as None and rejected. Now mirrors the factory default, with test.
• Misalignment swallowed by the empty-payload filter — write_csr dropped a non-empty locations vector attached to an empty payload before the length check ran. Now raises, with test.
• Non-uniform empty-cell contract — calculate_cell_statistics returned a bare [] for an empty located cell but a pair for populated ones; direct callers couldn't always unpack. _empty_cell_value now returns an empty (payload, locations) pair for located fields, with test.
• Located/unlocated dispatch on runtime shape — _aggregate_chunk_cells branched on isinstance(value, tuple) only; a producer drift could route a located field down the unlocated arm and surface as a distant CSR length error. The collect path is now unified with a fail-fast declared-vs-returned mismatch guard.
• Reader IO waste — read_locations read each shard's values array (the field's largest) and discarded it; it now opens only locations/offsets/cell_ids. The triplicated reader preamble was extracted into _shard_groups (used by all three readers).
• New codespell/mypy findings in the diff — the anc variable name tripped codespell (anc ==> and); renamed. New mypy errors in diff-added code (Optional narrowing in merge_tdigests, dict[str, ndarray] inference in calculate_cell_statistics, list-item in write_csr, NDArrayLike in the reader) fixed by proper narrowing/annotations — no # type: ignore added. Branch mypy is now 35 errors vs 39 on main (all remaining are pre-existing).
• Test-helper duplication — the two _point_words copies (one reading the private MortonIndexArray._data) collapsed into one tests/conftest.py helper that unwraps via the sanctioned zagg.grids.morton.morton_words adapter.
• Aliasing on empty-side merges — merge_tdigests' fast paths returned the caller's locations array by reference; now copies.

Verified clean (empirically, not just by reading)

• clip2order/mort2healpix on point-kind words are bit-identical to area words at all orders; no live path feeds assign output to encode_cell_ids (children stay area-kind); viz/catalog/shard-map unaffected.
• Scalar / 0-d / mixed / list lat-lon inputs to the new assign behave as before.
• All ragged consumers (runner, Lambda handler, sharded fanout) pass sink entries opaquely to write_ragged_to_zarr — no 2-tuple destructuring anywhere.
• The stable-argsort located sort produces byte-identical digests to np.sort (equal values sort by value either way).
• mortie 0.8.3 (the existing floor) carries both from_latlon(points=True) and common_ancestor — verified by installing 0.8.3.

Standing findings (not folded — for review)

1. Per-centroid common_ancestor is a Python-loop hot spot: ~1 µs/call puts a located build at +24% per cell (~1.4 s extra per fully-populated 4096-cell shard at delta=256). No batched primitive exists in mortie 0.8.4; a segmented reduce (common_ancestor(words, offsets), rust-side) would erase it. Suggest bundling into the same mortie ask as the phase-6 encode surface — side-effecting/cross-repo, so left as a suggestion.
2. Mixed located/value-only stores are unreadable by read_locations: a store where older shards predate the field's location: raises on the first value-only shard. Clear error, arguably correct; a skip_missing= escape hatch is possible if partial backfills are an expected workflow — scope call.
3. Polymorphic 2/3-tuple ragged sink entries kept deliberately: the 2-tuple is the documented pre-High-resolution morton location channel for ragged/CSR fields #87 ragged_out contract and byte-identity anchor; an always-3-tuple would change it for existing consumers. Flagging the alternative for the record.
4. config.py is over the ~1000-line guidance (1382 on main → ~1460 here). Pre-existing overage; this PR adds ~75 lines of validation to the module that owns validation. Raised here per the convention — if a split is wanted, it should be its own refactor PR.

---

Generated by Claude Code

[espg]

🤖 from Claude

Status: phases 1–5 complete, adversarial-review findings folded (d93efe8), CI green on each phase head, full suite 1146 passed / 24 skipped locally. Applying waiting — the two items left are yours:

(1) Phase 6 gate: the numpy-level geo2mort(..., points=True) mortie issue is drafted paste-ready on the #87 thread (this session couldn't file it — repo scope is englacial/zagg only). Once your mortie-scoped agent lands it in a release, phase 6 here is a one-call swap in HealpixGrid.assign. Consider also bundling the segmented common_ancestor(words, offsets) reduce into that same mortie ask (see the review summary, standing finding 1).

(2) Merge now vs hold (from "Questions for review"): (a) merge phases 1–5 now and track phase 6 as a follow-up when the mortie release lands — recommended, the swap is isolated to one call site; or (b) hold this PR open until then. Say the word (or just relabel) and the next run acts accordingly.

---

Generated by Claude Code

[espg]

🤖 from Claude

Reconciled with main after #70 merged: merge commit 2853844 (merge, no rebase). One conflict, tests/test_config.py, resolved additively — both new test blocks kept (#70's temporal/TestMerra2StormTemplate tests + this PR's TestLocationChannel matrix); tests/test_processing.py / tests/test_runner.py auto-merged cleanly and the _aggregate_chunk_cells 5-tuple stub still applies against the restructured runner tests. Full suite on the merged tree: 1271 passed, 24 skipped; ruff/format clean on touched files (registry.py N818 still the known pre-existing local-only hit).

---

Generated by Claude Code

[espg]

branch conflicts need to be resolved. mortie espg/mortie#100 landed and is included in mortie 0.8.5 , so phase 6 is unblocked.

[espg]

🤖 from Claude

Status note (recorded from espg's direction, for the next routine pass): phase 6 is unblocked — espg merged espg/mortie#100 (numpy-level point-kind geo2mort(..., points=True), closing mortie#96) and cut mortie 0.8.5 containing it, which is exactly the upstream surface this PR's deferred phase 6 waits on. Clearing waiting accordingly; this PR is actionable again.

For whoever advances it: (1) verify 0.8.5 has propagated to PyPI before bumping the floor (mortie>=0.8.5, base dep and the lambda-extra pin in lockstep); (2) phase-6 scope per the #87 plan — swap the interim MortonIndexArray-based point encode to the numpy-level surface and drop the unwrap path (the ~12x penalty measured in mortie#96); the new call must produce byte-identical order-29 point-kind words vs the interim path; (3) with phase 6 landed this PR's body should flip Refs #87 → Closes #87 per its own deferral note.

[espg]

🤖 from Claude

Both items from your comment are done:

Reconcile — merge commit d7a869c brings in main post-#152. One conflict, src/zagg/processing/worker.py, where #152's streaming/buffered branch met this PR's 5-tuple _aggregate_chunk_cells: resolved keeping main's buffered/pooled structure, with the pooled arm unpacking the 5-tuple and the buffered arm setting an empty location sink. Because #152's StreamingAggregator doesn't thread per-cell locations, validate_streaming now rejects located fields with a clear error rather than silently dropping the channel (test added in tests/test_streaming.py); located streaming support is a small follow-up if wanted — the located merge_tdigests law already exists.

Phase 6 — commit 6fc4794: mortie floor bumped to >=0.8.5 and HealpixGrid.assign swapped from the MortonIndexArray.from_latlon(points=True) + morton_words unwrap to the numpy-level geo2mort(lats, lons, order=29, points=True) that espg/mortie#100 shipped (explicit order= kept as the order-29 drift self-check). Output verified word-identical to the wrapper path. Benchmark delta (1M points): 59.5 → 53.0 ms (~11% faster; area baseline 47.1 ms, so the point encode is now 1.13x area with no pandas machinery and no private-attribute unwrap).

Full suite on the reconciled tree with mortie 0.8.5: 1336 passed, 26 skipped; ruff/format clean on touched files (registry.py N818 still the known pre-existing local-only hit). PR body updated: phase 6 ticked, Refs #87 → Closes #87, benchmark table refreshed. All phases complete — awaiting your review.

---

Generated by Claude Code