From 0c86ea03850d6ddc621350a310a3d5199c4d5d8a Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 13 Aug 2026 09:29:44 +0200 Subject: [PATCH] Replace Sankoff cost matrix with cardinality-composition decomposition Replaced the legacy Sankoff parsimony pipeline with a new cardinality-composition decomposition that constructs row-normalized transition probability matrices symmetrized via geometric mean. This ensures reversibility, reduces free parameters from 240 to 120, and guarantees a zero diagonal. Tallies are now explicitly restricted to variable families to align with +ASC-corrected alignment populations. Additionally, fixed `--exclude-genome` handling to re-scan surviving sequences and drop newly monomorphic columns, preventing silent data corruption in downstream tree inference tools. --- docmd/theory/evolutionary_distances.md | 139 +++++++- iqtree.model | 16 + sankoff_params.yaml | 180 +++++++++-- src/obikindex/src/cardcomp.rs | 242 ++++++++++++++ src/obikindex/src/lib.rs | 9 +- src/obikindex/src/sankoff.rs | 402 ------------------------ src/obikindex/src/siblings.rs | 255 +++++++++++++-- src/obikmer/src/cmd/distance/mod.rs | 95 ++++-- src/obikmer/src/cmd/distance/sankoff.rs | 112 ++++--- 9 files changed, 901 insertions(+), 549 deletions(-) create mode 100644 iqtree.model create mode 100644 src/obikindex/src/cardcomp.rs delete mode 100644 src/obikindex/src/sankoff.rs diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index c7e2ebfe..f6b20da1 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -1094,15 +1094,136 @@ identified. Gives 120 parameters, but derived from two small, well-estimated pieces (a 5×5 cardinality model, a 4×4 composition model) rather than fit or smoothed independently per pair. -**Status: designed, not implemented.** Would replace `build_cost_matrix` -(`obikindex/src/sankoff.rs`) and the single `c_ctx` scalar/parameter -entirely; `sub_cost`'s own calibration is untouched. Not yet decided -whether the "clean" parsimony-graph cost (`build_cost_matrix`'s current -output, still needed for `--tnt`/`--phyg`) should be replaced too, kept -as a separate simpler approximation, or derived as a special case of this -same decomposition (it likely can be — the same `shared`/`lost`/`gained` -logic, with the cardinality model reduced back to a single scalar `c_ctx` -substituted in, recovers exactly the current graph). +### Implemented (2026-08-12): `pairwise_cost_matrix` replaces `build_cost_matrix` entirely + +New module `obikindex/src/cardcomp.rs`, replacing `sankoff::build_cost_matrix` +and the `c_ctx`/`SankoffWeights`/`PHatEstimate`/`calibrate_p_hat`/ +`c_ctx_from_p_hat`/`mean_substitution_cost`/`substitution_costs_from_tally` +machinery it depended on outright — not kept in parallel as a fallback +(all now unreferenced outside their own tests; `sankoff.rs` itself is a +pending removal, not yet done). + +**New primitives, `obikindex/src/siblings.rs`:** +- `BasePairTally` gained a `same: [u64; 4]` field (diagonal — "both + genomes at the same single base", pooled from + [`base_pair_tally`](obikindex::KmerIndex::base_pair_tally)'s existing + traversal, extended to also tally the `bi == bj` case it previously + discarded). +- `CardinalityTally { counts: [[u64; 5]; 5] }` and + `KmerIndex::cardinality_tally`, a new traversal (same shape as + `snp_pseudo_alignment`'s — needs full per-genome presence masks, not + `scan_family_pairs`'s single-resolved-form view, since cardinality 2-4 + is exactly the signal being tallied, not noise to drop). Same + saturated/no-data pair exclusion as `base_pair_tally`. Restricted to + variable families (`family_size() >= 2`), matching + `snp_pseudo_alignment`'s own scope. Verified against the 20-genome + benchmark: counts match the earlier hand-rolled Python analysis exactly + (e.g. `c=0/c=0: 117,158,166`, `c=0/c=1: 13,707,223` — the same numbers + this whole investigation started from). + +**`cardcomp.rs`:** +- `cardinality_transition_probs`/`composition_transition_probs`: row- + normalise the two tallies into proper transition probability matrices, + diagonal included ("stay the same" is a real, calibrated outcome). +- `pairwise_cost_matrix`: for every pair of the 16 states, `shared = A∩B` + contributes `∏ P_composition(x→x)`; `lost = A\B`, `gained = B\A` are + parsimony-paired (`best_pairing_cost`, brute-force over the ≤4! + injections — small enough that hand-rolling beats a dependency) into + substitution events on `P_composition`, minimising total `-ln`; the + cardinality-difference leftover is priced once via + `P_cardinality(|A|→|B|)`, never chained. Row-normalised, `-ln`'d, then + **symmetrised**: `cost_sym(A,B) = (cost(A,B)+cost(B,A))/2` — equivalent + to taking the *geometric* mean of the two raw probabilities + (`-ln(√(P(A,B)·P(B,A))) = (-ln P(A,B) - ln P(B,A))/2`), not their + arithmetic mean. Required, not just convenient for IQ-TREE's + lower-triangular file format: Sankoff parsimony's score is independent + of where an *unrooted* tree (what TNT/PhyG actually search over) gets + rooted only if the cost matrix is symmetric — the discrete-parsimony + analogue of CTMC reversibility, established by direct reasoning, not + assumed. Bonus of the same decision: 120 free parameters instead of the + 240 a fully asymmetric matrix would need. + +**Verified on the 20-genome benchmark**: resulting matrix symmetric +(checked numerically, zero asymmetric cells), zero diagonal, no NaN/Inf. +`--tnt` output still loads into TNT with no triangle-inequality warning +(`scaled_metric_matrix`'s rounding-metric-closure step still needed and +still applied — nothing in the new construction guarantees the *rounded +integer* matrix stays a metric, even though the real-valued one is exact +by construction here, unlike the old Floyd-Warshall-closed matrix which +needed it for a different reason). IQ-TREE loads the new model file and +reports the same `π` as before (only `R` changed). + +### Two consistency bugs found and fixed post-implementation (2026-08-13) + +**`--exclude-genome` didn't drop columns that become monomorphic once the +excluded genome(s) are gone.** `snp_pseudo_alignment`'s "variable family" +test (`family_size() >= 2`) is a property of the annex computed over +*every* genome in the index — unaffected by the CLI-level exclusion, which +only dropped the excluded genome's *row*. A family variable only because +of the excluded genome stayed in the alignment as a now-constant column — +silently wrong data for TNT/PhyG, and a hard failure for IQ-TREE's `+ASC` +(verified: excluding 2 taxa on the benchmark left 116,351 such columns — +matches the manual `+ASC` failures hit earlier in this same investigation, +before `--exclude-genome` existed). Fixed in `drop_excluded` +(`obikmer/src/cmd/distance/mod.rs`): after dropping excluded rows, +re-scan each column among the *surviving* sequences and drop any that are +now constant. Verified: 908,723 → 792,372 sites after excluding 2 taxa, +zero monomorphic columns remain, `π` recomputed from the corrected +alignment matches an independent recount exactly. The compact-alphabet +renumbering (`iqtree::compact_alphabet`) needed no equivalent fix — it +already recomputes which of the 16 states occur fresh on every call, from +whatever alignment it's actually handed, so a symbol disappearing (e.g. +excluding every genome that carries `N`) is already handled correctly; +verified directly (excluded 13 genomes to force `N` out: "15 of 16 states" +reported, correctly-shaped model file). + +**`cardinality_tally`'s `family_size() >= 2` filter looked inconsistent +with `base_pair_tally` — removing it was tried, and was wrong; reverted.** +`cardinality_tally` (modelled after `snp_pseudo_alignment`) had the +filter; `base_pair_tally` didn't (it visits every family via +`scan_family_pairs` unconditionally, folding fully-invariant loci into its +own `same` diagonal). Read as `cardinality_tally` under-counting its +diagonal relative to `base_pair_tally`, and — independently — as another +angle on the `--exclude-genome` drift (`family_size()` being global-only +meant a family kept here post-exclusion could differ from what the +now-correctly-filtered alignment kept). First fix tried: drop +`cardinality_tally`'s filter entirely, matching `base_pair_tally`'s +whole-annex scope. + +**That fix was empirically wrong, confirmed by a real IQ-TREE run, not +just a hunch.** Log-likelihood dropped from the earlier correct run's +`-8,364,671`/`-8,371,082` to `-9,170,228` (worse fit, not better), with +repeated `NNI search needs unusual large number of steps (20) to +converge!` warnings — and the completed run's **total tree length came +out at 67.644**, roughly 30× the earlier correct runs' ~2.0, i.e. branches +blowing up/saturating. Root cause, only clear in hindsight: `+ASC` +("ascertainment bias correction") exists specifically because the +likelihood only ever sees *variable* sites — the alignment fed to +IQ-TREE, by construction, contains not one invariant column. Calibrating +`R` from a population overwhelmingly dominated by genome-wide invariant +background (family_size()<2 loci outnumber the ~908k variable ones by +orders of magnitude) describes a completely different population than the +one `+ASC` and the alignment actually model — the "consistency" argument +for matching `base_pair_tally`'s scope was real, but pointed the wrong +way: `base_pair_tally`'s own unrestricted `same` diagonal turned out to +have the *identical* latent bug (only unmasked once its diagonal existed +at all, which happened earlier the same day when `same` was added), not a +correct baseline to match `cardinality_tally` to. + +**Final fix**: restored `cardinality_tally`'s `family_size() >= 2` filter, +and gave `base_pair_tally`'s `same` diagonal the equivalent restriction — +`scan_family_pairs` (shared with `raw_snp_distance`, which legitimately +*does* want fully-invariant families counted as `shared`) now passes an +extra `variable: bool` (the family's own `family_size() >= 2`) to its +`on_pair` callback; `base_pair_tally` only increments `same` when +`variable` is true, `raw_snp_distance`'s callback ignores the new +argument. Both tallies now describe the same variable-families-only +population the `+ASC`-corrected alignment does. Verified: calibration +counts back to their original values exactly (`c=0/c=0`: 117,158,166, +matching the pre-regression run bit for bit), and a full IQ-TREE rerun +converged normally — log-likelihood `-8,389,106.273` (same order as the +two earlier correct runs), **total tree length 2.044** (was 67.644), no +NNI convergence warnings. ## Heterozygosity, ploidy, and consensus-assembly inputs diff --git a/iqtree.model b/iqtree.model new file mode 100644 index 00000000..09bd6ac5 --- /dev/null +++ b/iqtree.model @@ -0,0 +1,16 @@ +0.222641 +0.222347 0.005197 +0.206041 0.017442 0.019145 +0.222349 0.024185 0.003581 0.013346 +0.208691 0.017667 0.002921 0.000012 0.019168 +0.205784 0.013974 0.019121 0.000085 0.018901 0.000018 +0.027849 0.001121 0.001230 0.000109 0.001216 0.000109 0.000118 +0.222640 0.004745 0.023685 0.013831 0.005020 0.002986 0.013813 0.000889 +0.206337 0.017467 0.012967 0.000076 0.013365 0.000016 0.000059 0.000079 0.017590 +0.208662 0.003028 0.019388 0.000017 0.002764 0.000003 0.000018 0.000019 0.017788 0.000017 +0.027866 0.001122 0.001231 0.000109 0.000858 0.000016 0.000083 0.000009 0.001129 0.000100 0.000111 +0.206081 0.013994 0.012951 0.000059 0.018929 0.000017 0.000082 0.000085 0.017568 0.000078 0.000012 0.000076 +0.027872 0.001122 0.000833 0.000074 0.001217 0.000109 0.000080 0.000042 0.001130 0.000100 0.000017 0.000006 0.000108 +0.027846 0.000899 0.001230 0.000087 0.001216 0.000019 0.000118 0.000009 0.001129 0.000080 0.000111 0.000044 0.000108 0.000009 +0.009724 0.000172 0.000189 0.000044 0.000187 0.000044 0.000048 0.000047 0.000174 0.000041 0.000045 0.000043 0.000044 0.000043 0.000047 +0.904811 0.012806 0.013642 0.007640 0.013786 0.007933 0.004196 0.000140 0.012382 0.006685 0.007847 0.000195 0.007527 0.000202 0.000134 0.000076 diff --git a/sankoff_params.yaml b/sankoff_params.yaml index 6bb64d96..e4a0075c 100644 --- a/sankoff_params.yaml +++ b/sankoff_params.yaml @@ -1,27 +1,167 @@ ratio_ceiling: 0.5 -flank_length_m: 15 -p_hat: 0.010713940022064144 -p_hat_variance: 3.2110559785689563e-10 -n_pairs_included: 158 -n_loci_total: 33008305 -mean_sub_cost: 1.4764371767846869 -c_ctx: 1.7343663443925388 -substitutions: -- pair: A/C +cardinality_transitions: +- from: 0 + to: 0 + count: 117158166 + probability: 0.8249649657257814 +- from: 0 + to: 1 + count: 13707223 + probability: 0.09651891232567299 +- from: 0 + to: 2 + count: 10934100 + probability: 0.07699206755884405 +- from: 0 + to: 3 + count: 193315 + probability: 0.0013612205430842902 +- from: 0 + to: 4 + count: 23125 + probability: 0.00016283384661730445 +- from: 1 + to: 0 + count: 13707223 + probability: 0.8980473700324103 +- from: 1 + to: 1 + count: 1043692 + probability: 0.06837890181868832 +- from: 1 + to: 2 + count: 507953 + probability: 0.033279232106318904 +- from: 1 + to: 3 + count: 4270 + probability: 0.0002797548613631216 +- from: 1 + to: 4 + count: 225 + probability: 0.000014741181219368235 +- from: 2 + to: 0 + count: 10934100 + probability: 0.9551033885158036 +- from: 2 + to: 1 + count: 507953 + probability: 0.04437014765794788 +- from: 2 + to: 2 + count: 5241 + probability: 0.00045780602511512846 +- from: 2 + to: 3 + count: 690 + probability: 0.000060272115498843476 +- from: 2 + to: 4 + count: 96 + probability: 8.385685634621701e-6 +- from: 3 + to: 0 + count: 193315 + probability: 0.9743748708410829 +- from: 3 + to: 1 + count: 4270 + probability: 0.021522285898618442 +- from: 3 + to: 2 + count: 690 + probability: 0.0034778401100812 +- from: 3 + to: 3 + count: 98 + probability: 0.0004939541025912429 +- from: 3 + to: 4 + count: 26 + probability: 0.0001310490476262481 +- from: 4 + to: 0 + count: 23125 + probability: 0.9846291407647109 +- from: 4 + to: 1 + count: 225 + probability: 0.009580175423656646 +- from: 4 + to: 2 + count: 96 + probability: 0.004087541514093502 +- from: 4 + to: 3 + count: 26 + probability: 0.0011070424934003236 +- from: 4 + to: 4 + count: 14 + probability: 0.0005960998041386358 +composition_transitions: +- from: 'A' + to: 'A' + count: 162260 + probability: 0.47062167539692207 +- from: 'A' + to: 'C' count: 27810 - cost: 2.5439415226130238 -- pair: A/G + probability: 0.08066059899413536 +- from: 'A' + to: 'G' count: 130153 - cost: 1.0 -- pair: A/T + probability: 0.3774979842101294 +- from: 'A' + to: 'T' count: 24555 - cost: 2.6684722241428322 -- pair: C/G + probability: 0.07121974139881315 +- from: 'C' + to: 'A' + count: 27810 + probability: 0.07790832534920075 +- from: 'C' + to: 'C' + count: 184633 + probability: 0.5172401234879174 +- from: 'C' + to: 'G' count: 19637 - cost: 2.892062911764701 -- pair: C/T + probability: 0.05501207424963161 +- from: 'C' + to: 'T' count: 124878 - cost: 1.0413902163542994 -- pair: G/T + probability: 0.3498394769132503 +- from: 'G' + to: 'A' + count: 130153 + probability: 0.36057158212891627 +- from: 'G' + to: 'C' + count: 19637 + probability: 0.054401697680925745 +- from: 'G' + to: 'G' + count: 184557 + probability: 0.5112906308956874 +- from: 'G' + to: 'T' count: 26616 - cost: 2.5878424665170052 + probability: 0.07373608929447062 +- from: 'T' + to: 'A' + count: 24555 + probability: 0.07337692220342934 +- from: 'T' + to: 'C' + count: 124878 + probability: 0.3731689387464813 +- from: 'T' + to: 'G' + count: 26616 + probability: 0.07953574267426085 +- from: 'T' + to: 'T' + count: 158593 + probability: 0.4739183963758285 diff --git a/src/obikindex/src/cardcomp.rs b/src/obikindex/src/cardcomp.rs new file mode 100644 index 00000000..d3fb39bd --- /dev/null +++ b/src/obikindex/src/cardcomp.rs @@ -0,0 +1,242 @@ +//! Cardinality × composition decomposition of the 16-state transition cost +//! matrix — replaces `sankoff::build_cost_matrix`'s elementary-edge graph +//! and its Floyd-Warshall shortest-path closure, which double-counts +//! multi-hop transitions once IQ-TREE's own matrix exponential composes +//! them again. See `docmd/theory/evolutionary_distances.md`, "`R` via +//! `exp(-cost)` is wrong for a CTMC; cardinality/composition decomposition". +//! +//! 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 +//! `pairwise_transition_matrix`'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 crate::{BasePairTally, CardinalityTally}; + +/// Row-stochastic 5×5 cardinality transition probabilities (`0..=4`), +/// diagonal included ("stay at the same cardinality"), from +/// [`CardinalityTally`]'s pooled, symmetric co-occurrence counts. +pub 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 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, replacing +/// `sankoff::build_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). +pub fn pairwise_cost_matrix(p_card: &[[f64; 5]; 5], p_comp: &[[f64; 4]; 4]) -> [[f64; 16]; 16] { + let mut raw = [[0.0f64; 16]; 16]; + 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 log_p = 0.0; // accumulate ln(P), so 0.0 = probability 1 + 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]; + log_p += 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]; + log_p += if p > 0.0 { p.ln() } else { f64::NEG_INFINITY }; + } + } + log_p -= best_pairing_cost(&lost, &gained, p_comp); + + raw[a as usize][b as usize] = log_p.exp(); + } + } + + // Row-normalise to a proper transition probability matrix. + let mut p = [[0.0f64; 16]; 16]; + for a in 0..16 { + let row_sum: f64 = raw[a].iter().sum(); + if row_sum > 0.0 { + for b in 0..16 { + p[a][b] = raw[a][b] / row_sum; + } + } + } + + // Cost = -ln(P), then symmetrise (see doc comment: parsimony on an + // unrooted tree requires a symmetric cost matrix). + let mut cost = [[0.0f64; 16]; 16]; + for a in 0..16 { + for b in 0..16 { + cost[a][b] = if p[a][b] > 0.0 { -p[a][b].ln() } else { f64::INFINITY }; + } + } + 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); + 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]); + } + } + } +} diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index a48eb325..4bcc66dd 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -1,6 +1,7 @@ pub mod error; pub mod meta; pub mod state; +mod cardcomp; mod distance; mod dump; mod index; @@ -8,7 +9,6 @@ mod merge; mod numa; mod rebuild; mod reindex; -mod sankoff; mod select; mod siblings; mod stats; @@ -20,8 +20,5 @@ pub use merge::MergeMode; pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME}; pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED}; pub use stats::IndexBitsPerKmer; -pub use siblings::{BasePairTally, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; -pub use sankoff::{ - build_cost_matrix, c_ctx_from_p_hat, calibrate_p_hat, mean_substitution_cost, - substitution_costs_from_tally, PHatEstimate, SankoffWeights, -}; +pub use siblings::{BasePairTally, CardinalityTally, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; +pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix}; diff --git a/src/obikindex/src/sankoff.rs b/src/obikindex/src/sankoff.rs deleted file mode 100644 index 372ed21e..00000000 --- a/src/obikindex/src/sankoff.rs +++ /dev/null @@ -1,402 +0,0 @@ -//! Sankoff parsimony cost matrix for the 16-state (powerset of `{A,C,G,T}`) -//! family alphabet, and calibration of its parameters from real pairwise -//! SNP/shared counts. See `docmd/theory/evolutionary_distances.md`, -//! "Sankoff parsimony as the resolution of the 16-state model problem". - -use crate::{BasePairTally, RawSnpDistanceOutput}; - -/// Transition base pairs in the project's fixed bit convention (see -/// `siblings::central_base`: bit 0=A, 1=C, 2=G, 3=T). All other single-bit -/// swaps are transversions. -const TRANSITION_PAIRS: [(u8, u8); 2] = [(0, 2), (1, 3)]; // A<->G, C<->T - -/// Tunable weights for the 16-state Sankoff cost matrix — see "A concrete -/// Sankoff cost matrix for the 16-state alphabet" in the design doc. -/// -/// Only two costs, not three: an earlier version had a separate `c_gl` for -/// "gain/loss between two nonempty states" alongside `c_ctx` for "collapse -/// to `∅`", but presence/absence tracking never observes "the flanks" -/// independently of a whole 31-mer — losing one member of a family while -/// others remain (`{A,C}->{A}`) and losing the last one (`{A}->∅`) are the -/// same event: a specific, complete, homologous 31-mer that used to be -/// observed no longer is. Both are exactly what `c_ctx` (see -/// `c_ctx_from_p_hat`) prices, so it is used for every gain/loss uniformly -/// — including a compound family losing several members at once, `∅` -/// included, which costs `|X|*c_ctx` (no flat shortcut for `∅`: see -/// `build_cost_matrix`'s docs for why an earlier version's flat special -/// case there didn't hold up). -#[derive(Debug, Clone, Copy)] -pub struct SankoffWeights { - /// Cost of a single-base substitution `a <-> b` (`a != b`), symmetric - /// (`sub_cost[a][b] == sub_cost[b][a]`) and only meaningful off the - /// diagonal. The full 6-category (AC, AG, AT, CG, CT, GT) resolution a - /// symmetric cost matrix allows — see `substitution_costs_from_tally` - /// for calibrating it from real data, or `SankoffWeights::ts_tv` for - /// the coarser 2-category (transition/transversion) convenience - /// constructor. - pub sub_cost: [[f64; 4]; 4], - /// Cost of losing or gaining one family member — the same constant - /// everywhere a member is gained or lost, member count and target state - /// (`∅` or not) included: see the struct docs. - pub c_ctx: f64, -} - -impl SankoffWeights { - /// Convenience constructor for the coarser 2-category - /// (transition/transversion) substitution model: every transition pair - /// (A<->G, C<->T) costs `c_ts`, every transversion pair costs `c_tv`. - pub fn ts_tv(c_ts: f64, c_tv: f64, c_ctx: f64) -> Self { - let mut sub_cost = [[0.0; 4]; 4]; - for a in 0..4usize { - for b in 0..4usize { - if a == b { - continue; - } - let is_ts = TRANSITION_PAIRS.contains(&(a as u8, b as u8)); - sub_cost[a][b] = if is_ts { c_ts } else { c_tv }; - } - } - Self { sub_cost, c_ctx } - } -} - -/// Build the 16x16 Sankoff step-cost matrix, states indexed `0..=15` by -/// their bitmask (matching `siblings::iupac_code`'s convention — state `0` -/// is `∅`). -/// -/// All 16 states, `∅` included, form a single graph: an edge of weight -/// `c_ts`/`c_tv` between two states of equal cardinality differing by -/// exactly one element (classified by whether the swapped bases form a -/// transition or transversion pair), and an edge of weight `c_ctx` between -/// a state and one that is its strict subset plus exactly one element -/// (`∅` is a strict subset of every singleton state, so it connects -/// directly to each of them this way — no special case). `cost(X, Y)` is -/// the shortest-path distance in that graph (16 nodes — Floyd-Warshall, -/// trivial at this size), uniformly, including `X <-> ∅`: a compound -/// family losing several members at once, `∅` included, costs `|X|*c_ctx` -/// via that many single-element steps — no cheaper flat alternative for -/// `∅` specifically, since nothing distinguishes it from any other -/// multi-element loss (see design doc — an earlier flat special case here -/// was inconsistent with how every *other* multi-element transformation is -/// already priced, and had no principled justification once that was -/// noticed). -pub fn build_cost_matrix(w: &SankoffWeights) -> [[f64; 16]; 16] { - const INF: f64 = f64::INFINITY; - let mut dist = [[INF; 16]; 16]; - for (i, row) in dist.iter_mut().enumerate() { - row[i] = 0.0; - } - - for x in 0u8..16 { - for y in (x + 1)..16 { - let (xu, yu) = (x as usize, y as usize); - let diff = x ^ y; - let card_x = x.count_ones(); - let card_y = y.count_ones(); - let edge = if card_x == card_y && diff.count_ones() == 2 { - // Exactly one element swapped: the bit only in x is the base - // leaving, the bit only in y is the base entering. - let a = (x & diff).trailing_zeros() as u8; - let b = (y & diff).trailing_zeros() as u8; - Some(w.sub_cost[a as usize][b as usize]) - } else if card_x.abs_diff(card_y) == 1 && (x & y) == x.min(y) { - // True subset relationship: a pure gain/loss of one element. - Some(w.c_ctx) - } else { - None - }; - if let Some(cost) = edge { - dist[xu][yu] = cost; - dist[yu][xu] = cost; - } - } - } - - for k in 0..16 { - for i in 0..16 { - if dist[i][k].is_infinite() { - continue; - } - for j in 0..16 { - let via = dist[i][k] + dist[k][j]; - if via < dist[i][j] { - dist[i][j] = via; - } - } - } - } - - dist -} - -/// `p_hat` calibrated from real pairwise SNP/shared counts, restricted to -/// pairs below `ratio_ceiling`, and its variance as a pooled Bernoulli -/// proportion. -#[derive(Debug, Clone, Copy)] -pub struct PHatEstimate { - pub p_hat: f64, - /// `p_hat*(1-p_hat) / n_loci_total` — the pooled estimator's variance is - /// governed by the total number of loci across all included pairs, not - /// by any single pair's count (see design doc: this is why the - /// exclusion criterion below is the per-pair *ratio*, not a per-pair - /// minimum-count floor — a low-count pair barely moves the pool either - /// way, but a saturated pair with plenty of loci would bias it). - pub variance: f64, - pub n_pairs_included: usize, - pub n_loci_total: u64, -} - -/// Pool `snp`/`shared` counts across all genome pairs whose per-pair ratio -/// `snp/(snp+shared)` is at most `ratio_ceiling`, then estimate -/// `p_hat = sum(snp) / sum(snp+shared)` over the included pairs. -/// -/// A pair with no eligible locus at all (`snp+shared == 0`) is always -/// excluded (nothing to pool). `ratio_ceiling` should be well below 1.0 — -/// pairs at or near saturation (e.g. cross-domain comparisons, where nearly -/// every shared central-position family already differs) carry no -/// information about `p_hat` and bias it upward if pooled in. -pub fn calibrate_p_hat(raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> PHatEstimate { - let n = raw.snp.nrows(); - let mut snp_sum: u64 = 0; - let mut total_sum: u64 = 0; - let mut n_pairs = 0usize; - for i in 0..n { - for j in (i + 1)..n { - let snp = raw.snp[[i, j]]; - let shared = raw.shared[[i, j]]; - let total = snp + shared; - if total == 0 { - continue; - } - let ratio = snp as f64 / total as f64; - if ratio > ratio_ceiling { - continue; - } - snp_sum += snp; - total_sum += total; - n_pairs += 1; - } - } - let p_hat = if total_sum > 0 { snp_sum as f64 / total_sum as f64 } else { 0.0 }; - let variance = if total_sum > 0 { - p_hat * (1.0 - p_hat) / total_sum as f64 - } else { - 0.0 - }; - PHatEstimate { p_hat, variance, n_pairs_included: n_pairs, n_loci_total: total_sum } -} - -/// `c_ctx(p) = mean_sub_cost * [(2m*p) / (1 - (1-p)^(2m)) + p]` — see -/// "Context, detectability, and a 3-way ordinal distance per pair" in the -/// design doc for the bracketed term's derivation. `m` is the flank length -/// on *each* side of the central position (`k = 2m+1`); pass `(k-1)/2`, not -/// the project's minimizer-size parameter of the same name. -/// -/// The bracketed term is `E[mutation count | at least one occurred]` — a -/// *count* of mutations, not a cost. An earlier version used it directly as -/// the cost, implicitly pricing every mutation at a flat `1` regardless of -/// type. That stopped making sense once substitution costs were calibrated -/// per base-pair type (`substitution_costs_from_tally`): transitions and -/// transversions aren't equally likely (roughly 5-6x apart in this -/// project's own data) nor equally costly, so "one mutation" isn't worth a -/// flat unit — it's worth whatever the *empirical mix* of mutation types is -/// worth on average. `mean_sub_cost` (see `mean_substitution_cost`) is that -/// weighted average, turning the expected mutation *count* into an actual -/// expected cost. -pub fn c_ctx_from_p_hat(p_hat: f64, m: usize, mean_sub_cost: f64) -> f64 { - if p_hat <= 0.0 { - return 0.0; - } - let two_m = 2.0 * m as f64; - let expected_mutations = (two_m * p_hat) / (1.0 - (1.0 - p_hat).powf(two_m)) + p_hat; - expected_mutations * mean_sub_cost -} - -/// Mean substitution cost, weighted by each of the 6 base-pair categories' -/// observed frequency in `tally` — the empirical average cost of "one -/// mutation" under the calibrated substitution spectrum. Falls back to -/// `1.0` (a flat per-mutation cost) if `tally` has no signal at all (no -/// substitution ever observed), rather than dividing by zero. -pub fn mean_substitution_cost(tally: &BasePairTally, sub_cost: &[[f64; 4]; 4]) -> f64 { - let mut weighted_sum = 0.0; - let mut total = 0u64; - for a in 0..4 { - for b in (a + 1)..4 { - let count = tally.counts[a][b]; - weighted_sum += count as f64 * sub_cost[a][b]; - total += count; - } - } - if total == 0 { - return 1.0; - } - weighted_sum / total as f64 -} - -/// Derive the 4x4 symmetric substitution cost table from an observed -/// [`BasePairTally`]: `cost(a,b) = -ln(rate(a,b))`, normalised so the most -/// frequently observed substitution type costs exactly `1.0` — the standard -/// generalised-parsimony step-weighting heuristic (see design doc, -/// "Transition/transversion refinement"), generalised here from 2 -/// categories (Ts/Tv) to the full 6-category resolution a symmetric matrix -/// allows: a rarer substitution type is treated as less parsimonious to -/// invoke, and therefore costs more. -/// -/// A pair type never observed at all (`counts[a][b] == 0`) gets an infinite -/// cost — parsimony should never spend a mutation on something with no -/// empirical support in this data. -pub fn substitution_costs_from_tally(tally: &BasePairTally) -> [[f64; 4]; 4] { - let total: u64 = (0..4) - .flat_map(|a| (a + 1..4).map(move |b| (a, b))) - .map(|(a, b)| tally.counts[a][b]) - .sum(); - - let mut raw = [[0.0f64; 4]; 4]; - let mut min_cost = f64::INFINITY; - for a in 0..4 { - for b in (a + 1)..4 { - let count = tally.counts[a][b]; - let cost = if count > 0 && total > 0 { - -((count as f64 / total as f64).ln()) - } else { - f64::INFINITY - }; - raw[a][b] = cost; - raw[b][a] = cost; - if cost < min_cost { - min_cost = cost; - } - } - } - - let mut out = [[0.0f64; 4]; 4]; - for a in 0..4 { - for b in 0..4 { - if a != b { - out[a][b] = raw[a][b] / min_cost; - } - } - } - out -} - -#[cfg(test)] -mod tests { - use super::*; - use ndarray::Array2; - - fn weights() -> SankoffWeights { - SankoffWeights::ts_tv(1.0, 2.0, 10.0) - } - - #[test] - fn identity_is_zero() { - let m = build_cost_matrix(&weights()); - for row in m.iter().enumerate() { - assert_eq!(m[row.0][row.0], 0.0); - } - } - - #[test] - fn single_base_transition_costs_c_ts() { - let m = build_cost_matrix(&weights()); - // A = 0b0001 (state 1), G = 0b0100 (state 4): A<->G is a transition. - assert_eq!(m[1][4], 1.0); - assert_eq!(m[4][1], 1.0); - } - - #[test] - fn single_base_transversion_costs_c_tv() { - let m = build_cost_matrix(&weights()); - // A = 0b0001 (state 1), C = 0b0010 (state 2): A<->C is a transversion. - assert_eq!(m[1][2], 2.0); - } - - #[test] - fn gain_loss_costs_c_ctx() { - let m = build_cost_matrix(&weights()); - // A = 0b0001 (state 1) -> {A,C} = 0b0011 (state 3): pure gain of C. - assert_eq!(m[1][3], 10.0); - } - - #[test] - fn empty_set_costs_scale_with_cardinality() { - let m = build_cost_matrix(&weights()); - assert_eq!(m[0][1], 10.0); // {A} -> ∅: one member, one c_ctx step. - // N (all four) -> ∅: no flat shortcut — four single-element steps, - // same as losing four members down to a nonempty state would cost. - assert_eq!(m[0][15], 40.0); - assert_eq!(m[0][0], 0.0); - } - - #[test] - fn matrix_is_symmetric() { - let m = build_cost_matrix(&weights()); - for i in 0..16 { - for j in 0..16 { - assert_eq!(m[i][j], m[j][i], "asymmetry at ({i},{j})"); - } - } - } - - #[test] - fn calibrate_p_hat_excludes_saturated_pairs() { - // 3 genomes: pair (0,1) informative (10% SNP), pair (0,2) saturated - // (100% SNP) — must be excluded from the pooled estimate. - let mut snp = Array2::::zeros((3, 3)); - let mut shared = Array2::::zeros((3, 3)); - snp[[0, 1]] = 10; - shared[[0, 1]] = 90; - snp[[1, 0]] = 10; - shared[[1, 0]] = 90; - snp[[0, 2]] = 50; - shared[[0, 2]] = 0; - snp[[2, 0]] = 50; - shared[[2, 0]] = 0; - let raw = RawSnpDistanceOutput { snp, shared }; - - let est = calibrate_p_hat(&raw, 0.5); - assert_eq!(est.n_pairs_included, 1); - assert_eq!(est.n_loci_total, 100); - assert!((est.p_hat - 0.1).abs() < 1e-9); - } - - #[test] - fn c_ctx_from_p_hat_zero_is_zero() { - assert_eq!(c_ctx_from_p_hat(0.0, 15, 1.5), 0.0); - } - - #[test] - fn c_ctx_from_p_hat_matches_known_value() { - // p=0.1, m=15: [(30*0.1)/(1-(0.9)^30) + 0.1] * mean_sub_cost - let m = 15usize; - let p = 0.1; - let mean_sub_cost = 1.5; - let expected = ((2.0 * m as f64 * p) / (1.0 - (1.0 - p).powf(2.0 * m as f64)) + p) * mean_sub_cost; - assert!((c_ctx_from_p_hat(p, m, mean_sub_cost) - expected).abs() < 1e-12); - } - - #[test] - fn mean_substitution_cost_is_weighted_average() { - let mut counts = [[0u64; 4]; 4]; - counts[0][1] = 3; // A/C - counts[1][0] = 3; - counts[0][2] = 1; // A/G - counts[2][0] = 1; - let tally = BasePairTally { counts }; - let mut sub_cost = [[0.0f64; 4]; 4]; - sub_cost[0][1] = 2.0; - sub_cost[1][0] = 2.0; - sub_cost[0][2] = 1.0; - sub_cost[2][0] = 1.0; - // (3*2.0 + 1*1.0) / 4 = 1.75 - assert!((mean_substitution_cost(&tally, &sub_cost) - 1.75).abs() < 1e-12); - } - - #[test] - fn mean_substitution_cost_falls_back_to_one_with_no_data() { - let tally = BasePairTally { counts: [[0u64; 4]; 4] }; - let sub_cost = [[0.0f64; 4]; 4]; - assert_eq!(mean_substitution_cost(&tally, &sub_cost), 1.0); - } -} diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index dc33bf6a..ad8aea6b 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -660,10 +660,18 @@ impl KmerIndex { /// (tallied once, at its minorant) of every layer of the already-built /// sibling annex, resolves each genome's single observed form (`None` /// if absent or ambiguous/multi-copy), then calls `on_pair(acc, i, j, - /// bi, bj)` for every genome pair `(i, j)` where both are unambiguous - /// and single-copy (`bi == bj` means shared at that locus, `bi != bj` - /// means a SNP). Layers are processed in parallel (rayon); each gets - /// its own accumulator from `zero()`, combined pairwise via `combine`. + /// bi, bj, variable)` for every genome pair `(i, j)` where both are + /// unambiguous and single-copy (`bi == bj` means shared at that locus, + /// `bi != bj` means a SNP). `variable` is the family's own + /// `family_size() >= 2` (true if more than one member is observed + /// *anywhere* in the family, i.e. it isn't fully invariant across the + /// whole index) — `raw_snp_distance` ignores it (a fully-invariant + /// family is still legitimately "shared"), but callers whose diagonal + /// should only reflect genuine SNP-adjacent agreement, not the + /// genome-wide invariant background, need it (see + /// [`base_pair_tally`](Self::base_pair_tally)'s `same` field). Layers + /// are processed in parallel (rayon); each gets its own accumulator + /// from `zero()`, combined pairwise via `combine`. fn scan_family_pairs( &self, label: &str, @@ -673,7 +681,7 @@ impl KmerIndex { ) -> OKIResult where Acc: Send, - F: Fn(&mut Acc, usize, usize, u8, u8) + Sync, + F: Fn(&mut Acc, usize, usize, u8, u8, bool) + Sync, C: Fn(Acc, Acc) -> Acc, { let n_parts = self.n_partitions(); @@ -751,6 +759,7 @@ impl KmerIndex { if !is_minorant(kmer, mask, k) { continue; // family tallied once, at its minorant } + let variable = mask.family_size() >= 2; single_form.clear(); single_form.resize(n_cols, None); @@ -791,7 +800,7 @@ impl KmerIndex { continue; } let Some(bj) = single_form[j] else { continue }; - on_pair(&mut acc, i, j, bi, bj); + on_pair(&mut acc, i, j, bi, bj, variable); } } } @@ -816,7 +825,7 @@ impl KmerIndex { let (snp, shared) = self.scan_family_pairs( "raw_snp_distance", || (Array2::::zeros((n_genomes, n_genomes)), Array2::::zeros((n_genomes, n_genomes))), - |(snp, shared), i, j, bi, bj| { + |(snp, shared), i, j, bi, bj, _variable| { if bi == bj { shared[[i, j]] += 1; shared[[j, i]] += 1; @@ -837,9 +846,10 @@ impl KmerIndex { /// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG, /// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs /// whose overall SNP ratio in `raw` is at or below `ratio_ceiling` — - /// same saturation-exclusion discipline as `calibrate_p_hat`, for the - /// same reason: a saturated pair's observed base-pair mix trends toward - /// neutral base composition, not the true point-mutation spectrum. + /// same saturation-exclusion discipline as + /// [`cardinality_tally`](Self::cardinality_tally), for the same reason: + /// a saturated pair's observed base-pair mix trends toward neutral base + /// composition, not the true point-mutation spectrum. /// /// A second full pass over the annex, sharing /// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by @@ -858,25 +868,38 @@ impl KmerIndex { total > 0 && (snp as f64 / total as f64) <= ratio_ceiling }); - let counts = self.scan_family_pairs( + let (counts, same) = self.scan_family_pairs( "base_pair_tally", - || [[0u64; 4]; 4], - |counts, i, j, bi, bj| { - if bi != bj && included[[i, j]] { + || ([[0u64; 4]; 4], [0u64; 4]), + |(counts, same), i, j, bi, bj, variable| { + if !included[[i, j]] { + return; + } + if bi != bj { counts[bi as usize][bj as usize] += 1; counts[bj as usize][bi as usize] += 1; + } else if variable { + // Only count "stayed the same" from families that vary + // *somewhere* in the index — a fully invariant family + // (never varies anywhere) isn't a SNP-adjacent + // agreement, it's genome-wide background, and would + // otherwise swamp the diagonal (see + // `docmd/theory/evolutionary_distances.md`, the + // ascertainment-bias regression this was reverting). + same[bi as usize] += 1; } }, - |mut total, partial| { + |(mut counts, mut same), (partial_counts, partial_same)| { for a in 0..4 { + same[a] += partial_same[a]; for b in 0..4 { - total[a][b] += partial[a][b]; + counts[a][b] += partial_counts[a][b]; } } - total + (counts, same) }, )?; - Ok(BasePairTally { counts }) + Ok(BasePairTally { counts, same }) } } @@ -884,9 +907,15 @@ impl KmerIndex { pub 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` (0=A, 1=C, 2=G, 3=T). Diagonal always `0` (an `a == b` - /// locus is `shared`, not tallied here). + /// `a` and `b` (0=A, 1=C, 2=G, 3=T). 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` — the diagonal + /// `counts` omits, needed to build a proper row-stochastic composition + /// probability matrix (the "stay the same base" entries), not just the + /// substitution-cost off-diagonal. + pub same: [u64; 4], } /// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set @@ -918,6 +947,192 @@ fn iupac_code(mask: u8) -> u8 { } } +/// See [`KmerIndex::cardinality_tally`]. +pub struct CardinalityTally { + /// `counts[a][b] == counts[b][a]` = number of family sites, pooled over + /// included genome pairs, where one genome's family cardinality + /// (popcount of its presence mask, `0..=4`) 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], +} + +impl KmerIndex { + /// Cardinality co-occurrence, pooled only over genome pairs whose + /// overall SNP ratio in `raw` is at or below `ratio_ceiling` — same + /// saturation/no-data exclusion discipline as + /// [`base_pair_tally`](Self::base_pair_tally). Unlike + /// [`scan_family_pairs`](Self::scan_family_pairs) (which resolves each + /// genome to a single form and silently drops any genome carrying more + /// than one member of the family), this needs the *full* per-genome + /// presence mask — a family member count of 2, 3 or 4 is exactly the + /// signal being tallied, not noise to discard — so it re-implements the + /// traversal rather than reusing that helper. + /// + /// Restricted to variable families (`family_size() >= 2`), matching + /// `snp_pseudo_alignment`'s own scope — briefly removed, then + /// reinstated: without it, the diagonal is dominated by genome-wide + /// invariant background (family_size()<2 loci vastly outnumber the + /// ones that ever vary anywhere), which is inconsistent with the + /// `+ASC`-corrected alignment this matrix is ultimately used with — + /// `+ASC` exists specifically because the likelihood only ever sees + /// variable sites, so a rate model calibrated mostly from invariant + /// background sites doesn't describe the population it's applied to. + /// Verified empirically: removing the filter measurably worsened a + /// real IQ-TREE run (log-likelihood dropped, `NNI search needs + /// unusual large number of steps to converge` warnings appeared) — see + /// `docmd/theory/evolutionary_distances.md` for the full account. + /// [`base_pair_tally`](Self::base_pair_tally)'s own diagonal (`same`) + /// gets the matching restriction via `scan_family_pairs`'s new + /// `variable` flag, rather than a `family_size()` check of its own (it + /// doesn't have direct access to the family's mask). + pub fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult { + let n_parts = self.n_partitions(); + let n_genomes = self.meta.genomes.len(); + let with_counts = self.meta.config.with_counts; + let k = self.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; + + let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| { + if i == j { + return false; + } + let snp = raw.snp[[i, j]]; + let total = snp + raw.shared[[i, j]]; + total > 0 && (snp as f64 / total as f64) <= ratio_ceiling + }); + + let partition = KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + + let mut layer_dirs = Vec::new(); + for part in 0..n_parts { + let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + if !annex_path.exists() { + return Err(OKIError::InvalidInput(format!( + "no sibling annex at {} — run build_sibling_annex first", + annex_path.display() + ))); + } + layer_dirs.push(layer_dir); + } + } + + let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers"); + let partials: Vec<[[u64; 5]; 5]> = layer_dirs + .par_iter() + .map(|layer_dir| -> OKIResult<[[u64; 5]; 5]> { + let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); + let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + + let mut slot_kmer: Vec> = vec![None; annex.len()]; + let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) + .map_err(OKIError::Partition)?; + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if let Some(slot) = mphf.find(kmer) { + slot_kmer[slot] = Some(kmer); + } + } + + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(layer_dir)?) + }; + let n_cols = mat.n_cols().min(n_genomes); + + let mut counts = [[0u64; 5]; 5]; + let mut genome_mask: Vec = Vec::with_capacity(n_genomes); + + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + let Some(kmer) = slot_kmer[slot] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // family tallied once, at its minorant + } + if mask.family_size() < 2 { + // Fully invariant family (never varies anywhere in + // the index) — genome-wide background, not + // SNP-adjacent signal; would otherwise swamp the + // diagonal (`c=1/c=1` etc.), which needs to reflect + // the same variable-families-only population the + // `+ASC`-corrected alignment/likelihood actually + // models. See `base_pair_tally`'s `variable` gate + // on its own `same` diagonal for the matching fix. + continue; + } + + genome_mask.clear(); + genome_mask.resize(n_genomes, 0); + + for other in kmer.central_canonical_neighbors() { + let base = central_base(other, k); + if !mask.has(base) { + continue; + } + let presence: Option> = if other == kmer { + Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()) + } else { + let dest = partition_of(other, n_parts); + cache.find_presence(dest, other, n_genomes) + }; + let Some(presence) = presence else { continue }; + for (g, &present) in presence.iter().enumerate() { + if present { + genome_mask[g] |= 1 << base; + } + } + } + + for i in 0..n_genomes { + let card_i = genome_mask[i].count_ones() as usize; + for j in (i + 1)..n_genomes { + if !included[[i, j]] { + continue; + } + let card_j = genome_mask[j].count_ones() as usize; + counts[card_i][card_j] += 1; + if card_i != card_j { + counts[card_j][card_i] += 1; + } + } + } + } + + pb.inc(1); + Ok(counts) + }) + .collect::>>()?; + pb.finish_and_clear(); + + let mut total = [[0u64; 5]; 5]; + for partial in partials { + for a in 0..5 { + for b in 0..5 { + total[a][b] += partial[a][b]; + } + } + } + Ok(CardinalityTally { counts: total }) + } +} + /// A SNP-only pseudo-alignment: one row (byte sequence, IUPAC-coded) per /// genome, one column per variable family (`family_size() >= 2` — monomorphic /// families carry no signal and are skipped, unlike `raw_snp_distance`'s diff --git a/src/obikmer/src/cmd/distance/mod.rs b/src/obikmer/src/cmd/distance/mod.rs index c3b1035e..3d54cb44 100644 --- a/src/obikmer/src/cmd/distance/mod.rs +++ b/src/obikmer/src/cmd/distance/mod.rs @@ -10,9 +10,9 @@ use clap::Args; use kodama::{Method, linkage}; use obifastwrite::{JsonVal, write_record}; use obikindex::{ - DistanceMetric, KmerIndex, RawSnpDistanceOutput, SankoffWeights, - SiblingAnnexStats, SnpAlignment, build_cost_matrix, c_ctx_from_p_hat, calibrate_p_hat, - mean_substitution_cost, substitution_costs_from_tally, + DistanceMetric, KmerIndex, RawSnpDistanceOutput, + SiblingAnnexStats, SnpAlignment, + cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix, }; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; @@ -87,19 +87,20 @@ pub struct DistanceArgs { /// Exclude a genome (by its exact label) from every computation below /// that reads the sibling annex — `--raw-snp-distance`/`--raw-snp-counts`, - /// `--snp`, and `--sankoff` (and everything `--sankoff` implies: - /// `p_hat`, `sub_cost`, `c_ctx`, the exported matrix/alignment, - /// `--tnt`/`--phyg`/`--iqtree`). Repeatable. Does *not* affect the plain - /// `--metric` distance matrix/NJ/UPGMA path (a different, unrelated - /// computation). Applied by zeroing the excluded genome's row/column - /// after `raw_snp_distance` runs (a pair with zero counts is already - /// skipped by `calibrate_p_hat`/`base_pair_tally`, so this needs no - /// change to the underlying traversal) and by dropping its row from - /// `snp_pseudo_alignment`'s output — the annex is still built/scanned - /// for the excluded genome too, just not used afterward. For a genome - /// with almost no informative sites shared with anything else (see - /// `docmd/theory/evolutionary_distances.md`, the IQ-TREE/Mash rogue-taxon - /// discussion), its presence can otherwise silently bias `p_hat`/`R`. + /// `--snp`, and `--sankoff` (and everything `--sankoff` implies: the + /// cardinality/composition transition models, the exported + /// matrix/alignment, `--tnt`/`--phyg`/`--iqtree`). Repeatable. Does + /// *not* affect the plain `--metric` distance matrix/NJ/UPGMA path (a + /// different, unrelated computation). Applied by zeroing the excluded + /// genome's row/column after `raw_snp_distance` runs (a pair with zero + /// counts is already skipped by `base_pair_tally`/`cardinality_tally`, + /// so this needs no change to the underlying traversal) and by + /// dropping its row from `snp_pseudo_alignment`'s output — the annex + /// is still built/scanned for the excluded genome too, just not used + /// afterward. For a genome with almost no informative sites shared + /// with anything else (see `docmd/theory/evolutionary_distances.md`, + /// the IQ-TREE/Mash rogue-taxon discussion), its presence can + /// otherwise silently bias the transition models. #[arg(long = "exclude-genome", value_name = "LABEL")] pub exclude_genome: Vec, @@ -231,11 +232,12 @@ pub fn run(args: DistanceArgs) { // ── Genome exclusion (`--exclude-genome`) ─────────────────────────────── // Applied by zeroing a `RawSnpDistanceOutput`'s excluded rows/columns - // (`zero_excluded_pairs`) — `calibrate_p_hat`/`base_pair_tally` already - // skip any pair with zero total counts, so this needs no change to - // `obikindex`'s traversal — and by dropping the excluded genome's row - // from a `SnpAlignment` plus the matching label (`drop_excluded`), - // since an all-`∅` row for an "excluded" genome would otherwise still + // (`zero_excluded_pairs`) — `base_pair_tally`/`cardinality_tally` + // already skip any pair with zero total counts, so this needs no + // change to `obikindex`'s traversal — and by dropping the excluded + // genome's row from a `SnpAlignment` plus the matching label + // (`drop_excluded`), since an all-`∅` row for an "excluded" genome + // would otherwise still // reach TNT/PhyG/IQ-TREE as a real (empty) taxon. let exclude_mask: Vec = { let mut mask = vec![false; n]; @@ -263,15 +265,41 @@ pub fn run(args: DistanceArgs) { } } }; + // `snp_pseudo_alignment`'s "variable family" criterion + // (`mask.family_size() >= 2`) is a property of the annex, computed + // over *every* genome in the index — unaffected by `--exclude-genome`. + // So dropping excluded rows alone can leave columns that are variable + // only thanks to an excluded genome now monomorphic among the + // survivors — silently wrong data for TNT/PhyG, and a hard failure + // for IQ-TREE's `+ASC` (verified: excluding 2 taxa on the 20-genome + // benchmark left 116,351 such columns). Re-check variability among the + // *kept* genomes only, after dropping rows, and drop those columns too. let drop_excluded = |alignment: SnpAlignment| -> (SnpAlignment, Vec) { - let sequences = alignment.sequences.into_iter().enumerate() + let mut sequences: Vec> = alignment.sequences.into_iter().enumerate() .filter(|(i, _)| !exclude_mask[*i]) .map(|(_, seq)| seq) .collect(); - let kept_labels = labels.iter().enumerate() + let kept_labels: Vec = labels.iter().enumerate() .filter(|(i, _)| !exclude_mask[*i]) .map(|(_, l)| l.clone()) .collect(); + + if exclude_mask.iter().any(|&excluded| excluded) && !sequences.is_empty() { + let n_sites = sequences[0].len(); + let keep_col: Vec = (0..n_sites) + .map(|site| sequences.iter().any(|seq| seq[site] != sequences[0][site])) + .collect(); + for seq in &mut sequences { + let mut kept = Vec::with_capacity(seq.len()); + for (site, &b) in seq.iter().enumerate() { + if keep_col[site] { + kept.push(b); + } + } + *seq = kept; + } + } + (SnpAlignment { sequences }, kept_labels) }; @@ -330,21 +358,20 @@ pub fn run(args: DistanceArgs) { std::process::exit(1); }); zero_excluded_pairs(&mut raw); - let estimate = calibrate_p_hat(&raw, args.sankoff_ratio_ceiling); - let m = (idx.kmer_size() - 1) / 2; - let tally = idx.base_pair_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| { + let base_tally = idx.base_pair_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| { eprintln!("error computing base-pair tally: {e}"); std::process::exit(1); }); - let sub_cost = substitution_costs_from_tally(&tally); - let mean_sub_cost = mean_substitution_cost(&tally, &sub_cost); - let c_ctx = c_ctx_from_p_hat(estimate.p_hat, m, mean_sub_cost); - - let weights = SankoffWeights { sub_cost, c_ctx }; - let matrix = build_cost_matrix(&weights); - write_sankoff_matrix_csv(&matrix, &estimate, &weights, &args.output); - write_sankoff_params(&estimate, &tally, &weights, args.sankoff_ratio_ceiling, m, mean_sub_cost, &args.output); + let card_tally = idx.cardinality_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| { + eprintln!("error computing cardinality tally: {e}"); + std::process::exit(1); + }); + let p_card = cardinality_transition_probs(&card_tally); + let p_comp = composition_transition_probs(&base_tally); + let matrix = pairwise_cost_matrix(&p_card, &p_comp); + write_sankoff_matrix_csv(&matrix, &args.output); + write_sankoff_params(&card_tally, &p_card, &base_tally, &p_comp, args.sankoff_ratio_ceiling, &args.output); let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { eprintln!("error computing SNP pseudo-alignment: {e}"); diff --git a/src/obikmer/src/cmd/distance/sankoff.rs b/src/obikmer/src/cmd/distance/sankoff.rs index 5e02fda6..c8faaf2b 100644 --- a/src/obikmer/src/cmd/distance/sankoff.rs +++ b/src/obikmer/src/cmd/distance/sankoff.rs @@ -2,7 +2,7 @@ use std::io::{BufWriter, Write}; use std::path::PathBuf; use obifastwrite::{JsonVal, write_record}; -use obikindex::{BasePairTally, PHatEstimate, SankoffWeights, SnpAlignment}; +use obikindex::{BasePairTally, CardinalityTally, SnpAlignment}; use tracing::info; // ── Sankoff pseudo-alignment → FASTA ──────────────────────────────────────── @@ -65,19 +65,8 @@ pub(super) fn state_index_table() -> [u8; 128] { pub(super) fn write_sankoff_matrix_csv( matrix: &[[f64; 16]; 16], - estimate: &PHatEstimate, - weights: &SankoffWeights, output: &Option, ) { - info!( - p_hat = format_args!("{:.6}", estimate.p_hat), - variance = format_args!("{:.3e}", estimate.variance), - n_pairs_included = estimate.n_pairs_included, - n_loci_total = estimate.n_loci_total, - c_ctx = format_args!("{:.4}", weights.c_ctx), - "Sankoff matrix calibration" - ); - let path = output.as_ref() .map(|p| format!("{}_sankoff_matrix.csv", p.display())) .unwrap_or_else(|| "sankoff_matrix.csv".into()); @@ -109,60 +98,64 @@ pub(super) fn write_sankoff_matrix_csv( // rather than re-parsing free text. #[derive(serde::Serialize)] -struct SankoffSubstitution { - pair: String, +struct CardinalityTransition { + from: usize, + to: usize, count: u64, - cost: f64, + probability: f64, +} + +#[derive(serde::Serialize)] +struct CompositionTransition { + from: char, + to: char, + count: u64, + probability: f64, } #[derive(serde::Serialize)] struct SankoffParamsReport { ratio_ceiling: f64, - flank_length_m: usize, - p_hat: f64, - p_hat_variance: f64, - n_pairs_included: usize, - n_loci_total: u64, - /// Weighted-average substitution cost — see `c_ctx_from_p_hat`'s docs: - /// this is what turns the raw expected mutation *count* behind `c_ctx` - /// into an actual cost (not every mutation is worth a flat `1`). - mean_sub_cost: f64, - c_ctx: f64, - substitutions: Vec, + cardinality_transitions: Vec, + composition_transitions: Vec, } pub(super) fn write_sankoff_params( - estimate: &PHatEstimate, - tally: &BasePairTally, - weights: &SankoffWeights, + card_tally: &CardinalityTally, + p_card: &[[f64; 5]; 5], + base_tally: &BasePairTally, + p_comp: &[[f64; 4]; 4], ratio_ceiling: f64, - m: usize, - mean_sub_cost: f64, output: &Option, ) { const BASE_LETTER: [char; 4] = ['A', 'C', 'G', 'T']; - let mut substitutions = Vec::with_capacity(6); - for a in 0..4 { - for b in (a + 1)..4 { - substitutions.push(SankoffSubstitution { - pair: format!("{}/{}", BASE_LETTER[a], BASE_LETTER[b]), - count: tally.counts[a][b], - cost: weights.sub_cost[a][b], + let mut cardinality_transitions = Vec::with_capacity(25); + for a in 0..5 { + for b in 0..5 { + cardinality_transitions.push(CardinalityTransition { + from: a, + to: b, + count: card_tally.counts[a][b], + probability: p_card[a][b], }); } } - let report = SankoffParamsReport { - ratio_ceiling, - flank_length_m: m, - p_hat: estimate.p_hat, - p_hat_variance: estimate.variance, - n_pairs_included: estimate.n_pairs_included, - n_loci_total: estimate.n_loci_total, - mean_sub_cost, - c_ctx: weights.c_ctx, - substitutions, - }; + + let mut composition_transitions = Vec::with_capacity(16); + for a in 0..4 { + for b in 0..4 { + let count = if a == b { base_tally.same[a] } else { base_tally.counts[a][b] }; + composition_transitions.push(CompositionTransition { + from: BASE_LETTER[a], + to: BASE_LETTER[b], + count, + probability: p_comp[a][b], + }); + } + } + + let report = SankoffParamsReport { ratio_ceiling, cardinality_transitions, composition_transitions }; let path = output.as_ref() .map(|p| format!("{}_sankoff_params.yaml", p.display())) @@ -183,15 +176,18 @@ pub(super) fn write_sankoff_params( /// closure* of the result (Floyd-Warshall over the 16 states again, on the /// now-integer values). /// -/// `matrix` is already a metric in its real-valued form (it's a -/// shortest-path closure itself — see `obikindex::build_cost_matrix`), but -/// rounding each cell independently can still break the triangle -/// inequality: e.g. two real costs of `1.734` each round to `173`, summing -/// to `346`, while their own real sum `3.468` rounds to `347` — TNT then -/// reports "triangle inequality violated ... Fixed" and silently -/// substitutes its own corrected value. Re-closing after rounding makes -/// that correction explicit and reproducible here instead, rather than -/// left implicit and tool-version-dependent. +/// Unlike this project's earlier cost-matrix construction (a graph closed +/// by shortest path, guaranteeing a metric by construction), `matrix` here +/// comes from `obikindex::pairwise_cost_matrix`'s row-normalise-then-`-ln` +/// composition, which gives no such guarantee — so this closure isn't only +/// needed to correct integer-rounding artifacts (two real costs of `1.734` +/// each round to `173`, summing to `346`, while their own real sum `3.468` +/// rounds to `347` — TNT then reports "triangle inequality violated ... +/// Fixed" and silently substitutes its own corrected value), it may also +/// be the only thing making the *real-valued* matrix a metric in the first +/// place. Re-closing after rounding makes both corrections explicit and +/// reproducible here instead, rather than left implicit and +/// tool-version-dependent. pub(super) fn scaled_metric_matrix(matrix: &[[f64; 16]; 16], cost_scale: f64) -> [[i64; 16]; 16] { let mut m = [[0i64; 16]; 16]; for i in 0..16 {