Cohorts: run_cohort — N subjects per run

site_docs/guides/genotype.md

@@ -409,6 +409,62 @@ the allele frequencies are cartridge- or uniform-sourced — supplying an
 `model_id` / `source` / `version`, the plane's `model_checksum`, per-segment gene
 counts, the novel-allele count, and `source_field` (`"DataConfig.genotype_priors"`).
 
+## Cohorts
+
+A single genotype models one subject. To simulate a **cohort** — many subjects,
+each with their own genotype — use `run_cohort`. It runs the single-subject path
+once per subject and collects the results:
+
+```python
+import GenAIRR as ga
+import GenAIRR.data as gdata
+from GenAIRR.genotype import Genotype
+
+cfg = gdata.HUMAN_IGH_OGRDB
+# one sampled genotype per donor
+donors = [Genotype.sample(cfg, seed=s, subject_id=f"donor_{s}") for s in range(5)]
+
+cohort = ga.Experiment.on(cfg).recombine().run_cohort(
+    donors, n_per_subject=200, seed=0, expose_provenance=True)
+
+cohort.subject_ids            # ['donor_0', ..., 'donor_4']
+len(cohort)                   # 1000 total records
+cohort.result_for("donor_2")  # that donor's SimulationResult
+cohort.to_csv("cohort.csv")   # combined, subject-tagged, unique sequence_id
+```
+
+`run_cohort` returns a `CohortResult`:
+
+- `.subject_ids` / `.genotypes` / `.results` — per-subject, in input order.
+- `.result_for(sid)` / `.refdata_for(sid)` — one subject's `SimulationResult` and
+  the reference it ran against (preserved per subject, so `validate_records` and
+  novel-allele truth calls stay correct even when subjects differ).
+- `.records` — a fresh, subject-tagged, `sequence_id`-namespaced concatenation
+  (each `sequence_id` is `"{subject_id}_{...}"`, so combined AIRR/FASTA export
+  never collides).
+- `.to_dataframe()` / `.to_csv()` / `.to_fasta()` — combined export over a stable
+  union of columns.
+
+**Record counts.** `n_per_subject` applies to all subjects; pass `counts` (a
+parallel list, same length as the genotypes) to vary per-subject repertoire
+sizes. A count of `0` is allowed — that subject appears in the cohort with zero
+records.
+
+```python
+cohort = ga.Experiment.on(cfg).recombine().run_cohort(
+    donors, counts=[500, 200, 0, 1000, 300], seed=0)
+```
+
+**Subject IDs.** Taken from each genotype's `subject_id`; if none are set they are
+auto-assigned `subject_0..N-1`. Mixed (some set, some not) or duplicate IDs raise.
+
+**Determinism.** Each subject gets an independent sub-seed derived from `seed`, so
+a cohort is fully reproducible and subjects are independent.
+
+`run_cohort` is mutually exclusive with `with_genotype`, `restrict_alleles`, and
+`recombine(*_allele_weights=...)` (the genotype owns allele expression), and — in
+this release — is not combined with `receptor_revision` or clonal forks.
+
 ## Novel / private alleles
 
 Individuals carry germline alleles that aren't in any reference — *private* or
@@ -616,8 +672,6 @@ deletion calls.
 
 The genotype foundation is deliberately scoped. Deferred to later work:
 
-- **Cohorts** — many subjects, each with their own genotype, in one run
-  (`with_genotype` is single-subject; `result.genotypes` is a one-element list).
 - **External loaders** — importing genotypes from VDJbase / TIgGER / IgDiscover /
   partis output.
 - **Same-haplotype receptor revision** — `receptor_revision` with a genotype is

src/GenAIRR/cohort.py

