Add hhdc.block_rank and cell_index (issue #52)

[espg]

Closes #52.

Two reader primitives the zagg demo notebooks hand-roll today. Both are pure addressing — no t-digest algebra — so neither touches the issue #19 seam, and both run in the core (no-zagg-extra) leg.

Phases

• Phase 1 — moczarr.hhdc.block_rank + root-level moczarr.cell_index, tests, CHANGELOG.

Landed as one commit: the two functions share the spec §1/§8 addressing surface and one test module, and splitting them would have been file-splitting rather than phasing.

1. moczarr.hhdc.block_rank(words, block_order) -> (rank, order)

The block-local nested rank of each packed morton word, plus each word's own order. Pairs with the rank_to_rowcol that was already there:

rank, order = block_rank(located_words, block_order)
row, col = rank_to_rowcol(rank, int(order[0]) - block_order)

What it pins:

• Points normalize first. moczarr.convention.point_to_area29 re-bases the §1 suffix (48 + t28*4 + t29 → 28 + t28*5 + (t29 + 1)) before any digit is read; area words pass through. Skipping it is the notebook's bug — the level-28 digit comes out as (suffix - 28) // 5, which is 4..7 for every word in the point band.
• Mixed orders in one array. Not a constructed case: the committed strata fixture's own located companion carries order-29 points alongside order-22 and order-6 AREA fallbacks in a single array (tests/test_hhdc.py::TestBlockRank::test_a_real_located_companion_is_mixed_kind_and_mixed_order). Returning the per-word order is what lets a caller group by order - block_order and make one vectorized rank_to_rowcol call per depth.
• Vectorized over words. The decode loops over the ≤29 levels, never over the words, and breaks as soon as no word is that deep. No morton_decimal per word.
• Raises, never truncates, when block_order is finer than any word's own order (and on a negative block_order).
• Shape- and dtype-stable: the input's shape is preserved (a scalar word in yields 0-d arrays out), rank is uint64, order is int64 (deliberately not mortie's uint8 — order - block_order on a uint8 would underflow silently under NEP 50).

The packed-word geometry it decodes — [4-bit prefix | 54-bit body | 6-bit suffix], level L ≤ 27 at bit 6 + 2*(27 - L), levels 28/29 in the suffix's parent-first preorder band — is re-verified in TestPackedWordGeometry, not assumed. The order itself comes from mortie.orders_of rather than a second local read of the §1 suffix table.

2. moczarr.cell_index(store, field, morton_index, row, col)

Ported from zagg.readers.tdigest_tensor.cell_index — the reference implementation, which zagg's demo/06_paired.ipynb reaches across a package boundary for. Exported on the package root next to read_cell (and added to tests/test_surface.py's WORKFLOW roster), defined in moczarr/hhdc.py where _tensor_side/_chunk_word/rowcol_to_rank live — moczarr.ragged cannot import from moczarr.hhdc without a cycle.

Properties carried over verbatim from the reference:

• resolves the chunk start from the sibling morton coordinate and returns chunk_start + rowcol_to_rank(row, col, depth);
• searches only the array's stored spans (moczarr.ragged.stored_chunk_spans), one small slice of the morton coordinate per span — never the whole axis, never digest bytes (moczarr's slice goes through the span-cached _MortonWords, so the coordinate keeps its one-GET-per-stored-object posture);
• a coarser block_order block id names no single chunk and raises, as does a well-formed chunk id with nothing stored;
• (row, col) outside the chunk's (side, side) block raises.

One deliberate widening over the reference: morton_index accepts either currency — a packed area word or a decimal string — via moczarr.convention.morton_word, matching how subtree= already takes both in this package.

How it was tested

tests/test_hhdc.py, three new offline classes plus one probe-gated parity class (101 passed, 1 skipped in that module; the skip is the no-zagg-only TestMissingExtraHint):

