Add single-pass pairwise reduction pipeline for SNP alignment

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.
This commit is contained in:
Eric Coissac
2026-08-28 21:10:16 +02:00
parent e015362ce6
commit 9043868228
5 changed files with 683 additions and 2 deletions
@@ -65,8 +65,10 @@ pub(crate) fn snp_pseudo_alignment(
/// Appends one layer's worth of columns onto `sequences` — a "reduce" over /// Appends one layer's worth of columns onto `sequences` — a "reduce" over
/// [`sample_index`]'s per-layer batch, self-contained (no sampling/masking /// [`sample_index`]'s per-layer batch, self-contained (no sampling/masking
/// logic of its own, just character encoding) so it can be called /// logic of its own, just character encoding) so it can be called
/// independently of, and alongside, any other reduce over the same batch. /// independently of, and alongside, any other reduce over the same batch
fn reduce_alignment( /// `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], survivors: &[SurvivingFamily],
genome_indices: &[usize], genome_indices: &[usize],
free_loss: bool, free_loss: bool,
@@ -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::<u64>();
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<u8> = (0..4).filter(|&i| a & (1 << i) != 0 && b & (1 << i) == 0).collect();
let gained: Vec<u8> = (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::<f64>().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::<f64>() - 1.0).abs() < 1e-12);
assert!((p[1].iter().sum::<f64>() - 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::<f64>() - 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]
);
}
}
@@ -15,10 +15,13 @@
mod alignment; mod alignment;
mod annex; mod annex;
mod cardcomp;
mod entropy; mod entropy;
mod family_scan; mod family_scan;
mod masking; mod masking;
mod minorant_selection; mod minorant_selection;
mod pairwise;
mod sankoff;
mod stats; mod stats;
mod subsample; mod subsample;
@@ -30,8 +33,11 @@ pub use subsample::EntropyBias;
pub(crate) use alignment::snp_pseudo_alignment; pub(crate) use alignment::snp_pseudo_alignment;
pub(crate) use annex::build_layer_sibling_annex; 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 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 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 stats::{sibling_annex_stats, sibling_family_size_histogram};
pub(crate) use subsample::sample_index; pub(crate) use subsample::sample_index;
@@ -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<PairStats>,
}
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<i}(n_genomes-1-r) + (j-i-1)`.
fn flat_index(&self, i: usize, j: usize) -> 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::<u64>() / 2;
let shared: u64 = stats.same_all.iter().sum();
(snp, shared)
}
pub(crate) fn raw_snp_distance(&self) -> RawSnpDistanceOutput {
let mut snp = Array2::<u64>::zeros((self.n_genomes, self.n_genomes));
let mut shared = Array2::<u64>::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<bool> {
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<bool>) -> 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<u64>,
/// n×n count of eligible loci where the two genomes' single forms agree.
pub shared: Array2<u64>,
}
/// 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<u8> {
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;
}
}
}
}
}
}
@@ -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<EntropyBias>,
ratio_ceiling: f64,
) -> OKIResult<SankoffBundle> {
let n_genomes = cache.meta().genomes().len();
let genome_indices: Vec<usize> = (0..n_genomes)
.filter(|&g| !excluded.get(g).copied().unwrap_or(false))
.collect();
let mut sequences: Vec<Vec<u8>> = 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,
})
}