@@ -0,0 +1,157 @@
+"""Cohort orchestration results — N subjects each with their own diploid
+genotype, produced by :meth:`GenAIRR.Experiment.run_cohort`.
+
+``CohortResult`` is built from per-subject ``CohortSubjectResult`` entries and
+exposes per-subject access (``result_for`` / ``refdata_for``) plus combined
+export. It explicitly stores each subject's refdata because ``SimulationResult``
+does not preserve it (needed for ``validate_records`` and novel-allele subjects).
+See ``.private/specs/2026-06-17-genotype-cohorts-design.md``.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+
+def _resolve_subject_ids(raw_ids: List[Any]) -> List[str]:
+    """Resolve per-subject IDs: all-None -> subject_0..N-1; all-present ->
+    require unique after str() normalization; mixed -> raise."""
+    none_count = sum(1 for i in raw_ids if i is None)
+    n = len(raw_ids)
+    if none_count == n:
+        return [f"subject_{i}" for i in range(n)]
+    if none_count != 0:
+        raise ValueError(
+            "run_cohort: some genotypes have a subject_id and others don't; set "
+            "subject_id on all genotypes or none")
+    ids = [str(i) for i in raw_ids]
+    if len(ids) != len(set(ids)):
+        raise ValueError("run_cohort: duplicate subject_id after normalization")
+    return ids
+
+
+def _check_count(c, where: str) -> int:
+    if isinstance(c, bool) or not isinstance(c, int) or c < 0:
+        raise ValueError(f"{where} must be an int >= 0, got {c!r}")
+    return c
+
+
+def _resolve_counts(n_genotypes: int, n_per_subject, counts) -> List[int]:
+    """Resolve per-subject record counts. ``counts`` (a parallel sequence)
+    overrides ``n_per_subject`` when supplied; entries are validated as
+    non-bool ints >= 0."""
+    _check_count(n_per_subject, "n_per_subject")
+    if counts is None:
+        return [int(n_per_subject)] * n_genotypes
+    from collections.abc import Mapping
+    if isinstance(counts, (str, bytes, bytearray, Mapping)):
+        raise ValueError(
+            "run_cohort: counts must be a parallel list/tuple of ints, not a "
+            f"{type(counts).__name__}")
+    counts = list(counts)
+    if len(counts) != n_genotypes:
+        raise ValueError(
+            f"run_cohort: counts length {len(counts)} != number of genotypes "
+            f"{n_genotypes}")
+    return [_check_count(c, f"counts[{i}]") for i, c in enumerate(counts)]
+
+
+@dataclass(frozen=True)
+class CohortSubjectResult:
+    """One subject's slice of a cohort run."""
+
+    subject_id: str
+    genotype: Any            # GenAIRR.genotype.Genotype (snapshot)
+    result: Any              # GenAIRR.result.SimulationResult
+    refdata: Any             # the refdata this subject ran against (base or effective)
+    seed: int                # derived per-subject sub-seed
+    count: int               # records requested for this subject
+
+
+class CohortResult:
+    """Combined result of a :meth:`Experiment.run_cohort` run.
+
+    ``subjects`` is the single source of truth; ``subject_ids`` / ``genotypes`` /
+    ``results`` are derived so parallel lists can't drift."""
+
+    def __init__(self, subjects: List[CohortSubjectResult]):
+        self._subjects = list(subjects)
+        self._by_id = {s.subject_id: s for s in self._subjects}
+
+    @property
+    def subjects(self) -> List[CohortSubjectResult]:
+        return list(self._subjects)
+
+    @property
+    def subject_ids(self) -> List[str]:
+        return [s.subject_id for s in self._subjects]
+
+    @property
+    def genotypes(self) -> List[Any]:
+        return [s.genotype for s in self._subjects]
+
+    @property
+    def results(self) -> List[Any]:
+        return [s.result for s in self._subjects]
+
+    def result_for(self, subject_id: str):
+        return self._by_id[subject_id].result
+
+    def refdata_for(self, subject_id: str):
+        return self._by_id[subject_id].refdata
+
+    @property
+    def records(self) -> List[Dict[str, Any]]:
+        """A fresh concatenated list of every subject's record dicts (each
+        already ``subject_id``-tagged and ``sequence_id``-namespaced). Each dict
+        is a shallow copy, so neither the returned list nor mutating a record in
+        it affects the per-subject results — edit those via ``result_for``."""
+        out: List[Dict[str, Any]] = []
+        for s in self._subjects:
+            out.extend(dict(rec) for rec in s.result.records)
+        return out
+
+    def __len__(self) -> int:
+        return sum(len(s.result) for s in self._subjects)
+
+    def __repr__(self) -> str:
+        return f"<CohortResult subjects={len(self._subjects)} records={len(self)}>"
+
+    # ── combined export ─────────────────────────────────────────────
+    def to_dataframe(self, *, airr_strict: bool = False):
+        """Combined DataFrame over every subject's records. Column set is the
+        stable union across subjects (pandas fills missing keys with NaN)."""
+        import pandas as pd
+
+        from .result import _DEFAULT_COLUMN_ORDER
+
+        records = self.records
+        if not records:
+            # Empty cohort: preserve the default AIRR schema plus cohort-owned
+            # columns so an empty export still has a usable header (parity with
+            # SimulationResult.to_dataframe on an empty result).
+            cols = list(_DEFAULT_COLUMN_ORDER)
+            for extra in ("subject_id", "haplotype"):
+                if extra not in cols:
+                    cols.append(extra)
+            return pd.DataFrame(columns=cols)
+        if airr_strict:
+            from .result import _to_airr_strict
+            records = [_to_airr_strict(r) for r in records]
+        return pd.DataFrame(records)
+
+    def to_csv(self, path: str, *, airr_strict: bool = False) -> None:
+        """Write the combined records as CSV (union columns guaranteed by the
+        DataFrame)."""
+        self.to_dataframe(airr_strict=airr_strict).to_csv(path, index=False)
+
+    def to_fasta(self, path: str) -> None:
+        """Write one FASTA record per concatenated record. Header is exactly
+        ``>{sequence_id}`` (the namespaced id); body is the record's
+        ``sequence``. Unlike ``SimulationResult.to_fasta`` (whose headers come
+        from the enumerate index), this keeps cohort headers globally unique."""
+        with open(path, "w", encoding="utf-8") as fh:
+            for rec in self.records:
+                sid = rec.get("sequence_id", "")
+                seq = rec.get("sequence", "")
+                fh.write(f">{sid}\n{seq}\n")

