From 354f11893994cd4ef9a173794a35f912ff7eb290 Mon Sep 17 00:00:00 2001 From: d-laub Date: Sat, 26 Sep 2026 12:13:39 -0700 Subject: [PATCH 1/3] fix(rag): honor fixed axes between the outer axis and nested ragged axes Indexing and to_ak() on (b, p, ~v, ~w) arrays ignored the fixed p axis. Route R=2 arrays with rag_dim > 1 through the multidim segment gather, peel all-int leading keys on R=2, and wrap fixed leading dims in RegularArray in the numeric R=2 to_ak path. Closes #73 Co-Authored-By: Claude Opus 5.5 (1M context) --- python/seqpro/rag/_core.py | 53 +++++++---- python/seqpro/rag/_ingest.py | 7 +- skills/seqpro/SKILL.md | 9 +- tests/test_ragged_nested_fixed_axis.py | 122 +++++++++++++++++++++++++ 4 files changed, 166 insertions(+), 25 deletions(-) create mode 100644 tests/test_ragged_nested_fixed_axis.py diff --git a/python/seqpro/rag/_core.py b/python/seqpro/rag/_core.py index 37baf8c..4efa041 100644 --- a/python/seqpro/rag/_core.py +++ b/python/seqpro/rag/_core.py @@ -83,7 +83,12 @@ def _peel_string_row(rl: "RaggedLayout[Any]", lo: int, hi: int) -> "RaggedLayout class Ragged(NDArrayOperatorsMixin, Generic[RDTYPE_co]): - """A non-branching ragged array with a single ragged axis (Spec A).""" + """A non-branching ragged array with one ragged axis, or two adjacent ones. + + Shape is ``(*leading_int, None x R, *trailing_int)`` with ``R`` in ``{1, 2}``. + Fixed leading axes may sit between the outermost axis and the ragged axes, + e.g. ``(batch, ploidy, ~variants, ~window)``. + """ __slots__ = ("_layout",) @@ -657,6 +662,7 @@ def _getitem(self, where: Any) -> Any: and len(where) == 2 and not self._is_record and self._rl.n_ragged == 2 + and self.rag_dim == 1 and self._is_full_slice(where[0]) ): return self._getitem_inner(where[1]) @@ -670,7 +676,7 @@ def _getitem(self, where: Any) -> Any: # combining path would mis-handle the trailing fixed dim (e.g. a # (d0, None, K) array padded by to_padded), so fall through to the # sequential per-key path below. - if self._layout.shape.count(None) == 1 and self.rag_dim > 1: + if None in self._layout.shape and self.rag_dim > 1: return self._getitem_tuple_multidim(where) # Sequential per-key path. np.newaxis (None) keys are handled here: # track their positions, strip them from the tuple, apply the remaining @@ -719,6 +725,8 @@ def _getitem(self, where: Any) -> Any: if isinstance(self._layout, RecordLayout): return self._getitem_record(where) if self._layout.n_ragged == 2: + if self.rag_dim > 1: + return self._getitem_tuple_multidim((where,)) return self._getitem_r2(where) # Multi-dim leading shape: when rag_dim > 1 (e.g. shape (d0, d1, ..., None)) # index the first axis treating each "row" as a block of n_inner contiguous @@ -848,13 +856,19 @@ def _getitem_multidim(self, where: Any) -> "NDArray[Any] | Ragged[Any]": ) def _getitem_tuple_multidim(self, where: tuple[Any, ...]) -> Any: - """Handle a tuple index on a single-ragged Ragged (record or non-record) with rag_dim > 1. + """Handle a tuple index on a Ragged (record or non-record) with rag_dim > 1. Each key in ``where`` targets successive leading integer axes (like NumPy multi-dim indexing). All leading-dim keys are resolved simultaneously into a single flat segment selection so that e.g. ``rag[:, [0]]`` correctly selects from axis 1 (samples), not axis 0 (ranges). + For R=2 the leading keys select outer (O0) segments; O1 stays global, the + same lazy gather as ``_getitem_r2``. When every leading axis of an R=2 + array gets an int key, the one selected outer segment is peeled like + ``rag[i]`` on a ``rag_dim == 1`` array, so ``rag[i, j] == rag[i][j]``. + (R=1 keeps its 1-segment ``(None, ...)`` result.) + Keys that reach the ragged ``None`` axis or beyond are not handled here (fall through to normal single-key dispatch after leading dims are consumed). """ @@ -973,14 +987,17 @@ def _getitem_tuple_multidim(self, where: tuple[Any, ...]) -> Any: grid_combined = grid_combined + grids[j] * strides[j] flat_seg_idx = np.asarray(grid_combined.reshape(-1), dtype=np.int64) - # Gather segment starts/stops - starts, stops = self._starts_stops() + # Gather the outer segments' starts/stops; deeper levels stay global. + outer, *inner_offsets = self._layout.offsets + starts, stops = _level_bounds(outer) new_starts = np.ascontiguousarray(starts[flat_seg_idx], dtype=OFFSET_TYPE) new_stops = np.ascontiguousarray(stops[flat_seg_idx], dtype=OFFSET_TYPE) - new_offsets = np.stack([new_starts, new_stops], 0) + new_offsets = [np.stack([new_starts, new_stops], 0), *inner_offsets] - # Build new shape: scalar axes squeezed, array-selected axes kept - new_leading: tuple[int | None, ...] = tuple(out_leading) + # Build new shape: scalar axes squeezed, array-selected axes kept. An R=2 + # array with every leading axis squeezed keeps a size-1 axis to peel below. + all_int = not out_leading and len(new_offsets) == 2 + new_leading: tuple[int | None, ...] = (1,) if all_int else tuple(out_leading) new_shape: tuple[int | None, ...] = (*new_leading, *trailing) # Build result: record Ragged or non-record Ragged depending on layout type @@ -989,7 +1006,7 @@ def _getitem_tuple_multidim(self, where: tuple[Any, ...]) -> Any: new_fields = { name: RaggedLayout( data=fl.data, - offsets=[new_offsets], + offsets=new_offsets, shape=( *new_leading, *fl.shape[fl.shape.index(None) :], @@ -1000,7 +1017,7 @@ def _getitem_tuple_multidim(self, where: tuple[Any, ...]) -> Any: } result: Any = Ragged( RecordLayout( - offsets=[new_offsets], + offsets=new_offsets, shape=new_shape, fields=new_fields, ) @@ -1009,11 +1026,13 @@ def _getitem_tuple_multidim(self, where: tuple[Any, ...]) -> Any: result = Ragged( RaggedLayout( data=self._rl.data, - offsets=[new_offsets], + offsets=new_offsets, shape=new_shape, str_offsets=self._rl.str_offsets, ) ) + if all_int: + result = result[0] # Apply any remainder keys (reaching into or past the ragged axis) for k in remainder_keys: @@ -1031,15 +1050,9 @@ def _getitem_record(self, where: Any) -> Any: raise KeyError(where) return Ragged(field) # field.offsets[0] is the shared object (zero-copy) # numpy contract: a non-tuple key is treated as a 1-tuple (A[x] == A[(x,)]). - # For a single-ragged-axis record with >1 leading fixed axis, route through - # the multidim peel path directly (under the SAME guard it requires, so we - # never re-dispatch through __getitem__ and risk recursion). Records that are - # not single-None fall through to the flat record-rows path unchanged. - if ( - not isinstance(where, tuple) - and self.rag_dim > 1 - and rec.shape.count(None) == 1 - ): + # A record with >1 leading fixed axis routes through the multidim path + # directly (never re-dispatching through __getitem__, so no recursion). + if not isinstance(where, tuple) and self.rag_dim > 1: return self._getitem_tuple_multidim((where,)) return self._getitem_record_rows(where) diff --git a/python/seqpro/rag/_ingest.py b/python/seqpro/rag/_ingest.py index 3aac476..69e6875 100644 --- a/python/seqpro/rag/_ingest.py +++ b/python/seqpro/rag/_ingest.py @@ -223,7 +223,12 @@ def to_ak(rag: Any) -> ak.Array: ) else: leaf = NumpyArray(np.ascontiguousarray(data)) # type: ignore[arg-type] - return ak.Array(_wrap_list(o0, _wrap_list(o1, leaf))) + content = _wrap_list(o0, _wrap_list(o1, leaf)) + # Fixed dims between the outer axis and the first ragged axis (e.g. ploidy + # in (b, p, ~v, ~w)) partition O0's segments into regular groups. + for size in reversed(rag._layout.shape[1 : rag.rag_dim]): + content = RegularArray(content, size) + return ak.Array(content) # --- string-under-axis path (shape (L0, None), is_string=True) --- if rag._rl.str_offsets is not None and rag._layout.offsets: diff --git a/skills/seqpro/SKILL.md b/skills/seqpro/SKILL.md index 08490cc..ab59427 100644 --- a/skills/seqpro/SKILL.md +++ b/skills/seqpro/SKILL.md @@ -49,15 +49,15 @@ For exact signatures and kwargs, read the docstring directly (`sp.?` in a RE ## `Ragged` — variable-length sequence batches -`sp.Ragged` (backed by `python/seqpro/rag/_core.py`) is the canonical container for batches where sequences differ in length. It is a Rust-native class implementing `NDArrayOperatorsMixin` (NOT a subclass of `ak.Array`) with **exactly one ragged dimension**, plus zero-copy access to the underlying flat NumPy buffer and offsets. +`sp.Ragged` (backed by `python/seqpro/rag/_core.py`) is the canonical container for batches where sequences differ in length. It is a Rust-native class implementing `NDArrayOperatorsMixin` (NOT a subclass of `ak.Array`) with **one ragged dimension, or two adjacent ones (nested, e.g. `(batch, ploidy, None, None)`)**, plus zero-copy access to the underlying flat NumPy buffer and offsets. ### Mental model A `Ragged` has three things: - **`data`**: a flat contiguous `NDArray` of shape `(total_elements, *fixed_trailing_dims)`. Zero-copy access via `rag.data`. -- **`offsets`**: an `int64` array. Shape `(N+1,)` (contiguous, the common case) **or** `(2, N)` starts/stops (after some slices). Access via `rag.offsets`. -- **`shape`**: a tuple like `(batch, None, ohe_dim)` where exactly one entry is `None` — that's the ragged axis. `rag.rag_dim` gives its index. +- **`offsets`**: an `int64` array. Shape `(N+1,)` (contiguous, the common case) **or** `(2, N)` starts/stops (after some slices). Access via `rag.offsets`. A nested array has one such array per ragged axis, outermost first. +- **`shape`**: a tuple like `(batch, None, ohe_dim)` where one entry is `None` — that's the ragged axis — or two adjacent entries for a nested array. `rag.rag_dim` gives the index of the first. `rag.lengths` derives segment lengths from offsets (cheap, returns an `ndarray`). @@ -90,6 +90,7 @@ rag = sp.rag.Ragged.empty((10, None, 4), dtype=np.uint8) # batch of 10 OHE seq | Apply a `np.ufunc` | Just call it: `np.exp(rag)` — dispatched via `__array_ufunc__` (NDArrayOperatorsMixin) to return a `Ragged` | Manually unpack and rebuild | | Count top-level rows | `len(rag)` — returns `shape[0]` (raises if `shape[0]` is the ragged axis) | `rag.shape[0]` with manual int-cast | | Index one group of an opaque-string `Ragged` | `rag[i]` → a `Ragged` of `bytes`, one per string (`len(rag[i]) == rag.lengths[i]`); `rag[i][j]` is one `bytes` | Expect one concatenated `bytes` — that was the pre-0.22 behavior and it silently dropped the per-string boundaries | +| Index a nested array with fixed leading axes, e.g. `(b, p, ~v, ~w)` | NumPy-style keys on the fixed axes: `rag[i]` → `(p, ~v, ~w)`, `rag[:, j]` → `(b, ~v, ~w)`; `rag[i, j] == rag[i][j]` peels to `(n_v, ~w)` | Read `rag.offsets` by hand to find a row's ploid | | Insert a leading size-1 axis | `rag[np.newaxis]` — returns `Ragged` with shape `(1, *old_shape)` | Manual `from_offsets` rebuild | | Reinterpret bytes/dtype | `rag.view(np.uint8)` | `np.asarray(rag).view(...)` (loses ragged structure) | | Reshape non-ragged axes | `rag.reshape(batch, None, k, 4)` | Touch `rag.data.shape` directly | @@ -163,7 +164,7 @@ When in doubt, read `python/seqpro/rag/_core.py` — it's the live backend and t - **Offsets layout drifts after slicing.** `rag.offsets` may become `(2, N)` starts/stops instead of `(N+1,)`. Check `rag.is_contiguous` / call `rag.to_packed()` before any code that assumes `(N+1,)`. - **`rag.data` on a record layout is a dict.** Code like `rag.data.shape` will fail; branch on `isinstance(rag.data, dict)` or use `rag.parts` and inspect. -- **`Ragged` must have exactly one `None` in `shape`.** Constructing from data whose ragged structure doesn't match raises in `__init__`. Use `from_lengths` / `from_offsets` when in doubt. +- **`Ragged` must have one `None` in `shape`, or two adjacent ones (nested).** Constructing from data whose ragged structure doesn't match raises in `__init__`. Use `from_lengths` / `from_offsets` when in doubt. - **The Rust k-shuffle expects contiguous `uint8` with the last axis as sequence length.** `sp.k_shuffle` handles this for you; if calling `seqpro._k_shuffle` directly, ensure layout. ## Where to look (don't memorize — read the source) diff --git a/tests/test_ragged_nested_fixed_axis.py b/tests/test_ragged_nested_fixed_axis.py new file mode 100644 index 0000000..77265a7 --- /dev/null +++ b/tests/test_ragged_nested_fixed_axis.py @@ -0,0 +1,122 @@ +"""Issue #73: fixed axes between the leading axis and two ragged axes. + +Shape ``(b, p, ~v, ~w)`` is the layout genvarloader's ``_FlatWindow.to_ragged()`` +returns for variant windows. Indexing and ``to_ak()`` must honor the fixed ``p`` +axis. The oracle is awkward indexing on the equivalent nested Python list. +""" + +from __future__ import annotations + +from typing import Any + +import awkward as ak +import numpy as np +import pytest + +from seqpro.rag import Ragged + +# (b=2, p=2, ~v, ~w) +NESTED = [ + [[[0, 1, 2], [3, 4]], [[5]]], + [[], [[6, 7], [8], [9, 10, 11]]], +] +O0 = np.array([0, 2, 3, 3, 6]) +O1 = np.array([0, 3, 5, 6, 8, 9, 12]) + + +def _rag(scale: int = 1) -> Ragged[Any]: + data = np.arange(12, dtype=np.int64) * scale + return Ragged.from_offsets(data, (2, 2, None, None), [O0, O1]) + + +def _tolist(x: Any) -> Any: + if isinstance(x, Ragged): + return x.to_ak().tolist() + if isinstance(x, dict): + return {k: _tolist(v) for k, v in x.items()} + return np.asarray(x).tolist() + + +KEYS = [ + 0, + -1, + slice(1, None), + slice(None, None, -1), + [1, 0], + np.array([False, True]), + (0, 1), + (1, 0), + (-1, -1), + (slice(None), 1), + (slice(None), [1, 0]), + ([1, 0], 0), + (0, slice(None)), + ([0, 1], [1, 0]), + (np.array([True, False]), slice(None)), +] + + +def test_issue_repro(): + r = Ragged.from_offsets( + np.arange(9, dtype=np.uint8), + (2, 1, None, None), + [np.array([0, 2, 3]), np.array([0, 3, 5, 9])], + ) + assert r[0].shape == (1, None, None) + assert r[0, 0].shape == (2, None) + assert _tolist(r[0, 0]) == [[0, 1, 2], [3, 4]] + assert str(r.to_ak().type) == "2 * 1 * var * var * uint8" + + +def test_to_ak_type(): + assert str(_rag().to_ak().type) == "2 * 2 * var * var * int64" + assert _rag().to_ak().tolist() == NESTED + + +@pytest.mark.parametrize("key", KEYS, ids=repr) +def test_getitem_matches_awkward(key: Any): + expected = ak.Array(NESTED)[key].tolist() + assert _tolist(_rag()[key]) == expected + + +@pytest.mark.parametrize("key", KEYS, ids=repr) +def test_getitem_on_gathered_matches_awkward(key: Any): + # Gathered (2, N) outer offsets must index the same as packed ones. + gathered = _rag()[[1, 0]] + expected = ak.Array(NESTED)[[1, 0]][key].tolist() + assert _tolist(gathered[key]) == expected + + +@pytest.mark.parametrize( + ("key", "shape"), + [ + (0, (2, None, None)), + ((0, 1), (1, None)), + ((1, 0), (0, None)), + ((slice(None), 1), (2, None, None)), + ([1, 0], (2, 2, None, None)), + (([0, 1], [1, 0]), (2, None, None)), + ], + ids=repr, +) +def test_getitem_shape(key: Any, shape: tuple[int | None, ...]): + assert _rag()[key].shape == shape + + +def test_chained_equals_tuple(): + r = _rag() + assert _tolist(r[1][1]) == _tolist(r[1, 1]) + + +def test_record_getitem_matches_awkward(): + rec = Ragged.from_fields({"a": _rag(), "b": _rag(10)}) + oracle = {"a": ak.Array(NESTED), "b": ak.Array(NESTED) * 10} + for key in [0, [1, 0], (slice(None), 1), (0, slice(None))]: + got = rec[key] + assert isinstance(got, Ragged) + for name, arr in oracle.items(): + assert _tolist(got[name]) == arr[key].tolist(), (key, name) + peeled = rec[1, 1] + assert isinstance(peeled, dict) + for name, arr in oracle.items(): + assert _tolist(peeled[name]) == arr[1, 1].tolist(), name From 108b03063408d5dbab73e8b5603aa3819450c96d Mon Sep 17 00:00:00 2001 From: d-laub Date: Sat, 26 Sep 2026 12:15:25 -0700 Subject: [PATCH 2/3] fix(rag)!: emit leaf-level records from to_ak() for every rag_dim Record to_ak() zipped per row (depth_limit=1) when rag_dim <= 1 but at the leaf when rag_dim > 1, so rec.to_ak()[k] and rec[k].to_ak() could disagree. Always build leaf-level records, and stop wrapping trailing fixed dims in RegularArray: they already live in each field's leaf. BREAKING CHANGE: Ragged record to_ak() for arrays with one leading axis now returns records at the leaf (e.g. 2 * var * {a, b}) instead of one record of lists per row. Closes #75 Co-Authored-By: Claude Opus 5.5 (1M context) --- python/seqpro/rag/_ingest.py | 29 ++++++++++------------- skills/seqpro/SKILL.md | 1 + tests/test_ragged_core_records.py | 3 ++- tests/test_ragged_nested_diff.py | 2 +- tests/test_ragged_nested_fixed_axis.py | 32 ++++++++++++++++++++++++++ 5 files changed, 48 insertions(+), 19 deletions(-) diff --git a/python/seqpro/rag/_ingest.py b/python/seqpro/rag/_ingest.py index 69e6875..4ebcce9 100644 --- a/python/seqpro/rag/_ingest.py +++ b/python/seqpro/rag/_ingest.py @@ -155,18 +155,10 @@ def to_ak(rag: Any) -> ak.Array: if isinstance(rag._layout, RecordLayout): rec = rag._layout - # Fast path: shape has only one leading fixed dim (e.g. (L0, None) or - # (L0, None, None)). In this case per-field to_ak() + ak.zip works because all - # per-field arrays have the same outermost length and no extra RegularArray wrappers. - rag_dim = rec.shape.index(None) - if rag_dim <= 1: - return ak.zip({f: to_ak(rag[f]) for f in rag.fields}, depth_limit=1) - - # Multi-leading-axis path: shape like (b, p, ~v) where rag_dim >= 2 and there is - # exactly one ragged axis. Per-field to_ak() produces incompatible outer shapes: - # - numeric fields get RegularArray wrappers from _parts_to_content - # - string-under-axis fields do NOT (string path only wraps one level) - # Fix: build leaf contents from layout buffers and re-wrap uniformly. + # Records always sit at the leaf (the ak.zip default), whatever rag_dim is + # (#75). Per-field to_ak() + ak.zip can't do this in general: numeric + # fields get RegularArray wrappers but string-under-axis fields don't. + # So build leaf contents from the layout buffers and re-wrap uniformly. field_leaves: list[Any] = [] field_names: list[str] = [] for name, fl in rec.fields.items(): @@ -198,12 +190,15 @@ def to_ak(rag: Any) -> ak.Array: content: Any = RecordArray(field_leaves, field_names) - # Re-wrap by iterating reversed(shape[1:]), applying ListOffsetArray for None dims - # (innermost ragged first) and RegularArray for fixed dims. Shared offsets are - # consumed innermost-to-outermost so that the outer list wraps the inner list. - # The outermost dim (shape[0]) is implicit in the final array length. + # Re-wrap by iterating reversed(shape[1 : last ragged axis + 1]), applying + # ListOffsetArray for None dims (innermost ragged first) and RegularArray for + # fixed leading dims. Shared offsets are consumed innermost-to-outermost so + # that the outer list wraps the inner list. Trailing fixed dims already live + # in each field's leaf, and the outermost dim (shape[0]) is implicit in the + # final array length. + last_rag = len(rec.shape) - 1 - rec.shape[::-1].index(None) off_iter = iter(reversed(rec.offsets)) - for size in reversed(rec.shape[1:]): + for size in reversed(rec.shape[1 : last_rag + 1]): if size is None: off = next(off_iter) content = _wrap_list(off, content) diff --git a/skills/seqpro/SKILL.md b/skills/seqpro/SKILL.md index ab59427..10ade7c 100644 --- a/skills/seqpro/SKILL.md +++ b/skills/seqpro/SKILL.md @@ -129,6 +129,7 @@ The inputs **must share the same offsets object** (pass the same `shared_offsets - `rag["field"]` gives zero-copy single-field access and shares the parent's offsets object. Its `.data` is the flat NumPy buffer for that field. - `rag.to_numpy()` on a record layout returns a **dict `{field: dense ndarray}`** (raises if any field is still jagged — lengths must be uniform for a dense conversion). - `view` and `apply` are **not defined** on record layouts — operate per-field. +- `batch.to_ak()` puts the record at the leaf, like `ak.zip`'s default: `(2, None)` → `2 * var * {seq, score}`. Before 0.23 a record with one leading axis zipped per row (`depth_limit=1`) instead. - Peeling a row (`rag[i]` with an integer `i`) returns a **dict** whose entries all have the same length: numeric/char fields as `ndarray`, opaque-string fields as a `Ragged` of `bytes`. That's what makes `zip(row["start"], row["alt"])` correct. ### Hashing strings diff --git a/tests/test_ragged_core_records.py b/tests/test_ragged_core_records.py index 802ee89..4331070 100644 --- a/tests/test_ragged_core_records.py +++ b/tests/test_ragged_core_records.py @@ -654,7 +654,8 @@ def test_bridge_r2_record_roundtrip(): rec = Ragged(arr) assert rec._is_record and rec.shape == (2, None, None) assert rec["a"]._layout.offsets[0] is rec["b"]._layout.offsets[0] # shared O0 - assert rec.to_ak().to_list() == arr.to_list() + # to_ak() always emits leaf-level records, whatever depth was ingested (#75) + assert rec.to_ak().to_list() == ak.zip({"a": arr["a"], "b": arr["b"]}).to_list() # --------------------------------------------------------------------------- diff --git a/tests/test_ragged_nested_diff.py b/tests/test_ragged_nested_diff.py index c6d1c57..639ec11 100644 --- a/tests/test_ragged_nested_diff.py +++ b/tests/test_ragged_nested_diff.py @@ -231,7 +231,7 @@ def test_diff_record(t): ora_a = _oracle(outer, inner, data) ora_b = _oracle(outer, inner, data2) - ora_rec = ak.zip({"a": ora_a, "b": ora_b}, depth_limit=1) + ora_rec = ak.zip({"a": ora_a, "b": ora_b}) # leaf-level records (#75) assert rec.to_ak().to_list() == ora_rec.to_list() assert rec.to_packed().to_ak().to_list() == ak.to_packed(ora_rec).to_list() diff --git a/tests/test_ragged_nested_fixed_axis.py b/tests/test_ragged_nested_fixed_axis.py index 77265a7..891c8dc 100644 --- a/tests/test_ragged_nested_fixed_axis.py +++ b/tests/test_ragged_nested_fixed_axis.py @@ -120,3 +120,35 @@ def test_record_getitem_matches_awkward(): assert isinstance(peeled, dict) for name, arr in oracle.items(): assert _tolist(peeled[name]) == arr[1, 1].tolist(), name + + +@pytest.mark.parametrize("shape", [(2, None), (2, None, None), (2, 2, None, None)]) +def test_record_to_ak_is_leaf_level(shape: tuple[int | None, ...]): + # One record depth for every rag_dim (#75): records sit at the leaf, so + # indexing then converting equals converting then indexing. + if shape == (2, None): + a = Ragged.from_offsets(np.arange(5), shape, np.array([0, 3, 5])) + elif shape == (2, None, None): + a = Ragged.from_offsets( + np.arange(9), shape, [np.array([0, 2, 3]), np.array([0, 3, 5, 9])] + ) + else: + a = _rag() + rec = Ragged.from_fields({"a": a, "b": a * 10}) + out = rec.to_ak() + assert str(out.type).endswith("{a: int64, b: int64}") + assert out.tolist() == ak.zip({"a": a.to_ak(), "b": (a * 10).to_ak()}).tolist() + got = rec[[1, 0]] + assert isinstance(got, Ragged) + assert got.to_ak().tolist() == out[[1, 0]].tolist() + + +def test_record_to_ak_trailing_dim_stays_in_leaf(): + off = np.array([0, 2, 3]) + a = Ragged.from_offsets(np.arange(6).reshape(3, 2), (2, None, 2), off) + b = Ragged.from_offsets(np.arange(3), (2, None), off) + out = Ragged.from_fields({"a": a, "b": b}).to_ak() + assert str(out.type) == "2 * var * {a: 2 * int64, b: int64}" + assert ( + out.tolist() == ak.zip({"a": a.to_ak(), "b": b.to_ak()}, depth_limit=2).tolist() + ) From ff1f21a498c44369be7f36aa2d852cf7d38d1a3c Mon Sep 17 00:00:00 2001 From: d-laub Date: Sat, 26 Sep 2026 12:21:00 -0700 Subject: [PATCH 3/3] fix(rag)!: apply each tuple key to the next output axis A tuple key after a non-int key was applied to axis 0 of the intermediate result, so rag[:, 0] on (n, ~) returned row 0 instead of the first element of every row, and keys reaching ragged or trailing axes after a slice indexed the wrong axis. Track the output axis each key targets. Leading fixed axes are still indexed jointly; ragged axes take an int (per-row element) or a step-1 slice (Python semantics, negative bounds included); trailing axes index the data buffer. Unsupported combinations raise NotImplementedError instead of returning wrong data. Also vectorize the R=2 inner int-array gather and drop the now-dead _getitem_inner/_is_full_slice. BREAKING CHANGE: rag[:, k] and similar tuple keys now follow NumPy and awkward semantics; code relying on the old axis-0 behavior must index rows directly (rag[k]). Closes #74 Co-Authored-By: Claude Opus 5.5 (1M context) --- python/seqpro/rag/_core.py | 414 +++++++++++++++++++++---------- skills/seqpro/SKILL.md | 1 + tests/test_ragged_core.py | 8 +- tests/test_ragged_tuple_index.py | 203 +++++++++++++++ 4 files changed, 488 insertions(+), 138 deletions(-) create mode 100644 tests/test_ragged_tuple_index.py diff --git a/python/seqpro/rag/_core.py b/python/seqpro/rag/_core.py index 4efa041..f48db97 100644 --- a/python/seqpro/rag/_core.py +++ b/python/seqpro/rag/_core.py @@ -82,6 +82,64 @@ def _peel_string_row(rl: "RaggedLayout[Any]", lo: int, hi: int) -> "RaggedLayout ) +def _is_int(k: Any) -> bool: + return isinstance(k, (int, np.integer)) + + +def _last_ragged_axis(shape: tuple[int | None, ...]) -> int: + return len(shape) - 1 - shape[::-1].index(None) + + +def _elem_index(starts: NDArray[Any], stops: NDArray[Any], k: int) -> NDArray[np.int64]: + """Flat index of element ``k`` (negative counts from the end) of every row.""" + lens = stops - starts + if k >= 0: + bad, idx = lens <= k, starts + k + else: + bad, idx = lens < -k, stops + k + if np.any(bad): + raise IndexError(f"index {k} is out of range for some ragged rows") + return idx.astype(np.int64) + + +def _slice_bounds( + starts: NDArray[Any], stops: NDArray[Any], sl: slice +) -> tuple[NDArray[Any], NDArray[Any]]: + """Apply a step-1 slice to every row; bounds follow Python slice semantics.""" + if sl.step not in (None, 1): + raise NotImplementedError("stepped slices on a ragged axis are unsupported") + lens = stops - starts + + def resolve(v: Any, default: NDArray[Any]) -> NDArray[Any]: + if v is None: + return default + v = int(v) + return np.clip(lens + v if v < 0 else np.full_like(lens, v), 0, lens) + + lo = resolve(sl.start, np.zeros_like(lens)) + hi = np.maximum(resolve(sl.stop, lens), lo) + return (starts + lo).astype(OFFSET_TYPE), (starts + hi).astype(OFFSET_TYPE) + + +def _ranges(starts: NDArray[Any], counts: NDArray[Any]) -> NDArray[np.int64]: + """Concatenate ``range(s, s + c)`` for every ``(s, c)`` pair, vectorized.""" + counts = counts.astype(np.int64) + ends = np.cumsum(counts) + shift = np.repeat(starts.astype(np.int64) - (ends - counts), counts) + return np.arange(int(ends[-1]) if len(ends) else 0, dtype=np.int64) + shift + + +def _index_output_axis(obj: Any, p: int, k: Any) -> Any: + """Index output axis ``p`` of any ``Ragged.__getitem__`` result with ``k``.""" + if isinstance(obj, Ragged): + return obj._index_axis(p, k) + if isinstance(obj, dict): + return {name: _index_output_axis(v, p, k) for name, v in obj.items()} + if isinstance(obj, bytes) and p == 0: + return obj[k] + return np.asarray(obj)[(slice(None),) * p + (k,)] + + class Ragged(NDArrayOperatorsMixin, Generic[RDTYPE_co]): """A non-branching ragged array with one ragged axis, or two adjacent ones. @@ -657,76 +715,46 @@ def _getitem(self, where: Any) -> Any: str_offsets=self._rl.str_offsets, ) return Ragged(new_layout) - if ( - isinstance(where, tuple) - and len(where) == 2 - and not self._is_record - and self._rl.n_ragged == 2 - and self.rag_dim == 1 - and self._is_full_slice(where[0]) - ): - return self._getitem_inner(where[1]) if isinstance(where, tuple): - # Multi-dim tuple: each key in the tuple targets successive leading axes. - # We collapse all the integer-dim keys into one combined index so that - # e.g. rag[:, [0]] correctly selects sample axis 1, not range axis 0. - # Only meaningful when there are *multiple* leading int axes to combine - # (rag_dim > 1). When rag_dim == 1 there is a single leading axis and a - # tuple just indexes it then reaches into the ragged/trailing dims; the - # combining path would mis-handle the trailing fixed dim (e.g. a - # (d0, None, K) array padded by to_padded), so fall through to the - # sequential per-key path below. - if None in self._layout.shape and self.rag_dim > 1: - return self._getitem_tuple_multidim(where) - # Sequential per-key path. np.newaxis (None) keys are handled here: - # track their positions, strip them from the tuple, apply the remaining - # keys, then insert size-1 axes at the correct positions in the result. - has_none = any(k is None for k in where) - if has_none: - # Separate None positions from real indexing keys. - real_keys: list[Any] = [] - none_positions: list[int] = [] - pos = 0 # tracks the output axis number as we accumulate - for k in where: - if k is None: - none_positions.append(pos) - pos += 1 # new axis consumes one output position - else: - real_keys.append(k) - # An integer key collapses one axis (no output position); - # a slice/array key keeps one axis. - if not isinstance(k, (int, np.integer)): - pos += 1 - result: Any = self - for k in real_keys: - result = result[k] - # Insert size-1 leading axes at recorded positions. - if isinstance(result, Ragged): - for p in reversed(none_positions): - cur_shape = result._layout.shape - new_shape = cur_shape[:p] + (1,) + cur_shape[p:] - result = Ragged( - RaggedLayout( - data=result._rl.data, - offsets=list(result._layout.offsets), - shape=new_shape, - str_offsets=result._rl.str_offsets, - ) - ) - else: - # numpy array or scalar: use np.expand_dims - for p in reversed(none_positions): - result = np.expand_dims(result, p) - return result - result = self + # np.newaxis (None) keys: track their output positions, strip them, + # apply the remaining keys, then insert size-1 axes at those positions. + real_keys: list[Any] = [] + none_positions: list[int] = [] + pos = 0 # tracks the output axis number as we accumulate for k in where: - result = result[k] + if k is None: + none_positions.append(pos) + pos += 1 # new axis consumes one output position + else: + real_keys.append(k) + # An integer key collapses one axis (no output position); + # a slice/array key keeps one axis. + if not _is_int(k): + pos += 1 + result: Any = self._index_keys(tuple(real_keys)) + # Insert size-1 leading axes at recorded positions. + if isinstance(result, Ragged): + for p in reversed(none_positions): + cur_shape = result._layout.shape + new_shape = cur_shape[:p] + (1,) + cur_shape[p:] + result = Ragged( + RaggedLayout( + data=result._rl.data, + offsets=list(result._layout.offsets), + shape=new_shape, + str_offsets=result._rl.str_offsets, + ) + ) + else: + # numpy array or scalar: use np.expand_dims + for p in reversed(none_positions): + result = np.expand_dims(result, p) return result if isinstance(self._layout, RecordLayout): return self._getitem_record(where) if self._layout.n_ragged == 2: if self.rag_dim > 1: - return self._getitem_tuple_multidim((where,)) + return self._getitem_leading((where,))[0] return self._getitem_r2(where) # Multi-dim leading shape: when rag_dim > 1 (e.g. shape (d0, d1, ..., None)) # index the first axis treating each "row" as a block of n_inner contiguous @@ -855,8 +883,11 @@ def _getitem_multidim(self, where: Any) -> "NDArray[Any] | Ragged[Any]": ) ) - def _getitem_tuple_multidim(self, where: tuple[Any, ...]) -> Any: - """Handle a tuple index on a Ragged (record or non-record) with rag_dim > 1. + def _getitem_leading(self, where: tuple[Any, ...]) -> tuple[Any, int]: + """Index the leading fixed axes of a Ragged (record or not) with rag_dim > 1. + + Returns ``(result, n_kept)``: ``n_kept`` is how many output axes the keys + left in front of the first unindexed axis, i.e. where the next key goes. Each key in ``where`` targets successive leading integer axes (like NumPy multi-dim indexing). All leading-dim keys are resolved simultaneously into @@ -869,17 +900,15 @@ def _getitem_tuple_multidim(self, where: tuple[Any, ...]) -> Any: ``rag[i]`` on a ``rag_dim == 1`` array, so ``rag[i, j] == rag[i][j]``. (R=1 keeps its 1-segment ``(None, ...)`` result.) - Keys that reach the ragged ``None`` axis or beyond are not handled here - (fall through to normal single-key dispatch after leading dims are consumed). + ``where`` must not be longer than ``rag_dim``; ``_index_keys`` applies any + keys past the leading axes. """ rag_dim = self.rag_dim leading = self._layout.shape[:rag_dim] # e.g. (d0, d1, ..., dk) trailing = self._layout.shape[rag_dim:] # (None, *trailing_int) - # Split tuple: leading-dim keys vs. keys that reach or exceed the ragged axis. - n_leading_keys = min(len(where), rag_dim) - leading_keys = where[:n_leading_keys] - remainder_keys = where[n_leading_keys:] + n_leading_keys = len(where) + leading_keys = where # For each leading dim, resolve the key to an index array (or scalar). # We track whether each dim produces an int (scalar), a "slice" range, or a @@ -1032,14 +1061,177 @@ def _getitem_tuple_multidim(self, where: tuple[Any, ...]) -> Any: ) ) if all_int: - result = result[0] + return result[0], 0 + return result, len(new_leading) + + def _index_keys(self, keys: tuple[Any, ...]) -> Any: + """Apply a tuple of (non-newaxis) keys, each to the next output axis. + + A key after a non-int key must not go to axis 0 of the intermediate + result (issue #74): track the output axis ``p`` it targets instead. + """ + if not keys: + return self + if None in self._layout.shape and self.rag_dim > 1: + n_lead = min(len(keys), self.rag_dim) + result, p = self._getitem_leading(keys[:n_lead]) + rest = keys[n_lead:] + else: + result, p, rest = self, 0, keys - # Apply any remainder keys (reaching into or past the ragged axis) - for k in remainder_keys: - result = result[k] + def is_fancy(k: Any) -> bool: + return not _is_int(k) and not isinstance(k, slice) + if any(map(is_fancy, rest)) and sum(map(is_fancy, keys)) > 1: + # Leading fancy keys broadcast jointly in _getitem_leading; fancy keys + # applied one axis at a time would not follow NumPy's broadcasting. + raise NotImplementedError( + "multiple fancy keys are only supported on leading fixed axes" + ) + for k in rest: + result = _index_output_axis(result, p, k) + if not _is_int(k): + p += 1 return result + def _index_axis(self, p: int, k: Any) -> Any: + """Index output axis ``p`` with one key, leaving every other axis intact.""" + shape = self._layout.shape + if None not in shape: # flat opaque-string collection + if p == 0: + return self._getitem(k) + raise NotImplementedError("indexing into opaque strings is unsupported") + rag_dim = self.rag_dim + if p < rag_dim: + if p == 0: + return self._getitem(k) + return self._getitem_leading((*(slice(None),) * p, k))[0] + n_rag = shape.count(None) + if p < rag_dim + n_rag: + return self._index_level(p - rag_dim, k) + if p < len(shape): + return self._index_trailing(p, k) + if self._is_record or self._rl.is_string: + raise NotImplementedError("indexing into opaque strings is unsupported") + raise IndexError(f"too many indices for a Ragged of shape {shape}") + + def _index_level(self, j: int, k: Any) -> Any: + """Index ragged axis ``rag_dim + j``, i.e. offsets level ``j``, with ``k``.""" + layout = self._layout + offs = list(layout.offsets) + starts, stops = _level_bounds(offs[j]) + if isinstance(k, slice): + offs[j] = np.stack(_slice_bounds(starts, stops, k), 0) + return self._rebuild(offs, layout.shape) + if not _is_int(k): + if len(offs) == 2 and j == 0 and not self._is_record: + return self._getitem_inner_gather(k) + raise NotImplementedError( + "fancy or boolean keys on a ragged axis are unsupported" + ) + k = int(k) + axis = self.rag_dim + j + new_shape = layout.shape[:axis] + layout.shape[axis + 1 :] + if j < len(offs) - 1: + # Pick one sub-list per row: gather the next level's bounds. + idx = _elem_index(starts, stops, k) + nxt_starts, nxt_stops = _level_bounds(offs[j + 1]) + picked = np.stack([nxt_starts[idx], nxt_stops[idx]], 0) + return self._rebuild([*offs[:j], picked, *offs[j + 2 :]], new_shape) + if j == 0: + # Innermost level is the only one: one element per leading row. + return self._take(_elem_index(starts, stops, k), [], new_shape) + # Innermost of two levels: one element per middle row that O0 references. + # Gathering only those rows repacks O0 and skips unreferenced rows. + o0_starts, o0_stops = _level_bounds(offs[0]) + counts = o0_stops - o0_starts + rows = _ranges(o0_starts, counts) + idx = _elem_index(starts[rows], stops[rows], k) + return self._take(idx, [lengths_to_offsets(counts)], new_shape) + + def _index_trailing(self, p: int, k: Any) -> "Ragged[Any]": + """Index a trailing fixed axis: a plain index on the data buffer.""" + if self._is_record: + raise NotImplementedError( + "indexing a trailing axis of a record Ragged is unsupported; " + "index a field instead" + ) + rl = self._rl + last = _last_ragged_axis(rl.shape) + new_data = rl.data[(*(slice(None),) * (p - last), k)] + return Ragged( + RaggedLayout( + data=new_data, + offsets=list(rl.offsets), + shape=(*rl.shape[: last + 1], *new_data.shape[1:]), + str_offsets=rl.str_offsets, + ) + ) + + def _leaves(self) -> "dict[str, RaggedLayout[Any]] | None": + """Record fields, or None for a single leaf.""" + if isinstance(self._layout, RecordLayout): + return self._layout.fields + return None + + def _rebuild( + self, offsets: list[NDArray[Any]], shape: tuple[int | None, ...] + ) -> "Ragged[Any]": + """Same data, new offsets and shape (fields keep their own trailing dims).""" + prefix = shape[: _last_ragged_axis(shape) + 1] + fields = self._leaves() + if fields is None: + rl = self._rl + return Ragged(RaggedLayout(rl.data, offsets, shape, rl.str_offsets)) + new_fields = { + name: RaggedLayout( + fl.data, + offsets, + (*prefix, *fl.shape[_last_ragged_axis(fl.shape) + 1 :]), + fl.str_offsets, + ) + for name, fl in fields.items() + } + return Ragged(RecordLayout(offsets, shape, new_fields)) + + def _take( + self, + idx: NDArray[np.int64], + offsets: list[NDArray[Any]], + shape: tuple[int | None, ...], + ) -> Any: + """Gather elements ``idx`` of every leaf under new ``offsets``/``shape``. + + ``shape`` is the result shape including the record's (first field's) + trailing dims; with no ``offsets`` left the result is dense (a dict of + per-field results for a record), or a flat opaque-string collection. + """ + n_trailing = len(self._layout.shape) - 1 - _last_ragged_axis(self._layout.shape) + prefix = shape[: len(shape) - n_trailing] + + def leaf(fl: "RaggedLayout[Any]") -> Any: + if fl.str_offsets is not None: + str_starts, str_stops = _level_bounds(fl.str_offsets) + str_offsets = np.stack([str_starts[idx], str_stops[idx]], 0) + return Ragged(RaggedLayout(fl.data, offsets, prefix, str_offsets)) + data = fl.data[idx] + if not offsets: + dense = data.reshape( + (*cast("tuple[int, ...]", prefix), *data.shape[1:]) + ) + return dense[()] # a NumPy scalar when every axis was indexed + return Ragged(RaggedLayout(data, offsets, (*prefix, *data.shape[1:]))) + + fields = self._leaves() + if fields is None: + return leaf(self._rl) + out = {name: leaf(fl) for name, fl in fields.items()} + if not offsets: + return out + return Ragged( + RecordLayout(offsets, shape, {name: r._layout for name, r in out.items()}) + ) + def _getitem_record(self, where: Any) -> Any: rec = self._layout assert isinstance(rec, RecordLayout) @@ -1053,7 +1245,7 @@ def _getitem_record(self, where: Any) -> Any: # A record with >1 leading fixed axis routes through the multidim path # directly (never re-dispatching through __getitem__, so no recursion). if not isinstance(where, tuple) and self.rag_dim > 1: - return self._getitem_tuple_multidim((where,)) + return self._getitem_leading((where,))[0] return self._getitem_record_rows(where) def _gather_indices( @@ -1120,58 +1312,12 @@ def _row_gather(self, where: Any) -> "tuple[NDArray[Any], NDArray[Any]]": as contiguous OFFSET_TYPE arrays for the shared ragged axis.""" return self._gather_indices(where, *self._starts_stops()) - @staticmethod - def _is_full_slice(k: Any) -> bool: - return isinstance(k, slice) and k == slice(None) - - def _getitem_inner(self, sel: Any) -> "Ragged[Any]": - o0, o1 = self._layout.offsets - o0_starts, o0_stops = _level_bounds(o0) - o1_starts, o1_stops = _level_bounds(o1) - if isinstance(sel, (int, np.integer)): # k-th middle of each group -> R=1 - counts = o0_stops - o0_starts - if np.any(sel >= counts) or np.any(-sel > counts): - raise IndexError(f"inner index {sel} out of range for some group") - mid_idx = (o0_starts + (sel if sel >= 0 else counts + sel)).astype(np.int64) - ds = np.ascontiguousarray(o1_starts[mid_idx], dtype=OFFSET_TYPE) - de = np.ascontiguousarray(o1_stops[mid_idx], dtype=OFFSET_TYPE) - trailing = self._layout.shape[self.rag_dim + 2 :] - return Ragged( - RaggedLayout( - data=self._rl.data, - offsets=[np.stack([ds, de], 0)], - shape=(len(mid_idx), None, *trailing), - ) - ) - if isinstance(sel, slice): # local per-group slice -> R=2 - start, stop, step = sel.indices(1 << 62) - if step != 1: - raise NotImplementedError("step != 1 inner slices are unsupported") - if (sel.start is not None and sel.start < 0) or ( - sel.stop is not None and sel.stop < 0 - ): - raise NotImplementedError( - "negative inner-slice bounds (rag[:, -k:]) are unsupported" - ) - new_starts = np.minimum(o0_starts + start, o0_stops) - new_stops = np.maximum(np.minimum(o0_starts + stop, o0_stops), new_starts) - new_o0 = np.stack( - [new_starts.astype(OFFSET_TYPE), new_stops.astype(OFFSET_TYPE)], 0 - ) - trailing = self._layout.shape[self.rag_dim + 2 :] - return Ragged( - RaggedLayout( - data=self._rl.data, - offsets=[new_o0, o1], - shape=(len(new_starts), None, None, *trailing), - ) - ) - return self._getitem_inner_gather(sel) # mask / int-array -> Task 8 - def _getitem_inner_gather(self, sel: Any) -> "Ragged[Any]": + """Mask or int-array key on the first of two ragged axes.""" o0, o1 = self._layout.offsets o0_starts, o0_stops = _level_bounds(o0) o1_starts, o1_stops = _level_bounds(o1) + leading = self._layout.shape[: self.rag_dim] trailing = self._layout.shape[self.rag_dim + 2 :] if _where_is_bool(sel): # mask over the global middle axis -> R=2 from seqpro.seqpro import _ragged_nested_gather # type: ignore[missing-import] @@ -1188,23 +1334,23 @@ def _getitem_inner_gather(self, sel: Any) -> "Ragged[Any]": RaggedLayout( data=self._rl.data, offsets=[new_o0, np.stack([ds, de], 0)], - shape=(len(o0_starts), None, None, *trailing), + shape=(*leading, None, None, *trailing), ) ) - idx = np.atleast_1d( - np.asarray(sel, dtype=np.int64) - ) # uniform per-group int array - counts = o0_stops - o0_starts - if np.any(idx.max() >= counts) or np.any(idx.min() < -counts.min()): + # uniform per-group int array: middle (g, i) = element idx[i] of group g + idx = np.atleast_1d(np.asarray(sel, dtype=np.int64)) + counts = (o0_stops - o0_starts)[:, None] + if np.any(idx >= counts) or np.any(-idx > counts): raise IndexError("uniform inner index out of range for some group") - cols = [self._getitem_inner(int(k)) for k in idx] # each is (L0, ~K) R=1 - ds = np.stack([c._layout.offsets[0][0] for c in cols], 1).reshape(-1) - de = np.stack([c._layout.offsets[0][1] for c in cols], 1).reshape(-1) + mid = np.where(idx >= 0, o0_starts[:, None], o0_stops[:, None]) + idx + mid = mid.reshape(-1) + ds = np.ascontiguousarray(o1_starts[mid], dtype=OFFSET_TYPE) + de = np.ascontiguousarray(o1_stops[mid], dtype=OFFSET_TYPE) return Ragged( RaggedLayout( data=self._rl.data, offsets=[np.stack([ds, de], 0)], - shape=(len(counts), len(idx), None, *trailing), + shape=(*leading, len(idx), None, *trailing), ) ) diff --git a/skills/seqpro/SKILL.md b/skills/seqpro/SKILL.md index 10ade7c..82129bb 100644 --- a/skills/seqpro/SKILL.md +++ b/skills/seqpro/SKILL.md @@ -91,6 +91,7 @@ rag = sp.rag.Ragged.empty((10, None, 4), dtype=np.uint8) # batch of 10 OHE seq | Count top-level rows | `len(rag)` — returns `shape[0]` (raises if `shape[0]` is the ragged axis) | `rag.shape[0]` with manual int-cast | | Index one group of an opaque-string `Ragged` | `rag[i]` → a `Ragged` of `bytes`, one per string (`len(rag[i]) == rag.lengths[i]`); `rag[i][j]` is one `bytes` | Expect one concatenated `bytes` — that was the pre-0.22 behavior and it silently dropped the per-string boundaries | | Index a nested array with fixed leading axes, e.g. `(b, p, ~v, ~w)` | NumPy-style keys on the fixed axes: `rag[i]` → `(p, ~v, ~w)`, `rag[:, j]` → `(b, ~v, ~w)`; `rag[i, j] == rag[i][j]` peels to `(n_v, ~w)` | Read `rag.offsets` by hand to find a row's ploid | +| Index along the ragged or a trailing axis | NumPy/awkward tuple keys, each hitting the next output axis: `rag[:, 0]` → first element of every row (dense; `IndexError` if a row is too short), `rag[:, 1:]` / `rag[:, -2:]` → per-row slice (zero-copy), `rag[:, :, k]` on `(b, ~, K)` → trailing column. Works on nested arrays too (`rag[:, :, 0]` on `(b, ~v, ~w)` → first of each window) | Loop over rows; expect `NotImplementedError` for stepped slices, fancy/mask keys on a ragged axis (except the first of two ragged axes), and trailing keys on records | | Insert a leading size-1 axis | `rag[np.newaxis]` — returns `Ragged` with shape `(1, *old_shape)` | Manual `from_offsets` rebuild | | Reinterpret bytes/dtype | `rag.view(np.uint8)` | `np.asarray(rag).view(...)` (loses ragged structure) | | Reshape non-ragged axes | `rag.reshape(batch, None, k, 4)` | Touch `rag.data.shape` directly | diff --git a/tests/test_ragged_core.py b/tests/test_ragged_core.py index 6666763..1f40a0f 100644 --- a/tests/test_ragged_core.py +++ b/tests/test_ragged_core.py @@ -631,15 +631,15 @@ def test_r2_per_group_inner_slice(): np.testing.assert_array_equal(sub[1, 0], np.array([5, 6, 7])) -def test_r2_per_group_inner_slice_negative_raises(): +def test_r2_per_group_inner_slice_negative(): data = np.arange(10, dtype=np.int32) rag = Ragged.from_offsets( data, (2, None, None), - [np.array([0, 2, 4]), np.array([0, 3, 5, 8, 10])], + [np.array([0, 1, 4]), np.array([0, 3, 5, 8, 10])], ) - with pytest.raises(NotImplementedError, match="negative"): - rag[:, -2:] + # Python slice semantics per group (#74): group0 has 1 middle, group1 has 3 + assert rag[:, -2:].to_ak().tolist() == rag.to_ak()[:, -2:].tolist() # --------------------------------------------------------------------------- diff --git a/tests/test_ragged_tuple_index.py b/tests/test_ragged_tuple_index.py new file mode 100644 index 0000000..fe581c8 --- /dev/null +++ b/tests/test_ragged_tuple_index.py @@ -0,0 +1,203 @@ +"""Issue #74: tuple keys after the leading axes index the right axis. + +Each key in a tuple targets the next *output* axis, as in NumPy/awkward. Keys +after a non-int key used to be applied to axis 0 of the intermediate result. +The oracle is awkward indexing on the equivalent array. +""" + +from __future__ import annotations + +from typing import Any + +import awkward as ak +import numpy as np +import pytest + +from seqpro.rag import Ragged + +S = slice(None) + +A_LIST = [[0, 1, 2], [3, 4], [5, 6, 7, 8]] # (3, ~) +B_LIST = [[[0, 1, 2], [3, 4]], [[5, 6, 7, 8], [9, 10, 11]]] # (2, 2, ~) +C_LIST = [ + [[0, 1], [2, 3], [4, 5]], + [[6, 7], [8, 9]], + [[10, 11], [12, 13], [14, 15], [16, 17]], +] # (3, ~, 2) +N_LIST = [[[0, 1, 2], [3, 4]], [[5, 6, 7, 8]]] # (2, ~, ~) +F_LIST = [ # (2, 2, ~, ~) + [[[0, 1, 2], [3, 4]], [[5]]], + [[], [[6, 7], [8], [9, 10, 11]]], +] + + +def _a() -> Ragged[Any]: + return Ragged.from_offsets(np.arange(9), (3, None), np.array([0, 3, 5, 9])) + + +def _b() -> Ragged[Any]: + return Ragged.from_offsets(np.arange(12), (2, 2, None), np.array([0, 3, 5, 9, 12])) + + +def _c() -> Ragged[Any]: + data = np.arange(18).reshape(9, 2) + return Ragged.from_offsets(data, (3, None, 2), np.array([0, 3, 5, 9])) + + +def _n() -> Ragged[Any]: + return Ragged.from_offsets( + np.arange(9), (2, None, None), [np.array([0, 2, 3]), np.array([0, 3, 5, 9])] + ) + + +def _f() -> Ragged[Any]: + return Ragged.from_offsets( + np.arange(12), + (2, 2, None, None), + [np.array([0, 2, 3, 3, 6]), np.array([0, 3, 5, 6, 8, 9, 12])], + ) + + +FIXTURES = { + "a": (_a, A_LIST), + "b": (_b, B_LIST), + "c": (_c, C_LIST), + "n": (_n, N_LIST), + "f": (_f, F_LIST), +} + +CASES = [ + ("a", (S, 0)), + ("a", (S, -1)), + ("a", (S, slice(1, None))), + ("a", (S, slice(None, 2))), + ("a", (S, slice(-2, None))), + ("a", (S, slice(-5, -1))), + ("a", ([2, 0], 0)), + ("a", (slice(1, None), 1)), + ("a", (np.array([True, False, True]), -1)), + ("a", (0, 1)), + ("a", (S, 2)), # row 1 too short -> IndexError + ("b", (S, S, 0)), + ("b", (0, S, 0)), + ("b", (S, 1, -1)), + ("b", ([1, 0], S, slice(1, None))), + ("b", (1, 0, 1)), + ("b", (S, 0, slice(None, 2))), + ("c", (S, 0)), + ("c", (S, S, 1)), + ("c", (0, S, 1)), + ("c", (S, 0, 1)), + ("c", (S, slice(1, None), 0)), + ("c", (0, 1, 1)), + ("c", (S, -1, S)), + ("c", (S, S, [1, 0])), + ("n", (S, 0)), + ("n", (S, 0, 1)), + ("n", (S, S, 0)), + ("n", (S, slice(1, None), slice(1, None))), + ("n", (0, 1, 1)), + ("n", (S, -1, -1)), + ("n", (S, 0, slice(None, 2))), + ("n", (S, S, -1)), + ("f", (S, 1, 0)), + ("f", (S, 1, 0, -1)), + ("f", (S, S, S, 0)), + ("f", (0, S, S, slice(1, None))), + ("f", (S, S, S, slice(None, 1))), + ("f", (S, S, 0)), # b1p0 has no variants -> IndexError + ("f", (1, 1, S, 0)), +] + + +def _tolist(x: Any) -> Any: + if isinstance(x, Ragged): + return x.to_ak().tolist() + if isinstance(x, dict): + return {k: _tolist(v) for k, v in x.items()} + if isinstance(x, bytes): + return x + return np.asarray(x).tolist() + + +@pytest.mark.parametrize(("name", "key"), CASES, ids=repr) +def test_tuple_index_matches_awkward(name: str, key: tuple[Any, ...]): + make, nested = FIXTURES[name] + try: + expected = ak.Array(nested)[key] + except (IndexError, ValueError): + with pytest.raises(IndexError): + make()[key] + return + expected = expected.tolist() if isinstance(expected, ak.Array) else expected + assert _tolist(make()[key]) == expected + + +@pytest.mark.parametrize(("name", "key"), CASES, ids=repr) +def test_tuple_index_on_gathered(name: str, key: tuple[Any, ...]): + # Gathered (2, N) offsets must index the same as packed ones. + make, nested = FIXTURES[name] + perm = [1, 0] if name in {"b", "n", "f"} else [2, 0, 1] + try: + expected = ak.Array(nested)[perm][key] + except (IndexError, ValueError): + with pytest.raises(IndexError): + make()[perm][key] + return + expected = expected.tolist() if isinstance(expected, ak.Array) else expected + assert _tolist(make()[perm][key]) == expected + + +def _strings() -> Ragged[Any]: + # (3, ~) opaque strings: ["A", "CG"], ["T"], ["GA", "", "C"] + data = np.frombuffer(b"ACGTGAC", "S1") + str_off = np.array([0, 1, 3, 4, 6, 6, 7]) + return Ragged.from_offsets( + data, (3, None), np.array([0, 2, 3, 6]), str_offsets=str_off + ) + + +@pytest.mark.parametrize( + "key", [(S, 0), (S, -1), (S, slice(1, None)), ([2, 0], 0)], ids=repr +) +def test_string_under_axis(key: tuple[Any, ...]): + rag = _strings() + assert _tolist(rag[key]) == rag.to_ak()[key].tolist() + + +@pytest.mark.parametrize("name", ["a", "b", "n"]) +@pytest.mark.parametrize("key", [(S, 0), (S, slice(1, None)), (S, -1)], ids=repr) +def test_record(name: str, key: tuple[Any, ...]): + make, nested = FIXTURES[name] + rec = Ragged.from_fields({"x": make(), "y": make() * 10}) + oracle = {"x": ak.Array(nested), "y": ak.Array(nested) * 10} + got = rec[key] + for field, arr in oracle.items(): + assert _tolist(got[field]) == arr[key].tolist(), field + + +@pytest.mark.parametrize( + ("name", "key"), + [ + ("a", (S, slice(None, None, 2))), # stepped ragged slice + ("a", (S, [0, 1])), # fancy key on a ragged axis + ("a", (S, np.array([True, False]))), + ("n", (S, S, [0, 1])), + ], + ids=repr, +) +def test_unsupported_raises(name: str, key: tuple[Any, ...]): + with pytest.raises(NotImplementedError): + FIXTURES[name][0]()[key] + + +def test_record_trailing_key_raises(): + rec = Ragged.from_fields({"x": _c(), "y": _c()}) + with pytest.raises(NotImplementedError): + rec[:, :, 0] + + +def test_newaxis_with_ragged_key(): + # np.newaxis positions are still tracked around ragged-axis keys. + got = _a()[None, :, 0] + assert _tolist(got) == [[0, 3, 5]]