• TestPackedWordGeometry — the [prefix | body | suffix] layout re-verified against morton_word: levels 1/2/13/26/27 carry their 2-bit digit at bit 6 + 2*(27 - L), an order-27 id's suffix is its order, and the order-28/29 band is 28 + t28*5 + (t29 + 1) with the POINT twin at 48 + t28*4 + t29.
• TestBlockRank — the independent oracle is morton_decimal's digits (digit at level L, minus 1, is that level's rank), computed through the frozen decimal_rank, which shares no bit with the kernel under test:
  • test_matches_the_digit_oracle_for_every_order — 90 ids spanning every order 0..29, both signs, all six base cells, at block orders 0/1/4/6/27/28/29;
  • test_mixed_orders_in_one_array — orders 6/18/29 together, per-word order returned, one rank_to_rowcol round-trip per depth group;
  • test_point_words_normalize_to_their_area_twin — an order-29 area word and its point twin give identical (rank, order);
  • test_located_point_words_decode_in_band — the notebook's bug: the fixture's real point words all read (suffix - 28) // 5 > 3 un-normalized, and rank_to_rowcol(rank, 1) at the deepest level now stays in [0, 4);
  • test_located_words_rank_inside_the_cell_that_stores_them — an observation's shard-local rank starts with its cell's shard-local rank, at every order the companion mixes;
  • test_cell_words_reproduce_their_cells_axis_position — the tie to the tensor path: block_rank(cell_words, shard_order) reproduces np.flatnonzero(written) exactly, and one order finer reproduces the chunk-local rank the sweep reports;
  • the block_order > order and negative-block_order raises, and the shape/dtype contract.
• TestCellIndex — every populated cell of the fixture resolves to its axis index and read_cell finds a digest there; the bare rank would name a different cell; both id currencies agree; block_rank → rank_to_rowcol → cell_index round-trips each fixture cell to its recorded index; the coarser-shard-id, unwritten-chunk-id and out-of-block refusals.
• TestCellIndexParity — probe-gated side-by-side against zagg.readers.tdigest_tensor.cell_index (agrees on every populated cell; same refusal on a coarser block id). Gated like the existing subtree= leg because cell_index post-dates the extra's declared zagg>=0.40 floor; the offline class stays the enforcement wherever it skips.

Local state: ruff check src tests and ruff format --check src tests clean; pre-commit run clean on the touched files except pre-existing mypy errors in files this PR does not touch (ragged.py:290, coverage.py:452, moc_index.py:375, intersect.py:380/383). Full suite 1005 passed, 3 failed, 3 skipped — the three failures (test_moc_index.py::TestAlignment::test_inner_and_outer_parity, test_open.py::TestOpenHive::test_matches_manual_leaf_open, test_open_store.py::TestOpenStoreKwargsForwarding::test_laziness_preserved_per_node) reproduce identically on the branch point with the working tree stashed, so they are pre-existing in this local environment and untouched here.

Questions for review

