Genotypes: per-individual diploid germline (phased recombination + novel alleles)

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.<invalid>";
 pub const SAMPLE_ALLELE_UNSUPPORTED: &str = "sample_allele.<unsupported>";
+/// 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<ChoiceAddress> {
         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<ChoiceAddressPattern> {
         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",
+        );
     }
 }

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())
     }
 

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<GeneId, Vec<GeneCopy>>,
+    d: HashMap<GeneId, Vec<GeneCopy>>,
+    j: HashMap<GeneId, Vec<GeneCopy>>,
+}
+
+impl Haplotype {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    fn map(&self, seg: Segment) -> &HashMap<GeneId, Vec<GeneCopy>> {
+        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<GeneId, Vec<GeneCopy>> {
+        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<GeneCopy>) {
+        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<Item = GeneId> + '_ {
+        let mut genes: Vec<GeneId> = 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<F: Fn(GeneId) -> 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<AlleleId> {
+        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<String>,
+    source_refdata_hash: String,
+}
+
+impl Genotype {
+    pub fn new(
+        haplotypes: [Haplotype; 2],
+        chromosome_weights: [f32; 2],
+        subject_id: Option<String>,
+        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<GeneId> = 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)]);
+    }
+}

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;

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)]