diff --git a/engine_rs/src/passes/sample_genotype.rs b/engine_rs/src/passes/sample_genotype.rs index ba395cb..3b8b4b1 100644 --- a/engine_rs/src/passes/sample_genotype.rs +++ b/engine_rs/src/passes/sample_genotype.rs @@ -153,14 +153,19 @@ impl SampleGenotypePass { /// mode the feasibility-viable set is authoritative (empty → the /// caller raises). fn viable_set(&self, sim: &Simulation, ctx: &PassContext, strict: bool) -> Vec { + // A zero-weight chromosome is never expressed, so it can never make the + // genotype viable — exclude it from the candidate set entirely (rather + // than letting draw_haplotype fall back to a zero-weight chromosome). + let w = self.genotype.chromosome_weights(); + let expressible = |c: usize| w[c] as f64 > 0.0; let feasible: Vec = (0..2) - .filter(|&c| self.is_viable(c, sim, ctx, true)) + .filter(|&c| expressible(c) && self.is_viable(c, sim, ctx, true)) .collect(); if strict || !feasible.is_empty() { return feasible; } (0..2) - .filter(|&c| self.is_viable(c, sim, ctx, false)) + .filter(|&c| expressible(c) && self.is_viable(c, sim, ctx, false)) .collect() } diff --git a/site_docs/guides/genotype.md b/site_docs/guides/genotype.md index b3a8107..19c8957 100644 --- a/site_docs/guides/genotype.md +++ b/site_docs/guides/genotype.md @@ -264,6 +264,58 @@ for row in g.to_table(): print(row["gene"], row["zygosity"], row["haplotype_0"], row["haplotype_1"]) ``` +## Sampling from population priors + +Instead of specifying every gene, draw a plausible diploid genotype with +`Genotype.sample`. It uses an **independent per-gene, per-chromosome +Hardy-Weinberg model**: each gene on each chromosome is independently deleted +(`haplotype_deletion_prob`) or assigned an allele from that gene's frequencies, so +homozygous / heterozygous / hemizygous / deleted states emerge at the expected +rates. + +```python +g = Genotype.sample( + cfg, + seed=7, + allele_frequencies={ # {segment: {gene: {allele: weight}}} + "V": {"IGHVF1-G1": {"IGHVF1-G1*01": 8, "IGHVF1-G1*02": 2}}, + }, # unspecified genes -> uniform within gene + haplotype_deletion_prob=0.05, # 5% per-haplotype gene-absence + subject_id="DONOR_R1", +) +res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records(n=500, seed=1) +``` + +`Genotype.sample` always returns a **fully-specified, runnable** genotype. By +default `ensure_viable=True` re-draws (up to `max_resamples`, deterministically) +until at least one **expressible** (positive-weight) chromosome carries every +required segment, raising a clear error only if your deletion settings make that +impossible; pass `ensure_viable=False` to allow infeasible draws. + +!!! note "Default sampling is HW *conditioned on viability*" + With the default `ensure_viable=True`, draws that leave no complete usable + haplotype are rejected, so per-gene deletion/zygosity rates are Hardy-Weinberg + **conditioned on at least one viable chromosome** — not the unconditional HW + rates. This only matters at high `haplotype_deletion_prob` (e.g. a single-J + cartridge with a high J-deletion prob). Use `ensure_viable=False` for the raw, + unconditional draw (which may be infeasible and rejected at compile). + +Frequency priors accept the segment-aware +`{segment: {gene: {allele: weight}}}` shape (or a flat `{gene: …}` when the gene is +unambiguous); a supplied gene's listed alleles define its distribution (weight `0` +excludes an allele), and unspecified genes fall back to uniform. +`haplotype_deletion_prob` is a float or a per-gene/per-segment dict. + +!!! warning "What this model is — and isn't" + This is an **independent per-gene** sampler. It does **not** model linkage + disequilibrium, gene co-deletion blocks, ancestry, or donor-specific haplotype + structure, and it samples **catalogue alleles only** (no novel alleles) and + **deletion only** (no duplication). Supply explicit `allele_frequencies` from a + population source (e.g. VDJbase) for realistic per-gene frequencies; the default + is uniform within each gene. `allele_frequencies="usage_as_prior"` is an + explicit opt-in that reuses the cartridge's recombination `allele_usage` as a + frequency proxy — convenient but biologically approximate. + ## Novel / private alleles Individuals carry germline alleles that aren't in any reference — *private* or @@ -473,8 +525,6 @@ 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). -- **Population priors** — sampling a plausible diploid genotype from - allele/deletion frequencies (today genotypes are specified explicitly). - **External loaders** — importing genotypes from VDJbase / TIgGER / IgDiscover / partis output. - **Cartridge genotype plane** — persisting a population genotype model in a diff --git a/src/GenAIRR/genotype.py b/src/GenAIRR/genotype.py index 89213ea..47d3cfc 100644 --- a/src/GenAIRR/genotype.py +++ b/src/GenAIRR/genotype.py @@ -14,6 +14,8 @@ """ from __future__ import annotations +import math +from collections.abc import Mapping from typing import Dict, List, Optional, Set, Tuple _SEGMENTS = ("V", "D", "J") @@ -64,18 +66,19 @@ def with_subject(self, sid: str) -> "Genotype": self.subject_id = str(sid) return self - def chromosome_weights(self, w0: float, w1: float) -> "Genotype": - import math - - if not (math.isfinite(w0) and math.isfinite(w1)): - raise ValueError( - f"chromosome_weights must be finite, got {(w0, w1)}" - ) + @staticmethod + def _check_chromosome_weights(w0, w1) -> Tuple[float, float]: + for w in (w0, w1): + if isinstance(w, bool) or not isinstance(w, (int, float)) or not math.isfinite(w): + raise ValueError(f"chromosome_weights must be finite numbers, got {(w0, w1)}") if w0 < 0 or w1 < 0 or (w0 + w1) <= 0: raise ValueError( f"chromosome_weights must be non-negative and sum>0, got {(w0, w1)}" ) - self._chromosome_weights = (float(w0), float(w1)) + return (float(w0), float(w1)) + + def chromosome_weights(self, w0: float, w1: float) -> "Genotype": + self._chromosome_weights = self._check_chromosome_weights(w0, w1) return self def _check_allele(self, segment: str, gene: str, allele: str) -> None: @@ -390,6 +393,331 @@ def _snapshot(self) -> "Genotype": g._source_hash = self._source_hash return g + # ── population sampling ─────────────────────────────────────── + @staticmethod + def _required_segments(cfg) -> List[str]: + req = ["V", "J"] + if _alleles_by_gene(cfg, "D"): + req.insert(1, "D") # V, D, J + return req + + @classmethod + def _resolve_sample_segments(cls, cfg, segments_to_sample) -> List[str]: + required = cls._required_segments(cfg) + if segments_to_sample is None: + return required + segs = list(segments_to_sample) + seen = set() + for s in segs: + if s not in _SEGMENTS: + raise ValueError(f"unknown segment {s!r}; expected one of {_SEGMENTS}") + if s in seen: + raise ValueError(f"segments_to_sample contains duplicate segment {s!r}") + seen.add(s) + if not _alleles_by_gene(cfg, s): + raise ValueError(f"cartridge has no {s} segment") + missing = [r for r in required if r not in seen] + if missing: + raise ValueError( + f"segments_to_sample must cover the chain's required segments " + f"{required}; missing {missing}. Partial sampling is not supported " + f"by Genotype.sample (it must return a runnable genotype)." + ) + # canonical _SEGMENTS order so the result is independent of input order + return [s for s in _SEGMENTS if s in seen] + + @classmethod + def _gene_segment_index(cls, cfg, segs) -> Dict[str, List[str]]: + """gene name -> [segments it appears in] (for flat-shape disambiguation).""" + idx: Dict[str, List[str]] = {} + for seg in segs: + for gene in _alleles_by_gene(cfg, seg): + idx.setdefault(gene, []).append(seg) + return idx + + @staticmethod + def _weighted_pick(rng, pairs): + names = [n for (n, _w) in pairs] + weights = [w for (_n, w) in pairs] + return rng.choices(names, weights=weights, k=1)[0] + + @classmethod + def _normalize_freq_spec(cls, cfg, spec, segs): + """Return nested ``{seg: {gene: {allele: weight}}}`` from a nested or flat + ``allele_frequencies`` spec, fully validating segment/gene/allele + addressing and shapes (unknown names and malformed values raise).""" + if spec is None: + return {} + if not isinstance(spec, Mapping): + raise ValueError( + "allele_frequencies must be a mapping (or 'usage_as_prior' / None), " + f"got {type(spec).__name__}" + ) + nested: Dict[str, Dict[str, Dict[str, float]]] = {} + keys = set(spec) + if keys and keys <= set(_SEGMENTS): # segment-keyed (nested) shape + for seg, genes in spec.items(): + if seg not in segs: + raise ValueError( + f"allele_frequencies: segment {seg!r} is not being sampled " + f"(sampling {segs})" + ) + if not isinstance(genes, Mapping): + raise ValueError( + f"allele_frequencies[{seg!r}] must be a mapping of " + f"gene -> {{allele: weight}}, got {type(genes).__name__}" + ) + catalogue = _alleles_by_gene(cfg, seg) + for gene, alleles in genes.items(): + if gene not in catalogue: + raise ValueError( + f"allele_frequencies: {seg} gene {gene!r} is not in the cartridge" + ) + if not isinstance(alleles, Mapping) or not alleles: + raise ValueError( + f"allele_frequencies[{seg!r}][{gene!r}] must be a non-empty " + f"mapping of allele -> weight" + ) + nested.setdefault(seg, {})[gene] = alleles + else: # flat {gene: {...}} shape + gidx = cls._gene_segment_index(cfg, segs) + for gene, alleles in spec.items(): + segs_for = gidx.get(gene) + if not segs_for: + raise ValueError(f"allele_frequencies: unknown gene {gene!r}") + if len(segs_for) > 1: + raise ValueError( + f"allele_frequencies: gene {gene!r} is ambiguous across segments " + f"{segs_for}; use the {{segment: {{gene: ...}}}} shape" + ) + if not isinstance(alleles, Mapping) or not alleles: + raise ValueError( + f"allele_frequencies[{gene!r}] must be a non-empty mapping of " + f"allele -> weight" + ) + nested.setdefault(segs_for[0], {})[gene] = alleles + return nested + + @classmethod + def _usage_frequencies(cls, cfg, segs): + rm = getattr(cfg, "reference_models", None) + usage = getattr(rm, "allele_usage", None) if rm else None + if usage is None: + raise ValueError( + "allele_frequencies='usage_as_prior' requires a cartridge with a " + "typed reference_models.allele_usage; this cartridge has none" + ) + nested: Dict[str, Dict[str, Dict[str, float]]] = {} + seg_attr = {"V": "v", "D": "d", "J": "j"} + for seg in segs: + table = getattr(usage, seg_attr[seg], None) or {} + if not table: + raise ValueError( + f"allele_frequencies='usage_as_prior': cartridge allele_usage has no " + f"entries for requested segment {seg!r}" + ) + catalogue = _alleles_by_gene(cfg, seg) + for allele_name, w in table.items(): + gene = allele_name.split("*")[0] + if gene not in catalogue: + raise ValueError( + f"allele_frequencies='usage_as_prior': usage allele {allele_name!r} " + f"maps to {seg} gene {gene!r}, which is not in the cartridge catalogue" + ) + nested.setdefault(seg, {}).setdefault(gene, {})[allele_name] = w + return nested + + @classmethod + def _resolve_allele_frequencies(cls, cfg, spec, segs): + if spec == "usage_as_prior": + nested = cls._usage_frequencies(cfg, segs) + else: + nested = cls._normalize_freq_spec(cfg, spec, segs) + out: Dict[str, Dict[str, List[Tuple[str, float]]]] = {} + for seg in segs: + out[seg] = {} + for gene, alleles in _alleles_by_gene(cfg, seg).items(): + names = {a.name for a in alleles} + supplied = nested.get(seg, {}).get(gene) + if supplied is None: + out[seg][gene] = [(a.name, 1.0) for a in alleles] # uniform fallback + continue + pairs: List[Tuple[str, float]] = [] + total = 0.0 + for nm, w in supplied.items(): + if nm not in names: + raise ValueError(f"{seg} gene {gene}: {nm!r} is not a known allele") + if isinstance(w, bool) or not isinstance(w, (int, float)) or not math.isfinite(w) or w < 0: + raise ValueError( + f"{seg} gene {gene}: weight for {nm!r} must be finite and >= 0, got {w!r}" + ) + if w > 0: + pairs.append((nm, float(w))) + total += w + if total <= 0 or not pairs: + raise ValueError( + f"{seg} gene {gene}: at least one allele weight must be > 0" + ) + out[seg][gene] = pairs + return out + + @classmethod + def _resolve_haplotype_deletion(cls, cfg, spec, segs): + def _check(p, where): + if isinstance(p, bool) or not isinstance(p, (int, float)) or not math.isfinite(p): + raise ValueError(f"{where}: deletion probability must be a finite number, got {p!r}") + if not (0.0 <= p <= 1.0): + raise ValueError(f"{where}: deletion probability must be in [0, 1], got {p}") + return float(p) + + out = {seg: {gene: 0.0 for gene in _alleles_by_gene(cfg, seg)} for seg in segs} + if isinstance(spec, (int, float)) and not isinstance(spec, bool): + p = _check(spec, "haplotype_deletion_prob") + for seg in segs: + for gene in out[seg]: + out[seg][gene] = p + return out + if not isinstance(spec, Mapping): + raise ValueError( + f"haplotype_deletion_prob must be a float or a mapping, got {type(spec).__name__}" + ) + keys = set(spec) + if keys and keys <= set(_SEGMENTS): # nested {seg: {gene: prob}} + for seg, genes in spec.items(): + if seg not in segs: + raise ValueError(f"haplotype_deletion_prob: segment {seg!r} not being sampled") + if not isinstance(genes, Mapping): + raise ValueError( + f"haplotype_deletion_prob[{seg!r}] must be a mapping of " + f"gene -> probability, got {type(genes).__name__}" + ) + for gene, p in genes.items(): + if gene not in out[seg]: + raise ValueError(f"haplotype_deletion_prob: unknown {seg} gene {gene!r}") + out[seg][gene] = _check(p, f"haplotype_deletion_prob[{seg}][{gene}]") + else: # flat {gene: prob} + gidx = cls._gene_segment_index(cfg, segs) + for gene, p in spec.items(): + segs_for = gidx.get(gene) + if not segs_for: + raise ValueError(f"haplotype_deletion_prob: unknown gene {gene!r}") + if len(segs_for) > 1: + raise ValueError( + f"haplotype_deletion_prob: gene {gene!r} is ambiguous across " + f"segments {segs_for}; use the {{segment: {{gene: prob}}}} shape" + ) + out[segs_for[0]][gene] = _check(p, f"haplotype_deletion_prob[{gene}]") + return out + + @classmethod + def sample( + cls, + cfg, + *, + seed: int = 0, + allele_frequencies=None, + haplotype_deletion_prob=0.0, + segments_to_sample=None, + chromosome_weights: Tuple[float, float] = (0.5, 0.5), + subject_id: Optional[str] = None, + ensure_viable: bool = True, + max_resamples: int = 1000, + ) -> "Genotype": + """Sample a fully-specified diploid genotype from population priors. + + Independent per-gene, per-chromosome Hardy-Weinberg model: each gene on + each chromosome is independently deleted (prob + ``haplotype_deletion_prob``) or assigned one allele drawn from the gene's + allele frequencies. Homozygous/heterozygous/hemizygous/deleted states + emerge at Hardy-Weinberg rates. + + This is NOT a population haplotype model — no linkage disequilibrium, gene + co-deletion blocks, ancestry, or donor-specific haplotype structure. It + samples catalogue alleles only (no novel alleles) and deletion only (no + copy-number duplication). The default prior is uniform within each gene; + supply ``allele_frequencies`` for realistic per-gene frequencies. + + With ``ensure_viable=True`` (default), the draw is repeated (with a + deterministic sub-seed) up to ``max_resamples`` times until at least one + **positive-weight** chromosome carries every required segment, raising + ``ValueError`` if that is impossible under the given deletion settings. + + NOTE: with the default ``ensure_viable=True`` the result is Hardy-Weinberg + **conditioned on viability** (draws with no complete usable haplotype are + rejected), not the unconditional HW distribution. Use + ``ensure_viable=False`` for the raw (possibly infeasible) HW draw. + """ + import random + + if isinstance(max_resamples, bool) or not isinstance(max_resamples, int) or max_resamples < 1: + raise ValueError(f"max_resamples must be an int >= 1, got {max_resamples!r}") + cw = cls._check_chromosome_weights(*chromosome_weights) + segs = cls._resolve_sample_segments(cfg, segments_to_sample) + freqs = cls._resolve_allele_frequencies(cfg, allele_frequencies, segs) + delp = cls._resolve_haplotype_deletion(cfg, haplotype_deletion_prob, segs) + # Compute the cartridge content hash ONCE (it rebuilds refdata + hashes); + # reuse it across all draws instead of recomputing per attempt. + source_hash = cfg.cartridge_manifest()["hashes"]["refdata_content_hash"] + # Derive each attempt's sub-seed from a base RNG so a failed attempt at + # `seed` cannot collide with a direct draw at `seed + 1`. + base_rng = random.Random(seed) + + attempts = max_resamples if ensure_viable else 1 + for _attempt in range(attempts): + sub_seed = base_rng.getrandbits(63) + g = cls._draw_one(cfg, sub_seed, segs, freqs, delp, cw, subject_id, source_hash) + if not ensure_viable or g._is_viable(cfg, cw): + return g + raise ValueError( + f"could not sample a viable genotype after {max_resamples} attempts; " + f"haplotype_deletion_prob is too high to leave a complete, positive-weight " + f"haplotype for required segments {cls._required_segments(cfg)} " + f"(chromosome_weights={cw})" + ) + + @classmethod + def _draw_one(cls, cfg, seed, segs, freqs, delp, cw, subject_id, source_hash): + import random + + rng = random.Random(seed) + # Build a bare Genotype directly (avoid from_dataconfig, which recomputes + # the cartridge hash on every draw); reuse the precomputed source_hash. + g = cls.__new__(cls) + g._cfg = cfg + g._permissive = False + g.subject_id = subject_id + g._chromosome_weights = cw + g._slots = {s: {} for s in _SEGMENTS} + g._novel = {} + g._source_hash = source_hash + for seg in segs: + for gene in _alleles_by_gene(cfg, seg): + pdel = delp[seg][gene] + slots: List[List[Tuple[str, int, float]]] = [[], []] + for h in (0, 1): + if rng.random() < pdel: + continue # deleted on this chromosome + allele = cls._weighted_pick(rng, freqs[seg][gene]) + slots[h] = [(allele, 1, 1.0)] + g._slots[seg][gene] = slots + return g + + def _is_viable(self, cfg, chromosome_weights=None) -> bool: + """A genotype is viable iff at least one chromosome that can actually be + expressed (positive chromosome weight) carries every required segment. + A complete haplotype on a zero-weight chromosome does NOT count — the + engine would never draw it.""" + cw = chromosome_weights if chromosome_weights is not None else self._chromosome_weights + for c in (0, 1): + if cw[c] <= 0: + continue + if all( + any(self._slots[seg].get(gene, [[], []])[c] for gene in self._slots[seg]) + for seg in self._required_segments(cfg) + ): + return True + return False + # ── queries / export ────────────────────────────────────────── @property def is_permissive(self) -> bool: diff --git a/tests/test_genotype_population.py b/tests/test_genotype_population.py new file mode 100644 index 0000000..2ab9ea1 --- /dev/null +++ b/tests/test_genotype_population.py @@ -0,0 +1,207 @@ +"""Population-prior sampling: Genotype.sample(...).""" +import math + +import pytest + +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype + + +def _cfg(): + return gdata.HUMAN_IGH_OGRDB + + +def test_sample_is_fully_specified_and_deterministic(): + cfg = _cfg() + g1 = Genotype.sample(cfg, seed=5) + g2 = Genotype.sample(cfg, seed=5) + assert g1.to_table() == g2.to_table() # deterministic + for gene in cfg.v_alleles: + assert g1.is_specified("V", gene) + for gene in cfg.j_alleles: + assert g1.is_specified("J", gene) + for gene in cfg.d_alleles: + assert g1.is_specified("D", gene) + res = ga.Experiment.on(cfg).with_genotype(g1).recombine().run_records(n=20, seed=1) + assert len(res) == 20 + + +def test_different_seeds_differ(): + cfg = _cfg() + assert Genotype.sample(cfg, seed=1).to_table() != Genotype.sample(cfg, seed=2).to_table() + + +def test_segments_must_cover_required(): + cfg = _cfg() # VDJ + with pytest.raises(ValueError, match="required segment"): + Genotype.sample(cfg, seed=0, segments_to_sample=("V",)) # omits D, J + with pytest.raises(ValueError, match="no .* segment|required segment|unknown segment"): + Genotype.sample(cfg, seed=0, segments_to_sample=("V", "D", "J", "C")) + + +def test_frequencies_bias_homozygosity(): + cfg = _cfg() + vg = next(g for g, al in cfg.v_alleles.items() if len(al) >= 2) + a0, a1 = (a.name for a in cfg.v_alleles[vg][:2]) + freqs = {"V": {vg: {a0: 100.0, a1: 1.0}}} # heavily favour a0 + homo_a0 = sum( + 1 + for s in range(60) + if Genotype.sample(cfg, seed=s, allele_frequencies=freqs).carried_alleles("V", vg) == {a0} + ) + assert homo_a0 > 40 # dominant allele -> usually homozygous-common + + +def test_zero_weight_excludes_allele(): + cfg = _cfg() + vg = next(g for g, al in cfg.v_alleles.items() if len(al) >= 2) + a0, a1 = (a.name for a in cfg.v_alleles[vg][:2]) + freqs = {"V": {vg: {a0: 1.0, a1: 0.0}}} # a1 excluded + for s in range(40): + assert a1 not in Genotype.sample(cfg, seed=s, allele_frequencies=freqs).carried_alleles("V", vg) + + +def test_frequency_validation(): + cfg = _cfg() + vg = next(iter(cfg.v_alleles)) + a0 = cfg.v_alleles[vg][0].name + with pytest.raises(ValueError, match="not a known"): + Genotype.sample(cfg, seed=0, allele_frequencies={"V": {vg: {"NOPE*9": 1.0}}}) + with pytest.raises(ValueError, match="must be > 0"): + Genotype.sample(cfg, seed=0, allele_frequencies={"V": {vg: {a0: 0.0}}}) + with pytest.raises(ValueError, match="finite and >= 0"): + Genotype.sample(cfg, seed=0, allele_frequencies={"V": {vg: {a0: -1.0}}}) + + +def test_flat_gene_shape_accepted(): + cfg = _cfg() + vg = next(iter(cfg.v_alleles)) + a0 = cfg.v_alleles[vg][0].name + g = Genotype.sample(cfg, seed=0, allele_frequencies={vg: {a0: 1.0}}) # flat, unambiguous + assert g.carried_alleles("V", vg) <= {a0} + + +def test_usage_as_prior_requires_typed_usage(): + cfg = _cfg() # no typed allele_usage + with pytest.raises(ValueError, match="allele_usage|usage_as_prior"): + Genotype.sample(cfg, seed=0, allele_frequencies="usage_as_prior") + + +def test_deletion_prob_per_gene(): + cfg = _cfg() + drop = list(cfg.v_alleles)[0] + other = list(cfg.v_alleles)[1] + g = Genotype.sample(cfg, seed=0, haplotype_deletion_prob={"V": {drop: 1.0}}) + assert g.carried_alleles("V", drop) == set() # fully deleted + assert g.carried_alleles("V", other) # others still carried + + +def test_deletion_validation(): + cfg = _cfg() + with pytest.raises(ValueError, match=r"in \[0, 1\]"): + Genotype.sample(cfg, seed=0, haplotype_deletion_prob=1.5) + with pytest.raises(ValueError, match="unknown"): + Genotype.sample(cfg, seed=0, haplotype_deletion_prob={"V": {"NOPE": 0.5}}) + + +def test_ensure_viable_raises_when_all_deleted(): + cfg = _cfg() + with pytest.raises(ValueError, match="viable|deletion"): + Genotype.sample(cfg, seed=0, haplotype_deletion_prob=1.0, max_resamples=10) + g = Genotype.sample(cfg, seed=0, haplotype_deletion_prob=1.0, ensure_viable=False) + assert all(g.carried_alleles("V", gene) == set() for gene in cfg.v_alleles) + + +def test_max_resamples_validation(): + cfg = _cfg() + with pytest.raises(ValueError, match="max_resamples"): + Genotype.sample(cfg, seed=0, max_resamples=0) + + +def test_nested_freq_validation_no_silent_ignore(): + cfg = _cfg() + vg = next(iter(cfg.v_alleles)) + with pytest.raises(ValueError, match="not in the cartridge"): + Genotype.sample(cfg, seed=0, allele_frequencies={"V": {"NOPE": {"NOPE*01": 1.0}}}) + with pytest.raises(ValueError, match="must be a mapping"): + Genotype.sample(cfg, seed=0, allele_frequencies={"V": []}) + with pytest.raises(ValueError, match="non-empty mapping"): + Genotype.sample(cfg, seed=0, allele_frequencies={"V": {vg: None}}) + with pytest.raises(ValueError, match="non-empty mapping"): + Genotype.sample(cfg, seed=0, allele_frequencies={"V": {vg: {}}}) + with pytest.raises(ValueError, match="must be a mapping"): + Genotype.sample(cfg, seed=0, allele_frequencies=True) + + +def test_deletion_dict_non_mapping_raises(): + cfg = _cfg() + with pytest.raises(ValueError, match="must be a mapping"): + Genotype.sample(cfg, seed=0, haplotype_deletion_prob={"V": []}) + + +def test_segments_dedup_and_order_invariant(): + cfg = _cfg() + with pytest.raises(ValueError, match="duplicate segment"): + Genotype.sample(cfg, seed=0, segments_to_sample=("V", "V", "D", "J")) + # order-invariant: same seed, reordered segments -> identical genotype + a = Genotype.sample(cfg, seed=3, segments_to_sample=("V", "D", "J")) + b = Genotype.sample(cfg, seed=3, segments_to_sample=("J", "D", "V")) + assert a.to_table() == b.to_table() + + +def test_zero_weight_chromosome_never_expressed(): + cfg = _cfg() + # 100% express chromosome 1; chromosome 0 must never appear in records. + g = Genotype.sample(cfg, seed=4, chromosome_weights=(0.0, 1.0)) + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records( + n=60, seed=1, expose_provenance=True + ) + assert all(r["haplotype"] == 1 for r in res) + + +def test_zero_weight_chromosome_not_counted_viable(): + cfg = _cfg() + # Delete heavily on chromosome 1 only, but give all weight to chromosome 1: + # the only complete haplotype (chrom 0) has zero weight -> not viable -> raises. + with pytest.raises(ValueError, match="positive-weight|viable"): + Genotype.sample( + cfg, + seed=0, + haplotype_deletion_prob={"J": {gene: 1.0 for gene in cfg.j_alleles}}, + chromosome_weights=(1.0, 0.0), + max_resamples=5, + ) + + +def test_chromosome_weights_validation_in_sample(): + cfg = _cfg() + with pytest.raises(ValueError, match="finite"): + Genotype.sample(cfg, seed=0, chromosome_weights=(float("nan"), 1.0)) + with pytest.raises(ValueError, match="non-negative"): + Genotype.sample(cfg, seed=0, chromosome_weights=(-1.0, 1.0)) + + +def test_conditioned_vs_unconditioned_viability(): + cfg = _cfg() + # ensure_viable=False can yield an infeasible (no complete haplotype) draw at + # high deletion; ensure_viable=True never does. + g_uncond = Genotype.sample(cfg, seed=0, haplotype_deletion_prob=1.0, ensure_viable=False) + assert not g_uncond._is_viable(cfg) # all deleted -> not viable + with pytest.raises(ValueError): + Genotype.sample(cfg, seed=0, haplotype_deletion_prob=1.0, max_resamples=5) + + +def test_end_to_end_truth_calls_are_carried(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=11, haplotype_deletion_prob=0.1) + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records( + n=200, seed=2, expose_provenance=True + ) + for r in res: + for seg, col in (("V", "truth_v_call"), ("D", "truth_d_call"), ("J", "truth_j_call")): + call = r[col] + if not call: + continue + gene = call.split("*")[0] + assert call in g.carried_alleles(seg, gene), (seg, call)