diff --git a/engine_rs/src/assignment.rs b/engine_rs/src/assignment.rs index 9ba6134..18c707a 100644 --- a/engine_rs/src/assignment.rs +++ b/engine_rs/src/assignment.rs @@ -114,6 +114,13 @@ pub struct AlleleInstance { /// (today the DSL prevents chaining, but the field stays /// stable under any future relaxation). pub receptor_revision_original_id: Option, + /// The diploid rearrangement chromosome (0 or 1) this assignment was + /// sampled from, when phased genotype recombination ran; `None` on the + /// flat (no-genotype) path. Immutable provenance of the *original* + /// rearrangement — receptor revision preserves it even when a + /// cross-haplotype replacement (`same_haplotype=false`) comes from the + /// other chromosome. + pub haplotype: Option, } impl AlleleInstance { @@ -130,6 +137,7 @@ impl AlleleInstance { trim_3: 0, orientation: SegmentOrientation::Forward, receptor_revision_original_id: None, + haplotype: None, } } @@ -167,6 +175,18 @@ impl AlleleInstance { ..self } } + + /// Return a new instance with the rearrangement `haplotype` (0 or 1) + /// set; receiver unchanged. Panics if `hap` is not 0 or 1 — a diploid + /// genotype has exactly two haplotypes. + #[must_use] + pub fn with_haplotype(self, hap: u8) -> Self { + assert!(hap <= 1, "haplotype must be 0 or 1, got {}", hap); + Self { + haplotype: Some(hap), + ..self + } + } } // ────────────────────────────────────────────────────────────────── @@ -337,6 +357,21 @@ mod tests { assert_eq!(a.trim_3, 7); } + #[test] + fn allele_instance_haplotype_defaults_none_and_with_haplotype_isolates() { + let a = AlleleInstance::new(AlleleId::new(0)); + assert_eq!(a.haplotype, None); + let b = a.with_haplotype(1); + assert_eq!(a.haplotype, None); // receiver unchanged (persistent) + assert_eq!(b.haplotype, Some(1)); + } + + #[test] + #[should_panic(expected = "haplotype must be 0 or 1")] + fn with_haplotype_rejects_out_of_range() { + let _ = AlleleInstance::new(AlleleId::new(0)).with_haplotype(2); + } + #[test] fn assignments_starts_empty() { let a = AlleleAssignments::new(); diff --git a/engine_rs/src/passes/receptor_revision.rs b/engine_rs/src/passes/receptor_revision.rs index 25e8b7a..4b71b94 100644 --- a/engine_rs/src/passes/receptor_revision.rs +++ b/engine_rs/src/passes/receptor_revision.rs @@ -42,9 +42,11 @@ //! 2. `TrimChanged { V, Three, old: Some(prev.trim_3), new: derived_trim_3 }` //! 3. `SegmentReplaced { V, old_region, new_region, bytes_delta: 0 }` //! -//! Slice C ships the pass surface only — no public DSL method yet, -//! no AIRR `receptor_revision_applied` field yet. Those land in -//! Slices D / E. +//! The pass is wired through the `receptor_revision()` DSL method and +//! the AIRR record exposes `receptor_revision_applied` + `original_v_call` +//! (pre-revision V); `v_call` is the post-revision V. A genotype-aware +//! variant (see `GenotypeVConstraint`) restricts the replacement V to the +//! carried alleles on the drawn rearrangement chromosome. use crate::address; use crate::assignment::{AlleleInstance, TrimEnd}; @@ -55,9 +57,19 @@ use crate::pass::{Pass, PassContext, PassError, PassRequirement}; use crate::refdata::{Allele, AlleleId, RefDataConfig}; use crate::trace::ChoiceValue; +/// Per-haplotype carried-V candidates for genotype-aware receptor revision. +/// `per_hap[c]` is the list of `(allele, mass)` (mass = weight * copies) +/// carried on chromosome `c`. `same_haplotype` restricts the draw to the +/// rearrangement chromosome; otherwise both haplotypes are aggregated. +pub struct GenotypeVConstraint { + pub per_hap: [Vec<(AlleleId, f64)>; 2], + pub same_haplotype: bool, +} + pub struct ReceptorRevisionPass { prob: f64, v_distribution: Box>, + genotype_v: Option, } impl ReceptorRevisionPass { @@ -78,6 +90,7 @@ impl ReceptorRevisionPass { Self { prob, v_distribution, + genotype_v: None, } } @@ -86,6 +99,289 @@ impl ReceptorRevisionPass { self.prob } + /// Attach a genotype-V constraint, switching the pass to genotype-aware + /// candidate selection (carried alleles on the drawn chromosome, current + /// V excluded). Builder — returns `self`. + #[must_use] + pub fn with_genotype_constraint(mut self, constraint: GenotypeVConstraint) -> Self { + self.genotype_v = Some(constraint); + self + } + + /// Aggregated, current-excluded, same-length-eligible candidate pool for + /// genotype-aware revision. Aggregates mass by `AlleleId` (summing across + /// haplotypes when `!same_haplotype`, and collapsing duplicates within a + /// haplotype), drops the currently-assigned `current_v`, and keeps only + /// alleles whose length admits the retained slice. + fn genotype_eligible( + &self, + constraint: &GenotypeVConstraint, + c: usize, + current_v: AlleleId, + refdata: &RefDataConfig, + old_v_len: u32, + ) -> Vec<(AlleleId, f64)> { + use std::collections::BTreeMap; + let mut agg: BTreeMap = BTreeMap::new(); + let haps: &[Vec<(AlleleId, f64)>] = if constraint.same_haplotype { + std::slice::from_ref(&constraint.per_hap[c]) + } else { + &constraint.per_hap + }; + for list in haps { + for (id, mass) in list { + if *mass > 0.0 { + *agg.entry(id.index()).or_insert(0.0) += *mass; + } + } + } + agg.into_iter() + .filter(|(idx, _)| *idx != current_v.index()) + .filter(|(idx, _)| { + refdata + .get(Segment::V, AlleleId::new(*idx)) + .map(|a| a.len() >= old_v_len) + .unwrap_or(false) + }) + .map(|(idx, mass)| (AlleleId::new(idx), mass)) + .collect() + } + + fn execute_genotype_aware( + &self, + sim: &Simulation, + ctx: &mut PassContext, + strict: bool, + constraint: &GenotypeVConstraint, + ) -> Result { + if !sim.assignments.has(Segment::V) { + return Err(PassError::missing_assignment(self.name(), Segment::V)); + } + let refdata = ctx + .refdata + .ok_or_else(|| PassError::missing_refdata(self.name()))?; + let old_v_len = self.old_v_region_len(sim)?; + let current = sim.assignments.get(Segment::V).expect("V assigned"); + let current_v = current.allele_id; + let c = current.haplotype.ok_or_else(|| { + PassError::invalid_plan_state( + self.name(), + "genotype-aware receptor revision requires a haplotype-stamped V \ + assignment (SampleGenotypePass must run first)", + ) + })? as usize; + + let replaying = ctx.replay_cursor.is_some(); + let coin = if replaying { + ctx.replay_cursor + .as_deref_mut() + .expect("cursor") + .expect_bool(address::ChoiceAddress::ReceptorRevisionApplied) + .map_err(|r| PassError::replay(self.name(), r))? + } else { + ctx.rng.next_f64() < self.prob + }; + + if replaying { + ctx.trace.record_choice( + address::ChoiceAddress::ReceptorRevisionApplied, + ChoiceValue::Bool(coin), + ); + if !coin { + return Ok(sim.clone()); + } + let (new_id, derived_trim_3) = self.consume_replay_choices_genotype( + constraint, c, current_v, refdata, old_v_len, ctx, + )?; + return self.finish_apply(sim, refdata, new_id, derived_trim_3, c, strict, ctx); + } + + // Fresh: applied = coin AND an eligible alternate exists. + let eligible = self.genotype_eligible(constraint, c, current_v, refdata, old_v_len); + if coin && eligible.is_empty() { + if strict { + return Err(PassError::constraint_sampling( + self.name(), + address::ChoiceAddress::ReceptorRevisionVAllele.to_string(), + crate::dist::FilteredSampleError::EmptyAdmissibleSupport, + )); + } + ctx.trace.record_choice( + address::ChoiceAddress::ReceptorRevisionApplied, + ChoiceValue::Bool(false), + ); + return Ok(sim.clone()); + } + let applied = coin && !eligible.is_empty(); + ctx.trace.record_choice( + address::ChoiceAddress::ReceptorRevisionApplied, + ChoiceValue::Bool(applied), + ); + if !applied { + return Ok(sim.clone()); + } + let total: f64 = eligible.iter().map(|(_, m)| m).sum(); + let mut u = ctx.rng.next_f64() * total; + let mut new_id = eligible.last().expect("eligible non-empty").0; + for (id, m) in &eligible { + u -= *m; + if u <= 0.0 { + new_id = *id; + break; + } + } + let new_allele = refdata + .get(Segment::V, new_id) + .ok_or_else(|| PassError::missing_allele(self.name(), Segment::V, new_id.index()))?; + let derived_trim_3 = new_allele.len() - old_v_len; + self.finish_apply(sim, refdata, new_id, derived_trim_3, c, strict, ctx) + } + + /// Record the typed choices + commit the replacement (shared by the fresh + /// and replay genotype paths). Stamps the original rearrangement haplotype + /// `c` onto the replacement instance. + fn finish_apply( + &self, + sim: &Simulation, + refdata: &RefDataConfig, + new_id: AlleleId, + derived_trim_3: u32, + c: usize, + strict: bool, + ctx: &mut PassContext, + ) -> Result { + if derived_trim_3 > u16::MAX as u32 { + return Err(PassError::invalid_distribution_output( + self.name(), + address::ChoiceAddress::ReceptorRevisionVTrim3.to_string(), + derived_trim_3 as i64, + "trim_exceeds_u16", + )); + } + let new_allele = refdata + .get(Segment::V, new_id) + .ok_or_else(|| PassError::missing_allele(self.name(), Segment::V, new_id.index()))? + .clone(); + ctx.trace.record_choice( + address::ChoiceAddress::ReceptorRevisionVAllele, + ChoiceValue::AlleleId(new_id.index()), + ); + ctx.trace.record_choice( + address::ChoiceAddress::ReceptorRevisionVTrim3, + ChoiceValue::Int(derived_trim_3 as i64), + ); + let old_v_len = self.old_v_region_len(sim)?; + // Strict post-event contract arbitration — parity with the non-genotype + // path: build a hypothetical post-replacement IR (no trace/sink side + // effects) and ask the active contracts whether the receptor-revised + // state is admissible (e.g. productive_only). Permissive skips this. + if strict && ctx.contracts.is_some() { + let mut hypothetical_ctx = PassContext { + trace: ctx.trace, + rng: ctx.rng, + pass_index: ctx.pass_index, + refdata: ctx.refdata, + contracts: ctx.contracts, + feasibility: ctx.feasibility, + reference_index: ctx.reference_index, + replay_cursor: None, + event_log_sink: None, + }; + let hypothetical = self.commit_replacement( + sim.clone(), + &new_allele, + new_id, + derived_trim_3 as u16, + old_v_len, + Some(c as u8), + &mut hypothetical_ctx, + ); + self.validate_post_event_contracts(sim, &hypothetical, ctx)?; + } + Ok(self.commit_replacement( + sim.clone(), + &new_allele, + new_id, + derived_trim_3 as u16, + old_v_len, + Some(c as u8), + ctx, + )) + } + + /// Replay validation for genotype-aware revision: the recorded + /// `(v_allele, v_trim_3)` must be carried by the allowed pool (per + /// `same_haplotype`), differ from the current V, resolve in refdata, and + /// retain exactly `old_v_len` bytes. "Trace proposes, engine validates." + fn consume_replay_choices_genotype( + &self, + constraint: &GenotypeVConstraint, + c: usize, + current_v: AlleleId, + refdata: &RefDataConfig, + old_v_len: u32, + ctx: &mut PassContext, + ) -> Result<(AlleleId, u32), PassError> { + let id_index = ctx + .replay_cursor + .as_deref_mut() + .expect("cursor") + .expect_allele_id(address::ChoiceAddress::ReceptorRevisionVAllele) + .map_err(|r| PassError::replay(self.name(), r))?; + let trim_i64 = ctx + .replay_cursor + .as_deref_mut() + .expect("cursor") + .expect_int(address::ChoiceAddress::ReceptorRevisionVTrim3) + .map_err(|r| PassError::replay(self.name(), r))?; + let id = AlleleId::new(id_index); + + // Resolve the allele FIRST so an unresolvable (out-of-range) recorded id + // surfaces as `missing_allele` — matching the non-genotype replay path's + // diagnostic contract — before the eligible/current/length checks. + let allele = refdata + .get(Segment::V, id) + .ok_or_else(|| PassError::missing_allele(self.name(), Segment::V, id_index))?; + + let eligible = self.genotype_eligible(constraint, c, current_v, refdata, old_v_len); + if !eligible.iter().any(|(eid, _)| eid.index() == id_index) { + return Err(PassError::invalid_distribution_output( + self.name(), + address::ChoiceAddress::ReceptorRevisionVAllele.to_string(), + id_index as i64, + if id_index == current_v.index() { + "receptor_revision_allele_equals_current" + } else { + "receptor_revision_allele_not_carried_on_haplotype" + }, + )); + } + if trim_i64 < 0 || trim_i64 > u16::MAX as i64 { + return Err(PassError::invalid_distribution_output( + self.name(), + address::ChoiceAddress::ReceptorRevisionVTrim3.to_string(), + trim_i64, + "trim_out_of_range", + )); + } + let trim_3 = trim_i64 as u32; + let retained = (allele.len() as i64).checked_sub(trim_3 as i64).unwrap_or(-1); + if retained != old_v_len as i64 { + return Err(PassError::invalid_plan_state( + self.name(), + format!( + "replay length mismatch: allele {} length {} trim_3 {} retains {}, expected {}", + allele.name, + allele.len(), + trim_3, + retained, + old_v_len + ), + )); + } + Ok((id, trim_3)) + } + /// Resolve the single V region in `sim`. Receptor revision v1 /// requires exactly one — zero means assembly never ran, more /// than one violates the one-region-per-segment invariant the @@ -182,6 +478,7 @@ impl ReceptorRevisionPass { new_id: AlleleId, derived_trim_3: u16, old_v_len: u32, + haplotype: Option, ctx: &mut PassContext, ) -> Simulation { // Capture the pre-revision V identity **before** the builder @@ -215,11 +512,15 @@ impl ReceptorRevisionPass { // derived value but preserves every other field on the // instance (including the provenance) per the // `AlleleInstance::with_trim_3` contract. - builder.assign_allele( - Segment::V, - AlleleInstance::new(new_id) - .with_receptor_revision_original_id(preserved_original_id), - ); + // Preserve the rearrangement-chromosome provenance: receptor revision + // does not change which chromosome the receptor was rearranged on, even + // when a cross-haplotype (same_haplotype=false) replacement is drawn. + let mut new_inst = + AlleleInstance::new(new_id).with_receptor_revision_original_id(preserved_original_id); + if let Some(h) = haplotype { + new_inst = new_inst.with_haplotype(h); + } + builder.assign_allele(Segment::V, new_inst); // 2. Commit the derived 3' trim. builder.update_trim(Segment::V, TrimEnd::Three, derived_trim_3); // 3. Build the replacement nucleotides — the retained slice @@ -267,6 +568,12 @@ impl ReceptorRevisionPass { ctx: &mut PassContext, strict: bool, ) -> Result { + // Genotype-aware mode: restrict the replacement V to carried alleles on + // the drawn rearrangement chromosome (same-haplotype), excluding the + // current V. The non-genotype path below is unchanged (byte-identical). + if let Some(constraint) = self.genotype_v.as_ref() { + return self.execute_genotype_aware(sim, ctx, strict, constraint); + } // Pre-conditions: V must be assigned and refdata must be // available. The schedule analyser already orders us after // `AssembleSegmentPass(V)`, but a hand-built plan can skip @@ -397,6 +704,7 @@ impl ReceptorRevisionPass { new_id, derived_trim_3 as u16, old_v_len, + None, &mut hypothetical_ctx, ); self.validate_post_event_contracts(sim, &hypothetical, ctx)?; @@ -408,6 +716,7 @@ impl ReceptorRevisionPass { new_id, derived_trim_3 as u16, old_v_len, + None, ctx, )) } @@ -492,7 +801,27 @@ impl Pass for ReceptorRevisionPass { // The V replacement distribution is over the same V pool // covered by `refdata_content_hash`; skip it (same // rationale as `SampleAllelePass`) and pin only `prob`. - crate::passes::paramsig::fmt_prob("prob", self.prob) + let base = crate::passes::paramsig::fmt_prob("prob", self.prob); + match &self.genotype_v { + None => base, + Some(g) => { + // Genotype-aware mode pins the constraint state directly so + // two genotypes with different carried-V sets (or different + // same_haplotype) produce distinct plan signatures (replay- + // cache correctness), rather than relying on the genotype + // pass's own signature. + use std::fmt::Write; + let mut s = base; + let _ = write!(s, "|geno_rr|same={}", g.same_haplotype); + for (c, list) in g.per_hap.iter().enumerate() { + let mut rows: Vec<(u32, u64)> = + list.iter().map(|(id, m)| (id.index(), m.to_bits())).collect(); + rows.sort_unstable(); + let _ = write!(s, "|h{}={:?}", c, rows); + } + s + } + } } fn execute(&self, sim: &Simulation, ctx: &mut PassContext) -> Simulation { @@ -517,10 +846,12 @@ impl Pass for ReceptorRevisionPass { } fn requirements(&self) -> Vec { - // V must be assigned (and assembled) before receptor - // revision runs. RefData is required to resolve the - // replacement allele's bytes; the schedule analyser auto- - // derives the dependency from this declaration. + // The engine-enforced requirement is `AlleleAssignment(V)` (+ RefData to + // resolve the replacement allele's bytes); the schedule analyser derives + // the dependency from this declaration. The stronger "V region already + // assembled" guarantee is a lowering/schedule-position contract (Python + // inlines the pass after `push_assemble("J")`), backstopped at runtime by + // `old_v_region_len()` which errors if no V region is present. vec![ PassRequirement::RefData, PassRequirement::AlleleAssignment(Segment::V), @@ -666,6 +997,234 @@ mod tests { ); } + // ── genotype-aware candidate selection ────────────────────── + + fn constraint(per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> GenotypeVConstraint { + GenotypeVConstraint { per_hap, same_haplotype: same } + } + + fn geno_sim(v0: AlleleId) -> Simulation { + sim_v_assembled(v0) + .with_allele_assigned(Segment::V, AlleleInstance::new(v0).with_haplotype(0)) + } + + fn geno_pass(cfg: &RefDataConfig, per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> ReceptorRevisionPass { + ReceptorRevisionPass::new(1.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) + .with_genotype_constraint(constraint(per_hap, same)) + } + + #[test] + fn genotype_same_haplotype_excludes_current_and_restricts_to_chromosome() { + let (cfg, v0, v1) = two_v_refdata(); + // hap0 carries {V0, V1}; hap1 carries {V0}. Current is V0 on hap0. + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let (trace, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap(); + assert_eq!(trace.find("receptor_revision.applied").unwrap().value, ChoiceValue::Bool(true)); + // exclude-current: the only eligible alternate on hap0 is V1 + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); + assert_eq!(after.assignments.get(Segment::V).unwrap().receptor_revision_original_id, Some(v0)); + assert_eq!(after.assignments.get(Segment::V).unwrap().haplotype, Some(0)); + } + + fn run_permissive(pass: &ReceptorRevisionPass, cfg: &RefDataConfig, initial: Simulation) -> (Trace, Simulation) { + let mut trace = Trace::new(); + let mut rng = Rng::new(0xc0ff_ee); + let mut ctx = PassContext { + trace: &mut trace, + rng: &mut rng, + pass_index: 0, + refdata: Some(cfg), + contracts: None, + feasibility: None, + reference_index: None, + replay_cursor: None, + event_log_sink: None, + }; + let next = pass.execute(&initial, &mut ctx); + (trace, next) + } + + #[test] + fn genotype_no_eligible_alternate_permissive_applied_false() { + let (cfg, v0, _v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); + let (trace, after) = run_permissive(&pass, &cfg, geno_sim(v0)); + assert_eq!(trace.find("receptor_revision.applied").unwrap().value, ChoiceValue::Bool(false)); + assert!(trace.find("receptor_revision.v_allele").is_none()); + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v0); + } + + #[test] + fn genotype_no_eligible_alternate_strict_errors() { + let (cfg, v0, _v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap_err(); + assert!(matches!(err, PassError::ConstraintSampling { .. }), "got {err:?}"); + } + + #[test] + fn genotype_both_haplotypes_admits_other_chromosome_allele() { + let (cfg, v0, v1) = two_v_refdata(); + // V1 carried only on hap1; same_haplotype=false aggregates both. + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v1, 1.0)]], false); + let (_t, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap(); + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); + // haplotype provenance preserved as the original rearrangement chromosome (0) + assert_eq!(after.assignments.get(Segment::V).unwrap().haplotype, Some(0)); + } + + #[test] + fn genotype_missing_haplotype_stamp_errors() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + // sim_v_assembled assigns V0 WITHOUT a haplotype stamp. + let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), None, None).unwrap_err(); + assert!(matches!(err, PassError::InvalidPlanState { .. }), "got {err:?}"); + } + + #[test] + fn genotype_strict_post_event_contract_rejection_errors() { + use crate::contract::{Contract, ContractViolation}; + + // A contract whose verify() always rejects the post-event state. + struct RejectAll; + impl Contract for RejectAll { + fn name(&self) -> &str { + "reject_all_test" + } + fn verify( + &self, + _sim: &Simulation, + _refdata: Option<&RefDataConfig>, + ) -> Result<(), ContractViolation> { + Err(ContractViolation::new(self.name(), "rejected by test")) + } + } + + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let contracts = ContractSet::new().with(Box::new(RejectAll)); + // Strict mode (execute_checked) with a rejecting contract must surface a + // contract violation — parity with the non-genotype path. + let err = + run_with_ctx(&pass, &cfg, Some(&contracts), geno_sim(v0), None, None).unwrap_err(); + assert!(matches!(err, PassError::ContractViolation { .. }), "got {err:?}"); + } + + #[test] + fn genotype_replay_strict_post_event_contract_rejection_errors() { + use crate::contract::{Contract, ContractViolation}; + + struct RejectAll; + impl Contract for RejectAll { + fn name(&self) -> &str { + "reject_all_test" + } + fn verify( + &self, + _sim: &Simulation, + _refdata: Option<&RefDataConfig>, + ) -> Result<(), ContractViolation> { + Err(ContractViolation::new(self.name(), "rejected by test")) + } + } + + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let contracts = ContractSet::new().with(Box::new(RejectAll)); + // A valid applied=true replay (v1, trim 2) must STILL be rejected by the + // post-event contract in strict mode — replay parity with the fresh path. + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); + let err = run_with_ctx(&pass, &cfg, Some(&contracts), geno_sim(v0), Some(&mut cursor), None) + .unwrap_err(); + assert!(matches!(err, PassError::ContractViolation { .. }), "got {err:?}"); + } + + // ── genotype-aware replay validation ──────────────────────── + + fn replay_records(applied: bool, allele: Option, trim: Option) -> Vec { + let mut t = Trace::new(); + t.record_choice(address::ChoiceAddress::ReceptorRevisionApplied, ChoiceValue::Bool(applied)); + if let Some(a) = allele { + t.record_choice(address::ChoiceAddress::ReceptorRevisionVAllele, ChoiceValue::AlleleId(a)); + } + if let Some(tr) = trim { + t.record_choice(address::ChoiceAddress::ReceptorRevisionVTrim3, ChoiceValue::Int(tr)); + } + t.choices().to_vec() + } + + #[test] + fn replay_genotype_valid_replacement_reproduces() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); + let (_t, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap(); + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); + assert!(cursor.is_drained()); + } + + #[test] + fn replay_genotype_allele_not_carried_on_haplotype_errors() { + let (cfg, v0, v1) = two_v_refdata(); + // hap0 carries only V0; recorded V1 is not carried on the drawn chromosome. + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0), (v1, 1.0)]], true); + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); + assert!(matches!(err, PassError::InvalidDistributionOutput { .. }), "got {err:?}"); + } + + #[test] + fn replay_genotype_unresolvable_allele_id_reports_missing_allele() { + let (cfg, v0, _v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); + // out-of-range recorded v_allele -> missing_allele (not "not_carried"), + // matching the non-genotype replay diagnostic contract. + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(999), Some(0))); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); + match err { + PassError::MissingAllele { allele_id, .. } => assert_eq!(allele_id, 999), + other => panic!("expected MissingAllele, got {other:?}"), + } + } + + #[test] + fn replay_genotype_equals_current_allele_errors() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v0.index()), Some(0))); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); + assert!(matches!(err, PassError::InvalidDistributionOutput { .. }), "got {err:?}"); + } + + #[test] + fn replay_genotype_trim_length_mismatch_errors() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + // V1 len 8, old 6 -> trim must be 2; record 3 (retains 5) => mismatch. + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(3))); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); + assert!(matches!(err, PassError::InvalidPlanState { .. }), "got {err:?}"); + } + + #[test] + fn genotype_signature_differs_by_candidate_set_and_same_haplotype() { + let (cfg, v0, v1) = two_v_refdata(); + let no_geno = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) + .parameter_signature(); + let g_true = geno_pass_prob(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true).parameter_signature(); + let g_false = geno_pass_prob(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], false).parameter_signature(); + let g_other = geno_pass_prob(&cfg, [vec![(v0, 2.0)], vec![(v1, 1.0)]], true).parameter_signature(); + assert_ne!(g_true, no_geno); + assert_ne!(g_true, g_false); + assert_ne!(g_true, g_other); + } + + fn geno_pass_prob(cfg: &RefDataConfig, per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> ReceptorRevisionPass { + ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) + .with_genotype_constraint(constraint(per_hap, same)) + } + // ── prob=0: no replacement ────────────────────────────────── #[test] diff --git a/engine_rs/src/passes/sample_genotype.rs b/engine_rs/src/passes/sample_genotype.rs index 3b8b4b1..361ee1c 100644 --- a/engine_rs/src/passes/sample_genotype.rs +++ b/engine_rs/src/passes/sample_genotype.rs @@ -208,12 +208,12 @@ impl SampleGenotypePass { ) } - fn commit(&self, seg: Segment, sim: Simulation, id: AlleleId, ctx: &mut PassContext) -> Simulation { + fn commit(&self, seg: Segment, sim: Simulation, id: AlleleId, c: usize, 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)); + b.assign_allele(seg, AlleleInstance::new(id).with_haplotype(c as u8)); if let Some(sink) = ctx.event_log_sink.as_deref_mut() { sink.extend(b.seal_event_log_observer()); } @@ -325,7 +325,7 @@ impl SampleGenotypePass { } ctx.trace .record_choice(ChoiceAddress::SampleAllele(vseg), ChoiceValue::AlleleId(id.index())); - Ok(self.commit(seg, sim, id, ctx)) + Ok(self.commit(seg, sim, id, c, ctx)) } /// Replay (trace-injected) sampling of one segment within `c`. @@ -409,7 +409,7 @@ impl SampleGenotypePass { } ctx.trace .record_choice(ChoiceAddress::SampleAllele(vseg), ChoiceValue::AlleleId(id.index())); - Ok(self.commit(seg, sim, id, ctx)) + Ok(self.commit(seg, sim, id, c, ctx)) } /// Shared execute body. `strict` selects feasibility-filtered @@ -610,6 +610,30 @@ mod tests { } } + #[test] + fn commit_stamps_drawn_haplotype_on_assignments() { + // Only hap0 is viable (hap1 lacks J), so the drawn chromosome is + // always 0; the V and J assignments must carry haplotype Some(0). + 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); + let drawn = match outcome.trace.find("sample_haplotype").unwrap().value { + ChoiceValue::Haplotype(h) => h, + _ => panic!("expected Haplotype record"), + }; + assert_eq!(drawn, 0); + let sim = outcome.final_simulation(); + for seg in [Segment::V, Segment::J] { + assert_eq!( + sim.assignments.get(seg).unwrap().haplotype, + Some(drawn), + "{seg:?} assignment must carry the drawn haplotype", + ); + } + } + #[test] fn records_canonical_sample_allele_and_gene_addresses() { let g = Arc::new(test_support::geno_chrom1_deletes_j()); diff --git a/engine_rs/src/python/plan.rs b/engine_rs/src/python/plan.rs index 9886a56..18d8b80 100644 --- a/engine_rs/src/python/plan.rs +++ b/engine_rs/src/python/plan.rs @@ -1162,6 +1162,77 @@ impl PyPassPlan { Ok(()) } + /// Append a genotype-aware `ReceptorRevisionPass`. `v_rows` is the + /// genotype's carried V slots as `(haplotype, allele_index, copies, + /// weight)`; the pass restricts the replacement V to carried alleles on + /// the drawn chromosome (or both when `same_haplotype` is false), + /// excluding the current V. Validates prob, the V pool, and every row. + fn push_genotype_receptor_revision( + &mut self, + prob: f64, + same_haplotype: bool, + refdata: &PyRefDataConfig, + v_rows: Vec<(u8, u32, u8, f32)>, + ) -> PyResult<()> { + use crate::ir::Segment; + use crate::passes::receptor_revision::GenotypeVConstraint; + use crate::refdata::AlleleId; + if !prob.is_finite() || !(0.0..=1.0).contains(&prob) { + return Err(PyValueError::new_err(format!( + "prob must be a finite number in [0.0, 1.0], got {}", + prob + ))); + } + let cfg = refdata.inner(); + let pool = cfg.pool_for(Segment::V).ok_or_else(|| { + PyValueError::new_err("receptor_revision requires a V pool in refdata") + })?; + if pool.is_empty() { + return Err(PyValueError::new_err( + "receptor_revision requires a non-empty V pool", + )); + } + let pool_len = pool.len() as u32; + let mut per_hap: [Vec<(AlleleId, f64)>; 2] = [Vec::new(), Vec::new()]; + for (h, aid, copies, weight) in &v_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!( + "V allele id {} out of range (pool size {})", + aid, pool_len + ))); + } + if !(*weight as f64).is_finite() || *weight <= 0.0 { + return Err(PyValueError::new_err(format!( + "V candidate weight must be finite and > 0, got {}", + weight + ))); + } + if *copies == 0 { + return Err(PyValueError::new_err( + "V candidate copies must be >= 1", + )); + } + let mass = *weight as f64 * *copies as f64; + if !(mass > 0.0) { + return Err(PyValueError::new_err("V candidate mass must be > 0")); + } + per_hap[*h as usize].push((AlleleId::new(*aid), mass)); + } + let pass = ReceptorRevisionPass::new(prob, Box::new(AllelePoolDist::uniform(pool))) + .with_genotype_constraint(GenotypeVConstraint { + per_hap, + same_haplotype, + }); + self.inner_mut()?.push(Box::new(pass)); + Ok(()) + } + /// Append a `TrimPass` for `(segment, end)`. `end` is `"5"` or /// `"3"` (the prime end being trimmed). `length_pairs` defines /// the trim-amount distribution. diff --git a/site_docs/guides/genotype.md b/site_docs/guides/genotype.md index dedc749..92bb2a5 100644 --- a/site_docs/guides/genotype.md +++ b/site_docs/guides/genotype.md @@ -183,9 +183,11 @@ Notes and guard-rails: - `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)). + expression. `receptor_revision` **is** supported (see + [Receptor revision with a genotype](#receptor-revision-with-a-genotype)); the + clonal forks (`expand_clones` / `clonal_lineage` / `clonal_repertoire`) are still + rejected with a genotype in this release (see + [Limitations](#limitations-this-release)). ### Gene usage @@ -462,8 +464,40 @@ auto-assigned `subject_0..N-1`. Mixed (some set, some not) or duplicate IDs rais a cohort is fully reproducible and subjects are independent. `run_cohort` is mutually exclusive with `with_genotype`, `restrict_alleles`, and -`recombine(*_allele_weights=...)` (the genotype owns allele expression), and — in -this release — is not combined with `receptor_revision` or clonal forks. +`recombine(*_allele_weights=...)` (the genotype owns allele expression). It +**supports** `receptor_revision` (each subject's replacement V is restricted to +its own carried alleles); clonal forks are not combined with a cohort in this +release. + +## Receptor revision with a genotype + +[Receptor revision](../reference/experiment.md) models a post-recombination V +replacement. With a genotype attached, the replacement V is drawn from the +**carried** V alleles on the **drawn rearrangement chromosome** (the haplotype the +original V came from), excluding the current V — so the revised receptor stays +consistent with the individual's germline: + +```python +g = Genotype.sample(cfg, seed=0, subject_id="donor") +res = (ga.Experiment.on(cfg).with_genotype(g) + .recombine().receptor_revision(prob=0.2) # same_haplotype=True by default + .run_records(n=500, seed=1, expose_provenance=True)) +# revised records: original_v_call = pre-revision V; v_call / truth_v_call = the +# carried replacement; receptor_revision_applied = True +``` + +`same_haplotype=False` is a **synthetic control** that draws the replacement from +either chromosome's carried V alleles — useful for ablation studies, but not a +realistic model of secondary V rearrangement (which is a *cis*, same-chromosome +event). Either way the record's `haplotype` provenance keeps naming the original +rearrangement chromosome. + +This is **haplotype-aware V replacement**: it guarantees the replacement is an +allele the individual carries, but it does not model genomic V order, RSS +constraints, upstream-V availability, or deletion of intervening loci. Carried +**novel** alleles on the drawn chromosome are valid replacement targets. Receptor +revision works the same way inside [`run_cohort`](#cohorts) (per subject) and is +still not combined with the clonal forks. ## Novel / private alleles @@ -674,8 +708,6 @@ The genotype foundation is deliberately scoped. Deferred to later work: - **External loaders** — importing genotypes from VDJbase / TIgGER / IgDiscover / partis output. -- **Same-haplotype receptor revision** — `receptor_revision` with a genotype is - rejected for now. ## Backward compatibility diff --git a/src/GenAIRR/_compile.py b/src/GenAIRR/_compile.py index 9fa00e1..352e28a 100644 --- a/src/GenAIRR/_compile.py +++ b/src/GenAIRR/_compile.py @@ -154,9 +154,11 @@ def _extract_receptor_revision_prob(steps): giving the canonical "recombine → revise → mutate/corrupt" order the design doc §2 requires. - Returns ``(receptor_revision_prob, filtered_steps)``. + Returns ``(receptor_revision_prob, receptor_revision_same_haplotype, + filtered_steps)``. """ revision_prob = None + same_haplotype = True filtered = [] for step in steps: if isinstance(step, _ReceptorRevisionStep): @@ -168,9 +170,10 @@ def _extract_receptor_revision_prob(steps): "is a structural bug." ) revision_prob = step.prob + same_haplotype = step.same_haplotype continue filtered.append(step) - return revision_prob, filtered + return revision_prob, same_haplotype, filtered def _name_to_id(refdata, segment): @@ -285,6 +288,7 @@ def _lower_recombine( *, invert_d_prob=None, receptor_revision_prob=None, + receptor_revision_same_haplotype=True, genotype=None, ) -> None: chain = refdata.chain_type @@ -425,7 +429,21 @@ def _lower_recombine( # inlining instead of a standalone _ReceptorRevisionStep lower # path. if receptor_revision_prob is not None: - plan.push_receptor_revision(receptor_revision_prob, refdata) + if genotype is not None: + # Genotype-aware: restrict the replacement V to carried alleles + # on the drawn rearrangement chromosome. `refdata` here is the + # effective refdata (base + injected novels), so carried novel + # V alleles are valid candidates. `v_rows` is the genotype's + # carried V slots as (haplotype, allele_index, copies, weight). + v_rows = _genotype_segment_rows(genotype, refdata, "V") + plan.push_genotype_receptor_revision( + receptor_revision_prob, + receptor_revision_same_haplotype, + refdata, + v_rows, + ) + else: + plan.push_receptor_revision(receptor_revision_prob, refdata) else: # pragma: no cover — RefDataConfig only constructs vj/vdj. raise ValueError(f"unsupported chain_type {chain!r}") diff --git a/src/GenAIRR/_pipeline_ir.py b/src/GenAIRR/_pipeline_ir.py index ac7b6c7..6d900a1 100644 --- a/src/GenAIRR/_pipeline_ir.py +++ b/src/GenAIRR/_pipeline_ir.py @@ -301,9 +301,16 @@ class _ReceptorRevisionStep: biology of receptor revision is heavy-chain-only in v1. :meth:`Experiment.receptor_revision` rejects VJ at the DSL boundary before this dataclass is constructed. + + ``same_haplotype`` (default ``True``) only matters when a genotype + is attached: the replacement V is restricted to the carried V + alleles on the drawn rearrangement chromosome (the cis VH- + replacement model). ``False`` is a synthetic control that draws + from both chromosomes. Ignored on the no-genotype path. """ prob: float + same_haplotype: bool = True @dataclass(frozen=True) diff --git a/src/GenAIRR/experiment.py b/src/GenAIRR/experiment.py index 9633672..1bc95e4 100644 --- a/src/GenAIRR/experiment.py +++ b/src/GenAIRR/experiment.py @@ -2149,7 +2149,7 @@ def invert_d(self, *, prob: float = 0.05) -> "Experiment": self._steps.append(_InvertDStep(prob=prob_f)) return self - def receptor_revision(self, *, prob: float = 0.05) -> "Experiment": + def receptor_revision(self, *, prob: float = 0.05, same_haplotype: bool = True) -> "Experiment": """Append a receptor-revision step. Models post-recombination V-segment replacement: with @@ -2194,12 +2194,11 @@ def receptor_revision(self, *, prob: float = 0.05) -> "Experiment": the plan, giving the canonical "recombine → revise → mutate/corrupt" order the design doc §2 requires. - The DSL does **not** expose ``receptor_revision_applied`` - or ``original_v_call`` on the AIRR record yet — those are - the Slice E follow-up. End-to-end observability today is - via the trace (the three - ``receptor_revision.*`` records above) and the post-event - pool bytes. + AIRR records expose ``receptor_revision_applied`` (bool) and + ``original_v_call`` (the pre-revision V; empty when no + revision applied). ``v_call`` / ``truth_v_call`` are the + **post**-revision V. The three ``receptor_revision.*`` trace + records above remain available for replay. Returns ``self`` so the call chains fluently. """ @@ -2243,7 +2242,14 @@ def receptor_revision(self, *, prob: float = 0.05) -> "Experiment": raise ValueError( f"receptor_revision prob must be in [0.0, 1.0], got {prob_f}" ) - self._steps.append(_ReceptorRevisionStep(prob=prob_f)) + if not isinstance(same_haplotype, bool): + raise ValueError( + "receptor_revision same_haplotype must be a bool, got " + f"{same_haplotype!r}" + ) + self._steps.append( + _ReceptorRevisionStep(prob=prob_f, same_haplotype=same_haplotype) + ) return self def paired_end( @@ -2611,18 +2617,10 @@ 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)" - ) + # Receptor revision with a phased genotype is now haplotype-aware: + # the lowering builds a genotype-aware ReceptorRevisionPass that + # restricts the replacement V to carried alleles on the drawn + # rearrangement chromosome (see _lower_recombine). No rejection here. # Genotype provenance (subject_id / haplotype / result.genotypes) # is only threaded through the plain compiled path, not the @@ -2908,7 +2906,9 @@ def _build_simulator( # schedule edge or place the pass at the end of the plan # (after corruption), both of which break the design doc §2 # ordering. - receptor_revision_prob, steps = _extract_receptor_revision_prob(steps) + receptor_revision_prob, receptor_revision_same_haplotype, steps = ( + _extract_receptor_revision_prob(steps) + ) # Pull out the (at-most-one) paired-end step too. It must # land at the END of the plan, not inline with recombine — # see `_extract_paired_end_step` for the rationale on @@ -2932,6 +2932,7 @@ def _build_simulator( refdata, invert_d_prob=invert_d_prob, receptor_revision_prob=receptor_revision_prob, + receptor_revision_same_haplotype=receptor_revision_same_haplotype, genotype=self._genotype, ) else: @@ -2976,8 +2977,9 @@ def run_cohort( Mutually exclusive with :meth:`with_genotype`, :meth:`restrict_alleles`, and ``recombine(*_allele_weights=...)`` (the genotype owns allele - expression). Not supported with :meth:`receptor_revision` or clonal forks - in this release. + expression). :meth:`receptor_revision` is supported (each subject's + replacement V is restricted to its carried alleles on the drawn + chromosome); clonal forks are not supported in this release. """ import copy as _copy import random as _random @@ -3011,10 +3013,9 @@ def run_cohort( raise ValueError( "run_cohort() and recombine(*_allele_weights=...) are mutually " "exclusive: the genotype owns allele expression") - if any(isinstance(s, _ReceptorRevisionStep) for s in self._steps): - raise ValueError( - "run_cohort() is not supported with receptor_revision() in this " - "release (the revision pass is not haplotype-aware)") + # receptor_revision is supported per subject (each subject's compile + # builds a genotype-aware revision pass restricted to that subject's + # carried alleles on the drawn chromosome) — no rejection here. if self._has_clonal_fork(): raise ValueError( "run_cohort() is not supported together with expand_clones() / " diff --git a/tests/test_genotype_dsl.py b/tests/test_genotype_dsl.py index d0a6f3b..7ac2aac 100644 --- a/tests/test_genotype_dsl.py +++ b/tests/test_genotype_dsl.py @@ -46,11 +46,14 @@ def test_with_genotype_then_recombine_weights_raises(): ga.Experiment.on(_cfg()).with_genotype(g).recombine(v_allele_weights={a1: 2.0}) -def test_receptor_revision_with_genotype_raises_at_compile(): +def test_receptor_revision_with_genotype_compiles_and_runs(): + # Genotype-aware receptor revision is now supported: the replacement V is + # restricted to carried alleles on the drawn chromosome (no longer rejected). 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() + exp.compile() # no raise + res = exp.run_records(n=10, seed=1, expose_provenance=True) + assert len(res) == 10 def test_genotype_with_clonal_fork_raises_at_compile(): diff --git a/tests/test_genotype_receptor_revision.py b/tests/test_genotype_receptor_revision.py new file mode 100644 index 0000000..cef0f62 --- /dev/null +++ b/tests/test_genotype_receptor_revision.py @@ -0,0 +1,133 @@ +"""Genotype-aware (same-haplotype) receptor revision.""" +import pytest + +import GenAIRR as ga +import GenAIRR.data as gdata +from GenAIRR.genotype import Genotype + + +def _cfg(): + return gdata.HUMAN_IGH_OGRDB + + +def _two_allele_v_gene(cfg): + return next(g for g, al in cfg.v_alleles.items() if len(al) >= 2) + + +def test_receptor_revision_same_haplotype_must_be_bool(): + cfg = _cfg() + with pytest.raises(ValueError, match="same_haplotype"): + ga.Experiment.on(cfg).recombine().receptor_revision(prob=0.5, same_haplotype="yes") + + +def test_genotype_plus_receptor_revision_compiles_and_runs(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=0, subject_id="A") + res = (ga.Experiment.on(cfg).with_genotype(g) + .recombine().receptor_revision(prob=1.0) + .run_records(n=10, seed=1, expose_provenance=True)) + assert len(res) == 10 + + +def test_genotype_plus_clonal_fork_still_rejected(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=0, subject_id="A") + with pytest.raises(ValueError, match="clonal|expand_clones|clonal_lineage"): + (ga.Experiment.on(cfg).with_genotype(g) + .recombine().clonal_lineage(n_clones=2).compile()) + + +def test_run_cohort_plus_receptor_revision_runs(): + cfg = _cfg() + gs = [Genotype.sample(cfg, seed=s, subject_id=f"D{s}") for s in range(2)] + c = (ga.Experiment.on(cfg).recombine().receptor_revision(prob=1.0) + .run_cohort(gs, n_per_subject=5, seed=0)) + assert len(c) == 10 + + +def _carried_on_hap(g, seg, hap): + """All allele names carried on chromosome `hap` for a segment (union over + genes), read from the genotype's per-haplotype slots.""" + out = set() + for _gene, haps in g._slots[seg].items(): + out.update(a for (a, _c, _w) in haps[hap]) + return out + + +def test_same_haplotype_revision_truth_is_carried_changed_and_on_drawn_chromosome(): + cfg = _cfg() + g = Genotype.sample(cfg, seed=11, subject_id="A") + res = (ga.Experiment.on(cfg).with_genotype(g) + .recombine().receptor_revision(prob=1.0) + .run_records(n=200, seed=3, expose_provenance=True)) + applied = [r for r in res if r.get("receptor_revision_applied")] + assert applied, "prob=1 with eligible alternates must revise some records" + for r in applied: + tv, ov = r["truth_v_call"], r["original_v_call"] + assert tv and ov and tv != ov # revision changed the V + # revised V is carried on the SAME chromosome the rearrangement drew from + assert tv in _carried_on_hap(g, "V", r["haplotype"]) + + +def _equal_len_two_allele_v_gene(cfg): + for gene, al in cfg.v_alleles.items(): + if len(al) >= 2 and len(al[0].ungapped_seq) == len(al[1].ungapped_seq): + return gene, al[0].name, al[1].name + raise AssertionError("no V gene with two equal-length alleles") + + +def test_same_haplotype_false_admits_other_chromosome(): + # Heterozygous V gene (a0 on hap0, a1 on hap1, equal length); force the + # rearrangement to always draw hap0 via chromosome_weights=(1, 0). a1 is + # carried ONLY on hap1, so it can never be the original V — its appearance as + # a revised truth V proves same_haplotype=False pulled across chromosomes. + cfg = _cfg() + vg, a0, a1 = _equal_len_two_allele_v_gene(cfg) + g = (Genotype.from_dataconfig(cfg) + .heterozygous(vg, a0, a1, segment="V") + .complete_from_reference() + .chromosome_weights(1.0, 0.0) + .with_subject("A")) + res = (ga.Experiment.on(cfg).with_genotype(g) + .recombine().receptor_revision(prob=1.0, same_haplotype=False) + .run_records(n=400, seed=4, expose_provenance=True)) + applied = [r for r in res if r.get("receptor_revision_applied")] + assert applied + # every rearrangement drew hap0, so any revision target carried only on hap1 + # (here a1) is a genuine cross-chromosome replacement. + assert all(r["haplotype"] == 0 for r in res) + assert any(r["truth_v_call"] == a1 for r in applied), ( + "same_haplotype=False must admit the opposite chromosome's allele") + # sanity: a1 is indeed hap1-exclusive in this genotype + assert a1 in _carried_on_hap(g, "V", 1) and a1 not in _carried_on_hap(g, "V", 0) + + +def test_novel_v_can_be_revision_target(): + cfg = _cfg() + vg = _two_allele_v_gene(cfg) + base = cfg.v_alleles[vg][0].name + bs = next(a for a in cfg.v_alleles[vg] if a.name == base).ungapped_seq.upper() + novel_seq = None + for pos in range(len(bs)): + for nt in "ACGT": + if nt == bs[pos]: + continue + cand = bs[:pos] + nt + bs[pos + 1:] + try: + Genotype.from_dataconfig(cfg).add_novel_allele( + f"{vg}*97", base=base, sequence=cand, segment="V") + novel_seq = cand + break + except ValueError: + continue + if novel_seq: + break + novel = f"{vg}*97" + g = (Genotype.from_dataconfig(cfg) + .add_novel_allele(novel, base=base, sequence=novel_seq, segment="V") + .heterozygous(vg, base, novel).complete_from_reference().with_subject("A")) + res = (ga.Experiment.on(cfg).with_genotype(g) + .recombine().receptor_revision(prob=1.0) + .run_records(n=200, seed=5, expose_provenance=True)) + # the carried novel can appear as a revised truth V + assert any(r.get("truth_v_call") == novel for r in res) diff --git a/tests/test_receptor_revision_dsl.py b/tests/test_receptor_revision_dsl.py index 37830d8..fb75391 100644 --- a/tests/test_receptor_revision_dsl.py +++ b/tests/test_receptor_revision_dsl.py @@ -14,8 +14,9 @@ - Calling `.receptor_revision()` twice raises `ValueError`. - prob validation (NaN, out-of-range) raises at the DSL boundary. -No `receptor_revision_applied` / `original_v_call` AIRR field -assertions here — that surface lands in Slice E. +`receptor_revision_applied` / `original_v_call` AIRR-field behavior is +covered in the end-to-end / provenance tests, not here (this file pins the +DSL boundary). """ from __future__ import annotations