Push zunrplorkwkt #70

Merged
coissac merged 93 commits from push-zunrplorkwkt into main 2026-08-28 23:15:38 +00:00
5 changed files with 683 additions and 2 deletions
Showing only changes of commit 9043868228 - Show all commits
@@ -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,
@@ -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 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;
@@ -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,
})
}