diff --git a/engine_rs/src/address.rs b/engine_rs/src/address.rs index 9d3d9e4..55b189f 100644 --- a/engine_rs/src/address.rs +++ b/engine_rs/src/address.rs @@ -54,6 +54,8 @@ const SAMPLE_ALLELE_J: &str = "sample_allele.j"; const SAMPLE_ALLELE_D_INVERTED: &str = "sample_allele.d.inverted"; pub const SAMPLE_ALLELE_INVALID: &str = "sample_allele."; pub const SAMPLE_ALLELE_UNSUPPORTED: &str = "sample_allele."; +/// Per-rearrangement chromosome choice for a phased genotype. +const SAMPLE_HAPLOTYPE: &str = "sample_haplotype"; /// Pass name for `InvertDPass`. Used as the `name()` return value /// and as the pass-plan signature token, so external consumers @@ -458,6 +460,16 @@ pub enum ChoiceAddress { /// from `(allele, trim, orientation, length)` so only the /// length needs a trace address. PLength { end: PEnd }, + /// Haplotype (0/1): the chromosome drawn once per rearrangement by + /// `SampleHaplotypePass` for a phased genotype. V/D/J read it back. + SampleHaplotype, + /// GeneId: the gene chosen within the drawn chromosome for a segment + /// by `SampleGeneAllelePass`. + SampleGene(VdjSegment), + /// AlleleId: the within-slot allele draw, recorded only when a gene + /// slot carries more than one copy (single-copy slots are + /// deterministic and record nothing here). + SampleAlleleInSlot(VdjSegment), } impl ChoiceAddress { @@ -534,6 +546,11 @@ impl fmt::Display for ChoiceAddress { Self::PairedEndR2Length => f.write_str(PAIRED_END_R2_LENGTH), Self::PairedEndInsertSize => f.write_str(PAIRED_END_INSERT_SIZE), Self::PLength { end } => write!(f, "p.{}.length", end.suffix()), + Self::SampleHaplotype => f.write_str(SAMPLE_HAPLOTYPE), + Self::SampleGene(segment) => write!(f, "sample_gene.{}", segment.suffix()), + Self::SampleAlleleInSlot(segment) => { + write!(f, "sample_allele_in_slot.{}", segment.suffix()) + } } } } @@ -615,6 +632,13 @@ fn parse_choice_address(address: &str) -> Option { P_D5_LENGTH => Some(ChoiceAddress::PLength { end: PEnd::D5 }), P_D3_LENGTH => Some(ChoiceAddress::PLength { end: PEnd::D3 }), P_J5_LENGTH => Some(ChoiceAddress::PLength { end: PEnd::J5 }), + SAMPLE_HAPLOTYPE => Some(ChoiceAddress::SampleHaplotype), + "sample_gene.v" => Some(ChoiceAddress::SampleGene(VdjSegment::V)), + "sample_gene.d" => Some(ChoiceAddress::SampleGene(VdjSegment::D)), + "sample_gene.j" => Some(ChoiceAddress::SampleGene(VdjSegment::J)), + "sample_allele_in_slot.v" => Some(ChoiceAddress::SampleAlleleInSlot(VdjSegment::V)), + "sample_allele_in_slot.d" => Some(ChoiceAddress::SampleAlleleInSlot(VdjSegment::D)), + "sample_allele_in_slot.j" => Some(ChoiceAddress::SampleAlleleInSlot(VdjSegment::J)), _ => None, }; if exact.is_some() { @@ -747,6 +771,15 @@ pub enum ChoiceAddressPattern { /// [`ChoiceAddress::PLength`]. One pattern instance per /// `PEnd` — declared by `PAdditionPass`. PLength { end: PEnd }, + /// Singleton-family mirror of [`ChoiceAddress::SampleHaplotype`]. + /// Declared by `SampleHaplotypePass`. + SampleHaplotype, + /// Family mirror of [`ChoiceAddress::SampleGene`]. Declared by + /// `SampleGeneAllelePass`. + SampleGene(VdjSegment), + /// Family mirror of [`ChoiceAddress::SampleAlleleInSlot`]. Declared + /// by `SampleGeneAllelePass` (a potential draw on multi-copy slots). + SampleAlleleInSlot(VdjSegment), } impl ChoiceAddressPattern { @@ -793,6 +826,11 @@ impl fmt::Display for ChoiceAddressPattern { Self::PairedEndR2Length => f.write_str(PAIRED_END_R2_LENGTH), Self::PairedEndInsertSize => f.write_str(PAIRED_END_INSERT_SIZE), Self::PLength { end } => ChoiceAddress::PLength { end }.fmt(f), + Self::SampleHaplotype => f.write_str(SAMPLE_HAPLOTYPE), + Self::SampleGene(segment) => ChoiceAddress::SampleGene(segment).fmt(f), + Self::SampleAlleleInSlot(segment) => { + ChoiceAddress::SampleAlleleInSlot(segment).fmt(f) + } } } } @@ -889,6 +927,13 @@ fn parse_choice_address_pattern(address: &str) -> Option { P_D5_LENGTH => Some(ChoiceAddressPattern::PLength { end: PEnd::D5 }), P_D3_LENGTH => Some(ChoiceAddressPattern::PLength { end: PEnd::D3 }), P_J5_LENGTH => Some(ChoiceAddressPattern::PLength { end: PEnd::J5 }), + SAMPLE_HAPLOTYPE => Some(ChoiceAddressPattern::SampleHaplotype), + "sample_gene.v" => Some(ChoiceAddressPattern::SampleGene(VdjSegment::V)), + "sample_gene.d" => Some(ChoiceAddressPattern::SampleGene(VdjSegment::D)), + "sample_gene.j" => Some(ChoiceAddressPattern::SampleGene(VdjSegment::J)), + "sample_allele_in_slot.v" => Some(ChoiceAddressPattern::SampleAlleleInSlot(VdjSegment::V)), + "sample_allele_in_slot.d" => Some(ChoiceAddressPattern::SampleAlleleInSlot(VdjSegment::D)), + "sample_allele_in_slot.j" => Some(ChoiceAddressPattern::SampleAlleleInSlot(VdjSegment::J)), _ => None, }; @@ -1501,5 +1546,27 @@ mod tests { assert_pinned(ChoiceAddress::PLength { end: PEnd::D5 }, "p.d_5.length"); assert_pinned(ChoiceAddress::PLength { end: PEnd::D3 }, "p.d_3.length"); assert_pinned(ChoiceAddress::PLength { end: PEnd::J5 }, "p.j_5.length"); + + // Phased genotype (genotype-modeling PR1). New top-level + // `sample_haplotype` + `sample_gene.*` + `sample_allele_in_slot.*` + // namespaces; same additive policy as receptor revision / + // paired-end — old traces don't reference these strings, no + // ADDRESS_SCHEMA_VERSION bump. + assert_pinned(ChoiceAddress::SampleHaplotype, "sample_haplotype"); + assert_pinned(ChoiceAddress::SampleGene(VdjSegment::V), "sample_gene.v"); + assert_pinned(ChoiceAddress::SampleGene(VdjSegment::D), "sample_gene.d"); + assert_pinned(ChoiceAddress::SampleGene(VdjSegment::J), "sample_gene.j"); + assert_pinned( + ChoiceAddress::SampleAlleleInSlot(VdjSegment::V), + "sample_allele_in_slot.v", + ); + assert_pinned( + ChoiceAddress::SampleAlleleInSlot(VdjSegment::D), + "sample_allele_in_slot.d", + ); + assert_pinned( + ChoiceAddress::SampleAlleleInSlot(VdjSegment::J), + "sample_allele_in_slot.j", + ); } } diff --git a/engine_rs/src/feasibility.rs b/engine_rs/src/feasibility.rs index f1381d2..14d3f6b 100644 --- a/engine_rs/src/feasibility.rs +++ b/engine_rs/src/feasibility.rs @@ -219,6 +219,17 @@ impl VjProductiveFeasibility { .map(|instance| vec![instance.allele_id]); } + // Even at the *same* pass index, if the segment is already + // assigned in the partial simulation, treat it as committed. The + // consolidated `SampleGenotypePass` assigns V (and D) before + // sampling J within one pass index; without this, J feasibility + // would ignore the V already chosen and over-accept. The flat + // one-segment-per-pass path never assigns another segment at the + // same index, so this is a no-op there. + if let Some(instance) = sim.assignments.get(segment) { + return Some(vec![instance.allele_id]); + } + Some(domain.values.clone()) } diff --git a/engine_rs/src/genotype/mod.rs b/engine_rs/src/genotype/mod.rs new file mode 100644 index 0000000..7ff6bb7 --- /dev/null +++ b/engine_rs/src/genotype/mod.rs @@ -0,0 +1,179 @@ +//! Per-individual diploid genotype model (PR1: known reference alleles). +use std::collections::HashMap; + +use crate::ir::Segment; +use crate::refdata::{AlleleId, GeneId}; + +/// One carried allele in a haplotype gene slot. `copies` encodes +/// gene-copy multiplicity for the *same* allele; two different alleles +/// in a slot are two `GeneCopy` entries. `weight` is relative +/// within-slot expression. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct GeneCopy { + pub allele: AlleleId, + pub copies: u8, + pub weight: f32, +} + +/// One chromosome's carried alleles, per V/D/J gene. An absent or empty +/// slot means the gene is deleted on this chromosome. +#[derive(Clone, Debug, Default)] +pub struct Haplotype { + v: HashMap>, + d: HashMap>, + j: HashMap>, +} + +impl Haplotype { + pub fn new() -> Self { + Self::default() + } + + fn map(&self, seg: Segment) -> &HashMap> { + match seg { + Segment::V => &self.v, + Segment::D => &self.d, + Segment::J => &self.j, + _ => panic!("Haplotype: segment must be V/D/J, got {seg:?}"), + } + } + fn map_mut(&mut self, seg: Segment) -> &mut HashMap> { + match seg { + Segment::V => &mut self.v, + Segment::D => &mut self.d, + Segment::J => &mut self.j, + _ => panic!("Haplotype: segment must be V/D/J, got {seg:?}"), + } + } + + /// Set (replace) the copies carried for a gene on this chromosome. + /// An empty `copies` vec means the gene is deleted here. + pub fn set(&mut self, seg: Segment, gene: GeneId, copies: Vec) { + self.map_mut(seg).insert(gene, copies); + } + pub fn slot(&self, seg: Segment, gene: GeneId) -> &[GeneCopy] { + self.map(seg).get(&gene).map(Vec::as_slice).unwrap_or(&[]) + } + pub fn is_deleted(&self, seg: Segment, gene: GeneId) -> bool { + self.slot(seg, gene).is_empty() + } + /// Genes with at least one carried copy on this chromosome, in + /// ascending GeneId order (deterministic). + pub fn present_genes(&self, seg: Segment) -> impl Iterator + '_ { + let mut genes: Vec = self + .map(seg) + .iter() + .filter(|(_, v)| !v.is_empty()) + .map(|(g, _)| *g) + .collect(); + genes.sort_by_key(|g| g.index()); + genes.into_iter() + } + /// (GeneId, usage-weight) for each present gene, weight from `usage`. + pub fn gene_weights f64>(&self, seg: Segment, usage: &F) -> Vec<(GeneId, f64)> { + self.present_genes(seg).map(|g| (g, usage(g))).collect() + } + /// All carried allele ids for a segment across all present genes. + pub fn carried_alleles(&self, seg: Segment) -> Vec { + let mut out = Vec::new(); + for g in self.present_genes(seg) { + for c in self.slot(seg, g) { + out.push(c.allele); + } + } + out + } +} + +/// A diploid genotype: two chromosomes + draw weights + provenance. +#[derive(Clone, Debug)] +pub struct Genotype { + haplotypes: [Haplotype; 2], + chromosome_weights: [f32; 2], + subject_id: Option, + source_refdata_hash: String, +} + +impl Genotype { + pub fn new( + haplotypes: [Haplotype; 2], + chromosome_weights: [f32; 2], + subject_id: Option, + source_refdata_hash: String, + ) -> Self { + Self { + haplotypes, + chromosome_weights, + subject_id, + source_refdata_hash, + } + } + pub fn haplotype(&self, c: usize) -> &Haplotype { + &self.haplotypes[c] + } + pub fn chromosome_weights(&self) -> [f32; 2] { + self.chromosome_weights + } + pub fn subject_id(&self) -> Option<&str> { + self.subject_id.as_deref() + } + pub fn source_refdata_hash(&self) -> &str { + &self.source_refdata_hash + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::refdata::AlleleId; + + fn copy(id: u32) -> GeneCopy { + GeneCopy { + allele: AlleleId::new(id), + copies: 1, + weight: 1.0, + } + } + + #[test] + fn haplotype_reports_carried_alleles_per_gene_with_deletion_as_empty() { + let mut h = Haplotype::new(); + h.set(Segment::V, GeneId::new(0), vec![copy(10)]); // carried + h.set(Segment::V, GeneId::new(1), vec![]); // deleted + assert_eq!(h.slot(Segment::V, GeneId::new(0)).len(), 1); + assert!(h.is_deleted(Segment::V, GeneId::new(1))); + assert!(h.is_deleted(Segment::V, GeneId::new(2))); // absent == deleted + let genes: Vec = h.present_genes(Segment::V).collect(); + assert_eq!(genes, vec![GeneId::new(0)]); // only non-empty slots + } + + #[test] + fn genotype_carries_two_haplotypes_and_chromosome_weights() { + let mut h0 = Haplotype::new(); + let mut h1 = Haplotype::new(); + h0.set(Segment::V, GeneId::new(0), vec![copy(10)]); + h1.set(Segment::V, GeneId::new(0), vec![copy(11)]); // heterozygous + let g = Genotype::new([h0, h1], [0.5, 0.5], Some("S1".into()), "sha256:x".into()); + assert_eq!(g.chromosome_weights(), [0.5, 0.5]); + assert_eq!(g.subject_id(), Some("S1")); + assert_eq!( + g.haplotype(0).slot(Segment::V, GeneId::new(0))[0].allele, + AlleleId::new(10) + ); + assert_eq!( + g.haplotype(1).slot(Segment::V, GeneId::new(0))[0].allele, + AlleleId::new(11) + ); + } + + #[test] + fn gene_weights_restrict_to_present_genes_and_apply_usage() { + // chromosome 0 carries genes 0 and 1; usage favors gene 1. + let mut h = Haplotype::new(); + h.set(Segment::V, GeneId::new(0), vec![copy(10)]); + h.set(Segment::V, GeneId::new(1), vec![copy(20)]); + let usage = |g: GeneId| if g.index() == 1 { 3.0 } else { 1.0 }; + let w = h.gene_weights(Segment::V, &usage); + assert_eq!(w, vec![(GeneId::new(0), 1.0), (GeneId::new(1), 3.0)]); + } +} diff --git a/engine_rs/src/lib.rs b/engine_rs/src/lib.rs index 2c385ed..6d85e2a 100644 --- a/engine_rs/src/lib.rs +++ b/engine_rs/src/lib.rs @@ -44,6 +44,7 @@ pub mod contract; pub mod dist; pub mod event; pub mod feasibility; +pub mod genotype; pub mod ir; pub mod junction; pub mod lineage; diff --git a/engine_rs/src/passes/mod.rs b/engine_rs/src/passes/mod.rs index a059a45..3b90d4b 100644 --- a/engine_rs/src/passes/mod.rs +++ b/engine_rs/src/passes/mod.rs @@ -31,6 +31,7 @@ pub(crate) mod paramsig; pub mod receptor_revision; pub mod sample_allele; pub mod sample_base; +pub mod sample_genotype; pub mod trim; #[cfg(test)] diff --git a/engine_rs/src/passes/sample_genotype.rs b/engine_rs/src/passes/sample_genotype.rs new file mode 100644 index 0000000..ba395cb --- /dev/null +++ b/engine_rs/src/passes/sample_genotype.rs @@ -0,0 +1,740 @@ +//! `SampleGenotypePass` — phased, genotype-aware V(D)J allele sampling. +//! +//! Replaces the three flat `SampleAllelePass` passes when a genotype is +//! attached. In ONE pass it: +//! 1. draws the chromosome (haplotype) once, restricted to **viable** +//! haplotypes (those carrying a feasible allele for every required +//! segment under the active contracts + feasibility), and +//! 2. for each of V/(D)/J, samples a gene present on that chromosome +//! (usage-weighted) then the allele within the gene slot +//! (single-copy slots are deterministic), assigning the slot. +//! +//! The chromosome is a local variable — there is no cross-pass state to +//! share (each pass gets a fresh per-pass trace), which is exactly why +//! haplotype + gene + allele all live in one pass. The canonical +//! `sample_allele.{seg}` choice + the slot assignment are emitted just +//! like the flat path, so every downstream pass (assemble/trim/AIRR/ +//! replay) is unchanged. Gene + within-slot choices are recorded as +//! extra addresses for provenance and replay. +use std::sync::Arc; + +use crate::address::{self, ChoiceAddress, ChoiceAddressPattern}; +use crate::assignment::AlleleInstance; +use crate::contract::ChoiceContext; +use crate::dist::FilteredSampleError; +use crate::genotype::Genotype; +use crate::ir::{Segment, Simulation, SimulationBuilder}; +use crate::pass::{AlleleIdSupport, Pass, PassCompileFact, PassContext, PassEffect, PassError}; +use crate::refdata::{AlleleId, GeneId}; +use crate::rng::Rng; +use crate::trace::ChoiceValue; + +pub struct SampleGenotypePass { + genotype: Arc, + d_required: bool, + // Per-segment gene-usage weights (empty => uniform over present genes). + usage_v: Vec<(GeneId, f64)>, + usage_d: Vec<(GeneId, f64)>, + usage_j: Vec<(GeneId, f64)>, +} + +impl SampleGenotypePass { + pub fn new( + genotype: Arc, + d_required: bool, + usage_v: Vec<(GeneId, f64)>, + usage_d: Vec<(GeneId, f64)>, + usage_j: Vec<(GeneId, f64)>, + ) -> Self { + Self { + genotype, + d_required, + usage_v, + usage_d, + usage_j, + } + } + + fn segments(&self) -> Vec { + if self.d_required { + vec![Segment::V, Segment::D, Segment::J] + } else { + vec![Segment::V, Segment::J] + } + } + + fn usage(&self, seg: Segment) -> &[(GeneId, f64)] { + match seg { + Segment::V => &self.usage_v, + Segment::D => &self.usage_d, + Segment::J => &self.usage_j, + _ => &[], + } + } + + fn usage_of(&self, seg: Segment, g: GeneId) -> f64 { + let table = self.usage(seg); + if table.is_empty() { + 1.0 + } else { + table + .iter() + .find(|(gg, _)| *gg == g) + .map(|(_, w)| *w) + .unwrap_or(1.0) + } + } + + fn vseg(seg: Segment) -> address::VdjSegment { + seg.try_into().expect("V/D/J segment") + } + + fn allele_feasible( + &self, + seg: Segment, + id: AlleleId, + sim: &Simulation, + ctx: &PassContext, + ) -> bool { + let choice = ChoiceValue::AlleleId(id.index()); + let vseg = Self::vseg(seg); + let contract_ok = ctx.contracts.map_or(true, |k| { + k.admits_typed( + sim, + ctx.refdata, + ChoiceContext::none().with_address(ChoiceAddress::SampleAllele(vseg)), + &choice, + ) + .is_ok() + }); + let feasible_ok = ctx.feasibility.map_or(true, |f| { + f.admits( + ctx.pass_index, + sim, + ctx.refdata, + address::sample_allele_vdj(seg), + &choice, + ) + }); + contract_ok && feasible_ok + } + + /// Carried alleles on chromosome `c` for `seg`, optionally filtered + /// by the active contracts + feasibility. + fn admissible_alleles( + &self, + c: usize, + seg: Segment, + sim: &Simulation, + ctx: &PassContext, + filter: bool, + ) -> Vec { + let carried = self.genotype.haplotype(c).carried_alleles(seg); + if !filter { + return carried; + } + carried + .into_iter() + .filter(|id| self.allele_feasible(seg, *id, sim, ctx)) + .collect() + } + + fn is_viable(&self, c: usize, sim: &Simulation, ctx: &PassContext, filter: bool) -> bool { + self.segments() + .iter() + .all(|seg| !self.admissible_alleles(c, *seg, sim, ctx, filter).is_empty()) + } + + /// Viable chromosomes. The feasibility-filtered viable set is + /// preferred; in permissive mode, only if NO chromosome is + /// feasibility-viable do we fall back to presence-viability — exactly + /// mirroring `SampleAllelePass` (filter first, fall back to the + /// unconstrained draw only on empty admissible support). In strict + /// mode the feasibility-viable set is authoritative (empty → the + /// caller raises). + fn viable_set(&self, sim: &Simulation, ctx: &PassContext, strict: bool) -> Vec { + let feasible: Vec = (0..2) + .filter(|&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)) + .collect() + } + + fn draw_haplotype(&self, viable: &[usize], rng: &mut Rng) -> usize { + let w = self.genotype.chromosome_weights(); + let total: f64 = viable.iter().map(|&c| w[c] as f64).sum(); + if total <= 0.0 { + return *viable.first().expect("viable non-empty"); + } + let mut x = rng.next_f64() * total; + for &c in viable { + x -= w[c] as f64; + if x < 0.0 { + return c; + } + } + *viable.last().expect("viable non-empty") + } + + fn weighted_pick(items: &[(T, f64)], rng: &mut Rng) -> T { + let total: f64 = items.iter().map(|(_, w)| *w).sum(); + if total <= 0.0 { + return items.first().expect("non-empty").0; + } + let mut x = rng.next_f64() * total; + for (t, w) in items { + x -= *w; + if x < 0.0 { + return *t; + } + } + items.last().expect("non-empty").0 + } + + fn infeasible_error(&self) -> PassError { + PassError::constraint_sampling( + self.name(), + "sample_haplotype", + FilteredSampleError::EmptyAdmissibleSupport, + ) + } + + fn commit(&self, seg: Segment, sim: Simulation, id: AlleleId, ctx: &mut PassContext) -> Simulation { + let mut b = SimulationBuilder::from_simulation(sim); + if ctx.event_log_sink.is_some() { + b.attach_event_log_observer(); + } + b.assign_allele(seg, AlleleInstance::new(id)); + if let Some(sink) = ctx.event_log_sink.as_deref_mut() { + sink.extend(b.seal_event_log_observer()); + } + b.seal() + } + + /// Union of carried alleles across both haplotypes for a segment — + /// the support advertised to the feasibility/schedule analyzer. + fn union_support(&self, seg: Segment) -> Vec<(AlleleId, f64)> { + let mut ids: Vec = Vec::new(); + for c in 0..2 { + for id in self.genotype.haplotype(c).carried_alleles(seg) { + if !ids.contains(&id) { + ids.push(id); + } + } + } + ids.into_iter().map(|id| (id, 1.0)).collect() + } + + /// Candidate `(allele, mass)` copies in a gene slot on chromosome + /// `c`. `mass = weight * copies` (copy-number dosage). When `filter` + /// is set, copies are restricted to those admissible under the active + /// contracts + feasibility. + fn slot_candidates( + &self, + seg: Segment, + gene: GeneId, + c: usize, + sim: &Simulation, + ctx: &PassContext, + filter: bool, + ) -> Vec<(AlleleId, f64)> { + self.genotype + .haplotype(c) + .slot(seg, gene) + .iter() + .filter(|cp| !filter || self.allele_feasible(seg, cp.allele, sim, ctx)) + .map(|cp| (cp.allele, cp.weight as f64 * cp.copies as f64)) + .collect() + } + + /// Live (fresh-RNG) sampling of one segment within chromosome `c`. + /// Records in canonical order: `sample_gene` → `sample_allele_in_slot` + /// (only when the gene slot has >1 raw copy) → `sample_allele`. + fn sample_segment_live( + &self, + seg: Segment, + c: usize, + sim: Simulation, + ctx: &mut PassContext, + strict: bool, + ) -> Result { + let hap = self.genotype.haplotype(c); + let vseg = Self::vseg(seg); + // Filter-then-fallback, mirroring SampleAllelePass: prefer + // feasibility-admissible candidates; only when NONE are admissible + // do we fall back to the unfiltered carried set (permissive), or + // raise (strict). `filter` is true whenever feasible candidates + // exist on this chromosome+segment. + let feasible_exists = hap.present_genes(seg).any(|g| { + !self + .slot_candidates(seg, g, c, &sim, ctx, true) + .is_empty() + }); + if strict && !feasible_exists { + return Err(PassError::constraint_sampling( + self.name(), + address::sample_allele_vdj(seg), + FilteredSampleError::EmptyAdmissibleSupport, + )); + } + let filter = strict || feasible_exists; + // Genes present on this chromosome that have >=1 candidate copy, + // weighted by gene usage * total copy dosage (so a duplicated + // gene recombines more often). + let mut genes: Vec<(GeneId, f64)> = Vec::new(); + for g in hap.present_genes(seg) { + let cands = self.slot_candidates(seg, g, c, &sim, ctx, filter); + if cands.is_empty() { + continue; + } + let dosage: f64 = cands.iter().map(|(_, m)| *m).sum(); + genes.push((g, self.usage_of(seg, g) * dosage)); + } + if genes.is_empty() { + return Err(PassError::constraint_sampling( + self.name(), + address::sample_allele_vdj(seg), + FilteredSampleError::EmptyAdmissibleSupport, + )); + } + let gene = Self::weighted_pick(&genes, ctx.rng); + let cands = self.slot_candidates(seg, gene, c, &sim, ctx, filter); + let id = Self::weighted_pick(&cands, ctx.rng); + + // "Multi-copy slot" is decided by the RAW slot length (genotype + // structure), independent of feasibility filtering, so replay can + // reconstruct whether a within-slot record exists from the + // genotype alone. + let multi = hap.slot(seg, gene).len() > 1; + ctx.trace + .record_choice(ChoiceAddress::SampleGene(vseg), ChoiceValue::GeneId(gene.index())); + if multi { + ctx.trace.record_choice( + ChoiceAddress::SampleAlleleInSlot(vseg), + ChoiceValue::AlleleId(id.index()), + ); + } + ctx.trace + .record_choice(ChoiceAddress::SampleAllele(vseg), ChoiceValue::AlleleId(id.index())); + Ok(self.commit(seg, sim, id, ctx)) + } + + /// Replay (trace-injected) sampling of one segment within `c`. + /// Consumes records in the same order live emits them — `sample_gene`, + /// then `sample_allele_in_slot` (iff the raw slot has >1 copy), then + /// the canonical `sample_allele` (the assigned id, source of truth). + fn sample_segment_replay( + &self, + seg: Segment, + c: usize, + sim: Simulation, + ctx: &mut PassContext, + ) -> Result { + let vseg = Self::vseg(seg); + let gene_idx = ctx + .replay_cursor + .as_deref_mut() + .expect("replay cursor present") + .expect_gene_id(ChoiceAddress::SampleGene(vseg)) + .map_err(|r| PassError::replay(self.name(), r))?; + let gene = GeneId::new(gene_idx); + let slot = self.genotype.haplotype(c).slot(seg, gene); + // Replay validation: a recorded trace must be admissible against + // this genotype — the recorded gene must be carried on the drawn + // chromosome (non-empty slot). Mirrors SampleAllelePass's + // "trace proposes, engine validates" contract. + if slot.is_empty() { + return Err(PassError::invalid_distribution_output( + self.name(), + address::sample_allele_vdj(seg), + gene_idx as i64, + "genotype_gene_not_carried_on_haplotype", + )); + } + let multi = slot.len() > 1; + let slot_recorded = if multi { + Some( + ctx.replay_cursor + .as_deref_mut() + .expect("replay cursor present") + .expect_allele_id(ChoiceAddress::SampleAlleleInSlot(vseg)) + .map_err(|r| PassError::replay(self.name(), r))?, + ) + } else { + None + }; + let allele = ctx + .replay_cursor + .as_deref_mut() + .expect("replay cursor present") + .expect_allele_id(ChoiceAddress::SampleAllele(vseg)) + .map_err(|r| PassError::replay(self.name(), r))?; + // The canonical allele must be one the gene slot actually carries. + if !slot.iter().any(|cp| cp.allele.index() == allele) { + return Err(PassError::invalid_distribution_output( + self.name(), + address::sample_allele_vdj(seg), + allele as i64, + "genotype_allele_not_in_gene_slot", + )); + } + // When a within-slot record exists it must agree with the + // canonical allele (live writes the same id to both). + if let Some(slot_id) = slot_recorded { + if slot_id != allele { + return Err(PassError::invalid_distribution_output( + self.name(), + address::sample_allele_vdj(seg), + slot_id as i64, + "genotype_slot_record_disagrees_with_canonical_allele", + )); + } + } + let id = AlleleId::new(allele); + + ctx.trace + .record_choice(ChoiceAddress::SampleGene(vseg), ChoiceValue::GeneId(gene_idx)); + if let Some(slot_id) = slot_recorded { + ctx.trace + .record_choice(ChoiceAddress::SampleAlleleInSlot(vseg), ChoiceValue::AlleleId(slot_id)); + } + ctx.trace + .record_choice(ChoiceAddress::SampleAllele(vseg), ChoiceValue::AlleleId(id.index())); + Ok(self.commit(seg, sim, id, ctx)) + } + + /// Shared execute body. `strict` selects feasibility-filtered + /// viability/sampling (and structured errors) vs presence-based + /// permissive sampling (feasibility advisory, never errors given a + /// complete haplotype). + fn run( + &self, + sim: &Simulation, + ctx: &mut PassContext, + strict: bool, + ) -> Result { + let replaying = ctx.replay_cursor.is_some(); + let c = if replaying { + let recorded = ctx + .replay_cursor + .as_deref_mut() + .expect("replay cursor present") + .expect_haplotype(ChoiceAddress::SampleHaplotype) + .map_err(|r| PassError::replay(self.name(), r))?; + let viable = self.viable_set(sim, ctx, strict); + if !viable.contains(&(recorded as usize)) { + return Err(self.infeasible_error()); + } + ctx.trace + .record_choice(ChoiceAddress::SampleHaplotype, ChoiceValue::Haplotype(recorded)); + recorded as usize + } else { + let viable = self.viable_set(sim, ctx, strict); + if viable.is_empty() { + return Err(self.infeasible_error()); + } + let c = self.draw_haplotype(&viable, ctx.rng); + ctx.trace + .record_choice(ChoiceAddress::SampleHaplotype, ChoiceValue::Haplotype(c as u8)); + c + }; + + let mut current = sim.clone(); + for seg in self.segments() { + current = if replaying { + self.sample_segment_replay(seg, c, current, ctx)? + } else { + self.sample_segment_live(seg, c, current, ctx, strict)? + }; + } + Ok(current) + } +} + +impl Pass for SampleGenotypePass { + fn name(&self) -> &str { + "sample_genotype" + } + + /// Encode the attached genotype so two different genotypes produce + /// distinct plan signatures (replay-cache correctness): source + /// cartridge hash, subject, chromosome weights, and the full + /// per-haplotype/per-segment carried slots (allele id, copies, + /// weight bits). + fn parameter_signature(&self) -> String { + use std::fmt::Write; + let g = &self.genotype; + let mut s = String::new(); + let _ = write!( + s, + "geno|hash={}|subj={}|cw={:?}|d={}", + g.source_refdata_hash(), + g.subject_id().unwrap_or(""), + g.chromosome_weights(), + self.d_required, + ); + for c in 0..2usize { + let hap = g.haplotype(c); + for seg in self.segments() { + for gene in hap.present_genes(seg) { + let mut copies: Vec<(u32, u8, u32)> = hap + .slot(seg, gene) + .iter() + .map(|cp| (cp.allele.index(), cp.copies, cp.weight.to_bits())) + .collect(); + copies.sort_unstable(); + let _ = write!(s, "|h{}.{:?}.g{}={:?}", c, seg, gene.index(), copies); + } + } + } + s + } + + fn execute(&self, sim: &Simulation, ctx: &mut PassContext) -> Simulation { + // Permissive: viability is presence-based (feasibility advisory), + // so a genotype with >=1 complete haplotype (guaranteed by the + // compile-time presence check) never errors here. + self.run(sim, ctx, false) + .expect("SampleGenotypePass permissive execution must not error") + } + + fn execute_checked( + &self, + sim: &Simulation, + ctx: &mut PassContext, + ) -> Result { + self.run(sim, ctx, true) + } + + fn declared_choice_patterns(&self) -> Vec { + let mut patterns = vec![ChoiceAddressPattern::SampleHaplotype]; + for seg in self.segments() { + let vseg = Self::vseg(seg); + patterns.push(ChoiceAddressPattern::SampleGene(vseg)); + patterns.push(ChoiceAddressPattern::SampleAlleleInSlot(vseg)); + patterns.push(ChoiceAddressPattern::SampleAllele(vseg)); + } + patterns + } + + fn effects(&self) -> Vec { + self.segments() + .into_iter() + .map(PassEffect::AssignAllele) + .collect() + } + + fn compile_facts(&self) -> Vec { + self.segments() + .into_iter() + .map(|seg| PassCompileFact::AlleleSampleSupport { + segment: seg, + support: AlleleIdSupport::from_weighted_pairs(Some(self.union_support(seg))), + }) + .collect() + } +} + +#[cfg(test)] +pub(crate) mod test_support { + use super::*; + use crate::genotype::{GeneCopy, Haplotype}; + + pub fn copy(id: u32) -> GeneCopy { + GeneCopy { + allele: AlleleId::new(id), + copies: 1, + weight: 1.0, + } + } + + /// hap0 carries V (gene0 -> allele 0) + J (gene0 -> allele 100); + /// hap1 carries V but no J → only hap0 viable. + pub fn geno_chrom1_deletes_j() -> Genotype { + let mut h0 = Haplotype::new(); + h0.set(Segment::V, GeneId::new(0), vec![copy(0)]); + h0.set(Segment::J, GeneId::new(0), vec![copy(100)]); + let mut h1 = Haplotype::new(); + h1.set(Segment::V, GeneId::new(0), vec![copy(1)]); + Genotype::new([h0, h1], [0.5, 0.5], Some("S1".into()), "sha256:test".into()) + } + + /// Neither haplotype carries a J gene → no viable haplotype. + pub fn geno_both_delete_j() -> Genotype { + let mut h0 = Haplotype::new(); + h0.set(Segment::V, GeneId::new(0), vec![copy(0)]); + let mut h1 = Haplotype::new(); + h1.set(Segment::V, GeneId::new(0), vec![copy(1)]); + Genotype::new([h0, h1], [0.5, 0.5], Some("S1".into()), "sha256:test".into()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pass::testing::PassRuntime; + use crate::pass::PassPlan; + + #[test] + fn draws_only_the_viable_haplotype_and_assigns_carried_alleles() { + let g = Arc::new(test_support::geno_chrom1_deletes_j()); + let pass = SampleGenotypePass::new(g, false, vec![], vec![], vec![]); + let mut plan = PassPlan::new(); + plan.push(Box::new(pass)); + for seed in 0..30u64 { + let outcome = PassRuntime::execute(&plan, Simulation::new(), seed); + match outcome.trace.find("sample_haplotype").unwrap().value { + ChoiceValue::Haplotype(c) => assert_eq!(c, 0, "seed {seed}"), + _ => panic!("wrong variant"), + } + let sim = outcome.final_simulation(); + assert_eq!( + sim.assignments.get(Segment::V).unwrap().allele_id, + AlleleId::new(0) + ); + assert_eq!( + sim.assignments.get(Segment::J).unwrap().allele_id, + AlleleId::new(100) + ); + // single-copy slots => no within-slot record + assert!(outcome.trace.find("sample_allele_in_slot.v").is_none()); + } + } + + #[test] + fn records_canonical_sample_allele_and_gene_addresses() { + let g = Arc::new(test_support::geno_chrom1_deletes_j()); + let pass = SampleGenotypePass::new(g, false, vec![], vec![], vec![]); + let mut plan = PassPlan::new(); + plan.push(Box::new(pass)); + let outcome = PassRuntime::execute(&plan, Simulation::new(), 7); + match outcome.trace.find("sample_allele.v").unwrap().value { + ChoiceValue::AlleleId(id) => assert_eq!(id, 0), + _ => panic!("expected AlleleId at sample_allele.v"), + } + match outcome.trace.find("sample_gene.v").unwrap().value { + ChoiceValue::GeneId(g) => assert_eq!(g, 0), + _ => panic!("expected GeneId at sample_gene.v"), + } + } + + #[test] + fn errors_when_no_haplotype_viable() { + let g = Arc::new(test_support::geno_both_delete_j()); + let pass = SampleGenotypePass::new(g, false, vec![], vec![], vec![]); + let mut plan = PassPlan::new(); + plan.push(Box::new(pass)); + let result = + PassRuntime::execute_strict_with_context(&plan, Simulation::new(), 0, None, None); + assert!(result.is_err(), "expected genotype-infeasibility error"); + } + + #[test] + fn replay_reproduces_live_phased_choices() { + use crate::pass::PassContext; + use crate::replay::TraceCursor; + use crate::rng::Rng; + use crate::trace::Trace; + + let g = Arc::new(test_support::geno_chrom1_deletes_j()); + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleGenotypePass::new( + g.clone(), + false, + vec![], + vec![], + vec![], + ))); + // Live run — capture the full trace (haplotype + gene + allele). + let live = PassRuntime::execute(&plan, Simulation::new(), 13); + let records: Vec<_> = live.trace.choices().to_vec(); + let live_sim = live.final_simulation(); + let live_v = live_sim.assignments.get(Segment::V).unwrap().allele_id; + let live_j = live_sim.assignments.get(Segment::J).unwrap().allele_id; + + // Replay the captured trace through a fresh pass instance. + let replay_pass = SampleGenotypePass::new(g, false, vec![], vec![], vec![]); + let mut cursor = TraceCursor::from_owned(records); + let mut trace = Trace::new(); + let mut rng = Rng::new(999); + let sim = Simulation::new(); + let result; + { + let mut ctx = PassContext { + trace: &mut trace, + rng: &mut rng, + pass_index: 0, + refdata: None, + contracts: None, + feasibility: None, + reference_index: None, + replay_cursor: Some(&mut cursor), + event_log_sink: None, + }; + result = replay_pass.run(&sim, &mut ctx, true); + } + let replayed = result.expect("genotype replay must succeed"); + assert_eq!( + replayed.assignments.get(Segment::V).unwrap().allele_id, + live_v + ); + assert_eq!( + replayed.assignments.get(Segment::J).unwrap().allele_id, + live_j + ); + } + + #[test] + fn replay_rejects_allele_not_carried_in_genotype_slot() { + use crate::pass::PassContext; + use crate::replay::TraceCursor; + use crate::rng::Rng; + use crate::trace::Trace; + + let g = Arc::new(test_support::geno_chrom1_deletes_j()); + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleGenotypePass::new( + g.clone(), + false, + vec![], + vec![], + vec![], + ))); + let live = PassRuntime::execute(&plan, Simulation::new(), 13); + // Tamper: rewrite the canonical V allele to one not in the slot. + let mut records: Vec<_> = live.trace.choices().to_vec(); + for r in records.iter_mut() { + if r.address == "sample_allele.v" { + r.value = ChoiceValue::AlleleId(999); + } + } + let replay_pass = SampleGenotypePass::new(g, false, vec![], vec![], vec![]); + let mut cursor = TraceCursor::from_owned(records); + let mut trace = Trace::new(); + let mut rng = Rng::new(1); + let sim = Simulation::new(); + let result; + { + let mut ctx = PassContext { + trace: &mut trace, + rng: &mut rng, + pass_index: 0, + refdata: None, + contracts: None, + feasibility: None, + reference_index: None, + replay_cursor: Some(&mut cursor), + event_log_sink: None, + }; + result = replay_pass.run(&sim, &mut ctx, true); + } + assert!( + result.is_err(), + "replay must reject an allele not carried in the genotype slot" + ); + } +} diff --git a/engine_rs/src/python/plan.rs b/engine_rs/src/python/plan.rs index a45ff1f..9886a56 100644 --- a/engine_rs/src/python/plan.rs +++ b/engine_rs/src/python/plan.rs @@ -345,6 +345,137 @@ impl PyPassPlan { Ok(()) } + /// Append a single `SampleGenotypePass` that replaces the three flat + /// `SampleAllelePass` passes when a genotype is attached. The pass + /// draws the chromosome once, then V/(D)/J alleles from that + /// chromosome's carried set (phased). + /// + /// `v`/`d`/`j` are flat rows `(haplotype, allele_id, copies, weight)` + /// already resolved to this refdata's allele ids; rows are grouped by + /// the gene the allele belongs to (via the segment's `GeneIndex`). + /// + /// `*_weights` are optional pool-aligned allele-usage weight vectors + /// (from the cartridge's allele-usage model). They are aggregated to + /// gene-level usage (sum of an allele's weights per gene) so gene + /// choice reflects empirical usage instead of being uniform. + #[pyo3(signature = (refdata, chromosome_weights, subject_id, source_hash, v, d, j, d_required, v_weights=None, d_weights=None, j_weights=None))] + #[allow(clippy::too_many_arguments)] + fn push_genotype_recombine( + &mut self, + refdata: &PyRefDataConfig, + chromosome_weights: (f32, f32), + subject_id: Option, + source_hash: String, + v: Vec<(u8, u32, u8, f32)>, + d: Vec<(u8, u32, u8, f32)>, + j: Vec<(u8, u32, u8, f32)>, + d_required: bool, + v_weights: Option>, + d_weights: Option>, + j_weights: Option>, + ) -> PyResult<()> { + use crate::genotype::{GeneCopy, Genotype, Haplotype}; + use crate::ir::Segment; + use crate::refdata::{AlleleId, GeneId, GeneIndex}; + use std::collections::HashMap; + + let cfg = refdata.inner(); + + let build_index = |seg: Segment| -> PyResult { + let pool = cfg + .pool_for(seg) + .ok_or_else(|| PyValueError::new_err(format!("no pool for segment {:?}", seg)))?; + Ok(GeneIndex::build(pool)) + }; + let v_index = build_index(Segment::V)?; + let j_index = build_index(Segment::J)?; + let d_index = if d_required { + Some(build_index(Segment::D)?) + } else { + None + }; + + // Aggregate pool-aligned allele weights to gene-level usage. + let gene_usage = |idx: &GeneIndex, weights: &Option>| -> Vec<(GeneId, f64)> { + match weights { + None => Vec::new(), + Some(w) => idx + .genes() + .map(|(g, _)| { + let mass: f64 = idx + .alleles_of(g) + .iter() + .map(|a| w.get(a.as_usize()).copied().unwrap_or(0.0)) + .sum(); + (g, mass) + }) + .collect(), + } + }; + let usage_v = gene_usage(&v_index, &v_weights); + let usage_j = gene_usage(&j_index, &j_weights); + let usage_d = match &d_index { + Some(di) => gene_usage(di, &d_weights), + None => Vec::new(), + }; + + let mut haps = [Haplotype::new(), Haplotype::new()]; + let fill = |haps: &mut [Haplotype; 2], + seg: Segment, + idx: &GeneIndex, + rows: &[(u8, u32, u8, f32)]| + -> PyResult<()> { + let pool_len = cfg.pool_for(seg).map(|p| p.len() as u32).unwrap_or(0); + let mut grouped: HashMap<(u8, u32), Vec> = HashMap::new(); + for (h, aid, copies, weight) in rows { + if *h > 1 { + return Err(PyValueError::new_err(format!( + "haplotype index must be 0 or 1, got {}", + h + ))); + } + if *aid >= pool_len { + return Err(PyValueError::new_err(format!( + "{:?} allele id {} out of range (pool size {})", + seg, aid, pool_len + ))); + } + let allele = AlleleId::new(*aid); + let gene = idx.gene_of(allele); + grouped.entry((*h, gene.index())).or_default().push(GeneCopy { + allele, + copies: *copies, + weight: *weight, + }); + } + for ((h, gene_idx), copies) in grouped { + haps[h as usize].set(seg, GeneId::new(gene_idx), copies); + } + Ok(()) + }; + fill(&mut haps, Segment::V, &v_index, &v)?; + fill(&mut haps, Segment::J, &j_index, &j)?; + if let Some(di) = &d_index { + fill(&mut haps, Segment::D, di, &d)?; + } + + let genotype = std::sync::Arc::new(Genotype::new( + haps, + [chromosome_weights.0, chromosome_weights.1], + subject_id, + source_hash, + )); + self.inner_mut()? + .push(Box::new(crate::passes::sample_genotype::SampleGenotypePass::new( + genotype, + d_required, + usage_v, + usage_d, + usage_j, + ))); + Ok(()) + } + /// Append an `AssembleSegmentPass` for `segment`. The matching /// `SampleAllelePass` must already be earlier in the plan /// (otherwise the assembler will fail at execute time with a diff --git a/engine_rs/src/python/trace.rs b/engine_rs/src/python/trace.rs index b8b9cc9..03d369d 100644 --- a/engine_rs/src/python/trace.rs +++ b/engine_rs/src/python/trace.rs @@ -17,6 +17,8 @@ fn choice_value_to_py(py: Python<'_>, v: &ChoiceValue) -> PyObject { ChoiceValue::Bases(bs) => PyBytes::new_bound(py, bs).into_py(py), ChoiceValue::AlleleId(id) => id.into_py(py), ChoiceValue::Bool(b) => b.into_py(py), + ChoiceValue::Haplotype(h) => h.into_py(py), + ChoiceValue::GeneId(g) => g.into_py(py), } } diff --git a/engine_rs/src/refdata.rs b/engine_rs/src/refdata.rs index dafaec0..cb97b75 100644 --- a/engine_rs/src/refdata.rs +++ b/engine_rs/src/refdata.rs @@ -95,6 +95,89 @@ impl AlleleId { } } +/// Stable, refdata-local identifier for a gene within one segment's +/// pool. Assigned in first-appearance order over the pool's alleles, so +/// it is deterministic for a fixed cartridge (and therefore safe to +/// record in the trace for replay, gated by `refdata_content_hash`). +#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)] +pub struct GeneId(u32); + +impl GeneId { + pub const fn new(idx: u32) -> Self { + Self(idx) + } + pub const fn index(self) -> u32 { + self.0 + } + pub const fn as_usize(self) -> usize { + self.0 as usize + } +} + +/// Gene-level grouping over a single `AllelePool`, derived from each +/// allele's `gene` string. Built on demand; the pool itself is +/// unchanged, so `Allele` and `refdata_content_hash` are untouched. +#[derive(Clone, Debug)] +pub struct GeneIndex { + names: Vec, // GeneId.index() -> gene name + by_name: std::collections::HashMap, // gene name -> GeneId + alleles: Vec>, // GeneId.index() -> alleles, pool order + gene_of: Vec, // AlleleId.index() -> GeneId +} + +impl GeneIndex { + /// Build the gene grouping from a pool. Genes are numbered in the + /// order their first allele appears in the pool. + pub fn build(pool: &AllelePool) -> Self { + let mut names: Vec = Vec::new(); + let mut by_name: std::collections::HashMap = + std::collections::HashMap::new(); + let mut alleles: Vec> = Vec::new(); + let mut gene_of: Vec = Vec::with_capacity(pool.len()); + for (id, allele) in pool.iter() { + let gid = *by_name.entry(allele.gene.clone()).or_insert_with(|| { + let g = GeneId::new(names.len() as u32); + names.push(allele.gene.clone()); + alleles.push(Vec::new()); + g + }); + alleles[gid.as_usize()].push(id); + gene_of.push(gid); + } + Self { + names, + by_name, + alleles, + gene_of, + } + } + + pub fn len(&self) -> usize { + self.names.len() + } + pub fn is_empty(&self) -> bool { + self.names.is_empty() + } + pub fn gene_id(&self, name: &str) -> Option { + self.by_name.get(name).copied() + } + pub fn gene_name(&self, g: GeneId) -> &str { + &self.names[g.as_usize()] + } + pub fn alleles_of(&self, g: GeneId) -> &[AlleleId] { + &self.alleles[g.as_usize()] + } + pub fn gene_of(&self, a: AlleleId) -> GeneId { + self.gene_of[a.as_usize()] + } + pub fn genes(&self) -> impl Iterator { + self.names + .iter() + .enumerate() + .map(|(i, n)| (GeneId::new(i as u32), n.as_str())) + } +} + // ────────────────────────────────────────────────────────────────── // ChainType — VJ (light) vs VDJ (heavy) // ────────────────────────────────────────────────────────────────── @@ -1180,3 +1263,46 @@ mod tests { assert_eq!(cfg.j_pool.len(), 1); } } + +#[cfg(test)] +mod gene_index_tests { + use super::*; + + fn pool_with(genes: &[(&str, &str)]) -> AllelePool { + // genes: (allele_name, gene_name) + let mut p = AllelePool::new(); + for (name, gene) in genes { + let _ = p.push(Allele { + name: (*name).to_string(), + gene: (*gene).to_string(), + seq: vec![b'A'; 10], + segment: Segment::V, + anchor: Some(3), + functional_status: None, + subregions: Vec::new(), + }); + } + p + } + + #[test] + fn gene_index_groups_alleles_by_gene_in_first_appearance_order() { + let pool = pool_with(&[ + ("IGHV1-2*01", "IGHV1-2"), + ("IGHV1-2*02", "IGHV1-2"), + ("IGHV3-23*01", "IGHV3-23"), + ]); + let idx = GeneIndex::build(&pool); + + assert_eq!(idx.len(), 2); + let g12 = idx.gene_id("IGHV1-2").unwrap(); + let g323 = idx.gene_id("IGHV3-23").unwrap(); + assert_eq!(g12.index(), 0); // first appearance + assert_eq!(g323.index(), 1); + assert_eq!(idx.gene_name(g12), "IGHV1-2"); + assert_eq!(idx.alleles_of(g12), &[AlleleId::new(0), AlleleId::new(1)]); + assert_eq!(idx.alleles_of(g323), &[AlleleId::new(2)]); + assert_eq!(idx.gene_of(AlleleId::new(1)), g12); + assert!(idx.gene_id("nope").is_none()); + } +} diff --git a/engine_rs/src/replay.rs b/engine_rs/src/replay.rs index 560a4f9..70f9291 100644 --- a/engine_rs/src/replay.rs +++ b/engine_rs/src/replay.rs @@ -137,6 +137,8 @@ pub fn choice_value_kind(value: &ChoiceValue) -> &'static str { ChoiceValue::Bases(_) => "Bases", ChoiceValue::AlleleId(_) => "AlleleId", ChoiceValue::Bool(_) => "Bool", + ChoiceValue::Haplotype(_) => "Haplotype", + ChoiceValue::GeneId(_) => "GeneId", } } @@ -302,6 +304,24 @@ impl TraceCursor { other => Err(kind_mismatch(position, &address_str, "Bool", &other)), } } + + /// Consume the next record as a `ChoiceValue::Haplotype`. + pub fn expect_haplotype(&mut self, address: ChoiceAddress) -> Result { + let (position, address_str, value) = self.advance_with_address(address)?; + match value { + ChoiceValue::Haplotype(h) => Ok(h), + other => Err(kind_mismatch(position, &address_str, "Haplotype", &other)), + } + } + + /// Consume the next record as a `ChoiceValue::GeneId`. + pub fn expect_gene_id(&mut self, address: ChoiceAddress) -> Result { + let (position, address_str, value) = self.advance_with_address(address)?; + match value { + ChoiceValue::GeneId(g) => Ok(g), + other => Err(kind_mismatch(position, &address_str, "GeneId", &other)), + } + } } /// Build a `ValueKindMismatch` without re-borrowing the cursor. Free diff --git a/engine_rs/src/trace.rs b/engine_rs/src/trace.rs index ff61836..a15ee91 100644 --- a/engine_rs/src/trace.rs +++ b/engine_rs/src/trace.rs @@ -63,6 +63,12 @@ pub enum ChoiceValue { /// A boolean choice (e.g., D inversion: yes/no, receptor /// revision: yes/no, contaminant injection: yes/no). Bool(bool), + + /// A chromosome index for phased genotype sampling (0 or 1). + Haplotype(u8), + + /// A refdata-local gene identifier (see `refdata::GeneId`). + GeneId(u32), } /// On-disk discriminant tag for `ChoiceValue`. Lives as a separate @@ -78,6 +84,8 @@ enum ChoiceValueWire { Bases(String), AlleleId(u32), Bool(bool), + Haplotype(u8), + GeneId(u32), } impl Serialize for ChoiceValue { @@ -92,6 +100,8 @@ impl Serialize for ChoiceValue { ), ChoiceValue::AlleleId(id) => ChoiceValueWire::AlleleId(id), ChoiceValue::Bool(b) => ChoiceValueWire::Bool(b), + ChoiceValue::Haplotype(h) => ChoiceValueWire::Haplotype(h), + ChoiceValue::GeneId(g) => ChoiceValueWire::GeneId(g), }; wire.serialize(ser) } @@ -120,6 +130,8 @@ impl<'de> Deserialize<'de> for ChoiceValue { } ChoiceValueWire::AlleleId(id) => ChoiceValue::AlleleId(id), ChoiceValueWire::Bool(b) => ChoiceValue::Bool(b), + ChoiceValueWire::Haplotype(h) => ChoiceValue::Haplotype(h), + ChoiceValueWire::GeneId(g) => ChoiceValue::GeneId(g), }) } } @@ -314,6 +326,15 @@ mod tests { assert_eq!(run_trace.choices()[2].address, "third.choice"); } + #[test] + fn haplotype_and_gene_id_choice_values_round_trip_through_wire() { + for v in [ChoiceValue::Haplotype(1), ChoiceValue::GeneId(7)] { + let json = serde_json::to_string(&v).unwrap(); + let back: ChoiceValue = serde_json::from_str(&json).unwrap(); + assert_eq!(v, back); + } + } + #[test] fn trace_find_by_exact_address() { let mut t = Trace::new(); diff --git a/mkdocs.yml b/mkdocs.yml index d864b06..e8814ec 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -207,6 +207,7 @@ nav: - Clonal simulation overview: guides/clonal-families.md - Clonal lineage trees: guides/clonal-lineage.md - Clonal repertoires (TCR & abundance): guides/clonal-repertoire.md + - Genotypes (per-individual diploid): guides/genotype.md - Junction N/P additions: guides/junction-additions.md - Targeted SHM rates: guides/shm-targeting.md - Corruption + sequencing artefacts: guides/corruption-sequencing.md diff --git a/site_docs/assets/genotype-tigger-recovery.png b/site_docs/assets/genotype-tigger-recovery.png new file mode 100644 index 0000000..2d7a666 Binary files /dev/null and b/site_docs/assets/genotype-tigger-recovery.png differ diff --git a/site_docs/guides/genotype.md b/site_docs/guides/genotype.md new file mode 100644 index 0000000..b3a8107 --- /dev/null +++ b/site_docs/guides/genotype.md @@ -0,0 +1,490 @@ +# Genotypes: per-individual diploid germline + +

