fix(rag)!: nested fixed axes, tuple-key axes, and record to_ak depth

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)
@@ -223,7 +218,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:

skills/seqpro/SKILL.md

@@ -49,15 +49,15 @@ For exact signatures and kwargs, read the docstring directly (`sp.<fn>?` 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,8 @@ 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 |
+| 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 |
@@ -128,6 +130,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
@@ -163,7 +166,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)

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()
 
 
 # ---------------------------------------------------------------------------

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()
 
 
 # ---------------------------------------------------------------------------

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()

tests/test_ragged_nested_fixed_axis.py

@@ -0,0 +1,154 @@
+"""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
+
+
+@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()
+    )