From 0d8aa1101f7d06c2c244d7a59601792cf5b76e5c Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:29:58 +0300 Subject: [PATCH 1/9] feat(cohort): CohortSubjectResult + CohortResult core accessors --- src/GenAIRR/cohort.py | 74 ++++++++++++++++++++++++++++++++++ tests/test_genotype_cohorts.py | 35 ++++++++++++++++ 2 files changed, 109 insertions(+) create mode 100644 src/GenAIRR/cohort.py create mode 100644 tests/test_genotype_cohorts.py diff --git a/src/GenAIRR/cohort.py b/src/GenAIRR/cohort.py new file mode 100644 index 0000000..f900cdf --- /dev/null +++ b/src/GenAIRR/cohort.py @@ -0,0 +1,74 @@ +"""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 + + +@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). Mutating + the returned list does not affect the cohort.""" + out: List[Dict[str, Any]] = [] + for s in self._subjects: + out.extend(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"" diff --git a/tests/test_genotype_cohorts.py b/tests/test_genotype_cohorts.py new file mode 100644 index 0000000..db47109 --- /dev/null +++ b/tests/test_genotype_cohorts.py @@ -0,0 +1,35 @@ +"""Cohorts: Experiment.run_cohort + CohortResult.""" +import pytest + +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype +from GenAIRR.cohort import CohortResult, CohortSubjectResult + + +def _cfg(): + return gdata.HUMAN_IGH_OGRDB + + +def test_cohort_result_core_accessors(): + cfg = _cfg() + g0 = Genotype.sample(cfg, seed=0, subject_id="A") + g1 = Genotype.sample(cfg, seed=1, subject_id="B") + r0 = ga.Experiment.on(cfg).with_genotype(g0).recombine().run_records(n=3, seed=1) + r1 = ga.Experiment.on(cfg).with_genotype(g1).recombine().run_records(n=2, seed=2) + subjects = [ + CohortSubjectResult(subject_id="A", genotype=g0, result=r0, refdata=object(), seed=11, count=3), + CohortSubjectResult(subject_id="B", genotype=g1, result=r1, refdata=object(), seed=22, count=2), + ] + c = CohortResult(subjects) + assert c.subject_ids == ["A", "B"] + assert c.genotypes == [g0, g1] + assert c.results == [r0, r1] + assert c.result_for("B") is r1 + assert c.refdata_for("A") is subjects[0].refdata + assert len(c) == 5 # total records + assert len(c.records) == 5 # fresh concatenated list + c.records.append({"x": 1}) # mutating the returned list... + assert len(c.records) == 5 # ...does not affect the cohort + with pytest.raises(KeyError): + c.result_for("NOPE") From 74244ffe08d1015ce13f5eccbb5a7c4d93193619 Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:30:36 +0300 Subject: [PATCH 2/9] feat(cohort): combined to_dataframe/to_csv/to_fasta with union columns + unique headers --- src/GenAIRR/cohort.py | 28 ++++++++++++++++++++++++++++ tests/test_genotype_cohorts.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/src/GenAIRR/cohort.py b/src/GenAIRR/cohort.py index f900cdf..4e830ea 100644 --- a/src/GenAIRR/cohort.py +++ b/src/GenAIRR/cohort.py @@ -72,3 +72,31 @@ def __len__(self) -> int: def __repr__(self) -> str: return f"" + + # ── 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 + + records = self.records + 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") diff --git a/tests/test_genotype_cohorts.py b/tests/test_genotype_cohorts.py index db47109..760f676 100644 --- a/tests/test_genotype_cohorts.py +++ b/tests/test_genotype_cohorts.py @@ -33,3 +33,37 @@ def test_cohort_result_core_accessors(): assert len(c.records) == 5 # ...does not affect the cohort with pytest.raises(KeyError): c.result_for("NOPE") + + +def test_cohort_result_export_union_and_fasta(tmp_path): + cfg = _cfg() + g0 = Genotype.sample(cfg, seed=0, subject_id="A") + g1 = Genotype.sample(cfg, seed=1, subject_id="B") + # subject A WITH provenance columns, subject B WITHOUT -> schemas differ + rA = ga.Experiment.on(cfg).with_genotype(g0).recombine().run_records( + n=2, seed=1, expose_provenance=True) + rB = ga.Experiment.on(cfg).with_genotype(g1).recombine().run_records(n=2, seed=2) + # namespace ids the way run_cohort will (so headers are unique) + for rec in rA.records: + rec["sequence_id"] = "A_" + rec["sequence_id"] + for rec in rB.records: + rec["sequence_id"] = "B_" + rec["sequence_id"] + c = CohortResult([ + CohortSubjectResult("A", g0, rA, object(), 1, 2), + CohortSubjectResult("B", g1, rB, object(), 2, 2), + ]) + df = c.to_dataframe() + assert len(df) == 4 + # union of columns: truth_v_call exists (from A) even though B lacks it + assert "truth_v_call" in df.columns + + fasta = tmp_path / "cohort.fasta" + c.to_fasta(str(fasta)) + headers = [ln[1:].strip() for ln in fasta.read_text().splitlines() if ln.startswith(">")] + assert len(headers) == 4 + assert len(set(headers)) == 4 # globally unique + assert all(h.startswith("A_") or h.startswith("B_") for h in headers) + + csv = tmp_path / "cohort.csv" + c.to_csv(str(csv)) + assert csv.read_text().count("\n") >= 5 # header + 4 rows From d0fb0181e436e56524b3bec8cf7079f6bebd536e Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:31:19 +0300 Subject: [PATCH 3/9] feat(cohort): subject-id + counts resolution helpers --- src/GenAIRR/cohort.py | 38 ++++++++++++++++++++++++++++++++++ tests/test_genotype_cohorts.py | 32 +++++++++++++++++++++++++++- 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/src/GenAIRR/cohort.py b/src/GenAIRR/cohort.py index 4e830ea..cef2ee1 100644 --- a/src/GenAIRR/cohort.py +++ b/src/GenAIRR/cohort.py @@ -13,6 +13,44 @@ 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 + 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.""" diff --git a/tests/test_genotype_cohorts.py b/tests/test_genotype_cohorts.py index 760f676..d38a61a 100644 --- a/tests/test_genotype_cohorts.py +++ b/tests/test_genotype_cohorts.py @@ -4,7 +4,12 @@ import GenAIRR as ga import GenAIRR.data as gdata from GenAIRR.genotype import Genotype -from GenAIRR.cohort import CohortResult, CohortSubjectResult +from GenAIRR.cohort import ( + CohortResult, + CohortSubjectResult, + _resolve_counts, + _resolve_subject_ids, +) def _cfg(): @@ -67,3 +72,28 @@ def test_cohort_result_export_union_and_fasta(tmp_path): csv = tmp_path / "cohort.csv" c.to_csv(str(csv)) assert csv.read_text().count("\n") >= 5 # header + 4 rows + + +def test_resolve_subject_ids(): + assert _resolve_subject_ids([None, None, None]) == ["subject_0", "subject_1", "subject_2"] + assert _resolve_subject_ids(["A", "B"]) == ["A", "B"] + assert _resolve_subject_ids([1, 2]) == ["1", "2"] # normalized to str + with pytest.raises(ValueError, match="duplicate"): + _resolve_subject_ids(["A", "A"]) + with pytest.raises(ValueError, match="duplicate"): + _resolve_subject_ids([1, "1"]) # collide after str() + with pytest.raises(ValueError, match="some .* subject_id"): + _resolve_subject_ids(["A", None]) # mixed + + +def test_resolve_counts(): + assert _resolve_counts(3, 5, None) == [5, 5, 5] + assert _resolve_counts(3, 5, [1, 2, 0]) == [1, 2, 0] + with pytest.raises(ValueError, match="length"): + _resolve_counts(3, 5, [1, 2]) + for bad in (-1, True, "2", 1.5): + with pytest.raises(ValueError, match="n_per_subject"): + _resolve_counts(2, bad, None) + for bad_list in ([1, -1], [1, True], [1, "2"]): + with pytest.raises(ValueError, match="counts"): + _resolve_counts(2, 1, bad_list) From 9e23a0b94b25945dd1bd38c7f2a29aa0a4c87c38 Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:32:34 +0300 Subject: [PATCH 4/9] feat(cohort): Experiment.run_cohort happy path (clone-per-subject loop) --- src/GenAIRR/experiment.py | 114 +++++++++++++++++++++++++++++++++ tests/test_genotype_cohorts.py | 39 +++++++++++ 2 files changed, 153 insertions(+) diff --git a/src/GenAIRR/experiment.py b/src/GenAIRR/experiment.py index 7c9a61d..3400fb9 100644 --- a/src/GenAIRR/experiment.py +++ b/src/GenAIRR/experiment.py @@ -2948,6 +2948,120 @@ 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") + + # 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) + # Namespace sequence_id so combined export never collides. + 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, *, diff --git a/tests/test_genotype_cohorts.py b/tests/test_genotype_cohorts.py index d38a61a..bedaf4d 100644 --- a/tests/test_genotype_cohorts.py +++ b/tests/test_genotype_cohorts.py @@ -97,3 +97,42 @@ def test_resolve_counts(): for bad_list in ([1, -1], [1, True], [1, "2"]): with pytest.raises(ValueError, match="counts"): _resolve_counts(2, 1, bad_list) + + +def test_run_cohort_happy_path(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"D{s}") for s in range(3)] + c = (ga.Experiment.on(cfg).recombine() + .run_cohort(gs, n_per_subject=5, seed=0, expose_provenance=True)) + assert isinstance(c, CohortResult) + assert c.subject_ids == ["D0", "D1", "D2"] + assert len(c.genotypes) == 3 + assert len(c) == 15 + # every record is subject-tagged and its sequence_id is namespaced + unique + sids = [r["sequence_id"] for r in c.records] + assert len(set(sids)) == 15 + for r in c.records: + assert r["sequence_id"].startswith(r["subject_id"] + "_") + # each subject's result exposes its genotype + assert c.result_for("D1").genotypes[0].subject_id == "D1" + + +def test_run_cohort_deterministic_and_independent(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"D{s}") for s in range(2)] + a = ga.Experiment.on(cfg).recombine().run_cohort(gs, n_per_subject=4, seed=7) + b = ga.Experiment.on(cfg).recombine().run_cohort(gs, n_per_subject=4, seed=7) + assert [r["sequence"] for r in a.records] == [r["sequence"] for r in b.records] + d = ga.Experiment.on(cfg).recombine().run_cohort(gs, n_per_subject=4, seed=8) + assert [r["sequence"] for r in a.records] != [r["sequence"] for r in d.records] + + +def test_run_cohort_does_not_mutate_base_experiment(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"D{s}") for s in range(2)] + exp = ga.Experiment.on(cfg).recombine() + exp.run_cohort(gs, n_per_subject=2, seed=0) + assert exp._genotype is None + # still usable as a plain (no-genotype) experiment afterwards + res = exp.run_records(n=2, seed=1) + assert len(res) == 2 From abbaff13197215038282eb993eb4e765b97f6376 Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:33:08 +0300 Subject: [PATCH 5/9] test(cohort): auto-ids, counts override/zero, mutual-exclusion, cartridge-hash, mid-loop failure safety --- tests/test_genotype_cohorts.py | 84 ++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/tests/test_genotype_cohorts.py b/tests/test_genotype_cohorts.py index bedaf4d..1dd31da 100644 --- a/tests/test_genotype_cohorts.py +++ b/tests/test_genotype_cohorts.py @@ -136,3 +136,87 @@ def test_run_cohort_does_not_mutate_base_experiment(): # still usable as a plain (no-genotype) experiment afterwards res = exp.run_records(n=2, seed=1) assert len(res) == 2 + + +def test_run_cohort_auto_subject_ids_and_no_mutation_of_originals(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s) for s in range(3)] # no subject_id set + c = ga.Experiment.on(cfg).recombine().run_cohort(gs, n_per_subject=1, seed=0) + assert c.subject_ids == ["subject_0", "subject_1", "subject_2"] + # originals are untouched (resolution happened on snapshots) + assert all(g.subject_id is None for g in gs) + + +def test_run_cohort_counts_override_and_zero(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"D{s}") for s in range(3)] + c = ga.Experiment.on(cfg).recombine().run_cohort(gs, seed=0, counts=[4, 0, 2]) + assert [len(r) for r in c.results] == [4, 0, 2] + assert len(c) == 6 + # zero-count subject: empty records, but present with stamped genotype + refdata + zero = c.result_for("D1") + assert zero.records == [] + assert zero.genotypes[0].subject_id == "D1" + assert c.refdata_for("D1") is not None + + +def test_run_cohort_counts_length_mismatch_raises(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"D{s}") for s in range(2)] + with pytest.raises(ValueError, match="length"): + ga.Experiment.on(cfg).recombine().run_cohort(gs, counts=[1]) + + +def test_run_cohort_duplicate_subject_ids_raise(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id="DUP") for s in range(2)] + with pytest.raises(ValueError, match="duplicate"): + ga.Experiment.on(cfg).recombine().run_cohort(gs, n_per_subject=1) + + +def test_run_cohort_mutual_exclusions(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=0, subject_id="A") + with pytest.raises(ValueError, match="non-empty"): + ga.Experiment.on(cfg).recombine().run_cohort([]) + with pytest.raises(ValueError, match="with_genotype"): + ga.Experiment.on(cfg).with_genotype(g).recombine().run_cohort([g]) + with pytest.raises(ValueError, match="restrict_alleles"): + (ga.Experiment.on(cfg).recombine() + .restrict_alleles(v=cfg.v_alleles[next(iter(cfg.v_alleles))][0].name) + .run_cohort([g])) + + +def test_run_cohort_cartridge_hash_mismatch_raises(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=0, subject_id="A") + g._source_hash = "sha256:deadbeef" # forge a mismatch + with pytest.raises(ValueError, match="different cartridge"): + ga.Experiment.on(cfg).recombine().run_cohort([g]) + + +def test_run_cohort_failure_midway_leaves_base_experiment_clean(): + cfg = _cfg() + good = Genotype.sample(cfg, seed=0, subject_id="A") + bad = Genotype.sample(cfg, seed=1, subject_id="B") + import GenAIRR.genotype as _gmod + exp = ga.Experiment.on(cfg).recombine() + orig_snapshot = _gmod.Genotype._snapshot + + calls = {"n": 0} + + def boom(self): + calls["n"] += 1 + if calls["n"] == 2: # fail on the 2nd subject + raise RuntimeError("forced compile-time failure") + return orig_snapshot(self) + + _gmod.Genotype._snapshot = boom + try: + with pytest.raises(RuntimeError, match="forced"): + exp.run_cohort([good, bad], n_per_subject=1, seed=0) + finally: + _gmod.Genotype._snapshot = orig_snapshot + # base experiment is untouched and still runnable + assert exp._genotype is None + assert len(exp.run_records(n=2, seed=1)) == 2 From 7c0db1a80ef8ea53f23c9be26fe177f07eb40b93 Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:34:16 +0300 Subject: [PATCH 6/9] test(cohort): mixed novel/plain subjects + per-subject refdata/validate + e2e --- tests/test_genotype_cohorts.py | 67 ++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/test_genotype_cohorts.py b/tests/test_genotype_cohorts.py index 1dd31da..9b81e23 100644 --- a/tests/test_genotype_cohorts.py +++ b/tests/test_genotype_cohorts.py @@ -220,3 +220,70 @@ def boom(self): # base experiment is untouched and still runnable assert exp._genotype is None assert len(exp.run_records(n=2, seed=1)) == 2 + + +def _functional_novel_seq(cfg, base_name): + base = next(a for g in cfg.v_alleles.values() for a in g if a.name == base_name) + bs = base.ungapped_seq.upper() + for pos in range(len(bs)): + for nt in "ACGT": + if nt == bs[pos]: + continue + cand = bs[:pos] + nt + bs[pos + 1:] + try: + (Genotype.from_dataconfig(cfg) + .add_novel_allele(f"{base_name.split('*')[0]}*97", base=base_name, + sequence=cand, segment="V")) + return cand + except ValueError: + continue + raise AssertionError("no functional novel found") + + +def test_run_cohort_mixed_novel_and_plain_subjects(): + cfg = _cfg() + vg = next(iter(cfg.v_alleles)) + base = cfg.v_alleles[vg][0].name + novel = f"{vg}*97" + seq = _functional_novel_seq(cfg, base) + # subject A carries a novel; subject B is plain + gA = (Genotype.from_dataconfig(cfg) + .add_novel_allele(novel, base=base, sequence=seq, segment="V") + .homozygous(vg, novel).complete_from_reference().with_subject("A")) + gB = Genotype.sample(cfg, seed=2, subject_id="B") + c = (ga.Experiment.on(cfg).recombine() + .run_cohort([gA, gB], n_per_subject=40, seed=0, expose_provenance=True)) + # every record's truth call is carried by its subject's genotype + geno = {s.subject_id: s.genotype for s in c.subjects} + for r in c.records: + for seg, col in (("V", "truth_v_call"), ("D", "truth_d_call"), ("J", "truth_j_call")): + call = r.get(col) + if not call: + continue + assert call in geno[r["subject_id"]].carried_alleles(seg, call.split("*")[0]) + # the novel appears for subject A and never for subject B + assert any(r["truth_v_call"] == novel for r in c.result_for("A").records + if r.get("truth_v_call")) + assert all(r.get("truth_v_call") != novel for r in c.result_for("B").records) + + +def test_run_cohort_per_subject_validate_records(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"D{s}") for s in range(2)] + # validate_records=True runs each subject's validator against its own refdata + c = (ga.Experiment.on(cfg).recombine() + .run_cohort(gs, n_per_subject=10, seed=0, validate_records=True)) + assert len(c) == 20 + + +def test_run_cohort_end_to_end_from_sample(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"S{s}") for s in range(4)] + c = (ga.Experiment.on(cfg).recombine() + .run_cohort(gs, n_per_subject=25, seed=3, expose_provenance=True)) + assert len(c) == 100 + geno = {s.subject_id: s.genotype for s in c.subjects} + for r in c.records: + call = r.get("truth_v_call") + if call: + assert call in geno[r["subject_id"]].carried_alleles("V", call.split("*")[0]) From df8976cd73e413731b0ac9b6a41a8127263f1e74 Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:38:48 +0300 Subject: [PATCH 7/9] fix(cohort): CohortResult.records returns independent dict copies; test allele-weights exclusion --- src/GenAIRR/cohort.py | 7 ++++--- tests/test_genotype_cohorts.py | 8 ++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/GenAIRR/cohort.py b/src/GenAIRR/cohort.py index cef2ee1..545d688 100644 --- a/src/GenAIRR/cohort.py +++ b/src/GenAIRR/cohort.py @@ -98,11 +98,12 @@ def refdata_for(self, subject_id: str): @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). Mutating - the returned list does not affect the cohort.""" + 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(s.result.records) + out.extend(dict(rec) for rec in s.result.records) return out def __len__(self) -> int: diff --git a/tests/test_genotype_cohorts.py b/tests/test_genotype_cohorts.py index 9b81e23..ff971ca 100644 --- a/tests/test_genotype_cohorts.py +++ b/tests/test_genotype_cohorts.py @@ -36,6 +36,10 @@ def test_cohort_result_core_accessors(): assert len(c.records) == 5 # fresh concatenated list c.records.append({"x": 1}) # mutating the returned list... assert len(c.records) == 5 # ...does not affect the cohort + # mutating a record dict in the returned view must not corrupt the subject + recs = c.records + recs[0]["sequence_id"] = "TAMPERED" + assert c.records[0]["sequence_id"] != "TAMPERED" with pytest.raises(KeyError): c.result_for("NOPE") @@ -185,6 +189,10 @@ def test_run_cohort_mutual_exclusions(): (ga.Experiment.on(cfg).recombine() .restrict_alleles(v=cfg.v_alleles[next(iter(cfg.v_alleles))][0].name) .run_cohort([g])) + vname = cfg.v_alleles[next(iter(cfg.v_alleles))][0].name + with pytest.raises(ValueError, match="allele_weights"): + (ga.Experiment.on(cfg).recombine(v_allele_weights={vname: 100.0}) + .run_cohort([g])) def test_run_cohort_cartridge_hash_mismatch_raises(): From a4c1a2a452cec594e958dc5b1314652b51df659f Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:40:00 +0300 Subject: [PATCH 8/9] docs(cohort): run_cohort guide section + remove from limitations --- site_docs/guides/genotype.md | 58 ++++++++++++++++++++++++++++++++++-- 1 file changed, 56 insertions(+), 2 deletions(-) diff --git a/site_docs/guides/genotype.md b/site_docs/guides/genotype.md index 63b1b2a..dedc749 100644 --- a/site_docs/guides/genotype.md +++ b/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 From 2b104bd248ff3cc30d749ba92015857c005e1f64 Mon Sep 17 00:00:00 2001 From: thomas Date: Wed, 17 Jun 2026 16:56:20 +0300 Subject: [PATCH 9/9] fix(cohort): apply with_metadata per subject (+collision guard); reject mapping/bytes counts; empty-export schema --- src/GenAIRR/cohort.py | 16 +++++++++++++++ src/GenAIRR/experiment.py | 17 +++++++++++++++- tests/test_genotype_cohorts.py | 37 ++++++++++++++++++++++++++++++++++ 3 files changed, 69 insertions(+), 1 deletion(-) diff --git a/src/GenAIRR/cohort.py b/src/GenAIRR/cohort.py index 545d688..86a10a7 100644 --- a/src/GenAIRR/cohort.py +++ b/src/GenAIRR/cohort.py @@ -43,6 +43,11 @@ def _resolve_counts(n_genotypes: int, n_per_subject, counts) -> List[int]: _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( @@ -118,7 +123,18 @@ def to_dataframe(self, *, airr_strict: bool = False): 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] diff --git a/src/GenAIRR/experiment.py b/src/GenAIRR/experiment.py index 3400fb9..9633672 100644 --- a/src/GenAIRR/experiment.py +++ b/src/GenAIRR/experiment.py @@ -3019,6 +3019,14 @@ def run_cohort( 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() @@ -3053,7 +3061,14 @@ def run_cohort( n=count, seed=sub_seed, strict=strict, expose_provenance=expose_provenance, validate_records=validate_records) - # Namespace sequence_id so combined export never collides. + # 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( diff --git a/tests/test_genotype_cohorts.py b/tests/test_genotype_cohorts.py index ff971ca..71907ba 100644 --- a/tests/test_genotype_cohorts.py +++ b/tests/test_genotype_cohorts.py @@ -171,6 +171,43 @@ def test_run_cohort_counts_length_mismatch_raises(): ga.Experiment.on(cfg).recombine().run_cohort(gs, counts=[1]) +def test_run_cohort_counts_rejects_mapping_and_bytes(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"D{s}") for s in range(2)] + for bad in ({0: 5, 1: 7}, b"12", "12"): + with pytest.raises(ValueError, match="counts"): + ga.Experiment.on(cfg).recombine().run_cohort(gs, counts=bad) + + +def test_run_cohort_applies_with_metadata(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=0, subject_id="A") + c = (ga.Experiment.on(cfg).recombine().with_metadata(sample_id="S1") + .run_cohort([g], n_per_subject=3, seed=0)) + assert all(r["sample_id"] == "S1" for r in c.records) + # cohort-owned fields are still correct (not clobbered by metadata) + assert all(r["subject_id"] == "A" for r in c.records) + + +def test_run_cohort_metadata_collision_with_cohort_fields_raises(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=0, subject_id="A") + for bad_key in ("subject_id", "sequence_id", "haplotype"): + with pytest.raises(ValueError, match="cohort-owned|reserved"): + (ga.Experiment.on(cfg).recombine().with_metadata(**{bad_key: "X"}) + .run_cohort([g], n_per_subject=1)) + + +def test_run_cohort_empty_export_has_schema(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=0, subject_id="A") + c = ga.Experiment.on(cfg).recombine().run_cohort([g], counts=[0]) + df = c.to_dataframe() + assert len(df) == 0 + assert len(df.columns) > 0 # default AIRR columns preserved + assert "sequence_id" in df.columns + + def test_run_cohort_duplicate_subject_ids_raise(): cfg = _cfg() gs = [Genotype.sample(cfg, seed=s, subject_id="DUP") for s in range(2)]