Document and test RFD4's batching contract for several designs

No change to how write_cif indexes anything. RFD4 replicates at collate time --
datapipes/collate.py does `[sample for sample in samples for _ in
range(num_replicates)]` -- so every per-row quantity already has one row per
design by the time anything reads it, and the export loop in proteina.py reads
row i of each. write_cif already did the same thing, and keeping it that way is
the point: this package indexes the way RFD4 indexes.

What was missing was saying so, and testing it. The README example replicated
with `repeat`, which tiles; an input's designs sit together, so it must be
`repeat_interleave`, and the difference is invisible for one input and wrong
for more than one. It now also shows `num_designs_per_prompt` alongside the
replication, since a pipeline repeating its input internally still leaves the
caller to replicate the features it passes back in.

Tests cover one and three designs with replicated features, that unreplicated
features raise IndexError on the second design rather than quietly reading
another row, and that a row past the end raises.

An earlier draft of this added a helper deriving the design count from the
feature and output shapes, so that unreplicated features would work. That was
invention: it is not what RFD4 does, it would have been one more thing to keep
in step with upstream, and the first version of it silently read the wrong
input's mask for batches of more than one input. Following the reference is
cheaper than being clever next to it.

Verified against the Diffusers pipeline on branch rfd4-torch: prepare_cif's
features pass its check_inputs with num_designs_per_prompt=3, three designs
write from replicated features, and the unreplicated case raises.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

README.md

@@ -202,11 +202,25 @@ provide is the rest of a training run: the dataset loader, the task configuratio
 
 Source attribution and third-party licenses: [THIRD_PARTY_LICENSES.txt](THIRD_PARTY_LICENSES.txt).
 
-To generate several designs of one input in a batch, repeat the features and write each design with `row`:
+To generate several designs of one input, ask the pipeline for them and replicate the features to match, which is
+what RFD4 itself does: its collate step repeats each sample `num_replicates` times before batching, so every
+per-row quantity has one row per design and `row` indexes them all alike.
 
 ```python
-batch = {name: value.repeat(8, *[1] * (value.ndim - 1)) for name, value in features.items()}
-output = pipe(batch, generator=[torch.Generator("cuda:0").manual_seed(seed) for seed in range(8)])
-for row in range(8):
-    write_cif(output, batch, atoms, f"design_{row}.cif", row=row)
+designs = 8
+output = pipe(
+    features,
+    num_designs_per_prompt=designs,
+    generator=[torch.Generator("cuda:0").manual_seed(seed) for seed in range(designs)],
+)
+replicated = {
+    name: value.repeat_interleave(designs, dim=0) if torch.is_tensor(value) else value
+    for name, value in features.items()
+}
+for row in range(designs):
+    write_cif(output, replicated, atoms, f"design_{row}.cif", row=row)
 ```
+
+Use `repeat_interleave`, not `repeat`: an input's designs sit together, so output row `r` belongs to input
+`r // designs`. Passing the features unreplicated raises `IndexError` on the second design rather than writing the
+wrong row.

tests/test_write_cif_designs.py