src/GenAIRR/experiment.py

@@ -2948,6 +2948,135 @@ def _build_simulator(
             allow_curatable_refdata=allow_curatable_refdata,
         )
 
+    def run_cohort(
+        self,
+        genotypes,
+        *,
+        n_per_subject: int = 1,
+        seed: int = 0,
+        counts=None,
+        strict: bool = False,
+        expose_provenance: bool = False,
+        validate_records: bool = False,
+        allow_curatable_refdata: Optional[bool] = None,
+    ) -> "CohortResult":
+        """Run a cohort: N subjects, each with its own diploid genotype, in one
+        call. A Python loop around the single-subject genotype path — each
+        subject is compiled and run independently, records are tagged with
+        ``subject_id`` and given a namespaced ``sequence_id``, and the per-subject
+        ``SimulationResult`` (with its own refdata) is collected into a
+        :class:`~GenAIRR.cohort.CohortResult`.
+
+        ``n_per_subject`` applies to every subject; ``counts`` (a parallel
+        sequence, same length as ``genotypes``) overrides it per subject and may
+        contain ``0`` (that subject appears with zero records). Subject IDs are
+        taken from each genotype, or auto-assigned ``subject_0..N-1`` when all are
+        unset; mixed/duplicate IDs raise. Per-subject sub-seeds are derived
+        deterministically from ``seed``.
+
+        Mutually exclusive with :meth:`with_genotype`, :meth:`restrict_alleles`,
+        and ``recombine(*_allele_weights=...)`` (the genotype owns allele
+        expression). Not supported with :meth:`receptor_revision` or clonal forks
+        in this release.
+        """
+        import copy as _copy
+        import random as _random
+
+        from .cohort import (
+            CohortResult,
+            CohortSubjectResult,
+            _resolve_counts,
+            _resolve_subject_ids,
+        )
+        from .genotype import Genotype
+        from .result import SimulationResult
+
+        gts = list(genotypes)
+        if not gts:
+            raise ValueError("run_cohort: genotypes must be a non-empty sequence")
+        for g in gts:
+            if not isinstance(g, Genotype):
+                raise TypeError(
+                    f"run_cohort: every element must be a Genotype, got "
+                    f"{type(g).__name__}")
+
+        # Mutual exclusions — mirror with_genotype / compile.
+        if self._genotype is not None:
+            raise ValueError(
+                "run_cohort() and with_genotype() are mutually exclusive")
+        if any(v is not None for v in self._locks.values()):
+            raise ValueError(
+                "run_cohort() and restrict_alleles() are mutually exclusive")
+        if self._user_allele_weights_set:
+            raise ValueError(
+                "run_cohort() and recombine(*_allele_weights=...) are mutually "
+                "exclusive: the genotype owns allele expression")
+        if any(isinstance(s, _ReceptorRevisionStep) for s in self._steps):
+            raise ValueError(
+                "run_cohort() is not supported with receptor_revision() in this "
+                "release (the revision pass is not haplotype-aware)")
+        if self._has_clonal_fork():
+            raise ValueError(
+                "run_cohort() is not supported together with expand_clones() / "
+                "clonal_lineage() / clonal_repertoire() in this release")
+        # with_metadata() is applied per subject below, but it must not overwrite
+        # cohort-owned columns.
+        _cohort_owned = {"subject_id", "sequence_id", "haplotype"}
+        _md_collision = _cohort_owned & set(self._metadata)
+        if _md_collision:
+            raise ValueError(
+                f"run_cohort: with_metadata keys {sorted(_md_collision)} are "
+                f"cohort-owned (reserved); rename them")
+
+        # Cartridge-hash check (same as with_genotype).
+        live_hash = self._refdata.content_hash()
+        for g in gts:
+            if g._source_hash != live_hash:
+                raise ValueError(
+                    "run_cohort: a genotype was built against a different cartridge "
+                    f"(content hash {g._source_hash!r} != experiment {live_hash!r})")
+
+        resolved_counts = _resolve_counts(len(gts), n_per_subject, counts)
+        subject_ids = _resolve_subject_ids([g.subject_id for g in gts])
+        base_rng = _random.Random(seed)
+
+        subjects = []
+        for g, sid, count in zip(gts, subject_ids, resolved_counts):
+            sub_seed = base_rng.getrandbits(63)
+            snap = g._snapshot()
+            snap.subject_id = sid
+            # Clone the uncompiled experiment; never mutate self.
+            exp_i = _copy.copy(self)
+            exp_i._genotype = snap
+            compiled = exp_i.compile(allow_curatable_refdata=allow_curatable_refdata)
+            refdata_i = compiled.refdata
+            if count == 0:
+                # run() rejects n < 1; build an empty result and stamp the
+                # genotype manually (run_records would otherwise do it).
+                res = SimulationResult.from_outcomes(
+                    [], refdata_i, expose_provenance=expose_provenance)
+                res._genotypes = [snap]
+            else:
+                res = compiled.run_records(
+                    n=count, seed=sub_seed, strict=strict,
+                    expose_provenance=expose_provenance,
+                    validate_records=validate_records)
+            # Apply with_metadata() per subject (parity with run_records), then
+            # namespace sequence_id so combined export never collides. Metadata
+            # is stamped first; the cohort-owned sequence_id rewrite wins (and a
+            # collision was already rejected up front).
+            if self._metadata:
+                for rec in res.records:
+                    for key, value in self._metadata.items():
+                        rec[key] = value
+            for rec in res.records:
+                rec["sequence_id"] = f"{sid}_{rec.get('sequence_id', '')}"
+            subjects.append(CohortSubjectResult(
+                subject_id=sid, genotype=snap, result=res, refdata=refdata_i,
+                seed=sub_seed, count=count))
+
+        return CohortResult(subjects)
+
     def run_records(
         self,
         *,