A genotype in GenAIRR is one person's diploid germline +complement — which V/D/J alleles they carry, on which chromosome, in what copy +number. Attach a Genotype to an Experiment and V(D)J +recombination becomes haplotype-phased: the V, D and J of each +rearrangement are drawn from a single chromosome, honouring allele +presence/absence, zygosity, and gene deletion. With no genotype attached the +engine is byte-for-byte unchanged. This page explains exactly what a genotype is, +how the engine samples from it, how to build one, and how to use it to benchmark +genotype-inference tools — nothing here is a black box.

+ +## What a genotype is (and why it matters) + +Every person inherits two copies of the immunoglobulin heavy-chain locus — one on +each homologous chromosome (one **haplotype** from each parent). Across the +population the locus is extraordinarily polymorphic: a reference set may list +dozens of alleles per gene, but a *single individual* carries only a handful — +typically **one or two alleles per gene** — and may be **missing entire genes** +(deletion) or carry **extra copies** (duplication). That per-individual set is the +**genotype**. + +GenAIRR models four things a genotype encodes: + +| Concept | Meaning | In GenAIRR | +|---|---|---| +| **Allele presence/absence** | only carried alleles can rearrange | alleles not in the genotype are never sampled | +| **Diploid zygosity** | per gene: 1 allele (homozygous) or 2 (heterozygous) | `homozygous` / `heterozygous` | +| **Gene deletion / copy number** | a gene can be absent on one or both chromosomes, or duplicated | `delete_gene` / `duplicate_gene` | +| **Haplotype phasing** | V, D, J of one rearrangement come from one chromosome | drawn automatically, recorded per record | + +Phasing is what makes a genotype more than "a list of alleles to allow". The IGH +locus is physically on a chromosome, so a single recombination event splices a V, +a D and a J **from the same chromosome**. That linkage is exactly the signal +haplotype-inference methods exploit (e.g. the IGHJ6-anchor approach), and GenAIRR +reproduces it. + +!!! note "Supported loci and chains" + Genotypes work on any GenAIRR reference cartridge — BCR **and** TCR, heavy + **and** light/α/β chains. On **VDJ** loci (IGH, TRB, TRD) the genotype spans + V, D and J and each rearrangement draws all three from one chromosome. On + **VJ** loci (IGK, IGL, TRA, TRG) there is no D segment: genotype V and J, + D rows are simply not required and are ignored. The examples below use the + human IGH cartridge, but the same API applies to every locus; just use that + cartridge's gene/allele names. + +## Quick start + +```python +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype + +cfg = gdata.HUMAN_IGH_OGRDB + +# Build a diploid genotype: start from the reference, then edit specific genes. +g = ( + Genotype.from_dataconfig(cfg) + .complete_from_reference("homozygous_first_reference") # fill the rest + .heterozygous("IGHVF1-G1", "IGHVF1-G1*01", "IGHVF1-G1*02") # two alleles + .homozygous("IGHVF2-G4", "IGHVF2-G4*01") # one allele + .delete_gene("IGHVF3-G7", haplotype="both") # gene absent + .with_subject("DONOR01") +) + +result = ( + ga.Experiment.on(cfg) + .with_genotype(g) # recombination is now haplotype-phased + .recombine() + .run_records(n=1000, seed=7) +) + +result[0]["subject_id"] # 'DONOR01' — provenance on every record +result[0]["haplotype"] # 0 or 1 — which chromosome this read used +result.genotypes[0].to_table() # ground-truth genotype (per gene, per haplotype) +``` + +## How recombination samples from a genotype + +When a genotype is attached, recombination runs a single phased sampling pass per +rearrangement. The steps, in order: + +1. **Draw a chromosome.** One of the two haplotypes is chosen, weighted by the + `chromosome_weights` (default `[0.5, 0.5]`). This choice is made **once** and + shared by V, D and J — that is the phasing. +2. **Per segment, draw a gene then an allele.** Among the genes *present on the + chosen chromosome*, a gene is sampled (weighted by usage — see below), then the + allele follows from that chromosome's slot for the gene. A deleted gene is + simply not offered on the chromosome that lacks it. +3. **Assign and continue.** The chosen V/D/J alleles are assigned and the rest of + the pipeline (trimming, NP, assembly, SHM, corruption) runs unchanged. + +Every random choice (chromosome, gene, within-slot allele) is recorded to the +trace, so seeded runs are byte-stable and fully replayable. + +### Viability and `productive_only` + +A chromosome is only drawn if it is **viable** — it must carry at least one usable +allele for every required segment (V and J, plus D on heavy chains). This matters +with deletions: if one haplotype lacks a J gene entirely, only the other +chromosome is ever drawn. If **neither** chromosome can produce a rearrangement, +the genotype is rejected at compile time with a clear error (rather than failing +at run time). + +Under [`productive_only`](productive.md), viability also accounts for +productive-junction feasibility, and the V chosen earlier in the pass constrains +the J drawn later (the phased choices are evaluated together, not independently). + +### Strict vs permissive + +`Genotype.from_dataconfig(cfg)` is **strict**: any gene that could be used during +recombination but was never specified is an error when the experiment is compiled +(`compile()` / `run_records()`) — you must define +the whole genotype (use `complete_from_reference` to fill the genes you don't care +about). This guarantees a genuine diploid complement, which is what you want for a +ground-truth benchmark. + +`Genotype.permissive(cfg)` is a separate, explicitly **non-diploid** fallback: +unspecified genes are left to sample over *all* their reference alleles without +phasing. It exists for the "I only want to constrain a few genes" case — it is +**not** a biological genotype, and it is labelled as such in `to_table()` and +`repr`. + +In both modes, feasibility (e.g. `productive_only`) is applied the same way the +non-genotype path applies it: candidates are filtered to the feasible set, and the +unfiltered set is used only as a last resort when nothing is feasible — so a +genotype run never silently samples alleles a normal run would have avoided. + +## Building a genotype + +The `Genotype` builder is a fluent, validated editor over a `DataConfig`'s +reference alleles. Every method returns `self`, so calls chain. + +```python +g = Genotype.from_dataconfig(cfg) # strict (recommended) + +g.homozygous("IGHVF2-G4", "IGHVF2-G4*01") # 1 allele on both chromosomes +g.heterozygous("IGHVF1-G1", "IGHVF1-G1*01", "IGHVF1-G1*02") # different allele per chromosome +g.delete_gene("IGHVF3-G7", haplotype="both") # gene absent entirely (homozygous deletion) +g.homozygous("IGHVF3-G8", "IGHVF3-G8*01") # carried on both chromosomes... +g.delete_gene("IGHVF3-G8", haplotype=1) # ...then removed on chromosome 1 (hemizygous) +g.duplicate_gene("IGHVF1-G2", ["IGHVF1-G2*01", "IGHVF1-G2*02"], haplotype=0) # >1 copy on one chromosome +g.chromosome_weights(0.6, 0.4) # allelic-expression imbalance +g.with_subject("DONOR01") # provenance label + +g.complete_from_reference("homozygous_first_reference") # fill every unspecified gene +``` + +Every editing method takes a `segment` argument (`"V"` default, or `"D"` / `"J"`), +so genotype the D and J loci too — important since J anchors and D/J usage drive +haplotype-inference methods: + +| Method | Signature | Notes | +|---|---|---| +| `homozygous` | `(gene, allele, segment="V")` | one allele on both chromosomes | +| `heterozygous` | `(gene, allele0, allele1, segment="V")` | one allele per chromosome | +| `delete_gene` | `(gene, haplotype="both"\|0\|1, segment="V")` | whole-gene or one-chromosome (hemizygous) deletion | +| `duplicate_gene` | `(gene, alleles=[...], haplotype=0\|1, segment="V")` | >1 copy on one chromosome | +| `add_novel_allele` | `(name, *, base, mutations\|sequence, segment="V", allow_nonfunctional=False)` | define a private allele (see below) | +| `chromosome_weights` | `(w0, w1)` | allelic-expression imbalance (default 0.5/0.5) | +| `with_subject` | `(sid)` | provenance label stamped on every record | +| `complete_from_reference` | `(policy="homozygous_first_reference"\|"heterozygous_first_two")` | fill unspecified genes | + +```python +# Genotype the J locus too — e.g. heterozygous IGHJ6 + a homozygous IGHJ4: +g.heterozygous("IGHJ6", "IGHJ6*02", "IGHJ6*03", segment="J") +g.homozygous("IGHJ4", "IGHJ4*02", segment="J") +``` + +Notes and guard-rails: + +- **`delete_gene(..., haplotype=0|1)`** (one chromosome) requires the gene to be + specified first — deleting a single haplotype of an *unspecified* gene would + silently delete both, so it raises instead. +- **`complete_from_reference(policy=...)`** fills only genes you haven't touched. + `"homozygous_first_reference"` (default) makes each unspecified gene homozygous + for its first cartridge allele — note this is the *first listed* allele, **not** + a population-frequency-common one (GenAIRR has no frequency prior in this + release; the name says exactly what it does). `"heterozygous_first_two"` uses + the first two alleles. +- Unknown gene/allele names, NaN/inf chromosome weights, and segments left with no + usable allele are all rejected at build/attach with clear messages. +- `with_genotype` is **mutually exclusive** with + [`restrict_alleles`](../reference/experiment.md) and the + `recombine(*_allele_weights=...)` kwargs — the genotype owns allele presence and + expression. It is also rejected together with `receptor_revision` and with the + clonal forks (`expand_clones` / `clonal_lineage` / `clonal_repertoire`) in this + release (see [Limitations](#limitations-this-release)). + +### Gene usage + +Within a chromosome, which *gene* is used is weighted by the cartridge's typed +allele-usage model (`reference_models.allele_usage`), aggregated to the gene +level, and scaled by **copy-number dosage** (a duplicated gene recombines +proportionally more often). Cartridges that don't author a typed `allele_usage` +fall back to uniform-over-present-genes (× dosage). See +[Allele usage](v-usage.md) and [Estimate models from data](estimate-cartridge-models.md) +for authoring usage. + +### More genotype recipes + +**A richer diploid genotype** — several heterozygous genes, a homozygous gene, +a whole-gene (homozygous) deletion, a hemizygous deletion, and allelic-expression +imbalance, with everything else filled from the reference: + +```python +g = ( + Genotype.from_dataconfig(cfg) + .heterozygous("IGHVF1-G1", "IGHVF1-G1*01", "IGHVF1-G1*02") + .heterozygous("IGHVF1-G2", "IGHVF1-G2*01", "IGHVF1-G2*02") + .homozygous("IGHVF2-G4", "IGHVF2-G4*01") + .delete_gene("IGHVF3-G7", haplotype="both") # absent on both chromosomes + .homozygous("IGHVF3-G8", "IGHVF3-G8*01") # carried on both... + .delete_gene("IGHVF3-G8", haplotype=1) # ...then removed on chr 1 (hemizygous) + .chromosome_weights(0.65, 0.35) # chromosome 0 expressed more + .complete_from_reference() # the remaining genes + .with_subject("DONOR_A") +) +``` + +**Gene duplication** — one chromosome carries two alleles of the same gene +(specify the gene on both chromosomes first, then add the extra copy to one): + +```python +g = ( + Genotype.from_dataconfig(cfg) + .homozygous("IGHVF1-G3", "IGHVF1-G3*01") # both chromosomes carry *01 + .duplicate_gene("IGHVF1-G3", ["IGHVF1-G3*01", "IGHVF1-G3*02"], haplotype=0) # chr 0 now carries two copies + .complete_from_reference() + .with_subject("DONOR_DUP") +) +# chromosome 0 carries {*01, *02}, chromosome 1 carries {*01}; +# the extra copy raises this gene's recombination share (copy-number dosage). +``` + +**Build a fully-specified strict genotype programmatically** — drive the builder +from a per-gene plan (the natural shape if you load a genotype from a table or +generate many subjects): + +```python +plan = { + "IGHVF1-G1": ("IGHVF1-G1*01", "IGHVF1-G1*02"), # 2 alleles -> heterozygous + "IGHVF1-G2": ("IGHVF1-G2*01",), # 1 allele -> homozygous + "IGHVF3-G7": (), # 0 alleles -> deleted + # ... one entry per gene you want to pin +} + +g = Genotype.from_dataconfig(cfg) +for gene, alleles in plan.items(): + if not alleles: + g.delete_gene(gene, haplotype="both") + elif len(alleles) == 1: + g.homozygous(gene, alleles[0]) + else: + g.heterozygous(gene, alleles[0], alleles[1]) +g.complete_from_reference().with_subject("DONOR_B") +``` + +**Inspect the non-trivial genes** of any genotype: + +```python +for row in g.to_table(): + if row["zygosity"] != "homozygous": + print(row["gene"], row["zygosity"], row["haplotype_0"], row["haplotype_1"]) +``` + +## Novel / private alleles + +Individuals carry germline alleles that aren't in any reference — *private* or +*novel* alleles. Discovering them is a central task for IgDiscover, partis, and +TIgGER's `findNovelAlleles`. GenAIRR can plant them as ground truth. + +`add_novel_allele` derives a private allele from a reference **base** allele by +applying point `mutations` (or supplying an explicit `sequence` of the same +length), inheriting the base's gene, anchor, functional status and V sub-regions. +The novel allele is then placed like any allele, and at `compile()` time it is +injected into an **effective reference** (base catalogue + your private alleles) +so it flows through alignment and AIRR output as a genuine allele: + +```python +g = ( + Genotype.from_dataconfig(cfg) + .add_novel_allele("IGHVF1-G1*i01", base="IGHVF1-G1*01", + mutations=[(38, "C"), (41, "A")]) # two point variants + .complete_from_reference() + .heterozygous("IGHVF1-G1", "IGHVF1-G1*01", "IGHVF1-G1*i01") # one reference + one private + .with_subject("DONOR_N") +) + +result = ( + ga.Experiment.on(cfg).with_genotype(g).recombine() + .run_records(n=500, seed=3, expose_provenance=True) +) +# The private allele is sampled, assembled and reported like any allele — +# its name appears in v_call / truth_v_call and the reads carry its variants. +``` + +The novel allele's **gene is taken from its name** and must match the base +allele's gene; it must be a same-length (substitution-only) variant. The +synthesized coding sequence is **validated** — for V/J the conserved anchor codon +must still encode the conserved residue (Cys for V, Trp/Phe for J) and the coding +frame must be stop-free. A variant that breaks either is rejected unless you pass +`allow_nonfunctional=True` (then it is kept and marked non-functional). Novel +alleles are flagged in the ground truth: each `to_table()`/`to_tsv()` row carries +a `novel` list of the private alleles carried at that gene. + +**Benchmarking novel-allele discovery.** Plant a novel allele, simulate, then run +the discovery tool against the **base** germline (the cartridge *without* your +private alleles) so the tool must rediscover it from the reads — and score its +output against the planted novel sequence. (Write the base germline FASTA from +`cfg.v_alleles`; write the truth from `genotype.to_table()`.) + +## Ground truth and provenance + +A genotype experiment emits, by construction, everything an evaluation needs: + +- **Per-record fields:** `subject_id` and `haplotype` (`0`/`1`, the chromosome the + rearrangement used) are stamped on every AIRR record. Standard truth columns + (`truth_v_call`, …) are available with `expose_provenance=True`. +- **`result.genotypes`:** the list of attached `Genotype` objects (one per subject). +- **`Genotype.to_table()` / `to_tsv(path)`:** the ground-truth genotype as a table + — one row per (segment, gene) with `zygosity` + (`homozygous` / `heterozygous` / `hemizygous` / `deleted`), the carried alleles + per haplotype, and per-haplotype `allele:copies:weight` detail. This is the + reference a genotype-inference benchmark compares against. + +```python +for row in result.genotypes[0].to_table(): + if row["zygosity"] != "homozygous": # show the interesting genes + print(row["gene"], row["zygosity"], row["haplotype_0"], row["haplotype_1"]) +``` + +## Research workflow: benchmarking genotype inference + +The point of simulating from a *known* genotype is that you can run a +genotype-inference tool on the resulting repertoire and score it against the +planted truth — with no real-data uncertainty about what the right answer is. + +The recipe is the same for any tool: + +1. Build a `Genotype`, simulate a repertoire, write the AIRR table + (`result.to_tsv(...)`) and/or reads FASTA, and the ground truth + (`genotype.to_tsv(...)`). +2. Run the inference tool to recover the per-individual allele set. +3. Compare recovered vs planted: presence/absence, zygosity, and (for + discovery tools) any novel alleles. + +### Worked example: TIgGER and IgDiscover recover a planted genotype + +To show this end to end we planted a diploid IGH genotype in `human_igh` — +**3 heterozygous** V genes (two alleles each), **3 homozygous** (one allele), +and **3 fully deleted** genes — and filled the rest from the reference. We +simulated 4,000 reads with light SHM, then ran two independent AIRR +genotype-inference tools on the result: +[**TIgGER**](https://tigger.readthedocs.io) (Immcantation; consumes the AIRR +table) and [**IgDiscover**](https://igdiscover.se) (germline discovery from the +raw reads, with its own IgBLAST). + +Because GenAIRR already emits AIRR records with `v_call` **and** +`sequence_alignment`, TIgGER's `inferGenotype` consumes the rearrangement table +**directly — no separate IgBLAST step is needed**: + +```r +library(tigger); library(airr) +rep <- read_rearrangement("repertoire.tsv") # GenAIRR's AIRR output +germ_v <- readIgFasta("germline_V.fasta") # cartridge V germline (names match v_call) +geno <- inferGenotype(rep, germline_db = germ_v, find_unmutated = TRUE) +plotGenotype(geno) +``` + +TIgGER recovered the planted genotype **exactly**: every heterozygous gene → two +alleles, every homozygous gene → one, every deleted gene → **absent**. Across all +52 V genes, allele-presence **precision = 1.00**, **recall = 1.00**, and the +per-gene allele count matched the truth for **52/52** genes. + +**IgDiscover**, run on the raw reads with the cartridge as its starting database, +independently agreed: **precision = 1.00** (zero false-positive alleles), +**recall = 0.96** (50/52 carried alleles), with **all three deletions correct** +and **all heterozygous genes fully resolved** (both alleles recovered). The two +missed alleles were low-expression single-copy genes below IgDiscover's default +expression threshold — a tool-tuning matter, not a simulation artefact. + +![GenAIRR-simulated genotype recovered by TIgGER and IgDiscover: planted vs inferred allele counts agree for every gene](../assets/genotype-tigger-recovery.png) + +*(A) The nine study genes: both tools' inferred allele counts match the planted +zygosity for each (heterozygous → 2, homozygous → 1, deleted → 0). (B) All 52 V +genes fall on the agreement diagonal; presence precision = 1.00 for both tools, +recall 1.00 (TIgGER) / 0.96 (IgDiscover), zero false-positive alleles, all +deletions correct.* + +### Reproduce it + +This builds the **exact** genotype behind the figure — 3 heterozygous, 3 +homozygous, and 3 deleted study V genes, the rest filled from the reference — +simulates 4,000 reads with light SHM at `seed=7`, and writes every input the two +tools need plus the ground truth to score against: + +```python +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype + +cfg = gdata.HUMAN_IGH_OGRDB +HET = ["IGHVF1-G1", "IGHVF1-G2", "IGHVF1-G3"] # 2 alleles each +HOM = ["IGHVF2-G4", "IGHVF3-G5", "IGHVF3-G6"] # 1 allele +DEL = ["IGHVF3-G7", "IGHVF3-G8", "IGHVF3-G9"] # deleted (both chromosomes) + +g = Genotype.from_dataconfig(cfg).complete_from_reference("homozygous_first_reference") +for gene in HET: + a0, a1 = (a.name for a in cfg.v_alleles[gene][:2]) + g.heterozygous(gene, a0, a1) +for gene in HOM: + g.homozygous(gene, cfg.v_alleles[gene][0].name) +for gene in DEL: + g.delete_gene(gene, haplotype="both") +g.with_subject("DONOR01") + +res = ( + ga.Experiment.on(cfg).with_genotype(g).recombine() + .mutate(rate=0.004) # light SHM, as in real data + .run_records(n=4000, seed=7, expose_provenance=True) +) + +res.to_tsv("repertoire.tsv") # AIRR table → TIgGER +g.to_tsv("truth_genotype.tsv") # ground truth to score against + +with open("reads.fasta", "w") as fh: # raw reads → IgDiscover / partis + for r in res: + fh.write(f">{r['sequence_id']}\n{r['sequence'].upper()}\n") + +with open("germline_V.fasta", "w") as fh: # cartridge V germline (names match v_call) + for gene, alleles in cfg.v_alleles.items(): + for a in alleles: + fh.write(f">{a.name}\n{a.ungapped_seq.upper()}\n") +``` + +**Score it.** Run TIgGER (R snippet above) on `repertoire.tsv`, or IgDiscover on +`reads.fasta` with the cartridge as its starting database +(`igdiscover init --database db/ --single-reads reads.fasta project/ && cd project +&& igdiscover run`). Then compare each tool's per-gene allele set against +`g.to_table()` (the planted truth): allele-presence precision/recall, zygosity, +and deletion calls. With the genotype above this yields TIgGER precision/recall +1.00 (52/52 genes) and IgDiscover precision 1.00 / recall 0.96 — the figure. + +### Running other tools on the same data + +The only difference between tools is whether they consume the **AIRR table** +(TIgGER) or the **raw reads** (`reads.fasta`, which you can write from `result`), +running their own aligner: + +- **[IgDiscover](https://igdiscover.se)** — germline *discovery* from reads (its + own IgBLAST + iterative filtering). Initialise with the cartridge germline as + the starting database and the simulated reads, then run the pipeline; the + `final/database/V.fasta` expressed-allele set is the recovered genotype: + + ```bash + igdiscover init --database db/ --single-reads reads.fasta project/ + cd project && igdiscover run + ``` + +- **[partis](https://github.com/psathyrella/partis)** — HMM annotation with + per-sample germline inference (`partis cache-parameters --infname reads.fa + --initial-germline-dir db/`). partis also reports per-sample allele support and + novel alleles, scored the same way. + +Because the genotype is planted, every tool is scored identically: recovered +allele set vs `genotype.to_table()` — presence precision/recall, zygosity, and +deletion calls. + +## Limitations (this release) + +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 + cartridge. +- **Same-haplotype receptor revision** — `receptor_revision` with a genotype is + rejected for now. + +## Backward compatibility + +The genotype machinery is purely additive. An experiment with **no** genotype +attached produces byte-identical output to previous releases (pinned by a +checksum test). Attaching a genotype is the only thing that switches recombination +onto the phased path. diff --git a/src/GenAIRR/__init__.py b/src/GenAIRR/__init__.py index 9602c01..40aa089 100644 --- a/src/GenAIRR/__init__.py +++ b/src/GenAIRR/__init__.py @@ -27,6 +27,7 @@ # The simulation entry point. from .experiment import CompiledExperiment, Experiment, dataconfig_to_refdata +from .genotype import Genotype from .result import FamilyValidationReport, SimulationResult, ValidationReport from ._validation import FamilyValidationFailedError, RecordValidationFailedError diff --git a/src/GenAIRR/_compile.py b/src/GenAIRR/_compile.py index 7f4e96a..9fa00e1 100644 --- a/src/GenAIRR/_compile.py +++ b/src/GenAIRR/_compile.py @@ -173,6 +173,111 @@ def _extract_receptor_revision_prob(steps): return revision_prob, filtered +def _name_to_id(refdata, segment): + """Map allele name -> pool id for a segment against this refdata.""" + size = { + "V": refdata.v_pool_size, + "D": refdata.d_pool_size, + "J": refdata.j_pool_size, + }[segment]() + getter = { + "V": refdata.v_allele, + "D": refdata.d_allele, + "J": refdata.j_allele, + }[segment] + return {getter(i).name: i for i in range(size)} + + +def _genotype_segment_rows(genotype, refdata, segment): + """Resolve a genotype's per-haplotype slots for a segment into flat + ``(haplotype, allele_id, copies, weight)`` rows for the engine. + + Strict genotypes require every gene present in the cartridge to be + specified (call ``complete_from_reference()``); permissive genotypes + fill unspecified genes with all reference alleles single-copy on both + haplotypes (a NON-diploid fallback).""" + name_to_id = _name_to_id(refdata, segment) + cfg_by_gene = { + "V": genotype._cfg.v_alleles, + "D": genotype._cfg.d_alleles, + "J": genotype._cfg.j_alleles, + }[segment] or {} + rows = [] + for gene, allele_objs in cfg_by_gene.items(): + slot = genotype._slots[segment].get(gene) + if slot is None: + if genotype.is_permissive: + for h in (0, 1): + for a in allele_objs: + if a.name in name_to_id: + rows.append((h, name_to_id[a.name], 1, 1.0)) + continue + raise ValueError( + f"genotype is strict but {segment} gene {gene!r} is unspecified; " + f"call complete_from_reference() or specify it before with_genotype()" + ) + for h, copies in enumerate(slot): + for (allele_name, copy_count, weight) in copies: + rows.append((h, name_to_id[allele_name], int(copy_count), float(weight))) + return rows + + +def _genotype_presence_ok(v_rows, d_rows, j_rows, d_required): + """True iff at least one chromosome carries every required segment — + i.e. a phased rearrangement is possible. A genotype where (say) hap0 + has V-only and hap1 has J-only has non-empty union support but no + viable haplotype, which would otherwise panic at runtime.""" + def has(rows, h): + return any(r[0] == h for r in rows) + for h in (0, 1): + if has(v_rows, h) and has(j_rows, h) and (not d_required or has(d_rows, h)): + return True + return False + + +def _push_genotype_recombine(genotype, step, plan, refdata, *, d_required): + """Push the single phased ``SampleGenotypePass`` for an attached + genotype, replacing the flat per-segment allele sampling. Cartridge + allele-usage weights (resolved onto ``step``) are passed through and + aggregated to gene-level usage by the engine.""" + v_rows = _genotype_segment_rows(genotype, refdata, "V") + j_rows = _genotype_segment_rows(genotype, refdata, "J") + d_rows = _genotype_segment_rows(genotype, refdata, "D") if d_required else [] + if not _genotype_presence_ok(v_rows, d_rows, j_rows, d_required): + raise ValueError( + "genotype has no complete haplotype: every chromosome is missing at " + "least one of the required segments (V/" + + ("D/" if d_required else "") + + "J), so no phased rearrangement is possible" + ) + # Gene usage is driven by the cartridge's TYPED allele-usage plane + # (``reference_models.allele_usage``), resolved onto ``step.weights_*`` + # by recombine(). NOTE: bundled configs that don't author a typed + # allele_usage leave these ``None`` here, so gene choice for those is + # uniform-over-present-genes (× copy dosage). The legacy per-gene + # usage dict is intentionally NOT consulted (mirrors recombine()'s + # precedence chain). + v_weights = list(step.weights_v) if step.weights_v is not None else None + d_weights = list(step.weights_d) if step.weights_d is not None else None + j_weights = list(step.weights_j) if step.weights_j is not None else None + plan.push_genotype_recombine( + refdata, + ( + float(genotype._chromosome_weights[0]), + float(genotype._chromosome_weights[1]), + ), + genotype.subject_id, + genotype._source_hash, + v_rows, + d_rows, + j_rows, + d_required, + v_weights, + d_weights, + j_weights, + ) + + def _lower_recombine( step: _RecombineStep, plan: "_engine.PassPlan", @@ -180,6 +285,7 @@ def _lower_recombine( *, invert_d_prob=None, receptor_revision_prob=None, + genotype=None, ) -> None: chain = refdata.chain_type np1 = list(step.np1_lengths) @@ -238,8 +344,11 @@ def _lower_recombine( "receptor_revision is only valid for VDJ chains; the " "DSL boundary should have rejected this earlier." ) - plan.push_sample_allele("V", refdata, allowed_ids=v_ids, weights=v_weights) - plan.push_sample_allele("J", refdata, allowed_ids=j_ids, weights=j_weights) + if genotype is not None: + _push_genotype_recombine(genotype, step, plan, refdata, d_required=False) + else: + plan.push_sample_allele("V", refdata, allowed_ids=v_ids, weights=v_weights) + plan.push_sample_allele("J", refdata, allowed_ids=j_ids, weights=j_weights) if step.trim_v_3: plan.push_trim("V", "3", list(step.trim_v_3)) if step.trim_j_5: @@ -257,9 +366,12 @@ def _lower_recombine( plan.push_p_addition("J_5", p_j_5) plan.push_assemble("J") elif chain == "vdj": - plan.push_sample_allele("V", refdata, allowed_ids=v_ids, weights=v_weights) - plan.push_sample_allele("D", refdata, allowed_ids=d_ids, weights=d_weights) - plan.push_sample_allele("J", refdata, allowed_ids=j_ids, weights=j_weights) + if genotype is not None: + _push_genotype_recombine(genotype, step, plan, refdata, d_required=True) + else: + plan.push_sample_allele("V", refdata, allowed_ids=v_ids, weights=v_weights) + plan.push_sample_allele("D", refdata, allowed_ids=d_ids, weights=d_weights) + plan.push_sample_allele("J", refdata, allowed_ids=j_ids, weights=j_weights) if step.trim_v_3: plan.push_trim("V", "3", list(step.trim_v_3)) if step.trim_d_5: diff --git a/src/GenAIRR/_compiled.py b/src/GenAIRR/_compiled.py index 9c53d6d..e6684e0 100644 --- a/src/GenAIRR/_compiled.py +++ b/src/GenAIRR/_compiled.py @@ -39,7 +39,14 @@ class CompiledExperiment: time; ``run()`` only accepts execution parameters. """ - __slots__ = ("_simulator", "_refdata", "_steps", "_dataconfig", "_metadata") + __slots__ = ( + "_simulator", + "_refdata", + "_steps", + "_dataconfig", + "_metadata", + "_genotype", + ) def __init__( self, @@ -48,6 +55,7 @@ def __init__( steps: Sequence[Any] = (), dataconfig: Optional["DataConfig"] = None, metadata: Optional[Dict[str, Any]] = None, + genotype: Optional[Any] = None, ) -> None: self._simulator = simulator self._refdata = refdata @@ -58,6 +66,9 @@ def __init__( self._steps: Tuple[Any, ...] = tuple(steps) self._dataconfig = dataconfig self._metadata = dict(metadata) if metadata else {} + # Attached single-subject genotype (or None). When set, + # run_records stamps subject_id + haplotype provenance. + self._genotype = genotype @property def simulator(self) -> "_engine.CompiledSimulator": @@ -187,12 +198,25 @@ def run_records( result = SimulationResult.from_outcomes( outcomes, self._refdata, expose_provenance=expose_provenance ) + if self._genotype is not None: + self._stamp_genotype_provenance(outcomes, result) if validate_records: from ._validation import _raise_on_validation_failure _raise_on_validation_failure(result.validate_records(self._refdata)) return result + def _stamp_genotype_provenance(self, outcomes, result) -> None: + """Stamp per-record ``subject_id`` + ``haplotype`` (the chromosome + the rearrangement drew from) and expose the genotype on the + result. Used only when a genotype is attached.""" + subject = self._genotype.subject_id + for outcome, rec in zip(outcomes, result._records): + rec["subject_id"] = subject + hap = outcome.trace().find("sample_haplotype") + rec["haplotype"] = hap.value if hap is not None else None + result._genotypes = [self._genotype] + def stream( self, *, diff --git a/src/GenAIRR/experiment.py b/src/GenAIRR/experiment.py index 88cee7a..7c9a61d 100644 --- a/src/GenAIRR/experiment.py +++ b/src/GenAIRR/experiment.py @@ -429,6 +429,8 @@ class Experiment: "_metadata", "_contracts", "_allow_curatable_refdata", + "_genotype", + "_user_allele_weights_set", ) def __init__( @@ -474,6 +476,15 @@ def __init__( # catalogue (bundled mouse_igh / human_tcrb) that includes # pseudogene/ORF alleles. self._allow_curatable_refdata: bool = False + # Single-subject diploid genotype attached via ``with_genotype``. + # ``None`` => the flat (uniform/usage-weighted) allele path runs + # unchanged. When set, recombination lowers to the phased + # genotype path (one ``SampleGenotypePass``). + self._genotype = None + # True once the user passed an explicit ``*_allele_weights`` to + # ``recombine`` — distinct from cartridge-usage defaults. Used to + # enforce mutual exclusion with ``with_genotype``. + self._user_allele_weights_set: bool = False @classmethod def on(cls, source: ExperimentInput) -> "Experiment": @@ -1604,6 +1615,44 @@ def _is_tcr_refdata(self) -> bool: first_v_name = self._refdata.v_allele(0).name return first_v_name.upper().startswith("TR") + def with_genotype(self, genotype) -> "Experiment": + """Attach a single-subject diploid genotype. + + With a genotype attached, V(D)J recombination becomes + haplotype-phased: V, D and J of each rearrangement are drawn from + a single chromosome, honouring the genotype's allele + presence/absence, zygosity, and copy-number/deletion. With no + genotype, the flat (uniform / usage-weighted) path runs unchanged. + + Mutually exclusive with :meth:`restrict_alleles` and the + ``recombine(*_allele_weights=...)`` kwargs — the genotype owns + allele presence and within-gene expression. + + Raises ``ValueError`` if the genotype was built against a + different cartridge (content-hash mismatch), or if allele locks / + explicit allele weights are already set. + """ + live_hash = self._refdata.content_hash() + if genotype._source_hash != live_hash: + raise ValueError( + "genotype was built against a different cartridge (content hash " + f"{genotype._source_hash!r} != experiment {live_hash!r})" + ) + if any(v is not None for v in self._locks.values()): + raise ValueError( + "with_genotype() and restrict_alleles() are mutually exclusive" + ) + if self._user_allele_weights_set: + raise ValueError( + "with_genotype() and recombine(*_allele_weights=...) are mutually " + "exclusive: the genotype owns allele expression" + ) + # Snapshot the (mutable) builder so later edits to ``genotype`` + # cannot desync the compiled engine genotype from + # ``result.genotypes`` (review #8). + self._genotype = genotype._snapshot() + return self + def restrict_alleles( self, *, @@ -1640,6 +1689,11 @@ def restrict_alleles( a VJ chain. - ``TypeError`` if an unexpected input shape is passed. """ + if self._genotype is not None: + raise ValueError( + "restrict_alleles() and with_genotype() are mutually exclusive: " + "a genotype already owns allele presence and within-gene expression" + ) for segment, value in (("V", v), ("D", d), ("J", j)): if value is _UNSET: continue @@ -1770,6 +1824,21 @@ def recombine( ``ValueError`` for unknown allele names or non-positive weights. """ + # Explicit allele weights conflict with an attached genotype: + # the genotype owns allele presence + within-gene expression, and + # the phased lowering ignores recombine-step weights. Reject the + # combination instead of silently dropping the weights. + if any( + w is not None + for w in (v_allele_weights, d_allele_weights, j_allele_weights) + ): + self._user_allele_weights_set = True + if self._genotype is not None: + raise ValueError( + "recombine(*_allele_weights=...) and with_genotype() are " + "mutually exclusive: the genotype owns allele expression" + ) + # VJ chains have no NP2 region — surface user mistakes loudly # instead of silently dropping the argument. if np2_lengths is not None and self._refdata.chain_type != "vdj": @@ -2541,6 +2610,30 @@ def compile(self, *, allow_curatable_refdata: Optional[bool] = None): """ if allow_curatable_refdata is None: allow_curatable_refdata = self._allow_curatable_refdata + + # Receptor revision is not supported alongside a phased genotype + # in this release: the revision pass samples a replacement V from + # its own distribution with no chromosome/carried-allele + # awareness. Reject the combination (same-haplotype revision is a + # planned follow-on). + if self._genotype is not None and any( + isinstance(s, _ReceptorRevisionStep) for s in self._steps + ): + raise ValueError( + "receptor_revision() is not supported with with_genotype() in this " + "release (the revision pass is not haplotype-aware)" + ) + + # Genotype provenance (subject_id / haplotype / result.genotypes) + # is only threaded through the plain compiled path, not the + # clonal/lineage/repertoire forked classes. Reject the + # combination rather than silently dropping provenance (review + # #9); genotype + clonal cohorts are a planned follow-on. + if self._genotype is not None and self._has_clonal_fork(): + raise ValueError( + "with_genotype() is not supported together with expand_clones() / " + "clonal_lineage() / clonal_repertoire() in this release" + ) from dataclasses import replace as _replace # On raw RefDataConfig with default-on trim, warn at compile @@ -2744,19 +2837,37 @@ def compile(self, *, allow_curatable_refdata: Optional[bool] = None): metadata=self._metadata, ) + # When the attached genotype defines novel/private alleles, compile + # against an *effective* reference = base catalogue + injected novel + # alleles, so they become real pool entries the engine samples, + # assembles, and reports like any allele. No genotype, or a genotype + # without novel alleles, uses the base refdata unchanged. + effective_refdata = self._refdata + if self._genotype is not None and self._genotype.has_novel(): + if self._dataconfig is None: + raise ValueError( + "genotype with novel alleles requires a DataConfig-backed " + "experiment (Experiment.on(dataconfig), not a raw RefDataConfig)" + ) + effective_refdata = dataconfig_to_refdata( + self._genotype.effective_dataconfig() + ) + simulator = self._build_simulator( self._steps, contracts, any_lock, replace_fn=_replace, allow_curatable_refdata=allow_curatable_refdata, + refdata=effective_refdata, ) return CompiledExperiment( simulator, - self._refdata, + effective_refdata, steps=tuple(self._steps), dataconfig=self._dataconfig, metadata=self._metadata, + genotype=self._genotype, ) def _build_simulator( @@ -2767,10 +2878,16 @@ def _build_simulator( *, replace_fn, allow_curatable_refdata: bool = False, + refdata=None, ): """Compile a list of steps into a `GenAIRR._engine.CompiledSimulator`. Lifted out of `compile()` so the clonal-fork branch can build - two simulators from sub-step-lists with a shared body.""" + two simulators from sub-step-lists with a shared body. + + ``refdata`` overrides ``self._refdata`` — used when a genotype with + novel alleles compiles against an *effective* reference (base + + injected private alleles).""" + refdata = refdata if refdata is not None else self._refdata plan = _engine.PassPlan() # Pull the (at-most-one) `_InvertDStep` out of the step # sequence and thread its probability into the recombine @@ -2812,12 +2929,13 @@ def _build_simulator( _lower_recombine( step, plan, - self._refdata, + refdata, invert_d_prob=invert_d_prob, receptor_revision_prob=receptor_revision_prob, + genotype=self._genotype, ) else: - lower_step(step, plan, self._refdata) + lower_step(step, plan, refdata) # Paired-end is sequencing-stage / readout-stage: lower # it AFTER every biology + corruption pass so the trace # records land last. See `_extract_paired_end_step` for @@ -2825,7 +2943,7 @@ def _build_simulator( if paired_end_step is not None: _lower_paired_end(paired_end_step, plan) return plan.compile( - refdata=self._refdata, + refdata=refdata, respect=contracts, allow_curatable_refdata=allow_curatable_refdata, ) diff --git a/src/GenAIRR/genotype.py b/src/GenAIRR/genotype.py new file mode 100644 index 0000000..89213ea --- /dev/null +++ b/src/GenAIRR/genotype.py @@ -0,0 +1,481 @@ +"""Per-individual diploid genotype builder (PR1: known reference alleles). + +A :class:`Genotype` is an editable, narrowed view of a ``DataConfig``'s +reference alleles. Attach it to an experiment with +``Experiment.with_genotype(g)`` to make V(D)J recombination +haplotype-phased: V, D and J of each rearrangement are drawn from a single +chromosome, honouring presence/absence, zygosity, and copy-number/deletion. + +Default construction (:meth:`Genotype.from_dataconfig`) is **strict**: a gene +that is used during recombination but never assigned here is an error at +attach time. Use :meth:`complete_from_reference` to fill unspecified genes +with a valid diploid state. :meth:`Genotype.permissive` is a separate, +explicitly non-diploid fallback (see its docstring). +""" +from __future__ import annotations + +from typing import Dict, List, Optional, Set, Tuple + +_SEGMENTS = ("V", "D", "J") + + +def _alleles_by_gene(cfg, segment: str) -> Dict[str, List]: + return { + "V": cfg.v_alleles, + "D": cfg.d_alleles, + "J": cfg.j_alleles, + }[segment] or {} + + +class Genotype: + """A diploid genotype over a ``DataConfig``'s reference alleles.""" + + def __init__(self, cfg, *, permissive: bool = False): + self._cfg = cfg + self._permissive = bool(permissive) + self.subject_id: Optional[str] = None # plain attribute (read anywhere) + self._chromosome_weights: Tuple[float, float] = (0.5, 0.5) + # segment -> gene -> [hap0 list[(allele, copies, weight)], hap1 ...] + self._slots: Dict[str, Dict[str, List[List[Tuple[str, int, float]]]]] = { + s: {} for s in _SEGMENTS + } + # Private/novel alleles defined on this individual: + # name -> {"allele": , "gene": str, "segment": str, + # "base": str, "mutations": list} + self._novel: Dict[str, Dict] = {} + self._source_hash: str = cfg.cartridge_manifest()["hashes"]["refdata_content_hash"] + + # ── constructors ────────────────────────────────────────────── + @classmethod + def from_dataconfig(cls, cfg) -> "Genotype": + """A strict genotype: unspecified-but-used genes error at attach.""" + return cls(cfg, permissive=False) + + @classmethod + def permissive(cls, cfg) -> "Genotype": + """NOT a biological diploid genotype. A reference-wide fallback: + unspecified genes sample over ALL reference alleles WITHOUT + phasing for those genes. Use only to constrain a few genes; never + mistake this for 'complete genotype from refdata'.""" + return cls(cfg, permissive=True) + + # ── editing ─────────────────────────────────────────────────── + 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)}" + ) + 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 self + + def _check_allele(self, segment: str, gene: str, allele: str) -> None: + by_gene = _alleles_by_gene(self._cfg, segment) + if gene not in by_gene: + raise ValueError(f"{gene!r} is not a known {segment} gene in this cartridge") + names = {a.name for a in by_gene[gene]} + # also accept novel alleles defined on this genotype for this gene + names |= { + n + for n, info in self._novel.items() + if info["gene"] == gene and info["segment"] == segment + } + if allele not in names: + raise ValueError( + f"{allele!r} is not a known allele of {gene!r} " + f"(define novel alleles with add_novel_allele first)" + ) + + def homozygous(self, gene: str, allele: str, segment: str = "V") -> "Genotype": + self._check_allele(segment, gene, allele) + self._slots[segment][gene] = [[(allele, 1, 1.0)], [(allele, 1, 1.0)]] + return self + + def heterozygous( + self, gene: str, allele0: str, allele1: str, segment: str = "V" + ) -> "Genotype": + self._check_allele(segment, gene, allele0) + self._check_allele(segment, gene, allele1) + self._slots[segment][gene] = [[(allele0, 1, 1.0)], [(allele1, 1, 1.0)]] + return self + + def delete_gene(self, gene: str, haplotype="both", segment: str = "V") -> "Genotype": + if haplotype not in ("both", 0, 1): + raise ValueError( + f"haplotype must be 'both', 0, or 1, got {haplotype!r}" + ) + # One-haplotype (hemizygous) deletion requires the gene to be + # specified first, otherwise the *other* haplotype would also be + # empty — silently producing a full deletion that + # complete_from_reference() then skips. Full ("both") deletion of + # an unspecified gene is fine. + if haplotype != "both" and gene not in self._slots[segment]: + raise ValueError( + f"specify {segment} gene {gene!r} (homozygous/heterozygous) before " + f"deleting one haplotype; deleting a single haplotype of an " + f"unspecified gene would delete both" + ) + cur = self._slots[segment].get(gene, [[], []]) + # copy to avoid aliasing if the slot was shared + cur = [list(cur[0]), list(cur[1])] + if haplotype in ("both", 0): + cur[0] = [] + if haplotype in ("both", 1): + cur[1] = [] + self._slots[segment][gene] = cur + return self + + def duplicate_gene( + self, gene: str, alleles: List[str], haplotype: int, segment: str = "V" + ) -> "Genotype": + if haplotype not in (0, 1): + raise ValueError(f"haplotype must be 0 or 1, got {haplotype!r}") + for a in alleles: + self._check_allele(segment, gene, a) + cur = self._slots[segment].get(gene, [[], []]) + cur = [list(cur[0]), list(cur[1])] + cur[haplotype] = [(a, 1, 1.0) for a in alleles] + self._slots[segment][gene] = cur + return self + + # ── novel / private alleles ─────────────────────────────────── + def _find_ref_allele(self, segment: str, name: str): + for alleles in _alleles_by_gene(self._cfg, segment).values(): + for a in alleles: + if a.name == name: + return a + return None + + def add_novel_allele( + self, + name: str, + *, + base: str, + mutations: Optional[List[Tuple[int, str]]] = None, + sequence: Optional[str] = None, + segment: str = "V", + allow_nonfunctional: bool = False, + ) -> "Genotype": + """Define a private/novel allele not present in the reference. + + Derive it from a reference ``base`` allele by either applying point + ``mutations`` (a list of ``(0-based position, base)``) or supplying + an explicit same-length ``sequence`` (substitutions only — no + indels — so the gene's reading frame and the conserved-anchor + position are preserved). The novel allele's **gene is taken from + its name** and must equal the base allele's gene; sub-regions and + the anchor position are inherited (valid for same-length variants). + + The synthesized coding sequence is **validated**: for V/J the + conserved anchor codon must still encode the conserved residue + (Cys for V, Trp/Phe for J) and the coding frame must contain no + internal stop codon. A variant that breaks either is rejected + unless ``allow_nonfunctional=True`` (in which case it is kept and + marked non-functional). + + Registers the novel allele under ``name`` so it can then be placed + with :meth:`homozygous` / :meth:`heterozygous` / :meth:`duplicate_gene`. + At compile time it is injected as a real entry in an *effective* + reference, so it flows through alignment and AIRR output like a + catalogue allele. + """ + import copy as _copy + + from .utilities.misc import translate + + base_allele = self._find_ref_allele(segment, base) + if base_allele is None: + raise ValueError(f"base allele {base!r} not found in {segment} reference") + # Gene identity comes from the name; it must match the base's gene + # (no cross-gene synthesis — that would corrupt gene identity). + name_gene = name.split("*")[0] + if name_gene != base_allele.gene: + raise ValueError( + f"novel name {name!r} implies gene {name_gene!r} but base {base!r} " + f"belongs to gene {base_allele.gene!r}; a novel allele must belong " + f"to its base allele's gene" + ) + gene = base_allele.gene + # Name must be unique across the WHOLE catalogue (all segments) and + # all previously-defined novel alleles. + if name in self._novel: + raise ValueError(f"novel allele name {name!r} already defined") + for seg in _SEGMENTS: + if self._find_ref_allele(seg, name) is not None: + raise ValueError(f"novel allele name {name!r} collides with a catalogue allele") + if (mutations is None) == (sequence is None): + raise ValueError("provide exactly one of `mutations` or `sequence`") + + base_ungapped = base_allele.ungapped_seq.upper() + gapped = list(base_allele.gapped_seq or "") + # ungapped index -> gapped index (positions of non-gap characters). + # Some custom cartridges carry no (or inconsistent) gapped sequence; + # in that case we can't project onto gaps, so fall back to an + # ungapped novel sequence (no gap-derived metadata). + ung_to_gap = [i for i, ch in enumerate(base_allele.gapped_seq or "") if ch != "."] + project_gaps = len(ung_to_gap) == len(base_ungapped) + seq = list(base_ungapped) + if sequence is not None: + sequence = sequence.upper() + if len(sequence) != len(seq): + raise ValueError( + f"explicit sequence length {len(sequence)} != base length {len(seq)}; " + "novel alleles are substitution-only (same length as the base)" + ) + if any(b not in "ACGT" for b in sequence): + raise ValueError("sequence must contain only A/C/G/T") + seq = list(sequence) + else: + if not mutations: + raise ValueError("`mutations` must be a non-empty list of (position, base)") + for pos, b in mutations: + if not isinstance(pos, int) or isinstance(pos, bool): + raise ValueError(f"mutation position must be an int, got {pos!r}") + if not (isinstance(b, str) and len(b) == 1): + raise ValueError(f"mutation base must be a single character, got {b!r}") + if not (0 <= pos < len(seq)): + raise ValueError(f"mutation position {pos} out of range [0,{len(seq)})") + if b.upper() not in "ACGT": + raise ValueError(f"mutation base {b!r} must be A/C/G/T") + seq[pos] = b.upper() + new_ungapped = "".join(seq) + if new_ungapped == base_ungapped: + raise ValueError("novel allele is identical to its base allele") + # Project the substitutions onto the gapped sequence too, so + # gap-dependent metadata stays consistent. If the base has no + # usable gapped sequence, fall back to the ungapped form. + if project_gaps: + for k, b in enumerate(seq): + gapped[ung_to_gap[k]] = b + new_gapped = "".join(gapped) + else: + new_gapped = new_ungapped + + novel = _copy.deepcopy(base_allele) + novel.name = name + novel.gene = gene + novel.ungapped_seq = new_ungapped + novel.gapped_seq = new_gapped + if hasattr(novel, "ungapped_len"): + novel.ungapped_len = len(new_ungapped) + + # Functional validation (V/J have a conserved coding frame). + functional, reason = True, None + anchor = getattr(novel, "anchor", None) + if segment in ("V", "J") and anchor is not None: + conserved = {"V": {"C"}, "J": {"W", "F"}}[segment] + anchor_aa = translate(new_ungapped[anchor : anchor + 3]) + if anchor_aa not in conserved: + functional = False + reason = ( + f"conserved anchor codon now encodes {anchor_aa!r}, " + f"expected one of {sorted(conserved)}" + ) + coding = new_ungapped[:anchor] if segment == "V" else new_ungapped[anchor:] + if "*" in translate(coding): + reason = (reason + "; " if reason else "") + "internal stop codon in coding frame" + functional = False + if not functional and not allow_nonfunctional: + raise ValueError( + f"novel allele {name!r} is non-functional ({reason}); pass " + f"allow_nonfunctional=True to keep it anyway" + ) + if not functional: + try: + novel.functional_status = "pseudogene" + except Exception: + pass + + self._novel[name] = { + "allele": novel, + "gene": gene, + "segment": segment, + "base": base, + "mutations": list(mutations) if mutations else None, + "functional": functional, + } + return self + + def has_novel(self) -> bool: + return bool(self._novel) + + def novel_allele_names(self) -> Set[str]: + return set(self._novel) + + def _carried_allele_names(self) -> Set[str]: + """Allele names actually placed on a haplotype (across segments).""" + names: Set[str] = set() + for seg in _SEGMENTS: + for haps in self._slots[seg].values(): + for hap in haps: + names.update(a for (a, _c, _w) in hap) + return names + + def effective_dataconfig(self): + """Return a copy of the source ``DataConfig`` with this genotype's + **carried** novel alleles appended to their genes' allele lists — + the reference the engine actually runs against when novel alleles + are present. A novel allele that was defined but never placed on a + haplotype is NOT injected (it would otherwise pollute the aligner + reference and could surface in ``v_call`` despite being absent from + the ground truth).""" + import copy as _copy + + cfg = _copy.deepcopy(self._cfg) + by_seg = {"V": cfg.v_alleles, "D": cfg.d_alleles, "J": cfg.j_alleles} + carried = self._carried_allele_names() + for name, info in self._novel.items(): + if name not in carried: + continue + d = by_seg[info["segment"]] + existing = list(d.get(info["gene"], [])) + existing.append(_copy.deepcopy(info["allele"])) + d[info["gene"]] = existing + return cfg + + def complete_from_reference( + self, policy: str = "homozygous_first_reference" + ) -> "Genotype": + """Fill every UNspecified gene with a valid diploid state. + + ``policy``: + - ``"homozygous_first_reference"`` (default): each unspecified + gene becomes homozygous for its **first cartridge allele** (NOT + a population-frequency-common allele — there is no frequency + prior in PR1; the name says exactly what it does). + - ``"heterozygous_first_two"``: first two cartridge alleles, one + per haplotype (homozygous if the gene has a single allele). + """ + for seg in _SEGMENTS: + for gene, allele_objs in _alleles_by_gene(self._cfg, seg).items(): + if gene in self._slots[seg] or not allele_objs: + continue + names = [a.name for a in allele_objs] + if policy == "homozygous_first_reference": + self.homozygous(gene, names[0], segment=seg) + elif policy == "heterozygous_first_two": + if len(names) >= 2: + self.heterozygous(gene, names[0], names[1], segment=seg) + else: + self.homozygous(gene, names[0], segment=seg) + else: + raise ValueError(f"unknown policy {policy!r}") + return self + + # ── snapshot ────────────────────────────────────────────────── + def _snapshot(self) -> "Genotype": + """Return an independent copy for attachment to an experiment, so + that mutating the builder after ``with_genotype()``/``compile()`` + cannot desync ``result.genotypes`` from the compiled engine + genotype. Shares the (immutable) cartridge reference; deep-copies + the editable slot state.""" + import copy as _copy + + g = Genotype.__new__(Genotype) + g._cfg = self._cfg + g._permissive = self._permissive + g.subject_id = self.subject_id + g._chromosome_weights = self._chromosome_weights + g._slots = _copy.deepcopy(self._slots) + g._novel = _copy.deepcopy(self._novel) + g._source_hash = self._source_hash + return g + + # ── queries / export ────────────────────────────────────────── + @property + def is_permissive(self) -> bool: + return self._permissive + + def is_specified(self, segment: str, gene: str) -> bool: + return gene in self._slots[segment] + + def carried_alleles(self, segment: str, gene: str) -> Set[str]: + out: Set[str] = set() + for hap in self._slots[segment].get(gene, [[], []]): + out.update(a for (a, _c, _w) in hap) + return out + + @staticmethod + def _zygosity(h0: List, h1: List) -> str: + s0 = {a for (a, _, _) in h0} + s1 = {a for (a, _, _) in h1} + if not s0 and not s1: + return "deleted" + if bool(s0) != bool(s1): # exactly one haplotype carries the gene + return "hemizygous" + if s0 == s1 and len(s0) == 1: + return "homozygous" + return "heterozygous" + + def to_table(self) -> List[Dict]: + """One row per (segment, gene) with full diploid truth: zygosity + (incl. ``hemizygous`` / ``deleted``), the carried alleles per + haplotype, and per-haplotype copy/weight detail. Suitable as a + ground-truth genotype table for inference benchmarks.""" + rows = [] + for seg in _SEGMENTS: + for gene, haps in self._slots[seg].items(): + h0, h1 = haps[0], haps[1] + carried = {a for (a, _, _) in h0} | {a for (a, _, _) in h1} + novel_here = sorted(carried & set(self._novel)) + rows.append( + { + "subject_id": self.subject_id, + "segment": seg, + "gene": gene, + "zygosity": self._zygosity(h0, h1), + "haplotype_0": sorted(a for (a, _, _) in h0), + "haplotype_1": sorted(a for (a, _, _) in h1), + # per-haplotype (allele, copies, weight) detail + "haplotype_0_detail": sorted(h0), + "haplotype_1_detail": sorted(h1), + "novel": novel_here, # carried alleles that are private/novel + "permissive": self._permissive, + } + ) + return rows + + def to_tsv(self, path: str) -> None: + import csv + + def _fmt(detail): + # allele:copies:weight ; ... + return ";".join(f"{a}:{c}:{w}" for (a, c, w) in detail) + + rows = self.to_table() + with open(path, "w", newline="") as fh: + w = csv.writer(fh, delimiter="\t") + w.writerow( + [ + "subject_id", + "segment", + "gene", + "zygosity", + "haplotype_0", + "haplotype_1", + "novel", + "permissive", + ] + ) + for r in rows: + w.writerow( + [ + r["subject_id"], + r["segment"], + r["gene"], + r["zygosity"], + _fmt(r["haplotype_0_detail"]), + _fmt(r["haplotype_1_detail"]), + ";".join(r["novel"]), + r["permissive"], + ] + ) diff --git a/src/GenAIRR/result.py b/src/GenAIRR/result.py index b2baabc..05f4b69 100644 --- a/src/GenAIRR/result.py +++ b/src/GenAIRR/result.py @@ -419,7 +419,7 @@ class SimulationResult: inspection — most users won't need them. """ - __slots__ = ("_records", "_outcomes", "_parents") + __slots__ = ("_records", "_outcomes", "_parents", "_genotypes") def __init__( self, @@ -428,6 +428,10 @@ def __init__( parents: Optional[Sequence] = None, ) -> None: self._records: List[Dict[str, Any]] = list(records) + # Per-subject ground-truth ``Genotype`` objects, populated by + # ``CompiledExperiment.run_records`` when a genotype is attached. + # ``None`` for non-genotype results. + self._genotypes: Optional[List] = None # ``outcomes`` is optional: callers that built records by # other means (e.g. round-tripping a TSV) don't have the # underlying Outcome objects available. @@ -495,6 +499,12 @@ def outcomes(self) -> Optional[List]: directly (e.g. loaded from a TSV).""" return self._outcomes + @property + def genotypes(self) -> Optional[List]: + """Per-subject ground-truth ``Genotype`` objects when the + experiment had a genotype attached, else ``None``.""" + return self._genotypes + @property def parents(self) -> Optional[List]: """Per-clone parent ``Outcome`` objects for clonal results; diff --git a/tests/test_clonal_parent_contract.py b/tests/test_clonal_parent_contract.py index 7c48014..1f85598 100644 --- a/tests/test_clonal_parent_contract.py +++ b/tests/test_clonal_parent_contract.py @@ -366,14 +366,20 @@ def test_pin_scaffold_clonal_truth_calls_stable_within_clone_under_normal_fixtur def test_pin_scaffold_simulationresult_slots_documented_for_extension() -> None: """``SimulationResult.__slots__`` is the documented attribute surface. After Slice 2 it carries - ``("_records", "_outcomes", "_parents")``. A reviewer adding - a new slot must update this pin in lockstep so the slot-list + ``("_records", "_outcomes", "_parents")``; the genotype slice adds + ``"_genotypes"`` (per-subject ground-truth genotypes). A reviewer + adding a new slot must update this pin in lockstep so the slot-list change shows up as a deliberate, audited diff.""" - assert SimulationResult.__slots__ == ("_records", "_outcomes", "_parents"), ( + assert SimulationResult.__slots__ == ( + "_records", + "_outcomes", + "_parents", + "_genotypes", + ), ( f"SimulationResult.__slots__ drifted to {SimulationResult.__slots__}; " - "expected ('_records', '_outcomes', '_parents'). Either Slice 2 " - "regressed (parent accessor removed) or a new slot landed without " - "updating the lockstep pin." + "expected ('_records', '_outcomes', '_parents', '_genotypes'). Either " + "Slice 2 regressed (parent accessor removed) or a new slot landed " + "without updating the lockstep pin." ) diff --git a/tests/test_genotype_backward_compat.py b/tests/test_genotype_backward_compat.py new file mode 100644 index 0000000..f7a8fc8 --- /dev/null +++ b/tests/test_genotype_backward_compat.py @@ -0,0 +1,33 @@ +"""Backward-compat: the genotype machinery is purely additive, so a +run with NO genotype attached must be byte-identical to master. + +``MASTER_DIGEST`` was captured on ``master`` (pre-genotype) for +``Experiment.on(HUMAN_IGH_OGRDB).recombine().run_records(n=100, seed=12345)`` +and re-verified identical on this branch's no-genotype path. +""" +import hashlib + +import GenAIRR as ga +import GenAIRR.data as gdata + +MASTER_DIGEST = "3be8e5ea124e1dfff256f93b5ddbb925fdb738d4d6f2eeb5763a81e5b6213460" + + +def _digest(records): + h = hashlib.sha256() + for rec in records: + h.update(repr(sorted(rec.items())).encode()) + return h.hexdigest() + + +def test_no_genotype_output_matches_master_baseline(): + res = ga.Experiment.on(gdata.HUMAN_IGH_OGRDB).recombine().run_records( + n=100, seed=12345 + ) + assert _digest(res) == MASTER_DIGEST + + +def test_no_genotype_output_is_deterministic(): + a = ga.Experiment.on(gdata.HUMAN_IGH_OGRDB).recombine().run_records(n=100, seed=12345) + b = ga.Experiment.on(gdata.HUMAN_IGH_OGRDB).recombine().run_records(n=100, seed=12345) + assert _digest(a) == _digest(b) diff --git a/tests/test_genotype_builder.py b/tests/test_genotype_builder.py new file mode 100644 index 0000000..06bcbfa --- /dev/null +++ b/tests/test_genotype_builder.py @@ -0,0 +1,113 @@ +"""Builder-level tests for GenAIRR.genotype.Genotype (PR1).""" +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_homozygous_then_subject_builds_diploid_genotype(): + cfg = _cfg() + v_gene = next(iter(cfg.v_alleles)) + a1 = cfg.v_alleles[v_gene][0].name + g = Genotype.from_dataconfig(cfg).homozygous(v_gene, a1).with_subject("S1") + assert g.subject_id == "S1" + assert g.carried_alleles("V", v_gene) == {a1} + + +def test_unknown_allele_name_raises(): + cfg = _cfg() + v_gene = next(iter(cfg.v_alleles)) + with pytest.raises(ValueError, match="not a known"): + Genotype.from_dataconfig(cfg).homozygous(v_gene, "IGHV-NOPE*99") + + +def test_unknown_gene_raises(): + cfg = _cfg() + with pytest.raises(ValueError, match="not a known"): + Genotype.from_dataconfig(cfg).homozygous("NOSUCHGENE", "x*01") + + +def test_strict_genotype_reports_unspecified_gene(): + cfg = _cfg() + g = Genotype.from_dataconfig(cfg) + assert g.is_specified("V", next(iter(cfg.v_alleles))) is False + + +def test_complete_from_reference_specifies_every_gene(): + cfg = _cfg() + g = Genotype.from_dataconfig(cfg).complete_from_reference() + assert all(g.is_specified("V", gene) for gene in cfg.v_alleles) + assert all(g.is_specified("J", gene) for gene in cfg.j_alleles) + + +def test_heterozygous_to_table_reports_zygosity(): + cfg = _cfg() + v_gene = next(iter(cfg.v_alleles)) + names = [a.name for a in cfg.v_alleles[v_gene]] + if len(names) < 2: + pytest.skip("need >=2 alleles for heterozygous test") + g = Genotype.from_dataconfig(cfg).heterozygous(v_gene, names[0], names[1]) + rows = [r for r in g.to_table() if r["gene"] == v_gene] + assert rows and rows[0]["zygosity"] == "heterozygous" + assert g.carried_alleles("V", v_gene) == {names[0], names[1]} + + +def test_delete_gene_one_haplotype_keeps_the_other(): + cfg = _cfg() + v_gene = next(iter(cfg.v_alleles)) + a1 = cfg.v_alleles[v_gene][0].name + g = Genotype.from_dataconfig(cfg).homozygous(v_gene, a1).delete_gene(v_gene, haplotype=1) + # haplotype 0 still carries a1; haplotype 1 deleted + rows = [r for r in g.to_table() if r["gene"] == v_gene] + assert rows[0]["haplotype_0"] == [a1] + assert rows[0]["haplotype_1"] == [] + + +def test_permissive_is_flagged(): + cfg = _cfg() + g = Genotype.permissive(cfg) + assert g.is_permissive is True + assert Genotype.from_dataconfig(cfg).is_permissive is False + + +def test_delete_one_haplotype_of_unspecified_gene_raises(): + cfg = _cfg() + v_gene = next(iter(cfg.v_alleles)) + with pytest.raises(ValueError, match="before deleting one haplotype"): + Genotype.from_dataconfig(cfg).delete_gene(v_gene, haplotype=1) + + +def test_one_haplotype_deletion_is_labelled_hemizygous(): + cfg = _cfg() + v_gene = next(iter(cfg.v_alleles)) + a1 = cfg.v_alleles[v_gene][0].name + g = Genotype.from_dataconfig(cfg).homozygous(v_gene, a1).delete_gene(v_gene, haplotype=1) + row = next(r for r in g.to_table() if r["gene"] == v_gene) + assert row["zygosity"] == "hemizygous" + assert row["haplotype_0"] == [a1] + assert row["haplotype_1"] == [] + + +def test_chromosome_weights_rejects_nan(): + cfg = _cfg() + with pytest.raises(ValueError, match="finite"): + Genotype.from_dataconfig(cfg).chromosome_weights(float("nan"), 1.0) + + +def test_snapshot_decouples_from_later_mutation(): + cfg = _cfg() + import GenAIRR as ga + + v_gene = next(iter(cfg.v_alleles)) + a1 = cfg.v_alleles[v_gene][0].name + g = Genotype.from_dataconfig(cfg).complete_from_reference().with_subject("S1") + exp = ga.Experiment.on(cfg).with_genotype(g) + # Mutate the builder AFTER attach — must not affect the attached snapshot. + g.with_subject("MUTATED").delete_gene(v_gene, haplotype="both") + assert exp._genotype.subject_id == "S1" + assert exp._genotype.carried_alleles("V", v_gene) # still carried in snapshot diff --git a/tests/test_genotype_dsl.py b/tests/test_genotype_dsl.py new file mode 100644 index 0000000..d0a6f3b --- /dev/null +++ b/tests/test_genotype_dsl.py @@ -0,0 +1,73 @@ +"""DSL-level guards for Experiment.with_genotype (PR1).""" +import pytest + +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype + + +def _cfg(): + return gdata.HUMAN_IGH_OGRDB + + +def _full_genotype(): + return Genotype.from_dataconfig(_cfg()).complete_from_reference().with_subject("S1") + + +def test_with_genotype_then_restrict_alleles_raises(): + g = _full_genotype() + v_gene = next(iter(_cfg().v_alleles)) + a1 = _cfg().v_alleles[v_gene][0].name + with pytest.raises(ValueError, match="mutually exclusive"): + ga.Experiment.on(_cfg()).with_genotype(g).restrict_alleles(v=a1) + + +def test_restrict_alleles_then_with_genotype_raises(): + g = _full_genotype() + v_gene = next(iter(_cfg().v_alleles)) + a1 = _cfg().v_alleles[v_gene][0].name + with pytest.raises(ValueError, match="mutually exclusive"): + ga.Experiment.on(_cfg()).restrict_alleles(v=a1).with_genotype(g) + + +def test_recombine_weights_then_with_genotype_raises(): + g = _full_genotype() + v_gene = next(iter(_cfg().v_alleles)) + a1 = _cfg().v_alleles[v_gene][0].name + with pytest.raises(ValueError, match="mutually exclusive"): + ga.Experiment.on(_cfg()).recombine(v_allele_weights={a1: 2.0}).with_genotype(g) + + +def test_with_genotype_then_recombine_weights_raises(): + g = _full_genotype() + v_gene = next(iter(_cfg().v_alleles)) + a1 = _cfg().v_alleles[v_gene][0].name + with pytest.raises(ValueError, match="mutually exclusive"): + ga.Experiment.on(_cfg()).with_genotype(g).recombine(v_allele_weights={a1: 2.0}) + + +def test_receptor_revision_with_genotype_raises_at_compile(): + g = _full_genotype() + exp = ga.Experiment.on(_cfg()).with_genotype(g).recombine().receptor_revision(prob=0.5) + with pytest.raises(ValueError, match="receptor_revision"): + exp.compile() + + +def test_genotype_with_clonal_fork_raises_at_compile(): + g = _full_genotype() + exp = ( + ga.Experiment.on(_cfg()) + .with_genotype(g) + .recombine() + .clonal_lineage(n_clones=2) + ) + with pytest.raises(ValueError, match="not supported together with"): + exp.compile() + + +def test_cartridge_hash_mismatch_raises(): + # Genotype built on IGH, attached to a TCRB experiment → mismatch. + g = Genotype.from_dataconfig(_cfg()).complete_from_reference() + other = gdata.HUMAN_TCRB_IMGT + with pytest.raises(ValueError, match="different cartridge|content hash"): + ga.Experiment.on(other).with_genotype(g) diff --git a/tests/test_genotype_engine.py b/tests/test_genotype_engine.py new file mode 100644 index 0000000..e297ace --- /dev/null +++ b/tests/test_genotype_engine.py @@ -0,0 +1,145 @@ +"""End-to-end phased-genotype recombination tests (PR1).""" +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype + + +def _cfg(): + return gdata.HUMAN_IGH_OGRDB + + +def test_phased_recombine_only_emits_carried_allele_for_overridden_gene(): + cfg = _cfg() + # Pick a V gene with >=2 alleles and carry ONLY its second allele. + v_gene = next(g for g, al in cfg.v_alleles.items() if len(al) >= 2) + names = [a.name for a in cfg.v_alleles[v_gene]] + carried = names[1] + g = ( + Genotype.from_dataconfig(cfg) + .complete_from_reference("homozygous_first_reference") + .homozygous(v_gene, carried) + .with_subject("S1") + ) + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records( + n=300, seed=1, expose_provenance=True + ) + # Use the ground-truth call (truth_v_call) — the evidence-based v_call + # can be ambiguous (comma-joined) between similar alleles. + seen_for_gene = { + r["truth_v_call"] for r in res if r["truth_v_call"].startswith(v_gene + "*") + } + # Only the carried allele of that gene may appear (never names[0]). + assert seen_for_gene <= {carried}, seen_for_gene + + +def test_records_carry_subject_and_haplotype_and_result_exposes_genotype(): + cfg = _cfg() + g = Genotype.from_dataconfig(cfg).complete_from_reference().with_subject("S1") + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records(n=20, seed=3) + assert all(r["subject_id"] == "S1" for r in res) + assert all(r["haplotype"] in (0, 1) for r in res) + assert res.genotypes is not None + assert res.genotypes[0].subject_id == "S1" + + +def test_no_genotype_result_has_no_genotypes_and_no_haplotype_field(): + cfg = _cfg() + res = ga.Experiment.on(cfg).recombine().run_records(n=10, seed=3) + assert res.genotypes is None + assert "subject_id" not in res[0] + assert "haplotype" not in res[0] + + +def test_phased_run_is_deterministic_under_same_seed(): + cfg = _cfg() + g = Genotype.from_dataconfig(cfg).complete_from_reference().with_subject("S1") + exp = ga.Experiment.on(cfg).with_genotype(g).recombine() + a = exp.run_records(n=40, seed=77) + b = exp.run_records(n=40, seed=77) + assert [r["v_call"] for r in a] == [r["v_call"] for r in b] + assert [r["j_call"] for r in a] == [r["j_call"] for r in b] + + +def test_deleted_gene_is_never_sampled(): + cfg = _cfg() + drop = list(cfg.v_alleles)[1] + g = ( + Genotype.from_dataconfig(cfg) + .complete_from_reference() + .delete_gene(drop, haplotype="both", segment="V") + .with_subject("S1") + ) + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records( + n=300, seed=5, expose_provenance=True + ) + assert all(not r["truth_v_call"].startswith(drop + "*") for r in res) + + +def test_heterozygous_expression_is_roughly_balanced(): + cfg = _cfg() + v_gene = next(g for g, al in cfg.v_alleles.items() if len(al) >= 2) + names = [a.name for a in cfg.v_alleles[v_gene]] + a0, a1 = names[0], names[1] + g = Genotype.from_dataconfig(cfg).complete_from_reference() + for other in cfg.v_alleles: + if other != v_gene: + g.delete_gene(other, haplotype="both", segment="V") + g.heterozygous(v_gene, a0, a1).with_subject("S1") + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records( + n=400, seed=9, expose_provenance=True + ) + calls = [r["truth_v_call"] for r in res] + assert set(calls) <= {a0, a1}, set(calls) + frac0 = calls.count(a0) / len(calls) + assert 0.35 < frac0 < 0.65, frac0 + + +def test_vdj_phasing_links_v_and_j_to_one_chromosome(): + """The core phasing guarantee: with hap0 = {V:gv0, J:gj0} and + hap1 = {V:gv1, J:gj1}, every rearrangement's V and J must come from + the SAME chromosome — never a cross pairing.""" + cfg = _cfg() + vgenes, jgenes = list(cfg.v_alleles), list(cfg.j_alleles) + gv0, gv1 = vgenes[0], vgenes[1] + gj0, gj1 = jgenes[0], jgenes[1] + g = Genotype.from_dataconfig(cfg).complete_from_reference() + # Keep only gv0/gv1 (V) and gj0/gj1 (J); delete everything else. + for gene in cfg.v_alleles: + if gene not in (gv0, gv1): + g.delete_gene(gene, "both", segment="V") + for gene in cfg.j_alleles: + if gene not in (gj0, gj1): + g.delete_gene(gene, "both", segment="J") + # Phase: hap0 carries gv0 + gj0; hap1 carries gv1 + gj1. + g.delete_gene(gv0, haplotype=1, segment="V") + g.delete_gene(gv1, haplotype=0, segment="V") + g.delete_gene(gj0, haplotype=1, segment="J") + g.delete_gene(gj1, haplotype=0, segment="J") + g.with_subject("S1") + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records( + n=200, seed=4, expose_provenance=True + ) + for r in res: + v_is_0 = r["truth_v_call"].startswith(gv0 + "*") + j_is_0 = r["truth_j_call"].startswith(gj0 + "*") + assert v_is_0 == j_is_0, (r["truth_v_call"], r["truth_j_call"]) + assert (r["haplotype"] == 0) == v_is_0 + + +def test_one_dead_haplotype_uses_the_live_one_under_productive_only(): + cfg = _cfg() + g = Genotype.from_dataconfig(cfg).complete_from_reference() + # Delete every J gene on haplotype 1 → only haplotype 0 can produce + # a rearrangement; productive_only must still succeed via haplotype 0. + for j_gene in cfg.j_alleles: + g.delete_gene(j_gene, haplotype=1, segment="J") + g.with_subject("S1") + res = ( + ga.Experiment.on(cfg) + .productive_only() + .with_genotype(g) + .recombine() + .run_records(n=40, seed=2) + ) + assert len(res) == 40 + assert all(r["haplotype"] == 0 for r in res) diff --git a/tests/test_genotype_novel.py b/tests/test_genotype_novel.py new file mode 100644 index 0000000..d52d2fd --- /dev/null +++ b/tests/test_genotype_novel.py @@ -0,0 +1,195 @@ +"""Novel / private allele support on genotypes.""" +import pytest + +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype + + +def _cfg(): + return gdata.HUMAN_IGH_OGRDB + + +# Guaranteed-safe SNPs for IGHVF1-G1*01: wobble of non-T-starting codons in +# the framework — cannot create a stop, anchor untouched. +_SAFE = [(38, "C"), (41, "A")] + + +def test_add_novel_allele_synthesizes_from_base_and_mutations(): + cfg = _cfg() + base = cfg.v_alleles["IGHVF1-G1"][0] + g = Genotype.from_dataconfig(cfg).add_novel_allele( + "IGHVF1-G1*i01", base="IGHVF1-G1*01", mutations=_SAFE + ) + assert g.has_novel() + assert "IGHVF1-G1*i01" in g.novel_allele_names() + nv = g._novel["IGHVF1-G1*i01"]["allele"] + assert nv.ungapped_seq[38] == "C" and nv.ungapped_seq[41] == "A" + assert len(nv.ungapped_seq) == len(base.ungapped_seq) # substitution-only + assert nv.gene == "IGHVF1-G1" and nv.anchor == base.anchor # gene/anchor inherited + # gapped sequence projected consistently (ungapped derived from it) + assert nv.gapped_seq.replace(".", "") == nv.ungapped_seq + assert g._novel["IGHVF1-G1*i01"]["functional"] is True + + +def test_novel_name_gene_must_match_base_gene(): + cfg = _cfg() + with pytest.raises(ValueError, match="implies gene"): + Genotype.from_dataconfig(cfg).add_novel_allele( + "IGHVF1-G2*i01", base="IGHVF1-G1*01", mutations=_SAFE # name gene != base gene + ) + + +def test_novel_allele_basic_validation(): + cfg = _cfg() + G = lambda: Genotype.from_dataconfig(cfg) + with pytest.raises(ValueError, match="base allele"): + G().add_novel_allele("NOPE*i01", base="NOPE*01", mutations=[(1, "T")]) + with pytest.raises(ValueError, match="collides with a catalogue"): + G().add_novel_allele("IGHVF1-G1*01", base="IGHVF1-G1*01", mutations=_SAFE) + with pytest.raises(ValueError, match="out of range"): + G().add_novel_allele("IGHVF1-G1*i01", base="IGHVF1-G1*01", mutations=[(99999, "T")]) + with pytest.raises(ValueError, match="exactly one"): + G().add_novel_allele("IGHVF1-G1*i01", base="IGHVF1-G1*01") + with pytest.raises(ValueError, match="must be an int"): + G().add_novel_allele("IGHVF1-G1*i01", base="IGHVF1-G1*01", mutations=[(1.5, "A")]) + + +def test_cross_segment_name_collision_rejected(): + # A novel can't be named for a different gene/segment than its base, so + # naming a V-derived novel after a real J allele is rejected outright — + # the truth table can never mislabel the real J row as novel. + cfg = _cfg() + j_name = next(iter(cfg.j_alleles[next(iter(cfg.j_alleles))])).name + with pytest.raises(ValueError): + Genotype.from_dataconfig(cfg).add_novel_allele( + j_name, base="IGHVF1-G1*01", mutations=_SAFE + ) + + +def test_nonfunctional_novel_rejected_by_default_and_allowed_explicitly(): + cfg = _cfg() + base = cfg.v_alleles["IGHVF1-G1"][0] + # Force a stop codon at framework codon 13 (positions 39,40,41 -> TAA). + stop = [(39, "T"), (40, "A"), (41, "A")] + with pytest.raises(ValueError, match="non-functional.*stop codon"): + Genotype.from_dataconfig(cfg).add_novel_allele( + "IGHVF1-G1*i01", base="IGHVF1-G1*01", mutations=stop + ) + # Explicit override keeps it, marked non-functional. + g = Genotype.from_dataconfig(cfg).add_novel_allele( + "IGHVF1-G1*i01", base="IGHVF1-G1*01", mutations=stop, allow_nonfunctional=True + ) + assert g._novel["IGHVF1-G1*i01"]["functional"] is False + + +def test_broken_anchor_codon_rejected(): + cfg = _cfg() + base = cfg.v_alleles["IGHVF1-G1"][0] + a = base.anchor + # rewrite the conserved Cys anchor codon to GGG (Gly) + with pytest.raises(ValueError, match="conserved anchor codon"): + Genotype.from_dataconfig(cfg).add_novel_allele( + "IGHVF1-G1*i01", base="IGHVF1-G1*01", + mutations=[(a, "G"), (a + 1, "G"), (a + 2, "G")], + ) + + +def test_novel_allele_placed_and_sampled_appears_in_airr(): + cfg = _cfg() + gene = "IGHVF1-G1" + g = ( + Genotype.from_dataconfig(cfg) + .add_novel_allele(f"{gene}*i01", base=f"{gene}*01", mutations=_SAFE) + .complete_from_reference() + .homozygous(gene, f"{gene}*i01") # carry ONLY the novel allele for this gene + .with_subject("DONOR_N") + ) + res = ( + ga.Experiment.on(cfg) + .with_genotype(g) + .recombine() + .run_records(n=300, seed=3, expose_provenance=True) + ) + truth_for_gene = { + r["truth_v_call"] for r in res if r["truth_v_call"].startswith(gene + "*") + } + assert truth_for_gene == {f"{gene}*i01"}, truth_for_gene + novel_reads = [r for r in res if r["truth_v_call"] == f"{gene}*i01"] + assert novel_reads + assert any(r["sequence"][38].upper() == "C" for r in novel_reads) + + +def test_to_table_and_tsv_expose_novel(tmp_path): + cfg = _cfg() + gene = "IGHVF1-G1" + g = ( + Genotype.from_dataconfig(cfg) + .add_novel_allele(f"{gene}*i01", base=f"{gene}*01", mutations=_SAFE) + .heterozygous(gene, f"{gene}*01", f"{gene}*i01") + ) + row = next(r for r in g.to_table() if r["gene"] == gene) + assert row["novel"] == [f"{gene}*i01"] + p = tmp_path / "truth.tsv" + g.to_tsv(str(p)) + header = p.read_text().splitlines()[0].split("\t") + assert "novel" in header + + +def test_unplaced_novel_allele_not_injected_or_emitted(): + cfg = _cfg() + gene = "IGHVF1-G1" + # Define a novel allele but NEVER place it on a haplotype. + g = ( + Genotype.from_dataconfig(cfg) + .add_novel_allele(f"{gene}*unplaced", base=f"{gene}*01", mutations=_SAFE) + .complete_from_reference() + .with_subject("S1") + ) + # Effective reference must not contain the unplaced novel allele. + eff = g.effective_dataconfig() + eff_names = {a.name for alleles in eff.v_alleles.values() for a in alleles} + assert f"{gene}*unplaced" not in eff_names + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records( + n=200, seed=4, expose_provenance=True + ) + assert all(f"{gene}*unplaced" not in r["v_call"] for r in res) + assert all(f"{gene}*unplaced" not in r["truth_v_call"] for r in res) + + +def test_haplotype_argument_validation(): + cfg = _cfg() + gene = "IGHVF1-G1" + a0 = cfg.v_alleles[gene][0].name + g = Genotype.from_dataconfig(cfg).homozygous(gene, a0) + with pytest.raises(ValueError, match="haplotype must be"): + g.delete_gene(gene, haplotype=2) + with pytest.raises(ValueError, match="haplotype must be"): + g.delete_gene(gene, haplotype="x") + with pytest.raises(ValueError, match="haplotype must be 0 or 1"): + Genotype.from_dataconfig(cfg).duplicate_gene(gene, [a0], haplotype=-1) + with pytest.raises(ValueError, match="haplotype must be 0 or 1"): + Genotype.from_dataconfig(cfg).duplicate_gene(gene, [a0], haplotype=2) + + +def test_novel_synthesis_tolerates_missing_gapped_seq(): + import copy + + cfg = copy.deepcopy(_cfg()) + # Simulate a custom cartridge whose base allele has no gapped sequence. + cfg.v_alleles["IGHVF1-G1"][0].gapped_seq = "" + g = Genotype.from_dataconfig(cfg).add_novel_allele( + "IGHVF1-G1*i01", base="IGHVF1-G1*01", mutations=_SAFE + ) + nv = g._novel["IGHVF1-G1*i01"]["allele"] + # falls back to ungapped form (no crash) + assert nv.gapped_seq == nv.ungapped_seq + assert nv.ungapped_seq[38] == "C" + + +def test_genotype_without_novel_is_unaffected(): + cfg = _cfg() + g = Genotype.from_dataconfig(cfg).complete_from_reference().with_subject("S1") + assert g.has_novel() is False + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records(n=10, seed=1) + assert len(res) == 10 diff --git a/tests/test_genotype_usage.py b/tests/test_genotype_usage.py new file mode 100644 index 0000000..887bc3c --- /dev/null +++ b/tests/test_genotype_usage.py @@ -0,0 +1,39 @@ +"""Genotype gene-usage wiring is active when the cartridge authors a +typed ``reference_models.allele_usage`` plane (review #3). + +This pins that the Rust gene-level usage aggregation is *reachable* and +*effective* — not just code-complete. Bundled configs that don't author +allele_usage fall back to uniform-over-present-genes (× copy dosage), +which is documented, not tested here. +""" +import dataclasses + +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype +from GenAIRR.reference_models import AlleleUsageSpec, ReferenceEmpiricalModels + + +def _cfg_with_v_usage(target_allele_name, weight=1000.0): + base = gdata.HUMAN_IGH_OGRDB + rm = ReferenceEmpiricalModels(allele_usage=AlleleUsageSpec(v={target_allele_name: weight})) + return dataclasses.replace(base, reference_models=rm) + + +def test_typed_allele_usage_biases_genotype_gene_choice(): + base = gdata.HUMAN_IGH_OGRDB + # Heavily weight the first allele of one V gene. + target_gene = list(base.v_alleles)[10] + target_allele = base.v_alleles[target_gene][0].name + cfg = _cfg_with_v_usage(target_allele, weight=1000.0) + + g = Genotype.from_dataconfig(cfg).complete_from_reference().with_subject("S1") + res = ga.Experiment.on(cfg).with_genotype(g).recombine().run_records( + n=300, seed=1, expose_provenance=True + ) + frac_target = sum( + 1 for r in res if r["truth_v_call"].startswith(target_gene + "*") + ) / len(res) + # With ~52 V genes uniform would give ~0.02; a 1000x usage weight on + # this gene must dominate. + assert frac_target > 0.5, frac_target