1. Placement of cell_index. The issue says it "belongs next to read_cell", which is moczarr.ragged — but ragged.py importing rowcol_to_rank/_tensor_side/_chunk_word from hhdc.py is a cycle (hhdc imports ragged). It is defined in hhdc.py and exported next to read_cell on the package root, so mz.read_cell / mz.cell_index sit together as the issue asks. Moving the three helpers down into ragged.py would be the alternative — happy to do that instead if the module-path spelling matters.
2. block_rank is not on the package root. It matches its neighbours rank_to_rowcol / rowcol_to_rank, which are moczarr.hhdc-only, and the issue names it moczarr.hhdc.block_rank. But it is the call a located-companion reader makes, so it is arguably a workflow name for tests/test_surface.py's roster. Promote it?
3. order dtype. Returned as int64 rather than mortie's uint8 specifically so order - block_order cannot wrap. Reasonable, or should it mirror orders_of?
4. Zero/fill words. A 0 word decodes as order 0 (mortie's documented behaviour for the empty sentinel) and therefore raises for any block_order > 0 rather than being silently skipped. That seemed right for a function whose contract is "raise, don't truncate", but a located companion never carries fill, so no test exercises it — flagging the choice rather than pinning it.
5. No live-store leg. The committed suite is offline-by-posture, and the only live test in the repo is env-gated (MOCZARR_LIVE_TESTS). The fixture's located companion turned out to carry exactly the mixed point/area shape the issue describes, so nothing is untested for want of the network — but if a live atl03_tdigest_o9.zarr leg is wanted under the same env gate, say so.

[espg]

🤖 from Claude

Status: all phases in the checklist are done and CI is green on f9d4884 — ruff, build, notebook, and all three test legs (3.12/3.13 × xdggs, plus 3.12 + zagg, which is the one that runs the new TestCellIndexParity against zagg's reference implementation).

One correction to the PR body's "how it was tested" paragraph: the three local failures I noted there (test_moc_index.py::TestAlignment::test_inner_and_outer_parity, test_open.py::TestOpenHive::test_matches_manual_leaf_open, test_open_store.py::TestOpenStoreKwargsForwarding::test_laziness_preserved_per_node) are local-environment only — they reproduce on the branch point with the tree stashed, and CI's full suite passes on this PR. They are an artifact of the borrowed interpreter I ran in (a zagg venv, not a uv sync of this repo), not a real regression, and nothing here needs to act on them.

Deliberately left in draft and with no waiting/blocked state label: the five items under "Questions for review" are genuine design questions, not blockers, so the code is complete and reviewable as it stands — but items (1) and (2) (where cell_index is defined, and whether block_rank should be promoted to the package root) would change the public surface if answered the other way, so flipping to ready-for-review before they are settled would misstate how settled the API is.

[espg]

🤖 from Claude (review)

The read-posture claim is unpinned — a whole-axis scan passes the whole suite.

The docstring (L775-777), the CHANGELOG entry and the PR body all state that cell_index searches "only the array's STORED spans ... never the whole axis, and no digest bytes". I mutated exactly that line to

    for span_start, span_stop in [(0, int(arr.shape[0]))]:

and tests/test_hhdc.py is still 101 passed, 1 skipped. Nothing in TestCellIndex or TestCellIndexParity can tell the two apart, because the fixture's cells axis is 16 cells and the answers are identical — only the bytes read differ.

The claim is true as written; I measured it with the repo's own harness:

GET 6/h_tdigest_signal/zarr.json
GET 6/morton/zarr.json
GET 6/morton/c/0
signal data gets: 0 ; morton gets: 1

So this is a test gap, not a bug — but it is a gap in a property this module already pins elsewhere with the same tool: test_hhdc.py:676-683 and :1006-1028 both use CountingStore (tests/test_ragged.py:244) to assert the per-object GET counts of read_tensors. A three-line CountingStore assertion on cell_index (zero {SIGNAL}/c/ GETs, one {GROUP}/morton/c/ GET) would make the posture a regression-tested contract instead of prose.

[espg]

🤖 from Claude

Fixed in 67dbe25. The claim is now pinned, following the read_tensors precedent you pointed at: tests/test_hhdc.py::TestCellIndexReadPosture drives cell_index through the repo's own CountingStore and asserts the read posture — the morton coordinate is touched, the signal-data array is not, and no span outside stored_chunk_spans is read.

Verified it actually goes red under exactly the mutation you used: replacing stored_chunk_spans(arr) with [(0, int(arr.shape[0]))] gives

FAILED tests/test_hhdc.py::TestCellIndexReadPosture::test_an_unstored_span_is_never_searched
1 failed, 112 passed, 1 skipped

then restored. Suite is 117 passed / 1 skipped on the branch.

You were right that the claim was true but unenforced — which is the shape where it quietly stops being true later.

[espg]

🤖 from Claude (review)

Open question (4) — the 0 fill word is partly a defect, not just a choice. Two separate problems:

(a) At block_order == 0 there is no raise at all. block_rank(np.array([0], np.uint64), 0) returns (0, 0) — a fill sentinel silently reported as a legitimate order-0, rank-0 word. 0 is not a legal word: the prefix nibble encodes base cell + 1, so morton_word("1") == 1152921504606846976 (prefix 1) and prefix 0 is unreachable. mortie itself refuses it — morton_decimal(0) raises ValueError: morton_index array contains an empty or invalid word; only orders_of is lenient and answers 0. So the "raise, don't truncate" contract has a hole exactly where the sentinel is indistinguishable from data.

(b) The block_order > 0 message misdiagnoses the most likely caller. The array a reader is most likely to hand this is the sibling morton coordinate, which is fill-padded by construction — this module's own _cells_order and _chunk_word both open with words[words != 0] for that reason, and the PR's own test_cell_words_reproduce_their_cells_axis_position has to slice words[written] before calling. What that caller currently gets is:

block_order 4 is finer than 3 of the 4 word(s) given (shallowest order 0):
such a word lies at or above the block, so it has no position INSIDE it ...

which points at block geometry, never at the fill. Detecting packed == 0 and either raising a message that names the sentinel, or refusing it up front at every block_order, would close both halves. Cheap, and it is the one input a located-companion reader will hit by accident.

[espg]

🤖 from Claude

Fixed in de3698b — taken as a defect, not a preference. The 0 FILL word is now refused at every block_order, including 0, rather than passing through as a (0, 0) sentinel indistinguishable from data.

The check runs before the order check so the fill is diagnosed as the fill, and the message names it and the fix:

N of the M word(s) given is the 0 FILL word, not a morton index — the morton coordinate is fill-padded over its unwritten rows; mask them out (words[words != 0])

That matches what this module already does internally (_cells_order / _chunk_word both mask words != 0), so the public entry point now behaves like the private ones. Docstring and CHANGELOG updated to state the refusal; pinned in tests/test_hhdc.py.

[espg]

🤖 from Claude (review)

The headline usage example is wrong for the mixed-order array the function exists to serve.

int(order[0]) - block_order takes the first word's depth and applies it to the whole array — but the paragraph six lines below (L169-174) says the array is mixed-order and that the caller must "group by order - block_order and make one vectorized rank_to_rowcol call per depth". The two contradict each other, and it is the snippet a notebook author copies.

It does not degrade quietly. Run the docstring line on a genuinely mixed array:

words = np.array([morton_word(i) for i in ids], dtype=np.uint64)  # orders 5/17/28
rank, order = block_rank(words, 4)
rank_to_rowcol(rank, int(order[0]) - 4)
# ValueError: rank must lie in [0, 4)

That is the same rank must lie in [0, ...) failure the CHANGELOG entry says block_rank was written to eliminate. Suggest making the example the grouped form that TestBlockRank.test_mixed_orders_in_one_array already demonstrates, e.g.

rank, order = block_rank(located_words, block_order)
for depth in np.unique(order - block_order):
    sel = order - block_order == depth
    row, col = rank_to_rowcol(rank[sel], int(depth))

and keep the single-depth one-liner only if it is explicitly labelled "when every word shares one order".

[espg]

🤖 from Claude

Fixed in 18aa54c. The headline example is now the grouped form, so it works on the mixed-order array the function exists for instead of raising rank must lie in [0, 4**d) — the error the CHANGELOG says block_rank was written to eliminate. It no longer contradicts the paragraph six lines below it.

[espg]

🤖 from Claude (review)

The widening ships without the guards its cited precedent has.

The PR says morton_index "accepts either currency ... matching how subtree= already takes both". subtree= does not go through bare morton_word — it goes through moczarr.convention.normalize_subtree, which range-checks the int and refuses a POINT word with an actionable message:

raise ValueError(
    f"subtree {subtree!r} is outside the uint64 range, not a packed morton "
    f"word; parse a decimal id by passing it as a string instead"
)

morton_word is a plain int pass-through, so anything integral reaches the span scan and comes back blaming the store:

input	result
-5	no stored read chunk ... carries morton id -5 — ... (a coarser block_order block id names no single chunk)
2**70	same message, id 1180591620717411303424
433144 (a decimal id typed as an int)	same message

The third row is the failure mode the widening itself creates, and it is not hypothetical at the orders this PR names: a 19-digit decimal id (order 18 — "GEDI cell words are order 18") read as an int has a non-zero prefix nibble and so looks like a well-formed packed word. Nothing is silently wrong (it always raises), but the error names the wrong cause. Routing morton_index through the same uint64-range / point-word validation normalize_subtree performs — or just widening the final message to mention the currency — would keep the two seams honest with each other.

[espg]

🤖 from Claude

Fixed in fe10d0e. morton_index is now validated as a read-chunk id before the span scan, so a malformed input is diagnosed as a caller error rather than surfacing as "no stored read chunk … (a coarser block_order block id names no single chunk)" — which blamed the store for a caller mistake.

-5, 2**70, and a decimal id passed as an int now each get a message naming the actual problem. The genuine coarser-block-id case keeps its original message, since that one really is about the store's chunk grid.

[espg]

🤖 from Claude (review)

Open question (2) — there is concrete ground for promoting block_rank, in tests/test_surface.py's own rationale.

Not a defect; the roster docstring is explicit that it is "a hand-curated FLOOR ... a name absent from it is unpinned rather than ruled out". But the reason issue #49 created that roster was:

open_ragged was importable only by module path (moczarr.ragged.open_ragged) while every sibling it is used alongside sat on the package root, so a notebook mixing the two spellings was forced by the library, not chosen.

This PR ships the two halves of one issue on opposite spellings, so the located-companion reader the issue was filed for writes:

rank, order = moczarr.hhdc.block_rank(words, block_order)
index = mz.cell_index(store, field, chunk, row, col)

which is the asymmetry the roster exists to prevent. The counter-argument in the PR body is real — rank_to_rowcol/rowcol_to_rank are hhdc-only too — but those are kernels a reader composes through block_rank, whereas block_rank is itself the call the notebook makes. Recording the ground here so the answer is a decision rather than an omission; either way is defensible and neither blocks.

[espg]

🤖 from Claude

Left standing for @espg — promoting block_rank to the package root changes the public surface, so it is not mine to take on a review comment (CLAUDE.md §6).

The argument for it is the one you cite: issue #49's own rationale was that "a notebook mixing the two spellings was forced by the library, not chosen", and this PR ships the two halves of one issue on opposite spellings — mz.cell_index on the root, mz.hhdc.block_rank under the module. A reader doing the located-companion walk touches both in the same breath.

The argument against is that block_rank pairs with rank_to_rowcol, which is hhdc-only, so the current placement keeps the layout kernel's three pieces together and the root surface narrow.

No code change in this PR either way.

[espg]

🤖 from Claude

Resolved as leave it in hhdc (espg, 2026-08-26). The #49 analogy is sound in general but does not apply here — checked against the actual import lines in the two notebooks this PR exists to serve:

• hhdc_viewer.ipynb: from moczarr.hhdc import block_rank, rank_to_rowcol — one line, no mixing.
• waveform_viewer.ipynb: mixes from moczarr.hhdc import rowcol_to_rank with mz.cell_index / mz.read_cell — but that mixing is entirely from PRE-EXISTING surface (rowcol_to_rank module-only, read_cell root), and this notebook never calls block_rank.

So promotion would not remove a single mixed spelling in the deliverable, while it would split the layout kernel across two surfaces. Keeping block_rank next to rank_to_rowcol, which is what it pairs with. No code change.