@@ -0,0 +1,118 @@
+"""`write_cif` across several designs, following RFD4's own batching contract.
+
+RFD4 replicates at collate time -- `datapipes/collate.py` does
+`[sample for sample in samples for _ in range(num_replicates)]` -- so by the time anything is
+indexed, every per-row quantity already has one row per design, and the export loop in
+`proteina.py` reads row `i` of each. `write_cif` does the same: `row` indexes both the output and
+the features, so the caller passes as many feature rows as the output has rows.
+
+These cover that contract, and that passing too few rows fails loudly rather than reading some
+other row's mask.
+"""
+
+import tempfile
+import types
+from pathlib import Path
+
+import pytest
+
+ROOT = Path(__file__).resolve().parents[1]
+INPUT = ROOT / "examples/inputs/inverse_folding.cif"
+
+
+def _prepared():
+    pytest.importorskip("atomworks")
+    import torch
+
+    from rfd4_proteina_prep import prepare_cif
+
+    if not INPUT.is_file():
+        pytest.skip(f"{INPUT.name} is not present")
+    return prepare_cif(INPUT, seed=0, device=torch.device("cpu"))
+
+
+def _replicate(features, n_designs: int):
+    """What RFD4's collate does, and what a pipeline repeating its input leaves the caller to do.
+
+    `repeat_interleave` rather than `repeat`: an input's designs sit together, so output row `r`
+    belongs to input `r // n_designs`. Tiling would pair the rows differently, and for more than
+    one input, wrongly.
+    """
+    return {
+        name: value.repeat_interleave(n_designs, dim=0) if hasattr(value, "repeat_interleave") else value
+        for name, value in features.items()
+    }
+
+
+def _output_of(features, n_designs: int):
+    """The shape a pipeline returns for `n_designs` designs of one input.
+
+    Every design carries the input's own coordinates, so the written files must agree; what is
+    under test is which row is read, not what a model would have produced.
+    """
+    import torch
+
+    tokens = features["coords"].shape[1]
+    repeat = lambda tensor: tensor.repeat_interleave(n_designs, dim=0)  # noqa: E731 - matches the pipeline
+    return types.SimpleNamespace(
+        coords_nm=repeat(features["coords_nm"]),
+        residue_type=repeat(features["residue_type"]),
+        atom_mask=repeat(features["atom37_mask"]),
+        insertion_logits=torch.zeros((n_designs, tokens, tokens)),
+    )
+
+
+@pytest.mark.parametrize("n_designs", [1, 3])
+def test_every_design_is_written(n_designs):
+    """Features replicated to match the output, one file per design."""
+    from rfd4_proteina_prep import write_cif
+
+    features, atoms = _prepared()
+    assert features["padding_mask"].shape[0] == 1, "prepare_cif returns a single row"
+    replicated = _replicate(features, n_designs)
+    assert replicated["padding_mask"].shape[0] == n_designs
+    output = _output_of(features, n_designs)
+
+    with tempfile.TemporaryDirectory() as directory:
+        written = [
+            Path(write_cif(output, replicated, atoms, Path(directory) / f"design_{row}.cif", row=row))
+            for row in range(n_designs)
+        ]
+        assert all(path.is_file() for path in written)
+        # Same coordinates in every design, so the files agree apart from the lines the writer
+        # derives from the output path and the clock.
+        skip = ("_entry.time", "_entry.id", "data_")
+        bodies = {
+            tuple(line for line in path.read_text().splitlines() if not line.startswith(skip)) for path in written
+        }
+        assert len(bodies) == 1, "designs of one input differ beyond their file-derived lines"
+
+
+def test_unreplicated_features_fail_loudly():
+    """Too few feature rows raises rather than quietly reading another row.
+
+    This is what a caller sees when the pipeline repeated its input internally and they passed the
+    features straight through. Row 0 happens to work, because one feature row is the row it wants;
+    every later row must fail rather than succeed on the wrong data.
+    """
+    from rfd4_proteina_prep import write_cif
+
+    features, atoms = _prepared()
+    output = _output_of(features, 3)
+    with tempfile.TemporaryDirectory() as directory:
+        write_cif(output, features, atoms, Path(directory) / "design_0.cif", row=0)
+        with pytest.raises(IndexError):
+            write_cif(output, features, atoms, Path(directory) / "design_1.cif", row=1)
+
+
+def test_row_beyond_the_output_is_an_error():
+    """A row past the designs that exist fails rather than writing design 0 again."""
+    from rfd4_proteina_prep import write_cif
+
+    features, atoms = _prepared()
+    n_designs = 2
+    output = _output_of(features, n_designs)
+    replicated = _replicate(features, n_designs)
+    with tempfile.TemporaryDirectory() as directory:
+        with pytest.raises(IndexError):
+            write_cif(output, replicated, atoms, Path(directory) / "design_9.cif", row=9)