From 9043868228422b79de94ea4974b719291afb6c0a Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 27 Aug 2026 16:54:42 +0200 Subject: [PATCH] Add single-pass pairwise reduction pipeline for SNP alignment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the previous two-pass design with a single-pass reduction that computes per-genome-pair substitution counts, cardinality and base-composition tallies, and transition cost matrices. Introduces internal APIs to derive downstream n×n matrices from reduced data and applies configurable filtering based on SNP ratios and exclusion masks. Adds a fused pipeline entry point aggregating alignment and distance outputs for downstream phylogenetic commands. Adjusts visibility of existing reduction functions to restrict external access while permitting sibling module invocation. --- .../src/siblings/algorithms/alignment.rs | 6 +- .../src/siblings/algorithms/cardcomp.rs | 335 ++++++++++++++++++ src/obikphylo/src/siblings/algorithms/mod.rs | 6 + .../src/siblings/algorithms/pairwise.rs | 263 ++++++++++++++ .../src/siblings/algorithms/sankoff.rs | 75 ++++ 5 files changed, 683 insertions(+), 2 deletions(-) create mode 100644 src/obikphylo/src/siblings/algorithms/cardcomp.rs create mode 100644 src/obikphylo/src/siblings/algorithms/pairwise.rs create mode 100644 src/obikphylo/src/siblings/algorithms/sankoff.rs diff --git a/src/obikphylo/src/siblings/algorithms/alignment.rs b/src/obikphylo/src/siblings/algorithms/alignment.rs index 6624cc75..2a017c7a 100644 --- a/src/obikphylo/src/siblings/algorithms/alignment.rs +++ b/src/obikphylo/src/siblings/algorithms/alignment.rs @@ -65,8 +65,10 @@ pub(crate) fn snp_pseudo_alignment( /// Appends one layer's worth of columns onto `sequences` — a "reduce" over /// [`sample_index`]'s per-layer batch, self-contained (no sampling/masking /// logic of its own, just character encoding) so it can be called -/// independently of, and alongside, any other reduce over the same batch. -fn reduce_alignment( +/// independently of, and alongside, any other reduce over the same batch — +/// `pub(super)`: `algorithms::sankoff::sankoff_bundle` runs this one +/// alongside `pairwise::reduce_pairwise` over the same `on_layer` call. +pub(super) fn reduce_alignment( survivors: &[SurvivingFamily], genome_indices: &[usize], free_loss: bool, diff --git a/src/obikphylo/src/siblings/algorithms/cardcomp.rs b/src/obikphylo/src/siblings/algorithms/cardcomp.rs new file mode 100644 index 00000000..74248fab --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/cardcomp.rs @@ -0,0 +1,335 @@ +//! Cardinality × composition decomposition of the 16-state transition cost +//! matrix — see `DevDocMD/theory/evolutionary_distances.md`, "`R` via +//! `exp(-cost)` is wrong for a CTMC; cardinality/composition decomposition". +//! Pure functions, no dependency on `IndexCache`/the sampling pipeline — +//! everything here operates on already-computed [`super::pairwise::CardinalityTally`]/ +//! [`super::pairwise::BasePairTally`]. +//! +//! Two small, directly-estimated first-order Markov models (diagonals +//! included — "stay the same" is a real, calibrated outcome, not implicit): +//! a 5-state cardinality model (`0..=4` members of a family) and a 4-state +//! base-composition model (`A,C,G,T`, unambiguous/cardinality-1 loci only — +//! composition bias is a substitution phenomenon, only meaningful when +//! cardinality is conserved). Composed per pair of 16 states via a +//! parsimony-style pairing of the non-shared elements (see +//! [`best_pairing_cost`]'s doc), then symmetrised — Sankoff parsimony's +//! score is only independent of root placement (as required, since +//! TNT/PhyG search *unrooted* trees) if the cost matrix is symmetric, the +//! discrete-parsimony analogue of CTMC reversibility. + +use super::pairwise::{BasePairTally, CardinalityTally}; + +/// Row-stochastic 5×5 cardinality transition probabilities (`0..=4`), +/// diagonal included ("stay at the same cardinality"), from +/// [`CardinalityTally`]'s pooled co-occurrence counts. +pub(crate) fn cardinality_transition_probs(tally: &CardinalityTally) -> [[f64; 5]; 5] { + let mut p = [[0.0f64; 5]; 5]; + for a in 0..5 { + let row_sum: u64 = tally.counts[a].iter().sum(); + if row_sum == 0 { + continue; + } + for b in 0..5 { + p[a][b] = tally.counts[a][b] as f64 / row_sum as f64; + } + } + p +} + +/// Row-stochastic 4×4 base-composition transition probabilities +/// (`A,C,G,T`), diagonal included ("stay the same base"), from +/// [`BasePairTally`]'s pooled substitution (off-diagonal) and agreement +/// (`same`, diagonal) counts — unambiguous, cardinality-1 loci only. +pub(crate) fn composition_transition_probs(tally: &BasePairTally) -> [[f64; 4]; 4] { + let mut p = [[0.0f64; 4]; 4]; + for a in 0..4 { + let row_sum = tally.same[a] + (0..4).map(|b| tally.counts[a][b]).sum::(); + if row_sum == 0 { + continue; + } + p[a][a] = tally.same[a] as f64 / row_sum as f64; + for b in 0..4 { + if a != b { + p[a][b] = tally.counts[a][b] as f64 / row_sum as f64; + } + } + } + p +} + +/// Enumerate every injective mapping from `small` (indices `0..small.len()`) +/// into `large` (indices `0..large.len()`, `large.len() >= small.len()`), +/// calling `visit` with each full permutation of chosen `large` indices. +/// `small.len() <= 4` always here (at most 4 bases can be lost or gained), +/// so a plain recursive enumeration (at most `4! = 24` calls) is simplest — +/// not worth a dependency for. +fn for_each_injection(small_len: usize, large_len: usize, visit: &mut dyn FnMut(&[usize])) { + let mut chosen = vec![0usize; small_len]; + let mut used = vec![false; large_len]; + fn go(pos: usize, chosen: &mut [usize], used: &mut [bool], visit: &mut dyn FnMut(&[usize])) { + if pos == chosen.len() { + visit(chosen); + return; + } + for l in 0..used.len() { + if used[l] { + continue; + } + used[l] = true; + chosen[pos] = l; + go(pos + 1, chosen, used, visit); + used[l] = false; + } + } + if small_len == 0 { + visit(&[]); + return; + } + go(0, &mut chosen, &mut used, visit); +} + +/// Parsimony-style pairing cost for the non-shared elements of a +/// transition: pair every element of the smaller of `lost`/`gained` with +/// some element of the larger one, choosing the pairing that minimises +/// total `-ln(P_composition)` (maximises likelihood) — the discrete +/// analogue of "prefer a substitution over an independent loss+gain". +/// Leftover elements of the larger side (`|large|-|small|` of them) are +/// *not* charged here — they're pure cardinality change, already priced by +/// `P_cardinality(|A|→|B|)` in the caller. +fn best_pairing_cost(lost: &[u8], gained: &[u8], p_comp: &[[f64; 4]; 4]) -> f64 { + let (small, large) = if lost.len() <= gained.len() { (lost, gained) } else { (gained, lost) }; + if small.is_empty() { + return 0.0; + } + let mut best = f64::INFINITY; + for_each_injection(small.len(), large.len(), &mut |perm| { + let mut cost = 0.0; + for (i, &li) in perm.iter().enumerate() { + let (x, y) = (small[i], large[li]); + let p = p_comp[x as usize][y as usize]; + cost += if p > 0.0 { -p.ln() } else { f64::INFINITY }; + } + if cost < best { + best = cost; + } + }); + best +} + +/// The full 16-state transition cost matrix. For every pair of states `A`, +/// `B` (bitmasks, bit `0..4` = `A,C,G,T`): `shared = A∩B` contributes +/// `P_composition(x→x)` per shared base (diagonal — "stayed"); the +/// remaining `lost = A\B`, `gained = B\A` are parsimony-paired into +/// substitutions via [`best_pairing_cost`], and whatever's left over after +/// pairing (only possible on one side, since one side is fully consumed) +/// is *pure* cardinality change, priced once via `P_cardinality(|A|→|B|)` +/// — never chained through intermediate states. Row-normalised (`P(A→·)` +/// sums to `1` over all `B`), converted to a cost via `-ln`, then +/// symmetrised (`(cost(A,B)+cost(B,A))/2` — the row-normalised `P` is not +/// symmetric in general, but a Sankoff parsimony cost must be, so the +/// score is independent of where an unrooted tree gets rooted). +/// +/// `free_loss`: drop the `P_cardinality(|A|→|B|)` factor entirely (never +/// added to `log_p`) — the same low/incomplete-coverage argument that +/// justifies recoding whole-family non-detection as `?` applies one level +/// down: whether a genome shows 1 vs 2 (etc.) detected members of a +/// *present* family is exactly as vulnerable to sampling failure as +/// whether the family was detected at all. Without this, `∅`-involving +/// transitions are neutralised (via the `?` recoding, bypassing this +/// matrix's row/column 0 entirely) but cardinality changes *between two +/// otherwise-detected, non-empty* states (e.g. `{A} -> {A,C}`) still +/// carried the same calibrated `P_cardinality` penalty as any other +/// gain/loss — inconsistent with `--free-loss`'s own rationale. With the +/// factor dropped, cost is driven only by composition matching +/// (shared-base retention and paired substitutions), never by a state +/// pair's cardinality difference alone. +pub(crate) fn pairwise_cost_matrix( + p_card: &[[f64; 5]; 5], + p_comp: &[[f64; 4]; 4], + free_loss: bool, +) -> [[f64; 16]; 16] { + let mut log_p = [[0.0f64; 16]; 16]; // ln(P), *before* row-normalisation + for a in 0u8..16 { + for b in 0u8..16 { + let shared = a & b; + let lost: Vec = (0..4).filter(|&i| a & (1 << i) != 0 && b & (1 << i) == 0).collect(); + let gained: Vec = (0..4).filter(|&i| b & (1 << i) != 0 && a & (1 << i) == 0).collect(); + + let mut lp = 0.0; // accumulate ln(P), so 0.0 = probability 1 + if !free_loss { + let card_a = a.count_ones() as usize; + let card_b = b.count_ones() as usize; + let p_c = p_card[card_a][card_b]; + lp += if p_c > 0.0 { p_c.ln() } else { f64::NEG_INFINITY }; + } + + for i in 0..4u8 { + if shared & (1 << i) != 0 { + let p = p_comp[i as usize][i as usize]; + lp += if p > 0.0 { p.ln() } else { f64::NEG_INFINITY }; + } + } + lp -= best_pairing_cost(&lost, &gained, p_comp); + + log_p[a as usize][b as usize] = lp; + } + } + + // Row-normalise and convert to cost entirely in log-space + // (log-sum-exp), never exponentiating a raw `log_p` value directly — + // for a transition reachable only through several low-probability + // steps, `log_p.exp()` can underflow to exactly `0.0` (anything below + // roughly `-709` does, in `f64`), silently turning a real, if small, + // probability into a hard `+∞` cost. `log_sum_exp` never exponentiates + // anything above `0` (every term is shifted by the row's own max + // first), so it stays accurate across the full range `f64` can + // represent, not just what survives a direct `exp()`. + let mut cost = [[0.0f64; 16]; 16]; + for a in 0..16 { + let row_max = log_p[a].iter().cloned().fold(f64::NEG_INFINITY, f64::max); + if row_max == f64::NEG_INFINITY { + // Every transition out of this state has probability 0 in the + // calibration data — genuinely unreachable, not underflow. + cost[a] = [f64::INFINITY; 16]; + continue; + } + let log_sum = row_max + log_p[a].iter().map(|&lp| (lp - row_max).exp()).sum::().ln(); + for b in 0..16 { + // `log_sum - log_p[a][b]` is `+∞` automatically when + // `log_p[a][b] == -∞` (IEEE 754: finite − (−∞) = +∞) — no + // separate branch needed for a genuinely zero-probability + // transition. + cost[a][b] = log_sum - log_p[a][b]; + } + } + + // Symmetrise (see doc comment: parsimony on an unrooted tree requires + // a symmetric cost matrix). + let mut sym = [[0.0f64; 16]; 16]; + for a in 0..16 { + for b in 0..16 { + sym[a][b] = if a == b { 0.0 } else { (cost[a][b] + cost[b][a]) / 2.0 }; + } + } + sym +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cardinality_probs_row_normalised() { + let mut counts = [[0u64; 5]; 5]; + counts[0][0] = 10; + counts[0][1] = 5; + counts[1][0] = 5; + counts[1][1] = 2; + let tally = CardinalityTally { counts }; + let p = cardinality_transition_probs(&tally); + assert!((p[0].iter().sum::() - 1.0).abs() < 1e-12); + assert!((p[1].iter().sum::() - 1.0).abs() < 1e-12); + assert!((p[0][0] - 10.0 / 15.0).abs() < 1e-12); + } + + #[test] + fn composition_probs_include_diagonal() { + let mut tally = BasePairTally { counts: [[0u64; 4]; 4], same: [0u64; 4] }; + tally.same[0] = 90; // A stays A 90 times + tally.counts[0][1] = 10; // A -> C 10 times + tally.counts[1][0] = 10; + let p = composition_transition_probs(&tally); + assert!((p[0][0] - 0.9).abs() < 1e-12); + assert!((p[0][1] - 0.1).abs() < 1e-12); + assert!((p[0].iter().sum::() - 1.0).abs() < 1e-12); + } + + #[test] + fn pairwise_cost_matrix_is_symmetric_and_zero_diagonal() { + // A plausible-looking pair of models (not calibrated from real + // data — just needs every state reachable with nonzero probability + // for this structural test). + let p_card = [ + [0.5, 0.3, 0.15, 0.04, 0.01], + [0.3, 0.4, 0.2, 0.08, 0.02], + [0.15, 0.2, 0.4, 0.2, 0.05], + [0.04, 0.08, 0.2, 0.5, 0.18], + [0.01, 0.02, 0.05, 0.18, 0.74], + ]; + let p_comp = [ + [0.7, 0.1, 0.15, 0.05], + [0.1, 0.7, 0.05, 0.15], + [0.15, 0.05, 0.7, 0.1], + [0.05, 0.15, 0.1, 0.7], + ]; + let cost = pairwise_cost_matrix(&p_card, &p_comp, false); + for a in 0..16 { + assert_eq!(cost[a][a], 0.0); + for b in 0..16 { + assert!((cost[a][b] - cost[b][a]).abs() < 1e-9, "cost[{a}][{b}]={} cost[{b}][{a}]={}", cost[a][b], cost[b][a]); + } + } + } + + #[test] + fn free_loss_ignores_cardinality_transition_probs() { + // A skewed cardinality model (cardinality change made artificially + // expensive) must have zero effect on the cost matrix once + // `free_loss` is set — the whole point of the flag. + let p_card_uniform = [[0.2f64; 5]; 5]; + let p_card_skewed = [ + [0.96, 0.01, 0.01, 0.01, 0.01], + [0.01, 0.96, 0.01, 0.01, 0.01], + [0.01, 0.01, 0.96, 0.01, 0.01], + [0.01, 0.01, 0.01, 0.96, 0.01], + [0.01, 0.01, 0.01, 0.01, 0.96], + ]; + let p_comp = [ + [0.7, 0.1, 0.15, 0.05], + [0.1, 0.7, 0.05, 0.15], + [0.15, 0.05, 0.7, 0.1], + [0.05, 0.15, 0.1, 0.7], + ]; + let cost_uniform = pairwise_cost_matrix(&p_card_uniform, &p_comp, true); + let cost_skewed = pairwise_cost_matrix(&p_card_skewed, &p_comp, true); + for a in 0..16 { + for b in 0..16 { + assert!( + (cost_uniform[a][b] - cost_skewed[a][b]).abs() < 1e-9, + "cost[{a}][{b}] differs between cardinality models under free_loss: {} vs {}", + cost_uniform[a][b], cost_skewed[a][b], + ); + } + } + } + + #[test] + fn underflow_prone_transition_gets_finite_cost_not_infinite() { + // Every off-diagonal composition probability is tiny (1e-200) but + // not exactly zero — small enough that the naive `log_p.exp()` + // (exponentiate before row-normalising) hard-underflows to `0.0` + // in `f64` well before the probability is truly zero, silently + // turning a real, if minuscule, probability into `+∞` cost. The + // log-sum-exp normalisation in `pairwise_cost_matrix` must keep + // this finite instead. + let p_card = [[0.2f64; 5]; 5]; + let tiny = 1e-200; + let mut p_comp = [[tiny; 4]; 4]; + for i in 0..4 { + p_comp[i][i] = 1.0 - 3.0 * tiny; + } + + let cost = pairwise_cost_matrix(&p_card, &p_comp, false); + + // {A,C} -> {G,T}: no shared bases, both members substituted — the + // shape most prone to underflow (several tiny-probability factors + // multiplied together). + let a = 0b0011u8; // A, C + let b = 0b1100u8; // G, T + assert!( + cost[a as usize][b as usize].is_finite(), + "cost must stay finite for a merely tiny (not exactly zero) probability, got {}", + cost[a as usize][b as usize] + ); + } +} diff --git a/src/obikphylo/src/siblings/algorithms/mod.rs b/src/obikphylo/src/siblings/algorithms/mod.rs index 9c0d95d5..2de23b8a 100644 --- a/src/obikphylo/src/siblings/algorithms/mod.rs +++ b/src/obikphylo/src/siblings/algorithms/mod.rs @@ -15,10 +15,13 @@ mod alignment; mod annex; +mod cardcomp; mod entropy; mod family_scan; mod masking; mod minorant_selection; +mod pairwise; +mod sankoff; mod stats; mod subsample; @@ -30,8 +33,11 @@ pub use subsample::EntropyBias; pub(crate) use alignment::snp_pseudo_alignment; pub(crate) use annex::build_layer_sibling_annex; +pub(crate) use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix}; pub(crate) use entropy::{ensure_layer_entropy_annex, family_entropy, family_entropy_4, iter_full_entropy}; pub(crate) use family_scan::{Selection, scan_layer_families}; +pub(crate) use pairwise::{BasePairTally, CardinalityTally, RawSnpDistanceOutput}; +pub(crate) use sankoff::{SankoffBundle, sankoff_bundle}; pub(crate) use stats::{sibling_annex_stats, sibling_family_size_histogram}; pub(crate) use subsample::sample_index; diff --git a/src/obikphylo/src/siblings/algorithms/pairwise.rs b/src/obikphylo/src/siblings/algorithms/pairwise.rs new file mode 100644 index 00000000..5307002c --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/pairwise.rs @@ -0,0 +1,263 @@ +//! Per-genome-pair composition/cardinality tally — a single reduce over +//! [`super::subsample::SurvivingFamily`] batches, storing enough per pair +//! (`i < j`, upper triangle only — reciprocal by construction, half the +//! naive `n × n` storage) that every one of [`RawSnpDistanceOutput`], +//! [`PairwiseTally::included`]'s `ratio_ceiling` filter, +//! [`BasePairTally`] and [`CardinalityTally`] derives from it afterward as +//! cheap `O(n²)` post-processing — never a second scan of the index. This +//! replaces the old two-pass design (`raw_snp_distance` fully computed +//! before `base_pair_tally`/`cardinality_tally` could even start, because +//! only the *complete* aggregate told you which pairs were +//! `ratio_ceiling`-eligible): keeping the *unfiltered* per-pair detail +//! instead of collapsing straight to a pooled tally means the filter can be +//! applied after the fact, on data already in hand. +//! +//! Deliberately `O(n²)` in genome count (unlike everything else in this +//! crate, which is `O(n)` or bounded per layer/round) — accepted as a +//! one-time cost for the target scale this project runs at (large-memory +//! machines), not something to optimise away. + +use ndarray::Array2; + +use super::subsample::SurvivingFamily; + +/// One pair's own tallies, unfiltered — `subst` is populated only by +/// variable families (a substitution, `bi != bj`, is only possible when the +/// family varies *somewhere* in the index; see the module docs), so it +/// needs no separate `variable` gate of its own. `same_all` counts every +/// agreement (`bi == bj`, any family); `same_variable` is the subset from +/// variable families only — [`RawSnpDistanceOutput`] wants the former, +/// [`BasePairTally`] the latter (they are *not* the same population — see +/// `BasePairTally::same`'s own docs). `cardinality` is restricted to +/// variable families only, matching [`CardinalityTally`]'s own scope. +#[derive(Default, Clone, Copy)] +struct PairStats { + subst: [[u64; 4]; 4], + same_all: [u64; 4], + same_variable: [u64; 4], + cardinality: [[u64; 5]; 5], +} + +/// See the module docs. `pub(crate)`: consumed by +/// `algorithms::sankoff`'s orchestration, not yet exposed past this crate. +pub(crate) struct PairwiseTally { + n_genomes: usize, + /// Upper triangle only (`i < j`), flat-indexed via [`Self::flat_index`]. + pairs: Vec, +} + +impl PairwiseTally { + pub(crate) fn new(n_genomes: usize) -> Self { + let n_pairs = n_genomes.saturating_sub(1) * n_genomes / 2; + Self { n_genomes, pairs: vec![PairStats::default(); n_pairs] } + } + + /// Flat index of unordered pair `(i, j)`, `i != j`, into the upper + /// triangle — standard "row-major over `i < j`" packing: row `i` + /// contributes `n_genomes - 1 - i` entries (columns `i+1..n_genomes`), + /// so pair `(i, j)` sits at `sum_{r usize { + let (i, j) = if i < j { (i, j) } else { (j, i) }; + debug_assert!(j < self.n_genomes && i != j); + let n = self.n_genomes; + (i * (2 * n - i - 1)) / 2 + (j - i - 1) + } + + #[inline] + fn pair(&self, i: usize, j: usize) -> &PairStats { + &self.pairs[self.flat_index(i, j)] + } + + #[inline] + fn pair_mut(&mut self, i: usize, j: usize) -> &mut PairStats { + let idx = self.flat_index(i, j); + &mut self.pairs[idx] + } + + /// This pair's total eligible-locus count (`snp + shared`, unfiltered) + /// and its raw SNP ratio (`snp / (snp + shared)`, `None` if no eligible + /// locus at all) — the two quantities [`Self::included`] and + /// [`RawSnpDistanceOutput`] both derive from. + fn snp_shared(&self, i: usize, j: usize) -> (u64, u64) { + let stats = self.pair(i, j); + // Every substitution was recorded twice (`subst[bi][bj]` and + // `subst[bj][bi]`, matching `BasePairTally::counts`'s own + // reciprocal convention) — halve the flat sum back to a true count. + let snp: u64 = stats.subst.iter().flatten().sum::() / 2; + let shared: u64 = stats.same_all.iter().sum(); + (snp, shared) + } + + pub(crate) fn raw_snp_distance(&self) -> RawSnpDistanceOutput { + let mut snp = Array2::::zeros((self.n_genomes, self.n_genomes)); + let mut shared = Array2::::zeros((self.n_genomes, self.n_genomes)); + for i in 0..self.n_genomes { + for j in (i + 1)..self.n_genomes { + let (s, sh) = self.snp_shared(i, j); + snp[[i, j]] = s; + snp[[j, i]] = s; + shared[[i, j]] = sh; + shared[[j, i]] = sh; + } + } + RawSnpDistanceOutput { snp, shared } + } + + /// Whether pair `(i, j)` is eligible for `base_pair_tally`/ + /// `cardinality_tally`'s calibration pool: enough eligible loci to have + /// a real SNP ratio, that ratio at or below `ratio_ceiling` (a pair this + /// close to saturation carries no information about the true + /// substitution rate — see `--sankoff-ratio-ceiling`'s own docs), and + /// neither genome excluded (`excluded[g]`, `--exclude-genome`) — an + /// excluded genome must never influence the calibrated matrices, not + /// just be absent from the reported output. + pub(crate) fn included(&self, ratio_ceiling: f64, excluded: &[bool]) -> Array2 { + Array2::from_shape_fn((self.n_genomes, self.n_genomes), |(i, j)| { + if i == j { + return false; + } + if excluded.get(i).copied().unwrap_or(false) || excluded.get(j).copied().unwrap_or(false) { + return false; + } + let (snp, shared) = self.snp_shared(i, j); + let total = snp + shared; + total > 0 && (snp as f64 / total as f64) <= ratio_ceiling + }) + } + + pub(crate) fn base_pair_tally(&self, included: &Array2) -> BasePairTally { + let mut counts = [[0u64; 4]; 4]; + let mut same = [0u64; 4]; + for i in 0..self.n_genomes { + for j in (i + 1)..self.n_genomes { + if !included[[i, j]] { + continue; + } + let stats = self.pair(i, j); + for a in 0..4 { + same[a] += stats.same_variable[a]; + for b in 0..4 { + counts[a][b] += stats.subst[a][b]; + } + } + } + } + BasePairTally { counts, same } + } + + /// Unlike [`Self::base_pair_tally`], **not** gated by `ratio_ceiling`: + /// that filter exists because a substitution-saturated pair's base + /// composition trends toward random noise, no longer informative about + /// the true substitution spectrum — a fact about *sequence divergence* + /// between the two genomes. Cardinality (how many family members each + /// genome carries) reflects each genome's own coverage/assembly + /// completeness and duplication/paralogy structure, not the pair's + /// mutual divergence — reusing the SNP-ratio filter here would exclude + /// pairs for a reason that has no established bearing on cardinality + /// co-occurrence. Only `--exclude-genome` applies. + pub(crate) fn cardinality_tally(&self, excluded: &[bool]) -> CardinalityTally { + let mut counts = [[0u64; 5]; 5]; + for i in 0..self.n_genomes { + if excluded.get(i).copied().unwrap_or(false) { + continue; + } + for j in (i + 1)..self.n_genomes { + if excluded.get(j).copied().unwrap_or(false) { + continue; + } + let stats = self.pair(i, j); + for a in 0..5 { + for b in 0..5 { + counts[a][b] += stats.cardinality[a][b]; + } + } + } + } + CardinalityTally { counts } + } +} + +/// Raw p-distance restricted to loci that are single-copy in **both** +/// genomes of a pair — see `PairwiseTally`'s module docs. +pub(crate) struct RawSnpDistanceOutput { + /// n×n count of eligible loci where the two genomes' single forms differ. + pub snp: Array2, + /// n×n count of eligible loci where the two genomes' single forms agree. + pub shared: Array2, +} + +/// Symmetric 6-category base-pair substitution tally (indexed +/// `0=A,1=C,2=G,3=T`), pooled over [`PairwiseTally::included`] genome pairs. +pub(crate) struct BasePairTally { + /// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled over + /// included genome pairs, where the two genomes' single forms are `a` + /// and `b`. Diagonal always `0` — an `a == b` locus is counted in + /// `same`, not here. + pub counts: [[u64; 4]; 4], + /// `same[a]` = number of eligible loci, pooled over included genome + /// pairs, where both genomes' single forms are `a` **and** the family + /// varies somewhere in the index — a fully invariant family's + /// agreement is genome-wide background, not SNP-adjacent signal (see + /// `PairStats::same_all` vs `same_variable`). + pub same: [u64; 4], +} + +/// Cardinality (`0..=4` members of a family) co-occurrence, pooled over +/// every genome pair except `--exclude-genome` ones (see +/// [`PairwiseTally::cardinality_tally`]'s own docs for why this is *not* +/// gated by `ratio_ceiling` the way [`BasePairTally`] is), restricted to +/// variable families (`family_size() >= 2`) — matching +/// `snp_pseudo_alignment`'s own scope. +pub(crate) struct CardinalityTally { + /// `counts[a][b] == counts[b][a]` = number of family sites, pooled over + /// included genome pairs, where one genome's family cardinality is `a` + /// and the other's is `b`. Diagonal is real data here (both genomes at + /// the same cardinality), unlike [`BasePairTally::counts`]. + pub counts: [[u64; 5]; 5], +} + +/// One layer's worth of [`SurvivingFamily`] folded into `tally` — the +/// "reduce" half of the map/reduce [`super::subsample::sample_index`] +/// drives, self-contained (no sampling/masking logic of its own) so it can +/// run alongside any other reduce over the same batch (see +/// `algorithms::alignment::reduce_alignment`, called from the same +/// `on_layer` closure in `algorithms::sankoff::sankoff_bundle`). +pub(crate) fn reduce_pairwise(survivors: &[SurvivingFamily], tally: &mut PairwiseTally) { + let n_genomes = tally.n_genomes; + for family in survivors { + let variable = family.mask.family_size() >= 2; + let genome_mask = &family.genome_mask; + let single_form = |g: usize| -> Option { + let m = genome_mask[g]; + (m.count_ones() == 1).then(|| m.trailing_zeros() as u8) + }; + + for i in 0..n_genomes { + let bi = single_form(i); + let card_i = genome_mask[i].count_ones() as usize; + for j in (i + 1)..n_genomes { + let stats = tally.pair_mut(i, j); + + if let (Some(bi), Some(bj)) = (bi, single_form(j)) { + if bi != bj { + stats.subst[bi as usize][bj as usize] += 1; + stats.subst[bj as usize][bi as usize] += 1; + } else { + stats.same_all[bi as usize] += 1; + if variable { + stats.same_variable[bi as usize] += 1; + } + } + } + + if variable { + let card_j = genome_mask[j].count_ones() as usize; + stats.cardinality[card_i][card_j] += 1; + if card_i != card_j { + stats.cardinality[card_j][card_i] += 1; + } + } + } + } + } +} diff --git a/src/obikphylo/src/siblings/algorithms/sankoff.rs b/src/obikphylo/src/siblings/algorithms/sankoff.rs new file mode 100644 index 00000000..85ac3062 --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/sankoff.rs @@ -0,0 +1,75 @@ +//! Fused entry point for the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` +//! pipeline: one [`sample_index`] call (one Bernoulli draw, shared by +//! everyone — the same family selection feeds the alignment and every +//! tally, never two independent samples of the same index) driving two +//! reduces over the same `on_layer` batches — [`reduce_alignment`] and +//! [`reduce_pairwise`] — so [`SnpAlignment`], [`RawSnpDistanceOutput`], +//! [`BasePairTally`] and [`CardinalityTally`] all come out of a single +//! scan. The latter three are cheap `O(n²)` post-processing over +//! [`PairwiseTally`] once the scan is done (see its own module docs for why +//! that removes the old two-pass `raw_snp_distance`-before-`base_pair_tally` +//! dependency entirely). + +use obikidxcache::index_cache::IndexCache; +use obikindex::OKIResult; + +use super::alignment::{SnpAlignment, reduce_alignment}; +use super::pairwise::{BasePairTally, CardinalityTally, PairwiseTally, RawSnpDistanceOutput, reduce_pairwise}; +use super::subsample::{EntropyBias, sample_index}; + +/// Every output the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` pipeline needs, +/// computed together from one shared, possibly-subsampled/entropy-biased +/// selection — see the module docs. +pub(crate) struct SankoffBundle { + pub alignment: SnpAlignment, + pub raw: RawSnpDistanceOutput, + pub base_pair_tally: BasePairTally, + pub cardinality_tally: CardinalityTally, +} + +/// See [`crate::siblings::extensions::SiblingExt::sankoff_bundle`] for the +/// public-facing docs — this is its implementation. `ratio_ceiling` gates +/// [`BasePairTally`] only (see [`PairwiseTally::cardinality_tally`]'s own +/// docs for why [`CardinalityTally`] uses a different — `excluded`-only — +/// inclusion rule). +pub(crate) fn sankoff_bundle( + cache: &IndexCache, + n: usize, + free_loss: bool, + no_ambiguity: bool, + excluded: &[bool], + entropy_bias: Option, + ratio_ceiling: f64, +) -> OKIResult { + let n_genomes = cache.meta().genomes().len(); + let genome_indices: Vec = (0..n_genomes) + .filter(|&g| !excluded.get(g).copied().unwrap_or(false)) + .collect(); + let mut sequences: Vec> = vec![Vec::new(); genome_indices.len()]; + let mut tally = PairwiseTally::new(n_genomes); + + sample_index( + cache, + n, + free_loss, + no_ambiguity, + excluded, + entropy_bias, + |_partition, _layer, survivors| { + reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences); + reduce_pairwise(&survivors, &mut tally); + }, + )?; + + let raw = tally.raw_snp_distance(); + let included = tally.included(ratio_ceiling, excluded); + let base_pair_tally = tally.base_pair_tally(&included); + let cardinality_tally = tally.cardinality_tally(excluded); + + Ok(SankoffBundle { + alignment: SnpAlignment { sequences, genome_indices }, + raw, + base_pair_tally, + cardinality_tally, + }) +}