From 55d7fa206762b38713bfb1c05287a414a9562ae8 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 12 Aug 2026 15:27:39 +0200 Subject: [PATCH 01/46] feat: add Sankoff parsimony model and directory locking Implements a calibrated 16-state Sankoff substitution cost matrix and CLI pipeline for evolutionary distance computation, including empirical calibration via saturation-filtered SNP counts. Refactors the sibling scanning stage to use batched transforms for improved synchronization efficiency. Introduces an OS-level advisory directory lock across all index-modifying commands to prevent concurrent write corruption. Updates dependencies and exposes new Sankoff utilities in the public API. --- docmd/theory/evolutionary_distances.md | 172 ++++++--- nj.nwk | 1 + sankoff_params.yaml | 27 ++ src/Cargo.lock | 154 +++++++- src/obikindex/src/lib.rs | 7 +- src/obikindex/src/sankoff.rs | 402 ++++++++++++++++++++ src/obikindex/src/siblings.rs | 306 ++++++++++----- src/obikmer/Cargo.toml | 8 +- src/obikmer/src/cmd/distance.rs | 505 ++++++++++++++++++++++++- src/obikmer/src/cmd/filter.rs | 7 + src/obikmer/src/cmd/index.rs | 10 + src/obikmer/src/cmd/merge.rs | 7 + src/obikmer/src/cmd/pack.rs | 7 + src/obikmer/src/cmd/reindex.rs | 7 + src/obikmer/src/cmd/select.rs | 10 + src/obisys/Cargo.toml | 1 + src/obisys/src/lib.rs | 43 +++ 17 files changed, 1494 insertions(+), 180 deletions(-) create mode 100644 nj.nwk create mode 100644 sankoff_params.yaml create mode 100644 src/obikindex/src/sankoff.rs diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index f927381e..b1fd239f 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -379,84 +379,156 @@ per pair loses. Not yet run. ### A concrete Sankoff cost matrix for the 16-state alphabet +**Why `c_gl` and `c_ctx` are the same constant, not two.** An earlier version +of this design used two independent parameters: `c_gl` for an ordinary +gain/loss step between two nonempty states (e.g. `{A,C}->{A}`), and `c_ctx` +only for a state collapsing all the way to `∅`. That distinction doesn't +survive contact with what the index actually observes. Presence/absence +tracking never sees "the flanks" separately from "the centre" — it sees +whole, distinct, homologous 31-mers (two 31-mers sharing >=30 bases are +homologous by construction, though orthology vs. paralogy is undecidable +from that alone — settled earlier, see "Locus eligibility"). Losing the `A` +form of a family while `C` remains (`{A,C}->{C}`) and losing the last +remaining form (`{A}->∅`) are the *same kind of event*: a specific, complete +31-mer that used to be observed no longer is, because at least one of its +31 positions mutated. There is no separate "gain/loss of a still-recognised +allele" mechanism distinct from "loss of context" — both are exactly the +event `c_ctx(p)` (below) already computes the expected cost of. So: one +constant, `c_ctx`, used everywhere a member is gained or lost — `X -> ∅` +included, and (see next point) not even as a special case there. + **Set-edit-distance formula.** For two states `X, Y ⊆ {A,C,G,T}`, split into `seulement_X = X \ Y` (size `a`) and `seulement_Y = Y \ X` (size `b`). Elements present in both cost nothing. Pair up to `min(a,b)` of the remaining elements as **substitutions** (cheaper than treating them as an -unrelated loss + gain whenever `c_sub < 2*c_gl`, which any sane parameter +unrelated loss + gain whenever `c_sub < 2*c_ctx`, which any sane parameter choice satisfies); whatever is left over after pairing is a pure **gain/loss**: ``` -cost(X,Y) = min(a,b)*c_sub + |a-b|*c_gl +cost(X,Y) = min(a,b)*c_sub + |a-b|*c_ctx ``` -Worked examples (`c_sub = c_gl = 1`): `{A}->{C}` = 1 (one substitution); +Worked examples (`c_sub = c_ctx = 1`): `{A}->{C}` = 1 (one substitution); `{A}->{A,C}` = 1 (one gain, no substitution pair available since nothing is only-in-Y that matches an only-in-X element after the shared `A` is excluded); `{A,C}->{G,T}` = 2 (two substitution pairs, `A/C` vs `G/T`, both same-size sets share nothing); `{A,C,G}->{A,C,T}` = 1 (`G`/`T` is the only mismatched pair, `A,C` shared). -**`∅` is a flat-cost special case, not a instance of the general formula.** -Applying the formula naively to e.g. `{A,C,G} -> ∅` would charge `3*c_gl` -(three independent losses). Reject that: total context disappearance -(flank-breaking, or true structural loss) is plausibly **one** event, not -`|X|` of them, so: +**`∅` is *not* a flat-cost special case — a corrected position, reversing +an earlier draft of this design.** An earlier version charged +`cost(X, ∅) = c_ctx` flat, independent of `|X|`, on the grounds that every +member of a family shares identical flanking sequence, so one mutation +breaking that context should unrecognise all of them at once rather than +`|X|` separate times. That argument doesn't survive comparison with how the +rest of the matrix already works: `{A,C,T} -> {A}` (losing two members, +one remaining) is charged `2*c_ctx` via the general formula above, with no +equivalent "maybe it was one shared event" discount — nothing distinguishes +that case from `{A,C} -> ∅` (losing two members, none remaining) other than +which state happens to be the target. Singling out `∅` for special +treatment was arbitrary, not principled: the uncertainty about "one event +or several" is identical in both cases, and the model already resolves it +uniformly elsewhere by simply counting elements. So `∅` is now an ordinary +node in the state graph below like any other, with no patch applied +afterwards: ``` -cost(X, ∅) = cost(∅, X) = c_ctx (constant, independent of |X|) +cost(X, ∅) = cost(∅, X) = |X| * c_ctx (via |X| single-element graph edges) cost(∅, ∅) = 0 ``` -`c_ctx` should not be guessed — it is exactly the value already derived for -the `>1` bucket in "Context, detectability, and a 3-way ordinal distance per -pair" above (`(2m*p_hat)/(1-(1-p_hat)^(2m)) + p_hat`), reused rather than -invented. +**Computing `c_ctx`.** Not guessed — built from the value already derived +for the `>1` bucket in "Context, detectability, and a 3-way ordinal +distance per pair" above: + +``` +c_ctx(p) = mean_sub_cost * [(2m*p) / (1 - (1-p)^(2m)) + p] +``` + +`p` is `p_hat`, the calibrated per-site mutation probability (see +`calibrate_p_hat` / the "Experiment" section below); `m = (k-1)/2`, the +flank length on each side of the centre (`k=31` → `m=15`). Read the +bracketed term, well, term by term: `(1-p)^(2m)` is the probability that +*none* of the `2m` flanking positions mutated, so `1-(1-p)^(2m)` is the +probability that *at least one* did — i.e. the probability that this +specific 31-mer stops being observable at all, which is exactly the event +`c_ctx` prices. `2m*p` is the unconditional expected mutation count over +those `2m` positions. Their ratio is the conditional expectation +`E[mutations | at least one occurred]` — provable directly: for any +nonnegative integer random variable `X`, `X · 1{X>=1} = X` always (both +sides are `0` when `X=0`, both are `X` otherwise), so +`E[X | X>=1] = E[X·1{X>=1}]/P(X>=1) = E[X]/P(X>=1)`, which is +`(2m*p)/(1-(1-p)^(2m))` here. The trailing `+p` adds the (much smaller, +first-order, not itself a conditional expectation) marginal contribution +of the centre position's own mutation probability. + +**The `mean_sub_cost` multiplier — a correction, not part of the original +derivation.** The bracketed term is a *count* of expected mutations, not a +cost — an earlier version used it directly as `c_ctx`, implicitly pricing +every one of those mutations at a flat `1` regardless of type. That stopped +being defensible once substitution costs were calibrated per base-pair +category (`substitution_costs_from_tally`): transitions are markedly more +frequent than transversions in this project's own data (roughly 5-6x, e.g. +130k vs. 20-28k observed instances per category in the eubacteria run +below) and correspondingly cheaper, so "one mutation" isn't worth a flat +unit. `mean_sub_cost` (`mean_substitution_cost`) is the empirical average +substitution cost, weighted by each category's observed frequency — it +converts the bracketed term's expected mutation *count* into an actual +expected *cost*. Concretely, in the eubacteria run: the bracketed term +alone is `~1.17` (close to `1`, i.e. context loss is *usually* attributable +to a single mutation, since `p_hat` is small); `mean_sub_cost ~= 1.48` +(pulled up from the cheapest transition cost of `1.0` by the substantial +minority of transversions in the mix); final `c_ctx ~= 1.73` — sensibly +between the cheapest transition (`~1.0`) and the transversion costs +(`~2.5`-`2.9`), rather than coinciding almost exactly with the cheapest +transition purely by construction accident, which is what the unmultiplied +version did and which is what exposed this gap in the first place. **Better construction method: shortest path in a small state graph, not the -closed-form formula directly.** Build a graph on the 16 states with two edge -types — substitution edges between same-cardinality sets differing by one -element (weight `c_sub`, or `c_ts`/`c_tv` if split further below), and -gain/loss edges between sets whose cardinality differs by one (weight -`c_gl`) — then define `cost(X,Y)` as shortest-path distance in that graph, -precomputed once (16 nodes, trivial) into a dense 16x16 matrix before -feeding it to Sankoff/TNT. Verified equivalent to the closed-form formula -above in the uniform-cost case (checked by hand on `{A}->{C,G,T}`: both give -`c_sub + 2*c_gl`). The graph construction is not just a reformulation for +closed-form formula directly.** Build a graph on all 16 states, `∅` +included, with two edge types — substitution edges between same-cardinality +sets differing by one element (weight `c_sub`, or `c_ts`/`c_tv` if split +further below), and gain/loss edges between sets whose cardinality differs +by one (weight `c_ctx`, per the unification above; `∅` connects to each +singleton state this way, being a strict subset of it) — then define +`cost(X,Y)` as shortest-path distance in that graph, uniformly for every +`X,Y` including `∅`, precomputed once (16 nodes, trivial) into a dense +16x16 matrix before feeding it to Sankoff/TNT. Verified equivalent to the +closed-form formula above in the +uniform-cost case (checked by hand on `{A}->{C,G,T}`: both give +`c_sub + 2*c_ctx`). The graph construction is not just a reformulation for its own sake: it is the version that generalises correctly once substitution costs stop being uniform (next point) — the closed-form's `min(a,b)` counting silently assumes *any* pairing costs the same, which breaks the moment `c_sub` depends on which two bases are involved. -**Transition/transversion refinement.** Split `c_sub` into `c_ts` (A<->G or -C<->T) and `c_tv` (the other four pairs) — already well-defined in canonical -space (see "Canonical invariance" above: a transition maps to a transition, -a transversion to a transversion, regardless of orientation). This turns +**Substitution refinement — implemented as the full 6-category symmetric +matrix, not just Ts/Tv.** `c_sub` was originally going to split into just +`c_ts` (A<->G or C<->T) and `c_tv` (the other four pairs) — already +well-defined in canonical space (see "Canonical invariance" above: a +transition maps to a transition, a transversion to a transversion, +regardless of orientation). A symmetric cost matrix allows finer +resolution than that 2-category split, though: each of the 6 distinct +undirected base pairs (AC, AG, AT, CG, CT, GT) can be costed +independently, with no grouping at all — `SankoffWeights::sub_cost` is a +full `[[f64;4];4]` table (`SankoffWeights::ts_tv` remains as a convenience +constructor for the coarser 2-category case, if ever wanted). This turns "pick `min(a,b)` substitution pairs" into a genuine (tiny, <=4 elements per side, trivially enumerable) minimum-cost bipartite matching problem instead of a plain count — the state-graph shortest-path construction handles this -automatically, no separate logic needed. Calibration is not a guess either: -the "Sufficient statistic: 4x4 base-pair tally" section below already plans -to collect the joint `(centre_i, centre_j)` distribution over resolved (`0` -or `1`) sites — that tally directly gives the empirical Ts/Tv ratio, from -which `c_ts`/`c_tv` follow (e.g. `cost ∝ -log(observed rate)`, the standard -generalised-parsimony step-weighting heuristic), reusing a statistic already -planned rather than adding a new one. - -**`gamma`/`mu` (i.e. `c_gl`/`c_sub`) sensitivity sweep before calibration.** -No strong prior on whether gain/loss events should cost more or less than a -point substitution — duplication/deletion rates are not a priori equal to -point-mutation rates, but the direction isn't obvious, and setting it too -high risks the parsimony search *eliminating* exactly the -heterozygosity/paralogy signal the design is meant to tolerate (see -"Heterozygosity, ploidy..." below). Cheap first step: run the topology at a -handful of ratios (`0.5, 1, 2, 5`) and check whether it's stable — a robust -topology across that range is far more trustworthy than one built on a -single, unvalidated guess. Empirical calibration of `c_gl` from the -family-size distribution already available (`sibling_annex_stats`) is a -natural follow-up once the sensitivity sweep shows the topology is worth -refining further. +automatically, no separate logic needed. Calibration is not a guess +either: `KmerIndex::base_pair_tally` collects the pooled `(centre_i, +centre_j)` distribution over resolved (SNP) sites — restricted to the same +ratio-ceiling-included genome pairs as `p_hat`'s own calibration, for the +same saturation-exclusion reason — and `substitution_costs_from_tally` +turns that into `cost(a,b) = -ln(rate(a,b))`, normalised so the most +frequent category costs `1.0` (the standard generalised-parsimony +step-weighting heuristic, generalised from 2 categories to 6). A second +full pass over the annex is required for this (see `base_pair_tally`'s own +docs for why it can't share `raw_snp_distance`'s single pass) — both now +share one traversal helper, `scan_family_pairs`, rather than duplicating +the per-family reconciliation logic. **Caveat carried over from the `D_F = min(a,b)` rejection earlier:** this cost matrix is only valid as **Sankoff step-cost input**, re-evaluated for @@ -469,7 +541,7 @@ specific genomes. **Feasibility confirmed:** TNT's `costs` command accepts custom step matrices for multistate characters, so this whole construction (16x16 -matrix derived from the state graph, `c_ts`/`c_tv`/`c_gl`/`c_ctx` as +matrix derived from the state graph, `c_ts`/`c_tv`/`c_ctx` as the three tunable parameters) is directly usable there — no new tooling required before testing it. @@ -491,7 +563,9 @@ emits a complete TNT script (`xread` matrix + `smatrix` step-matrix + `hold`/ `write_raw_snp_distance_csv` to emit raw counts, `c_ctx` was approximated as the unweighted mean of the pairwise `snp/(snp+shared)` ratios already present in `rawsnp.csv`, restricted to the relevant taxon subset. Used -values: `mu = gamma = 10`, `c_ctx = 18` (ratio `c_ctx/mu = 1.8`, consistent +values: `mu = gamma = 10`, `c_ctx = 18` (this run predates the `c_gl`/`c_ctx` +unification above — `gamma` here is what `c_gl` was called before it turned +out to just be `c_ctx`; ratio `c_ctx/mu = 1.8`, consistent with the observed pairwise ratios for genomes within Enterobacteriaceae/ Eubacteria, which cluster around 1.5-2 and barely move as the taxon set widens — the context signal is stable, not sensitive to which subset is diff --git a/nj.nwk b/nj.nwk new file mode 100644 index 00000000..47b92b3a --- /dev/null +++ b/nj.nwk @@ -0,0 +1 @@ +(((((((((((((Salmonella_enterica--P125109:0.009029074806027838,Salmonella_enterica--LT2:0.010374535686866677):0.00299392644685007,Salmonella_enterica--AKU_12601:0.010129090864056669):0.008889984754373365,Salmonella_enterica--CT18:0.029563522158784489):0.05178212639099835,(Escherichia_coli--CFT073:0.018265758526676266,(Escherichia_coli--EDL933:0.01586939152374907,(Escherichia_coli--K-12_MG1655:0.0043321525684933419,Escherichia_coli--K-12_W3110:0.004373581717010194):0.007835004272817602):0.005522415854494232):0.05395833543383322):0.007199364670842333,(Klebsiella_pneumoniae--HS11286:0.027306681424389149,(Klebsiella_pneumoniae--ATCC_13883:0.015176538442299506,Klebsiella_pneumoniae--MGH_78578:0.026462018414958045):0.0015885400482290244):0.060738462006352359):0.029481404183583916,Yersinia_ruckeri--YRB:0.10269826524749153):0.004316815340436764,Proteus_mirabilis--HI4320:0.11745645944715785):0.04439101156727093,Wolbachia_endosymbiont--GCF_000306885.1_ASM30688v1:0.17151356406519997):0.011344480816758762,Acidobacterium_capsulatum--ATCC_51196:0.17121217909071463):0.01869709592355495,Shouchella_clausii--KSM-K16:0.1554729916092201):0.04397682726757515,Bacillus_subtilis--168:0.10403127799132669):0.08828869671555731,Opitutus_terrae--PB90-1:0.06380512639846864):0.41655625366738377,Candidozyma_auris--GCF_003013715.1_ASM301371v2:0.22735026269091164,Saccharolobus_islandicus--M.16.4:0.21891918239548398); \ No newline at end of file diff --git a/sankoff_params.yaml b/sankoff_params.yaml new file mode 100644 index 00000000..6bb64d96 --- /dev/null +++ b/sankoff_params.yaml @@ -0,0 +1,27 @@ +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 + count: 27810 + cost: 2.5439415226130238 +- pair: A/G + count: 130153 + cost: 1.0 +- pair: A/T + count: 24555 + cost: 2.6684722241428322 +- pair: C/G + count: 19637 + cost: 2.892062911764701 +- pair: C/T + count: 124878 + cost: 1.0413902163542994 +- pair: G/T + count: 26616 + cost: 2.5878424665170052 diff --git a/src/Cargo.lock b/src/Cargo.lock index 72bc2355..4acaca3b 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -48,6 +48,15 @@ dependencies = [ "as-slice", ] +[[package]] +name = "aligned-vec" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b" +dependencies = [ + "equator", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -497,6 +506,18 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "console" +version = "0.16.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c" +dependencies = [ + "encode_unicode", + "libc", + "unicode-width", + "windows-sys 0.61.2", +] + [[package]] name = "const-oid" version = "0.10.2" @@ -820,6 +841,26 @@ dependencies = [ "syn", ] +[[package]] +name = "equator" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc" +dependencies = [ + "equator-macro", +] + +[[package]] +name = "equator-macro" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -907,6 +948,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs4" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8c6b3bd49c37d2aa3f3f2220233b29a7cd23f79d1fe70e5337d25fb390793de" +dependencies = [ + "rustix 0.38.44", + "windows-sys 0.52.0", +] + [[package]] name = "fs_at" version = "0.2.1" @@ -1228,13 +1279,26 @@ version = "0.17.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" dependencies = [ - "console", + "console 0.15.11", "number_prefix", "portable-atomic", "unicode-width", "web-time", ] +[[package]] +name = "indicatif" +version = "0.18.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c" +dependencies = [ + "console 0.16.4", + "portable-atomic", + "unicode-width", + "unit-prefix", + "web-time", +] + [[package]] name = "infer" version = "0.19.0" @@ -1326,12 +1390,9 @@ dependencies = [ [[package]] name = "kodama" -version = "0.2.3" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a44f3a71a44fbf49ce38152db7dc9adf959d4fe5c29344cd1858bdbda8d9091" -dependencies = [ - "num-traits", -] +checksum = "f868feceed4703c842925c14232667c5f29c8059d0354154b2d06ec3bf9daf82" [[package]] name = "lazy_static" @@ -1365,6 +1426,12 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -1695,7 +1762,7 @@ version = "0.1.0" dependencies = [ "crossbeam-channel", "hwlocality", - "indicatif", + "indicatif 0.17.11", "ndarray", "obicompactvec", "obikpartitionner", @@ -1720,7 +1787,7 @@ version = "1.1.44" dependencies = [ "clap", "csv", - "indicatif", + "indicatif 0.18.6", "kodama", "obidebruinj", "obifastwrite", @@ -1737,7 +1804,9 @@ dependencies = [ "obitaxonomy", "pprof", "rayon", + "serde", "serde_json", + "serde_yaml", "speedytree", "tracing", "tracing-subscriber", @@ -1749,7 +1818,7 @@ version = "0.1.0" dependencies = [ "cacheline-ef", "epserde", - "indicatif", + "indicatif 0.17.11", "memmap2", "niffler 3.0.0", "obicompactvec", @@ -1853,7 +1922,7 @@ dependencies = [ "memmap2", "niffler 3.0.0", "obikseq", - "rustix", + "rustix 1.1.4", "serde", "serde_json", "tempfile", @@ -1863,7 +1932,8 @@ dependencies = [ name = "obisys" version = "0.1.0" dependencies = [ - "indicatif", + "fs4", + "indicatif 0.17.11", "libc", "sysinfo", "tracing", @@ -2017,10 +2087,11 @@ dependencies = [ [[package]] name = "pprof" -version = "0.13.0" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef5c97c51bd34c7e742402e216abdeb44d415fbe6ae41d56b114723e953711cb" +checksum = "38a01da47675efa7673b032bf8efd8214f1917d89685e07e395ab125ea42b187" dependencies = [ + "aligned-vec", "backtrace", "cfg-if", "findshlibs", @@ -2028,15 +2099,15 @@ dependencies = [ "log", "nix 0.26.4", "once_cell", - "parking_lot", "prost", "prost-build", "prost-derive", "sha2", "smallvec", + "spin", "symbolic-demangle", "tempfile", - "thiserror 1.0.69", + "thiserror 2.0.18", ] [[package]] @@ -2386,6 +2457,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags 2.11.1", + "errno", + "libc", + "linux-raw-sys 0.4.15", + "windows-sys 0.59.0", +] + [[package]] name = "rustix" version = "1.1.4" @@ -2395,7 +2479,7 @@ dependencies = [ "bitflags 2.11.1", "errno", "libc", - "linux-raw-sys", + "linux-raw-sys 0.12.1", "windows-sys 0.61.2", ] @@ -2521,6 +2605,19 @@ dependencies = [ "zmij", ] +[[package]] +name = "serde_yaml" +version = "0.9.34+deprecated" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" +dependencies = [ + "indexmap", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + [[package]] name = "sha2" version = "0.10.9" @@ -2588,6 +2685,15 @@ dependencies = [ "rb_tree", ] +[[package]] +name = "spin" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3" +dependencies = [ + "lock_api", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -2736,7 +2842,7 @@ dependencies = [ "fastrand", "getrandom 0.3.4", "once_cell", - "rustix", + "rustix 1.1.4", "windows-sys 0.61.2", ] @@ -2920,6 +3026,18 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unit-prefix" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3" + +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" + [[package]] name = "untrusted" version = "0.9.0" @@ -3421,7 +3539,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156" dependencies = [ "libc", - "rustix", + "rustix 1.1.4", ] [[package]] diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 48daf548..a48eb325 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -8,6 +8,7 @@ mod merge; mod numa; mod rebuild; mod reindex; +mod sankoff; mod select; mod siblings; mod stats; @@ -19,4 +20,8 @@ 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::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; +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, +}; diff --git a/src/obikindex/src/sankoff.rs b/src/obikindex/src/sankoff.rs new file mode 100644 index 00000000..372ed21e --- /dev/null +++ b/src/obikindex/src/sankoff.rs @@ -0,0 +1,402 @@ +//! 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 90379b2c..dc33bf6a 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -14,18 +14,25 @@ //! mask by callers, not stored (see `FamilyMask` and //! [`sibling_annex_stats`](KmerIndex::sibling_annex_stats) below). //! -//! Per layer, an `obipipeline` `Flat` stage (throttled — see -//! `obipipeline::throttle`) generates each k-mer's 3 central variants, -//! interleaved across many in-flight k-mers by the scheduler's shared -//! worker pool rather than processed on a single thread. The actual -//! cross-partition lookup, though, reuses -//! `KmerPartition::query_partition_with` — the same partition-batching -//! mechanism `obikmer query` already uses (open a partition's files once, -//! answer a whole batch of queries against it) — rather than a per-item -//! pipeline stage: an earlier per-item design reopened/re-mmap'd every -//! target partition's files on every single lookup, which was fine at -//! toy scale but manifested as ~90% system time against a real index. -//! See `docmd/theory/evolutionary_distances.md`, Step 2b, "Mechanism". +//! Per layer, an `obipipeline` batch transform (throttled — see +//! `obipipeline::throttle`) generates a whole batch's central variants at +//! once (`BATCH_SIZE` source k-mers in, that batch's variants out as one +//! pipeline message), interleaved across many in-flight batches by the +//! scheduler's shared worker pool rather than processed on a single +//! thread. The actual cross-partition lookup reuses a `PartitionCache` of +//! every partition's already-open MPHF layers, built once for the whole +//! `build_sibling_annex` run, rather than reopening files per lookup or +//! per source layer. Two earlier, coarser-grained designs were tried and +//! measured (not guessed) to be worse, in order: (1) reopening/re-mmap'ing +//! every target partition's files on every single lookup — fine at toy +//! scale, ~90% system time against a real index; (2) a `Flat` pipeline +//! stage pushing one message per generated *variant* (up to 3 per source +//! k-mer) — cheaper than reopening files, but sampling a real run showed +//! most wall-clock time going into per-message channel send/notify +//! syscalls rather than the lookup itself, because a single k-mer's ≤3 +//! variants is far too fine a granularity to amortise a pipeline's +//! synchronisation cost over. See `docmd/theory/evolutionary_distances.md`, +//! Step 2b, "Mechanism". use std::path::Path; use std::sync::atomic::{AtomicU8, Ordering}; @@ -107,34 +114,30 @@ fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize { // ── obipipeline data types ───────────────────────────────────────────────── -/// One distinct k-mer of the layer currently being processed, at its local -/// MPHF slot — the pipeline's source item. Carries a throttle slot (shared, -/// `Arc`-wrapped so it can be cloned into each of the (up to 3) items this -/// one fans out into via the `Flat` stage) that is only released once every -/// one of those descendants has been fully processed — see the module docs' -/// "Throttling" note for why this is required, not optional, once a `Flat` -/// stage is in the pipeline. -struct SourceItem { - slot: usize, - kmer: CanonicalKmer, - _permit: Arc, +/// A batch of this layer's distinct k-mers (local MPHF slot + k-mer), the +/// pipeline's source item — batched, not one k-mer per item, so that +/// pipeline messages and their synchronisation cost stay amortised over +/// thousands of lookups (see `build_layer_sibling_annex`'s comment on +/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved +/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform, +/// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing. +struct SourceBatch { + items: Vec<(usize, CanonicalKmer)>, + _permit: ThrottleGuard, } -/// One of a source k-mer's (up to 3) central-substitution variants, already -/// routed to its destination partition and carrying its own central base -/// (0=A/1=C/2=G/3=T) — the mask bit it will set on a hit. Carries a clone of -/// the source item's throttle permit. -struct VariantQuery { - source_slot: usize, - dest_partition: usize, - variant: CanonicalKmer, - base: u8, - _permit: Arc, +/// One batch's worth of central-substitution variants (up to 3 per source +/// k-mer), each already routed to its destination partition and carrying +/// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a +/// hit. `(dest_partition, variant, source_slot, base)` per entry. +struct VariantBatch { + items: Vec<(usize, CanonicalKmer, usize, u8)>, + _permit: ThrottleGuard, } enum SibData { - Item(SourceItem), - Query(VariantQuery), + Batch(SourceBatch), + Variants(VariantBatch), } /// Every partition's already-open MPHF layers, built **once** for the whole @@ -341,57 +344,72 @@ impl KmerIndex { mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed); } - // ── obipipeline: Flat stage generates variants only — the actual - // cross-partition lookup reuses `KmerPartition::query_partition_with` - // (the same batching mechanism `obikmer query` already uses: open a - // partition's files once, answer a whole batch of queries against - // it) instead of one lookup per pipeline item. A per-item lookup - // (tried first) reopened/re-mmap'd every target partition's files on - // every single variant — fine at the scale of a handful of test - // k-mers, but with billions of lookups against a real index this - // manifested as ~90% system time, observed in practice. ─────────── + // ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one — + // the actual cross-partition lookup reuses + // `KmerPartition::query_partition_with` (the same batching mechanism + // `obikmer query` already uses: open a partition's files once, + // answer a whole batch of queries against it) instead of one lookup + // per pipeline item. A per-item lookup (tried first) reopened/ + // re-mmap'd every target partition's files on every single variant + // — fine at toy scale, but ~90% system time against a real index, + // observed in practice. A *later* attempt still pushed one pipeline + // message per generated variant (a `Flat` stage, `SourceItem` => + // `VariantQuery`, one k-mer in => up to 3 variants out as separate + // messages) — cheaper than reopening files, but sampling a real run + // showed most wall-clock time going into per-message channel + // send/notify syscalls instead of the lookup itself: the pipeline's + // whole point is amortising synchronisation over a batch, and a + // single k-mer's ≤3 variants is far too fine a granularity for + // that. Batching `BATCH_SIZE` source k-mers into one pipeline item + // — a plain 1-to-1 (`|`, not `||`) transform, batch in, batch of + // variants out, one message either way — keeps the per-message + // synchronisation cost amortised over thousands of lookups instead + // of one to three. ────────────────────────────────────────────── + const BATCH_SIZE: usize = 4096; let n_workers = obisys::effective_parallelism(); let capacity = 256; - // Throttling is not optional once a `Flat` stage is in the pipeline - // (see `obipipeline::throttle`'s docs): without it, every worker can - // become a simultaneous `Flat` producer, saturate the shared output - // channel, and deadlock against the scheduler's own dispatch loop — - // also observed in practice. The permit acquired here for a source - // k-mer is held (via the `Arc`-shared guard carried through - // `SourceItem` -> `VariantQuery`) until every one of its (up to 3) - // descendants has been read out of the pipeline by the accumulation - // loop below, not just until the `Flat` stage itself returns. + // Throttling limits how many *batches* are in flight at once — the + // permit is acquired per batch (not per k-mer) in the source + // thread, and released once its `VariantBatch` has been read out of + // the pipeline by the accumulation loop below. See + // `obipipeline::throttle`'s docs for why this is required, not + // optional, once a `Flat`-style stage sits in the pipeline. let sources: Vec<(usize, CanonicalKmer)> = slot_kmer .iter() .enumerate() .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer))) .collect(); - let throttled = obipipeline::throttle(sources.into_iter(), n_workers).map(|t| SourceItem { - slot: t.item.0, - kmer: t.item.1, - _permit: Arc::new(t.guard), + let batches: Vec> = sources + .chunks(BATCH_SIZE) + .map(|chunk| chunk.to_vec()) + .collect(); + let throttled = obipipeline::throttle(batches.into_iter(), n_workers).map(|t| SourceBatch { + items: t.item, + _permit: t.guard, }); let pipe = obipipeline::make_pipe! { - SibData : SourceItem => VariantQuery, - || { - move |item: SourceItem| -> Vec { - let kmer = item.kmer; - let permit = item._permit; - kmer.central_canonical_neighbors() - .into_iter() - .filter(|variant| *variant != kmer) - .map(|variant| VariantQuery { - source_slot: item.slot, - dest_partition: partition_of(variant, n_parts), - variant, - base: central_base(variant, k), - _permit: Arc::clone(&permit), - }) - .collect::>() + SibData : SourceBatch => VariantBatch, + | { + move |batch: SourceBatch| -> VariantBatch { + let mut items = Vec::with_capacity(batch.items.len() * 3); + for (slot, kmer) in batch.items { + for variant in kmer.central_canonical_neighbors() { + if variant == kmer { + continue; + } + items.push(( + partition_of(variant, n_parts), + variant, + slot, + central_base(variant, k), + )); + } + } + VariantBatch { items, _permit: batch._permit } } - } : Item => Query, + } : Batch => Variants, }; // ── Group generated variants by destination partition. `cache` @@ -400,11 +418,13 @@ impl KmerIndex { // lookup is not free just because the file isn't reopened. Grouping // keeps one partition's pages hot while its whole batch is resolved, // instead of faulting pages in and out as lookups jump between - // partitions in whatever order the `Flat` stage happens to produce - // them. The throttle permit drops here, once accumulated. ───────── + // partitions in whatever order the pipeline happens to produce + // them. Each batch's throttle permit drops here, once accumulated. let mut outgoing: Vec> = (0..n_parts).map(|_| Vec::new()).collect(); - for vq in pipe.apply(throttled, n_workers, capacity) { - outgoing[vq.dest_partition].push((vq.variant, vq.source_slot, vq.base)); + for vb in pipe.apply(throttled, n_workers, capacity) { + for (dest_partition, variant, source_slot, base) in vb.items { + outgoing[dest_partition].push((variant, source_slot, base)); + } } // ── Resolve each partition's batch against the cache in one @@ -635,9 +655,27 @@ pub struct RawSnpDistanceOutput { } impl KmerIndex { - /// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex - /// (run [`build_sibling_annex`](Self::build_sibling_annex) first). - pub fn raw_snp_distance(&self) -> OKIResult { + /// Shared traversal behind [`raw_snp_distance`](Self::raw_snp_distance) + /// and [`base_pair_tally`](Self::base_pair_tally): for every family + /// (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`. + fn scan_family_pairs( + &self, + label: &str, + zero: impl Fn() -> Acc + Sync, + on_pair: F, + combine: C, + ) -> OKIResult + where + Acc: Send, + F: Fn(&mut Acc, usize, usize, u8, u8) + Sync, + C: Fn(Acc, Acc) -> Acc, + { let n_parts = self.n_partitions(); let n_genomes = self.meta.genomes.len(); let with_counts = self.meta.config.with_counts; @@ -673,12 +711,11 @@ impl KmerIndex { } } - let pb = progress_bar("raw_snp_distance", layer_dirs.len() as u64, "layers"); - let partials: Vec<(Array2, Array2)> = layer_dirs + let pb = progress_bar(label, layer_dirs.len() as u64, "layers"); + let partials: Vec = layer_dirs .par_iter() - .map(|layer_dir| -> OKIResult<(Array2, Array2)> { - let mut snp = Array2::::zeros((n_genomes, n_genomes)); - let mut shared = Array2::::zeros((n_genomes, n_genomes)); + .map(|layer_dir| -> OKIResult { + let mut acc = zero(); 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)?; @@ -754,31 +791,102 @@ impl KmerIndex { continue; } let Some(bj) = single_form[j] else { continue }; - if bi == bj { - shared[[i, j]] += 1; - shared[[j, i]] += 1; - } else { - snp[[i, j]] += 1; - snp[[j, i]] += 1; - } + on_pair(&mut acc, i, j, bi, bj); } } } pb.inc(1); - Ok((snp, shared)) + Ok(acc) }) .collect::>>()?; pb.finish_and_clear(); - let mut snp = Array2::::zeros((n_genomes, n_genomes)); - let mut shared = Array2::::zeros((n_genomes, n_genomes)); - for (s, sh) in partials { - snp += &s; - shared += &sh; + let mut total = zero(); + for partial in partials { + total = combine(total, partial); } + Ok(total) + } + + /// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex + /// (run [`build_sibling_annex`](Self::build_sibling_annex) first). + pub fn raw_snp_distance(&self) -> OKIResult { + let n_genomes = self.meta.genomes.len(); + 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| { + if bi == bj { + shared[[i, j]] += 1; + shared[[j, i]] += 1; + } else { + snp[[i, j]] += 1; + snp[[j, i]] += 1; + } + }, + |(mut snp, mut shared), (s, sh)| { + snp += &s; + shared += &sh; + (snp, shared) + }, + )?; Ok(RawSnpDistanceOutput { snp, shared }) } + + /// 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. + /// + /// A second full pass over the annex, sharing + /// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by + /// it, not a blind re-scan) — needed because `raw_snp_distance` only + /// keeps aggregate SNP/shared counts per genome pair, not which bases + /// were actually involved at each locus, and the ratio-ceiling filter + /// can only be evaluated once the aggregate counts are known. + pub fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult { + let n_genomes = self.meta.genomes.len(); + 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 counts = self.scan_family_pairs( + "base_pair_tally", + || [[0u64; 4]; 4], + |counts, i, j, bi, bj| { + if bi != bj && included[[i, j]] { + counts[bi as usize][bj as usize] += 1; + counts[bj as usize][bi as usize] += 1; + } + }, + |mut total, partial| { + for a in 0..4 { + for b in 0..4 { + total[a][b] += partial[a][b]; + } + } + total + }, + )?; + Ok(BasePairTally { counts }) + } +} + +/// See [`KmerIndex::base_pair_tally`]. +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). + pub counts: [[u64; 4]; 4], } /// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index 9f3932d6..7a16f63c 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -22,15 +22,17 @@ obikindex = { path = "../obikindex", default-features = false } obitaxonomy = { path = "../obitaxonomy" } obilayeredmap = { path = "../obilayeredmap" } clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } serde_json = "1" +serde_yaml = "0.9.33" csv = "1" -kodama = "0.2" +kodama = "0.3.0" speedytree = "0.1" rayon = "1" -indicatif = "0.17" +indicatif = "0.18" tracing = "0.1.44" tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } -pprof = { version = "0.13", features = ["prost-codec"], optional = true } +pprof = { version = "0.15", features = ["prost-codec"], optional = true } [features] default = ["numa"] diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index 88524786..bd7a4e37 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -4,7 +4,11 @@ use std::path::PathBuf; use clap::Args; use kodama::{Method, linkage}; use obifastwrite::{JsonVal, write_record}; -use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; +use obikindex::{ + BasePairTally, DistanceMetric, KmerIndex, PHatEstimate, RawSnpDistanceOutput, SankoffWeights, + SiblingAnnexStats, SnpAlignment, build_cost_matrix, c_ctx_from_p_hat, calibrate_p_hat, + mean_substitution_cost, substitution_costs_from_tally, +}; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; @@ -94,9 +98,61 @@ pub struct DistanceArgs { #[arg(long)] pub snp: bool, + /// Calibrate a 16-state Sankoff cost matrix and its matching + /// pseudo-alignment from an already-built sibling annex (run with + /// `--sibling-annex` first, in this invocation or an earlier one), for + /// use with TNT/PhyG. See `docmd/theory/evolutionary_distances.md`, + /// "Sankoff parsimony as the resolution of the 16-state model problem". + #[arg(long)] + pub sankoff: bool, + + /// Exclude genome pairs whose raw SNP ratio exceeds this value from the + /// `p_hat` calibration pooled by `--sankoff` — a pair this close to + /// saturation carries no information about `p_hat` and would bias it + /// upward if pooled in (unlike a low eligible-loci count, which barely + /// moves the pooled estimate either way — see design doc). + #[arg(long, default_value = "0.5")] + pub sankoff_ratio_ceiling: f64, + + /// Also write _sankoff.tnt, a ready-to-run TNT script (`proc + /// ;`) for the same matrix/alignment `--sankoff` computes — + /// recoded to TNT's default xread alphabet (0-9A-F only; TNT rejects + /// the wider IUPAC set `--sankoff`'s own output uses unless `nstates + /// dna` is set, which imposes TNT's own incompatible DNA encoding + /// instead) with integer-scaled costs (TNT's smatrix/cost commands + /// reject decimals). Implies `--sankoff`. + #[arg(long)] + pub tnt: bool, + + /// Also write _sankoff.tcm and _sankoff.pg, a + /// custom-alphabet cost matrix and a ready-to-run PhyG script (`read`/ + /// `search`/`report`) for the same matrix/alignment `--sankoff` + /// computes. Reuses `--sankoff`'s own `_sankoff.fasta` directly — PhyG's + /// `tcm:` alphabet is read from the matrix file itself, so the IUPAC+`0` + /// alphabet needs no recoding here, unlike `--tnt`. Implies `--sankoff`. + #[arg(long)] + pub phyg: bool, + + /// Scale factor applied before rounding real-valued costs to the + /// integers both `--tnt`'s smatrix/cost commands and `--phyg`'s `tcm:` + /// matrix require. Keep this small: the total tree score is this scale + /// times the sum of per-character costs across every character (908k+ + /// for a typical run here), and there are hints in TNT's own manual + /// that at least some of its internal accumulators are 32-bit — a large + /// scale risks a silent integer overflow (undetectable, not just a + /// crash) far more costly than the resolution a bigger factor would + /// buy. Shared between `--tnt` and `--phyg` rather than split into two + /// flags: both scale the same calibrated matrix for the same reason + /// (integer-only cost commands), and no PhyG-specific accumulator-width + /// constraint has actually been found to justify a different default. + #[arg(long, default_value = "100")] + pub sankoff_cost_scale: f64, + /// Output prefix: _dist.csv, _shared.csv, /// _siblings.csv, _rawsnp.csv, _snp.fasta, - /// _nj.nwk, _upgma.nwk. + /// _sankoff_matrix.csv, _sankoff_params.yaml, + /// _sankoff.fasta, _sankoff.tnt, _sankoff.tcm, + /// _sankoff.pg, _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -117,6 +173,13 @@ pub fn run(args: DistanceArgs) { // deliberately decoupled: the annex is meant to be (re)built routinely, // the distribution only occasionally, on demand. if args.sibling_annex { + // Writes into the index directory — hold an exclusive lock for the + // duration so a second, concurrent `--sibling-annex` run on the same + // index can't corrupt these writes (see obisys::DirLock). + let _lock = obisys::DirLock::acquire(&args.index).unwrap_or_else(|e| { + eprintln!("error locking index directory {}: {e}", args.index.display()); + std::process::exit(1); + }); info!("building sibling-count/minorant annex"); idx.build_sibling_annex().unwrap_or_else(|e| { eprintln!("error building sibling annex: {e}"); @@ -144,15 +207,56 @@ pub fn run(args: DistanceArgs) { }); write_snp_fasta(&alignment, &labels, &args.output); } + if args.sankoff || args.tnt || args.phyg { + let raw = idx.raw_snp_distance().unwrap_or_else(|e| { + eprintln!("error computing raw SNP distance: {e}"); + std::process::exit(1); + }); + let estimate = calibrate_p_hat(&raw, args.sankoff_ratio_ceiling); + let m = (idx.kmer_size() - 1) / 2; - // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp` are - // their own operation, not a modifier on top of a distance-metric - // computation — a metric was never requested by asking for any of them, - // so there is nothing for the rest of this function to compute. Not a - // historical accident to keep: stop here rather than always also - // running a Jaccard (or whichever `--metric` defaults to) pass and - // printing an unrequested matrix. - if args.sibling_annex || args.sibling_stats || args.raw_snp_distance || args.snp { + let 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 alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { + eprintln!("error computing SNP pseudo-alignment: {e}"); + std::process::exit(1); + }); + write_sankoff_alignment_fasta(&alignment, &labels, &args.output); + + if args.tnt { + write_sankoff_tnt(&matrix, &alignment, &labels, &args.output, args.sankoff_cost_scale); + } + if args.phyg { + write_sankoff_phyg(&matrix, &args.output, args.sankoff_cost_scale); + } + } + + // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp`/ + // `--sankoff`/`--tnt` are their own operation, not a modifier on top of + // a distance-metric computation — a metric was never requested by + // asking for any of them, so there is nothing for the rest of this + // function to compute. Not a historical accident to keep: stop here + // rather than always also running a Jaccard (or whichever `--metric` + // defaults to) pass and printing an unrequested matrix. + if args.sibling_annex + || args.sibling_stats + || args.raw_snp_distance + || args.snp + || args.sankoff + || args.tnt + || args.phyg + { return; } @@ -364,6 +468,387 @@ fn write_snp_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option< if n_sites == 1 { "" } else { "s" }); } +// ── Sankoff pseudo-alignment → FASTA ──────────────────────────────────────── +// +// Same data as `--snp`'s pseudo-alignment (`SnpAlignment`/ +// `snp_pseudo_alignment`), re-coded so its symbols match the accompanying +// `--sankoff-matrix` output exactly: `0` for the empty/absent state instead +// of `-`, which TNT/PhyG would otherwise read as their own gap character +// rather than our "family absent" state. + +fn write_sankoff_alignment_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_sankoff.fasta", p.display())) + .unwrap_or_else(|| "sankoff.fasta".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); + for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { + let recoded: Vec = seq.iter().map(|&b| if b == b'-' { b'0' } else { b }).collect(); + write_record(&recoded, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| { + eprintln!("error writing {path}: {e}"); + std::process::exit(1); + }); + } + info!("Sankoff pseudo-alignment → {path} ({n_sites} site{})", + if n_sites == 1 { "" } else { "s" }); +} + +// ── Sankoff cost matrix → CSV ──────────────────────────────────────────────── +// +// 16 states indexed by bitmask (bit 0=A, 1=C, 2=G, 3=T; state 0 is `∅`), +// matching the convention already used for `--snp`'s IUPAC-coded output and +// for the external TNT/PhyG scripts this feeds. Calibration report (p_hat, +// its variance, how many pairs/loci went into it, the resulting c_ctx) goes +// to the log, not the CSV, since it's a run-level fact, not per-cell data. + +// IUPAC ambiguity code per state (same mapping as `siblings::iupac_code`, +// already used for `--snp`'s pseudo-alignment — a biologist reads "R" as +// "A or G" without needing this file's convention explained), with `0` +// standing in for the empty state (`-` would collide with TNT/PhyG's own +// gap/range syntax). Bit order: 0=A, 1=C, 2=G, 3=T. This project's +// canonical alphabet — `write_sankoff_tnt` below recodes it to TNT's own +// default alphabet at the adapter boundary, rather than using it here. +const STATE_SYMBOL: [char; 16] = [ + '0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N', +]; + +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()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + write!(f, "state").unwrap(); + for sym in STATE_SYMBOL { write!(f, ",{sym}").unwrap(); } + writeln!(f).unwrap(); + for (s, row) in matrix.iter().enumerate() { + write!(f, "{}", STATE_SYMBOL[s]).unwrap(); + for cost in row { write!(f, ",{cost:.4}").unwrap(); } + writeln!(f).unwrap(); + } + info!("Sankoff cost matrix → {path}"); +} + +// ── Sankoff calibration parameters → YAML report ──────────────────────────── +// +// Everything `--sankoff` estimates from real data, in one durable, +// machine-readable file: `p_hat` and its variance (with how many pairs/loci +// went into it), the derived `c_ctx`, and the base-pair substitution tally +// (raw counts, not just the derived costs) — kept for the same reason raw +// counts are kept anywhere else in this project: costs are a modelling +// choice built *from* the counts, and reproducing/re-deriving them later +// needs the counts, not just their current derived value. Structured (YAML, +// not an ad hoc key=value text file) so R/Python/etc. can load it directly +// rather than re-parsing free text. + +#[derive(serde::Serialize)] +struct SankoffSubstitution { + pair: String, + count: u64, + cost: 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, +} + +fn write_sankoff_params( + estimate: &PHatEstimate, + tally: &BasePairTally, + weights: &SankoffWeights, + 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 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 path = output.as_ref() + .map(|p| format!("{}_sankoff_params.yaml", p.display())) + .unwrap_or_else(|| "sankoff_params.yaml".into()); + let f = std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + }); + serde_yaml::to_writer(f, &report).unwrap_or_else(|e| { + eprintln!("error writing {path}: {e}"); + std::process::exit(1); + }); + info!("Sankoff calibration parameters → {path}"); +} + +// ── Sankoff cost matrix + alignment → ready-to-run TNT script ────────────── +// +// TNT's *default* xread reader only accepts its own 0-9A-F alphabet (see +// its manual: "up to 16 states are allowed by xread, using symbols 0-9 ... +// and A-F") — the wider IUPAC set `STATE_SYMBOL` uses is rejected as an +// "alien symbol" unless `nstates dna` is set, which imposes TNT's own fixed +// DNA encoding instead, incompatible with a custom smatrix. And TNT's +// `smatrix`/`cost` commands reject decimal costs ("found symbol . when +// reading transformation costs"). So: recode to TNT's alphabet and +// integer-scale the costs here, at this adapter's boundary, rather than +// degrading the project's own canonical (IUPAC, real-valued) output. + +const TNT_STATE_SYMBOL: [char; 16] = [ + '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', +]; + +fn write_sankoff_tnt( + matrix: &[[f64; 16]; 16], + alignment: &SnpAlignment, + labels: &[String], + output: &Option, + cost_scale: f64, +) { + let path = output.as_ref() + .map(|p| format!("{}_sankoff.tnt", p.display())) + .unwrap_or_else(|| "sankoff.tnt".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + + // IUPAC-ish symbol -> bitmask, to translate the alignment (which uses + // `STATE_SYMBOL`, `-` already normalised to `0` by `snp_pseudo_alignment` + // callers) into TNT's alphabet without re-deriving state indices. + let mut iupac_to_state = [0u8; 128]; + for (state, &sym) in STATE_SYMBOL.iter().enumerate() { + iupac_to_state[sym as usize] = state as u8; + } + + let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); + writeln!(f, "xread").unwrap(); + writeln!(f, "'obikmer central-position SNP families, calibrated Sankoff 16-state encoding'").unwrap(); + writeln!(f, "{n_sites} {}", labels.len()).unwrap(); + for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { + write!(f, "{label} ").unwrap(); + for &b in seq { + let b = if b == b'-' { b'0' } else { b }; + let state = iupac_to_state[b as usize]; + write!(f, "{}", TNT_STATE_SYMBOL[state as usize]).unwrap(); + } + writeln!(f).unwrap(); + } + writeln!(f, ";\n").unwrap(); + + let scaled_matrix = scaled_metric_matrix(matrix, cost_scale); + + writeln!(f, "smatrix =0 (family16)").unwrap(); + for i in 0..16 { + for j in (i + 1)..16 { + writeln!(f, "{}/{} {}", TNT_STATE_SYMBOL[i], TNT_STATE_SYMBOL[j], scaled_matrix[i][j]).unwrap(); + } + } + writeln!(f, ";\n").unwrap(); + + writeln!(f, "ccode ( 0.{} ;", n_sites - 1).unwrap(); + writeln!(f, "smatrix +0 0.{} ;", n_sites - 1).unwrap(); + writeln!(f).unwrap(); + + // Basename only (not the full `path`/`output` prefix): TNT's natural + // workflow is to `cd` into the output directory before `proc`-ing the + // script, and an absolute path here would break if that directory is + // later moved or copied elsewhere. + let tre_name = output.as_ref() + .and_then(|p| p.file_name()) + .map(|n| format!("{}_sankoff.tre", n.to_string_lossy())) + .unwrap_or_else(|| "sankoff.tre".into()); + + // TNT's plain command parser has no comment syntax of its own — `/* */` + // and `[ ]` are only recognised inside the (separately-enabled) macro + // scripting language, and fail with "No command!" here otherwise + // (verified against this file with the local TNT binary). `quote` is + // the closest working equivalent: it prints free text and does not + // otherwise affect parsing, so it doubles as an explanation of the + // defaults below when the script is run. `;` ends a `quote` block like + // any other TNT command, so the text itself must avoid semicolons. + writeln!(f, "quote").unwrap(); + writeln!(f, "Default search below (edit or delete this block to run your own strategy):").unwrap(); + writeln!(f, " hold N : size of TNT's tree buffer (how many equally-parsimonious").unwrap(); + writeln!(f, " trees it keeps in memory at once), 20 is a small, fast").unwrap(); + writeln!(f, " default, raise it if mult reports it had to drop trees.").unwrap(); + writeln!(f, " mult : traditional search (random addition sequences followed by").unwrap(); + writeln!(f, " TBR branch-swapping, TNT's own default replication count),").unwrap(); + writeln!(f, " a reasonable first-pass strategy on this data's memory").unwrap(); + writeln!(f, " footprint, xmult's ratchet/drift/tree-fusion buffers ran").unwrap(); + writeln!(f, " this out of RAM at TNT's default mxram on this dataset.").unwrap(); + writeln!(f, " export - F : write the trees held in the buffer to file F, in").unwrap(); + writeln!(f, " TNT/Hennig86 format ('-' means trees, as opposed to data).").unwrap(); + writeln!(f, ";").unwrap(); + writeln!(f, "hold 20;").unwrap(); + writeln!(f, "mult;").unwrap(); + writeln!(f, "export - {tre_name};").unwrap(); + + info!( + "TNT script → {path} (costs scaled x{cost_scale:.0}, runs a default `hold 20; mult;` \ + search and writes trees to {tre_name} in TNT's working directory — edit the trailing \ + comment block in the script to change this)\n\ + Run it with:\n \ + printf 'proc {path};\\nquit;\\n' | tnt\n\ + (or start `tnt` interactively and type `proc {path};`)" + ); +} + +/// Scale `matrix` by `cost_scale` and round to integers (TNT's smatrix/cost +/// commands reject decimals), then take the *metric 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 `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 +/// TNT-version-dependent. +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 { + for j in 0..16 { + m[i][j] = (matrix[i][j] * cost_scale).round() as i64; + } + } + for k in 0..16 { + for i in 0..16 { + for j in 0..16 { + let via = m[i][k] + m[k][j]; + if via < m[i][j] { + m[i][j] = via; + } + } + } + } + m +} + +// ── Sankoff cost matrix → PhyG custom-alphabet TCM + ready-to-run script ──── +// +// PhyG's `tcm:STRING` format needs no alphabet recoding, unlike `--tnt`: +// its parser reads the alphabet straight from the tcm file's own first +// line, so `--sankoff`'s own `_sankoff.fasta` (already IUPAC+`0`) is reused +// as-is via `prefasta:`. PhyG auto-adds its own indel/gap state as an +// (n+1)-th row/column of the tcm — inert here since the alignment already +// encodes absence as an ordinary state (`0`), never as `-` (see +// `write_sankoff_alignment_fasta`'s own comment on why, and the RAxML-era +// bug that motivated it). The gap row/column below reuses `matrix[i][0]`/ +// `matrix[0][j]` (cost to/from `∅`) as the closest principled value for a +// state that, in practice, is never actually triggered. + +fn write_sankoff_phyg(matrix: &[[f64; 16]; 16], output: &Option, cost_scale: f64) { + let scaled_matrix = scaled_metric_matrix(matrix, cost_scale); + + let basename = |suffix: &str| -> String { + output.as_ref() + .and_then(|p| p.file_name()) + .map(|n| format!("{}{suffix}", n.to_string_lossy())) + .unwrap_or_else(|| format!("sankoff{suffix}")) + }; + let full_path = |suffix: &str| -> String { + output.as_ref() + .map(|p| format!("{}{suffix}", p.display())) + .unwrap_or_else(|| format!("sankoff{suffix}")) + }; + + let tcm_path = full_path("_sankoff.tcm"); + let mut f = BufWriter::new(std::fs::File::create(&tcm_path).unwrap_or_else(|e| { + eprintln!("error creating {tcm_path}: {e}"); + std::process::exit(1); + })); + let alphabet_line = STATE_SYMBOL.iter().map(|c| c.to_string()).collect::>().join(" "); + writeln!(f, "{alphabet_line}").unwrap(); + for i in 0..16 { + let mut row: Vec = (0..16).map(|j| scaled_matrix[i][j]).collect(); + row.push(scaled_matrix[i][0]); // gap column: same cost as to/from ∅ + writeln!(f, "{}", row.iter().map(|v| v.to_string()).collect::>().join(" ")).unwrap(); + } + let mut gap_row: Vec = (0..16).map(|j| scaled_matrix[0][j]).collect(); + gap_row.push(0); + writeln!(f, "{}", gap_row.iter().map(|v| v.to_string()).collect::>().join(" ")).unwrap(); + info!("PhyG TCM → {tcm_path}"); + + let pg_path = full_path("_sankoff.pg"); + let mut f = BufWriter::new(std::fs::File::create(&pg_path).unwrap_or_else(|e| { + eprintln!("error creating {pg_path}: {e}"); + std::process::exit(1); + })); + let fasta_name = basename("_sankoff.fasta"); + let tcm_name = basename("_sankoff.tcm"); + let tre_name = basename("_sankoff.tre"); + writeln!(f, "read(prefasta:\"{fasta_name}\", tcm:\"{tcm_name}\")").unwrap(); + writeln!(f, "search(seconds:300, instances:4)").unwrap(); + writeln!(f, "report(\"{tre_name}\", graphs, newick, overwrite)").unwrap(); + + let pg_dir = std::path::Path::new(&pg_path).parent() + .filter(|d| !d.as_os_str().is_empty()) + .map(|d| d.display().to_string()) + .unwrap_or_else(|| ".".into()); + let pg_name = std::path::Path::new(&pg_path).file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_else(|| pg_path.clone()); + info!( + "PhyG script → {pg_path} (costs scaled x{cost_scale:.0}, runs a default 300s/4-instance \ + search and writes trees to {tre_name})\n\ + Run it with:\n \ + cd {pg_dir} && phyg {pg_name}\n\ + (`phyg` must run from that directory — `read()`/`report()` in the script use relative \ + file names)" + ); +} + // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── fn upgma_to_newick(dendro: &kodama::Dendrogram, names: &[String]) -> String { diff --git a/src/obikmer/src/cmd/filter.rs b/src/obikmer/src/cmd/filter.rs index ae3f45fa..ed0f407b 100644 --- a/src/obikmer/src/cmd/filter.rs +++ b/src/obikmer/src/cmd/filter.rs @@ -78,6 +78,13 @@ pub fn run(args: FilterCmdArgs) { filters.push(Box::new(MinComplexity { level_max: args.complexity_level_max, theta })); } + // Source is opened read-only above and needs no lock; only the + // destination is written. + let _lock = obisys::DirLock::acquire(&args.output).unwrap_or_else(|e| { + eprintln!("error locking output directory {}: {e}", args.output.display()); + std::process::exit(1); + }); + let mut rep = Reporter::new(); KmerIndex::rebuild(&args.output, &src, &filters, mode, args.force, &mut rep) .unwrap_or_else(|e| { diff --git a/src/obikmer/src/cmd/index.rs b/src/obikmer/src/cmd/index.rs index 31102fbd..e83f7910 100644 --- a/src/obikmer/src/cmd/index.rs +++ b/src/obikmer/src/cmd/index.rs @@ -156,6 +156,16 @@ pub fn run(args: IndexArgs) { let output = args.output.clone(); let mut rep = Reporter::new(); + // Locked for the whole build (including a possible --force removal + + // recreation below): a second `index` run resuming/overwriting the same + // output directory concurrently would otherwise corrupt it. Unlinking + // the lock file via --force's remove_dir_all is safe — the held file + // descriptor keeps the lock regardless of the directory entry. + let _lock = obisys::DirLock::acquire(&output).unwrap_or_else(|e| { + eprintln!("error locking output directory {}: {e}", output.display()); + std::process::exit(1); + }); + // ── Resolve evidence kind ──────────────────────────────────────────────── let (evidence, effective_kmer_size) = if args.approx { let (z, b, fp) = resolve_approx_params(args.findere_z, args.evidence_bits, args.fp); diff --git a/src/obikmer/src/cmd/merge.rs b/src/obikmer/src/cmd/merge.rs index f0ce2dc0..c5fbef8c 100644 --- a/src/obikmer/src/cmd/merge.rs +++ b/src/obikmer/src/cmd/merge.rs @@ -64,6 +64,13 @@ pub fn run(args: MergeArgs) { sources.len(), n_genomes, args.output.display() ); + // Only the destination is written; sources are opened read-only above + // and need no lock. + let _lock = obisys::DirLock::acquire(&args.output).unwrap_or_else(|e| { + eprintln!("error locking output directory {}: {e}", args.output.display()); + std::process::exit(1); + }); + let mut rep = Reporter::new(); KmerIndex::merge(&args.output, &source_refs, mode, args.force, args.rename_duplicates, args.budget_fraction, &mut rep).unwrap_or_else(|e| { eprintln!("error merging: {e}"); diff --git a/src/obikmer/src/cmd/pack.rs b/src/obikmer/src/cmd/pack.rs index f3c0d3da..4f34c52d 100644 --- a/src/obikmer/src/cmd/pack.rs +++ b/src/obikmer/src/cmd/pack.rs @@ -12,6 +12,13 @@ pub struct PackArgs { } pub fn run(args: PackArgs) { + // Modifies the index in place; acquired before opening so a concurrent + // writer can't slip in between the open and the pack below. + let _lock = obisys::DirLock::acquire(&args.index).unwrap_or_else(|e| { + eprintln!("error locking index directory {}: {e}", args.index.display()); + std::process::exit(1); + }); + let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| { eprintln!("error opening index: {e}"); std::process::exit(1); diff --git a/src/obikmer/src/cmd/reindex.rs b/src/obikmer/src/cmd/reindex.rs index 5c0bab76..3fd2c650 100644 --- a/src/obikmer/src/cmd/reindex.rs +++ b/src/obikmer/src/cmd/reindex.rs @@ -47,6 +47,13 @@ pub fn run(args: ReindexArgs) { IndexMode::Exact }; + // Modifies the index in place; acquired before opening so a concurrent + // writer can't slip in between the open and the reindex below. + let _lock = obisys::DirLock::acquire(&args.index).unwrap_or_else(|e| { + eprintln!("error locking index directory {}: {e}", args.index.display()); + std::process::exit(1); + }); + let mut idx = KmerIndex::open(&args.index).unwrap_or_else(|e| { eprintln!("error opening index: {e}"); std::process::exit(1); diff --git a/src/obikmer/src/cmd/select.rs b/src/obikmer/src/cmd/select.rs index 35719e8e..e74761b2 100644 --- a/src/obikmer/src/cmd/select.rs +++ b/src/obikmer/src/cmd/select.rs @@ -215,6 +215,16 @@ pub fn run(args: SelectArgs) { std::process::exit(1); } + // Lock whichever directory actually gets written: the source itself in + // --in-place mode, otherwise the (distinct) --output directory. Acquired + // before opening the source so a concurrent writer can't slip in between + // the open and the write below. + let lock_target = if args.in_place { &args.source } else { args.output.as_ref().unwrap() }; + let _lock = obisys::DirLock::acquire(lock_target).unwrap_or_else(|e| { + eprintln!("error locking {}: {e}", lock_target.display()); + std::process::exit(1); + }); + let mut src = KmerIndex::open(&args.source).unwrap_or_else(|e| { eprintln!("error opening source index: {e}"); std::process::exit(1); diff --git a/src/obisys/Cargo.toml b/src/obisys/Cargo.toml index da677db1..966e22ed 100644 --- a/src/obisys/Cargo.toml +++ b/src/obisys/Cargo.toml @@ -8,3 +8,4 @@ libc = "0.2" sysinfo = "0.33" indicatif = "0.17" tracing = "0.1" +fs4 = "0.9" diff --git a/src/obisys/src/lib.rs b/src/obisys/src/lib.rs index 21d4987b..65d5d4f4 100644 --- a/src/obisys/src/lib.rs +++ b/src/obisys/src/lib.rs @@ -8,6 +8,49 @@ use tracing::{debug, info, warn}; const BRAILLE: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]; +// ── DirLock ────────────────────────────────────────────────────────────────── + +/// Exclusive advisory lock on an index directory, held for the duration of +/// any command that writes into an already-existing index (building the +/// sibling annex, merging into a destination, filtering/selecting in place, +/// ...). Two such commands racing on the same directory can otherwise +/// corrupt each other's writes with no error from either side. +/// +/// Only the directory actually being *written to* needs a lock — a command +/// like `merge` that reads several source indexes to build one destination +/// only needs to lock the destination. +/// +/// Uses the OS's advisory file lock (`flock` on Unix, `LockFileEx` on +/// Windows) via `fs4`, not a hand-rolled PID file: the OS releases it +/// automatically on process exit, including a crash — no stale-lock cleanup +/// logic needed. +pub struct DirLock { + _file: std::fs::File, +} + +impl DirLock { + /// Block until the exclusive lock on `dir` is acquired (creating `dir` + /// and the lock file within it if needed). Logs once if the wait is + /// non-trivial, so a blocked command doesn't look silently hung. + pub fn acquire(dir: &std::path::Path) -> std::io::Result { + use fs4::fs_std::FileExt; + + std::fs::create_dir_all(dir)?; + let lock_path = dir.join(".obikmer.lock"); + let file = std::fs::OpenOptions::new() + .create(true) + .truncate(false) + .write(true) + .open(&lock_path)?; + + if file.try_lock_exclusive().is_err() { + info!(dir = %dir.display(), "waiting for another obikmer process to release this index"); + file.lock_exclusive()?; + } + Ok(Self { _file: file }) + } +} + // ── TracedBar ────────────────────────────────────────────────────────────────── /// Wrapper around `ProgressBar` that emits `tracing` events when stderr is not -- 2.54.0 From adf5b52dc7ebe5241e0234ef7d3f38b4d5105807 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 12 Aug 2026 18:23:22 +0200 Subject: [PATCH 02/46] feat(distance): implement native Sankoff calibration and backends Replaces external Python glue with native Rust modules for Sankoff model calibration, exporting calibrated cost matrices, FASTA alignments, and YAML parameters. Adds dedicated writers for TNT and PhyG that apply integer scaling and Floyd-Warshall metric closure to enforce triangle inequality. Integrates these exporters into the distance command pipeline to streamline downstream tree inference workflows, while updating theory documentation to reflect IQ-TREE integration and state renumbering improvements. --- docmd/theory/evolutionary_distances.md | 279 +++++++++++++ .../src/cmd/{distance.rs => distance/mod.rs} | 391 +----------------- src/obikmer/src/cmd/distance/phyg.rs | 81 ++++ src/obikmer/src/cmd/distance/sankoff.rs | 202 +++++++++ src/obikmer/src/cmd/distance/tnt.rs | 119 ++++++ 5 files changed, 690 insertions(+), 382 deletions(-) rename src/obikmer/src/cmd/{distance.rs => distance/mod.rs} (55%) create mode 100644 src/obikmer/src/cmd/distance/phyg.rs create mode 100644 src/obikmer/src/cmd/distance/sankoff.rs create mode 100644 src/obikmer/src/cmd/distance/tnt.rs diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index b1fd239f..6fd21d8b 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -620,6 +620,285 @@ simply does not survive domain-level divergence. Cross-domain placement would need conserved-marker characters (rRNA, ribosomal proteins), not this estimator. +### Native `--sankoff --tnt`/`--phyg` export (2026-08-12), superseding the external scripts above + +The ad hoc Python glue from the previous section is superseded: `obikmer +distance --sankoff` now calibrates the matrix natively (`p_hat`, 6-category +substitution costs, `c_ctx` weighted by `mean_sub_cost` — see the worked +example above) and `--tnt`/`--phyg` each write a ready-to-run script from +it, no external script needed. `∅` is an ordinary 16th state throughout +(never `-`), specifically to avoid gap-semantics confusion in downstream +tools — see "A concrete Sankoff cost matrix" above for why. + +**TNT (`--tnt`).** `write_sankoff_tnt` (`obikmer/src/cmd/distance.rs`) +recodes to TNT's own `0-9A-F` xread alphabet (its default reader rejects +the wider IUPAC set otherwise), scales and rounds costs to integers +(`smatrix`/`cost` reject decimals), then re-runs integer Floyd-Warshall on +the rounded matrix (`scaled_metric_matrix`) — independently rounding each +cell of an already-metric real-valued matrix can break the triangle +inequality (e.g. two real costs of `1.734` round to `173` each, summing to +`346`, while their real sum `3.468` rounds to `347`), which TNT otherwise +silently "fixes" itself with an unreproducible correction. Verified against +the real 20-genome benchmark index: zero triangle-inequality violations +after the fix, TNT loads the file without its "triangle inequality +violated... Fixed" warning. + +Two syntax facts worth recording because they're wrong in intuitive +guesses and contradicted actual TNT behavior when tested: TNT's plain +command stream has **no comment syntax** of its own — `/* */` and `[ ]` +only work inside the (separately-enabled, off by default) macro scripting +language, and error with "No command!" otherwise. The working substitute +is `quote TEXT ;` (prints the text, doesn't affect parsing) — but the text +itself can't contain a literal `;` (TNT's universal terminator); the +manual's own escape (`.,`) exists but the script here just avoids +semicolons in the text instead. + +The default search command embedded in the script is `mult` (traditional: +random addition sequences + TBR), not `xmult` (New Technology search: +ratchet/drift/tree-fusion). `xmult` with TNT's default `mxram` (16 MB, +must be set *before* `xread` if changed) ran out of RAM on the real +908k-character dataset ("`xmult - out of ram`"); `mult` does not, matching +what had already been validated by hand outside this session. + +**PhyG (`--phyg`).** `write_sankoff_phyg` writes a `tcm:` custom-alphabet +matrix (same scale+round+metric-closure treatment as TNT) and reuses +`--sankoff`'s own `_sankoff.fasta` as-is via `prefasta:` — PhyG's `tcm:` +alphabet is read from the matrix file's own first line, so (unlike TNT) no +recoding is needed. PhyG auto-adds its own indel/gap state as an +`(n+1)`-th row/column of the tcm; inert here since the alignment encodes +absence as `0`, never `-`. + +`report("file", newick, overwrite)` — exactly as shown in PhyG's own +manual — triggers `Unrecognized/missing report option ... defaulting to +'graphs'` on the locally installed binary (1.3, commit `3c1a1fa`); the +working form adds `graphs` explicitly: `report("file", graphs, newick, +overwrite)`. Manual/binary mismatches like this (also true of `criterion:` +— the binary accepts `parsimony`/`ml`/`pmdl`, the manual instead documents +`mapa`/`ncm`/`parsimony`/`pmdl`/`si`) mean command syntax against this PhyG +build should be verified empirically, not trusted from the PDF alone. +`instances:N` (not a separate CPU flag) is what actually parallelises the +search across cores — PhyG uses all physical cores by default but only +across as many instances as are running, so raise it to the physical core +count to use them all (the CLI-level `+RTS -NX -RTS` flag also exists but +controls something else: capping/limiting cores, not requesting more). + +Both scripts share one `--sankoff-cost-scale` (default `100`), not two +separate flags — they scale the same calibrated matrix for the same +reason (integer-only cost commands) and no PhyG-specific +accumulator-width constraint was ever found to justify a different +default from TNT's (TNT: hinted 32-bit accumulators in its own manual; +PhyG: no such hint found — Haskell's native `Int` is typically 64-bit). + +**Open problem: PhyG reports all branch lengths as `0.0`.** The graph-level +parsimony cost is correct (`3.3286×10⁸` on the real dataset, consistent +with TNT's `328574911` on the same calibrated matrix), but every individual +edge in the exported Newick shows `:0.0`, with the total cost only ever +shown as a whole-tree annotation (`[3.32860377e8]`). Not fixed, not fully +diagnosed — PhyG's manual describes per-edge branch length as computed by +ancestral-state (HTU) backtracking, well documented for sequence/standard +character types, but nothing found (the term "Sankoff" doesn't even appear +in the manual) confirming this backtracking is wired up for a custom +`tcm:` matrix character. Switching `criterion:` to a likelihood-family +option (`ml`, or the manual's `mapa`/`ncm`/`si`) was considered as a +possible fix but is very unlikely to be one: those criteria are +information-theoretic reparametrisations of the *same* step-counting +machinery as parsimony (`ncm` in particular is known in the literature to +be numerically equivalent to weighted parsimony), not classical +continuous-time-Markov ML with a real rate matrix — so they wouldn't +change how branch length is attributed per edge either. + +**Export format note (not a bug in the generator).** Neither script's +`.tre` output opens in PearTree (FigTree's successor) via File > Open — +association/Launch-Services quirks were ruled out (the file was opened +directly through the app, not by double-click). Likely cause, not yet +confirmed: TNT's export is a minimal NEXUS `begin trees;` block with no +preceding `Taxa` block and bare numeric (untranslated) leaf labels; PhyG's +is multiple raw Newick trees concatenated with no NEXUS wrapper at all +plus a trailing `[cost]` bracket tag after the root label. Both differ +from a "normal" single-tree, fully-declared NEXUS file; this is PhyG/TNT's +own export format, not something `write_sankoff_tnt`/`write_sankoff_phyg` +could fix without post-processing the *other* program's output after the +fact. + +### Next direction: genuine ML branch lengths, not parsimony (open, 2026-08-12) + +Decided: parsimony (the whole `--sankoff`/`--tnt`/`--phyg` pipeline above) +is a stopgap, not the destination. The goal is maximum likelihood with +real, calibrated branch lengths (expected substitutions/site), which +parsimony step-counts were never going to give directly (see the open +"branch lengths are `0.0`" problem above — even if fixed, TNT/PhyG-style +parsimony branch length is a step count, not a continuous ML estimate). + +**Model choices, settled:** +- **The exchangeability `R` is symmetric; the rate matrix `Q` is not.** + (Superseded an earlier, wrong framing here that treated "symmetric + model" as one thing — see the resolution below on `R` vs `Q` vs `π` for + the full reasoning.) `R(a,b) = R(b,a)` because `BasePairTally` never + captured direction — a fact about the data, not a modelling choice. + `Q(a,b) = R(a,b)·π_b` is asymmetric whenever the real state frequencies + `π` are (which they are, empirically) — biology drives this via `π`, + not via `R`. +- **`∅` stays an ordinary 16th state**, as already established for + TNT/PhyG — same reasoning applies to any ML tool: encode as a real + alphabet symbol, never as `-`/gap, or the RAxML-era failure (empty set + silently treated as missing data) repeats. + +**Stationary frequencies for the 16 states — resolved (2026-08-12).** +A CTMC needs a rate matrix `Q`, generally asymmetric. `Q(i,j) = R(i,j) · +π_j`, where `R` (exchangeability) is symmetric and `π` (stationary +frequencies) need not be — this factoring is what makes `Q` reversible +(satisfies detailed balance, `π_i·Q(i,j) = π_j·Q(j,i)`) for *any* `π`, not +just uniform, as long as `R` is symmetric. Two separate, both-easy +quantities, not one hard inverse problem: +- **`R` is already calibrated**: `sub_cost` (`-ln(observed rate)` per pair, + from `BasePairTally`) *is* `R` up to a log transform — recover it as + `R(a,b) = exp(-sub_cost(a,b))`. Symmetric by construction, because the + tally itself never distinguished direction (unordered-pair counts only) + — not a modelling choice, a fact about what the data can say. +- **`π` is a direct count**: empirical marginal frequency of each of the + 16 states across the whole alignment (same kind of scan already used to + confirm `N` occurs 1383 times in the real 20-genome benchmark). With + ~908k sites × 20 genomes, the counts are large enough that this is + precise on its own — no need to spend ML degrees of freedom + re-estimating it via IQ-TREE's `+FO`. + +Checked and ruled out along the way: IQ-TREE's `+F` (empirical, "compute +from the alignment") does **not** work as a shortcut for this — for a +custom-file morphology model, `readParameters` always requires the file's +own frequency line unconditionally; omitting it and passing `+F` instead +just fails (`ERROR: State frequencies could not be read`). `π` has to be +computed by `obikmer` and written into the file, not left to IQ-TREE. + +Net effect: no free-rate ML estimation needed for this piece at all (the +mistaken assumption that motivated most of this discussion — that +building an asymmetric `Q` a priori would require solving a linear system +from `Q` itself — doesn't apply, because `R`, the only piece that's +genuinely hard to get directionally, is symmetric and already in hand). + +**Candidate tool: IQ-TREE**, because it supports user-defined multistate +models (unlike RAxML's `MULTI` data type, which is limited to the +equal-rate Mk model and can't take a custom rate matrix at all — a genuine +tool limitation, not a gap-symbol encoding problem this time). IQ-TREE 3 +(3.0.1) is now installed locally (Homebrew, `iqtree3`). + +**IQ-TREE custom-model format — verified empirically against the local +binary (2026-08-12).** The web docs' `-mdef` NEXUS `begin models; +frequency NAME = ...; model NAME = ...; end;` mechanism (initially assumed +to apply directly, see history below) turned out to be for **named +components used inside `MIX{...}`/`FMIX{...}` mixture models only** — it +does **not** apply to a single, non-mixture custom morphology matrix, and +using it that way fails (`ERROR: File not found ` — traced in +IQ-TREE 3's own source, `model/modelmorphology.cpp`: any `-m` string that +isn't `MK`/`ORDERED`/`GTR`/`GTRX` is passed straight to +`ModelMarkov::readParameters()`, which opens it **as a literal file path**, +never consulting the `-mdef` models block at all for this data type). + +**The confirmed working recipe** (built a tiny 5-taxon/3-state toy dataset +and rate file, ran it end to end with `iqtree3`, got a real ML tree with +non-zero branch lengths and an optimized log-likelihood — ground truth, +not documentation): +- No `-mdef` needed. Write one plain file (any name) containing, as + whitespace/newline-separated numbers, in order: the **lower-triangular + rate matrix** (`N(N-1)/2` values, PAML row-major order — for 16 states, + 120 values, the same count and layout already produced for TNT's + `smatrix`), immediately followed by the **N state frequencies** on the + same stream (no header, no blank line required — confirmed by reading + `ModelMorphology::readRates`/`ModelMarkov::readStateFreq` directly, which + just pull tokens off the stream in sequence). +- Invoke with `-m +ASC` (`+ASC` for the no-constant-site + correction, as before). An explicit `+F{f1,...,fN}` on the command line + overrides the file's own frequency line if given (confirmed in + `ModelMorphology::init`) — useful once real calibrated stationary + frequencies exist, a placeholder equal-frequency line works meanwhile + (the still-open gap noted above). +- `--seqtype MORPH` (alphabet `0`-`9`,`A`-`V`, ≤31 states) — reuse the same + `0-9A-F` recoding already built for TNT (`TNT_STATE_SYMBOL`). + +**Risk, confirmed, precisely characterised, and resolved by design +(2026-08-12).** +`--seqtype MORPH{16}` **does not force the state count** for real ML +analysis — tested directly (`--seqtype MORPH{4}` on the 3-symbol toy +alignment gave the byte-for-byte identical 3-state result as no `{4}` at +all) and confirmed in source: the value it sets +(`params.alisim_num_states_morph`, `utils/tools.cpp`) is consumed only by +the `--alisim` simulator; the main analysis path always calls +`getDataBlockMorphStates`/an equivalent scan (`alignment/alignment.cpp`), +for both FASTA/PHYLIP and NEXUS input (a NEXUS `symbols=` declaration +doesn't change this either — checked, same code path). No CLI flag or +NEXUS declaration overrides it. + +The precise rule (from `getDataBlockMorphStates`, `alignment.cpp:1058`): +`N` = **one plus the highest state ordinal actually observed anywhere in +the alignment**, ordinal being the symbol's position in IQ-TREE's own +fixed table `"0123456789ABCDEFGHIJKLMNOPQRSTUV"` — not a count of distinct +symbols seen. So the risk is narrower than "any missing symbol breaks it": +concretely, it's whether the symbol mapped to state index 15 (`F` in the +`0-9A-F` recoding already used for TNT, i.e. our `N` = "all four bases +ambiguous") occurs **at least once anywhere** in the real alignment — if +it does, `N` correctly comes out to 16 regardless of which lower-index +symbols (including `0`/`∅`) are rare or absent; if it doesn't, `N` silently +undercounts and misaligns every value in a 16-entry rate/frequency file, +with no error to catch it. **Checked against the real biological alignment** +(`/tmp/msg_test/eub_sankoff.fasta`, the 20-genome benchmark index): the +symbol `N` (IUPAC "all four bases ambiguous," state index 15) occurs 1383 +times across 13 of the 20 sequences — present, so this specific real +dataset is not at risk. Still worth a real, general presence check inside +`obikmer` before this is wired in, rather than assuming every future +dataset will have `N` too (nothing in IQ-TREE would catch it if not). + +**Resolution: subset + compact-renumber, not rely on all 16 appearing.** +Since IQ-TREE always infers `N` from the alignment's own content and +nothing overrides that, the fix is to make the file `obikmer` writes match +that inference *by construction*, for every run, rather than hope the +16th (or any particular) state happens to occur: +1. Scan the real alignment for which of the 16 canonical states actually + occur anywhere (not per-column — anywhere in the whole alignment). +2. Renumber the occurring states to a **compact, consecutive** `0..k-1` + range, preserving their relative order (the original bitmask/`STATE_ + SYMBOL` ordering) — not just filtering, since a *gap* in the ordinal + sequence (e.g. keeping states `{0,1,2,4}` numbered as-is instead of + `{0,1,2,3}`) reproduces the exact same "highest observed ordinal" + miscount this was meant to fix. +3. Recode the alignment itself with this new compact `k`-symbol alphabet + (same recoding mechanism already used for TNT's `0-9A-F`, just over a + possibly-smaller symbol set). +4. Extract the matching `k×k` submatrix (rows/columns for the kept + states only) from the full calibrated 16×16 cost matrix, in the same + lower-triangular order the rate-matrix file needs — and, later, the + matching `k`-length subset of stationary frequencies once those are + calibrated (still the open gap noted earlier in this section). + +Consequence, and why nothing is lost: a state that never occurs in a given +alignment can, by definition, never contribute a transition to score in +that same alignment — dropping it from that run's matrix costs nothing. +The subset (and therefore `k`) can differ from one dataset/run to the +next; this has to be done freshly per alignment, not computed once and +reused. + +*(Superseded reasoning, kept for the record: the `-mdef` NEXUS route +below was the original plan, based on IQ-TREE's own web documentation for +protein mixture models, before the local install allowed testing it — +`GTRX` combined with a `-mdef`-referenced custom model, `+Fname` frequency +reference. Both pieces exist and parse without error individually, but +`GTRX`/`GTR` are IQ-TREE's own fixed built-in equal-structure multistate +model, not a hook for an arbitrary custom matrix; a custom matrix is a +file path in `-m` directly, no `-mdef` or `GTRX` involved.)* + +Source for the empirical findings above: `model/modelmorphology.cpp` and +`model/modelmarkov.cpp` in the local `iqtree/iqtree3` source (cloned to +inspect the exact parsing logic after documentation didn't resolve the +`+Fname` reference error) — more reliable here than the PDF/web manual, +which (like TNT/PhyG) doesn't always match this specific binary. Original +(partially superseded) sources: [Substitution +Models](https://iqtree.github.io/doc/Substitution-Models), [Complex +Models](https://iqtree.github.io/doc/Complex-Models). + +**Relation to the existing calibration.** `sub_cost[a][b] = -ln(rate)` +(see "A concrete Sankoff cost matrix" above) is already a log-rate — a +genuine CTMC rate matrix `Q` could plausibly be recovered as +`rate(a,b) = exp(-cost(a,b))`, renormalised so each row sums to zero, once +the stationary-frequency gap above is closed. Not yet attempted. + ## Heterozygosity, ploidy, and consensus-assembly inputs A within-genome multiplicity signal (more than one of the 4 central forms diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance/mod.rs similarity index 55% rename from src/obikmer/src/cmd/distance.rs rename to src/obikmer/src/cmd/distance/mod.rs index bd7a4e37..bd698caa 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance/mod.rs @@ -1,3 +1,7 @@ +mod phyg; +mod sankoff; +mod tnt; + use std::io::{self, BufWriter, Write}; use std::path::PathBuf; @@ -5,13 +9,17 @@ use clap::Args; use kodama::{Method, linkage}; use obifastwrite::{JsonVal, write_record}; use obikindex::{ - BasePairTally, DistanceMetric, KmerIndex, PHatEstimate, RawSnpDistanceOutput, SankoffWeights, + DistanceMetric, KmerIndex, RawSnpDistanceOutput, SankoffWeights, SiblingAnnexStats, SnpAlignment, build_cost_matrix, c_ctx_from_p_hat, calibrate_p_hat, mean_substitution_cost, substitution_costs_from_tally, }; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; +use phyg::write_sankoff_phyg; +use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params}; +use tnt::write_sankoff_tnt; + #[derive(clap::ValueEnum, Clone, Copy, Debug)] pub enum MetricArg { Jaccard, @@ -468,387 +476,6 @@ fn write_snp_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option< if n_sites == 1 { "" } else { "s" }); } -// ── Sankoff pseudo-alignment → FASTA ──────────────────────────────────────── -// -// Same data as `--snp`'s pseudo-alignment (`SnpAlignment`/ -// `snp_pseudo_alignment`), re-coded so its symbols match the accompanying -// `--sankoff-matrix` output exactly: `0` for the empty/absent state instead -// of `-`, which TNT/PhyG would otherwise read as their own gap character -// rather than our "family absent" state. - -fn write_sankoff_alignment_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option) { - let path = output.as_ref() - .map(|p| format!("{}_sankoff.fasta", p.display())) - .unwrap_or_else(|| "sankoff.fasta".into()); - let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { - eprintln!("error creating {path}: {e}"); - std::process::exit(1); - })); - let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); - for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { - let recoded: Vec = seq.iter().map(|&b| if b == b'-' { b'0' } else { b }).collect(); - write_record(&recoded, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| { - eprintln!("error writing {path}: {e}"); - std::process::exit(1); - }); - } - info!("Sankoff pseudo-alignment → {path} ({n_sites} site{})", - if n_sites == 1 { "" } else { "s" }); -} - -// ── Sankoff cost matrix → CSV ──────────────────────────────────────────────── -// -// 16 states indexed by bitmask (bit 0=A, 1=C, 2=G, 3=T; state 0 is `∅`), -// matching the convention already used for `--snp`'s IUPAC-coded output and -// for the external TNT/PhyG scripts this feeds. Calibration report (p_hat, -// its variance, how many pairs/loci went into it, the resulting c_ctx) goes -// to the log, not the CSV, since it's a run-level fact, not per-cell data. - -// IUPAC ambiguity code per state (same mapping as `siblings::iupac_code`, -// already used for `--snp`'s pseudo-alignment — a biologist reads "R" as -// "A or G" without needing this file's convention explained), with `0` -// standing in for the empty state (`-` would collide with TNT/PhyG's own -// gap/range syntax). Bit order: 0=A, 1=C, 2=G, 3=T. This project's -// canonical alphabet — `write_sankoff_tnt` below recodes it to TNT's own -// default alphabet at the adapter boundary, rather than using it here. -const STATE_SYMBOL: [char; 16] = [ - '0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N', -]; - -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()); - let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { - eprintln!("error creating {path}: {e}"); - std::process::exit(1); - })); - write!(f, "state").unwrap(); - for sym in STATE_SYMBOL { write!(f, ",{sym}").unwrap(); } - writeln!(f).unwrap(); - for (s, row) in matrix.iter().enumerate() { - write!(f, "{}", STATE_SYMBOL[s]).unwrap(); - for cost in row { write!(f, ",{cost:.4}").unwrap(); } - writeln!(f).unwrap(); - } - info!("Sankoff cost matrix → {path}"); -} - -// ── Sankoff calibration parameters → YAML report ──────────────────────────── -// -// Everything `--sankoff` estimates from real data, in one durable, -// machine-readable file: `p_hat` and its variance (with how many pairs/loci -// went into it), the derived `c_ctx`, and the base-pair substitution tally -// (raw counts, not just the derived costs) — kept for the same reason raw -// counts are kept anywhere else in this project: costs are a modelling -// choice built *from* the counts, and reproducing/re-deriving them later -// needs the counts, not just their current derived value. Structured (YAML, -// not an ad hoc key=value text file) so R/Python/etc. can load it directly -// rather than re-parsing free text. - -#[derive(serde::Serialize)] -struct SankoffSubstitution { - pair: String, - count: u64, - cost: 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, -} - -fn write_sankoff_params( - estimate: &PHatEstimate, - tally: &BasePairTally, - weights: &SankoffWeights, - 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 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 path = output.as_ref() - .map(|p| format!("{}_sankoff_params.yaml", p.display())) - .unwrap_or_else(|| "sankoff_params.yaml".into()); - let f = std::fs::File::create(&path).unwrap_or_else(|e| { - eprintln!("error creating {path}: {e}"); - std::process::exit(1); - }); - serde_yaml::to_writer(f, &report).unwrap_or_else(|e| { - eprintln!("error writing {path}: {e}"); - std::process::exit(1); - }); - info!("Sankoff calibration parameters → {path}"); -} - -// ── Sankoff cost matrix + alignment → ready-to-run TNT script ────────────── -// -// TNT's *default* xread reader only accepts its own 0-9A-F alphabet (see -// its manual: "up to 16 states are allowed by xread, using symbols 0-9 ... -// and A-F") — the wider IUPAC set `STATE_SYMBOL` uses is rejected as an -// "alien symbol" unless `nstates dna` is set, which imposes TNT's own fixed -// DNA encoding instead, incompatible with a custom smatrix. And TNT's -// `smatrix`/`cost` commands reject decimal costs ("found symbol . when -// reading transformation costs"). So: recode to TNT's alphabet and -// integer-scale the costs here, at this adapter's boundary, rather than -// degrading the project's own canonical (IUPAC, real-valued) output. - -const TNT_STATE_SYMBOL: [char; 16] = [ - '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', -]; - -fn write_sankoff_tnt( - matrix: &[[f64; 16]; 16], - alignment: &SnpAlignment, - labels: &[String], - output: &Option, - cost_scale: f64, -) { - let path = output.as_ref() - .map(|p| format!("{}_sankoff.tnt", p.display())) - .unwrap_or_else(|| "sankoff.tnt".into()); - let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { - eprintln!("error creating {path}: {e}"); - std::process::exit(1); - })); - - // IUPAC-ish symbol -> bitmask, to translate the alignment (which uses - // `STATE_SYMBOL`, `-` already normalised to `0` by `snp_pseudo_alignment` - // callers) into TNT's alphabet without re-deriving state indices. - let mut iupac_to_state = [0u8; 128]; - for (state, &sym) in STATE_SYMBOL.iter().enumerate() { - iupac_to_state[sym as usize] = state as u8; - } - - let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); - writeln!(f, "xread").unwrap(); - writeln!(f, "'obikmer central-position SNP families, calibrated Sankoff 16-state encoding'").unwrap(); - writeln!(f, "{n_sites} {}", labels.len()).unwrap(); - for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { - write!(f, "{label} ").unwrap(); - for &b in seq { - let b = if b == b'-' { b'0' } else { b }; - let state = iupac_to_state[b as usize]; - write!(f, "{}", TNT_STATE_SYMBOL[state as usize]).unwrap(); - } - writeln!(f).unwrap(); - } - writeln!(f, ";\n").unwrap(); - - let scaled_matrix = scaled_metric_matrix(matrix, cost_scale); - - writeln!(f, "smatrix =0 (family16)").unwrap(); - for i in 0..16 { - for j in (i + 1)..16 { - writeln!(f, "{}/{} {}", TNT_STATE_SYMBOL[i], TNT_STATE_SYMBOL[j], scaled_matrix[i][j]).unwrap(); - } - } - writeln!(f, ";\n").unwrap(); - - writeln!(f, "ccode ( 0.{} ;", n_sites - 1).unwrap(); - writeln!(f, "smatrix +0 0.{} ;", n_sites - 1).unwrap(); - writeln!(f).unwrap(); - - // Basename only (not the full `path`/`output` prefix): TNT's natural - // workflow is to `cd` into the output directory before `proc`-ing the - // script, and an absolute path here would break if that directory is - // later moved or copied elsewhere. - let tre_name = output.as_ref() - .and_then(|p| p.file_name()) - .map(|n| format!("{}_sankoff.tre", n.to_string_lossy())) - .unwrap_or_else(|| "sankoff.tre".into()); - - // TNT's plain command parser has no comment syntax of its own — `/* */` - // and `[ ]` are only recognised inside the (separately-enabled) macro - // scripting language, and fail with "No command!" here otherwise - // (verified against this file with the local TNT binary). `quote` is - // the closest working equivalent: it prints free text and does not - // otherwise affect parsing, so it doubles as an explanation of the - // defaults below when the script is run. `;` ends a `quote` block like - // any other TNT command, so the text itself must avoid semicolons. - writeln!(f, "quote").unwrap(); - writeln!(f, "Default search below (edit or delete this block to run your own strategy):").unwrap(); - writeln!(f, " hold N : size of TNT's tree buffer (how many equally-parsimonious").unwrap(); - writeln!(f, " trees it keeps in memory at once), 20 is a small, fast").unwrap(); - writeln!(f, " default, raise it if mult reports it had to drop trees.").unwrap(); - writeln!(f, " mult : traditional search (random addition sequences followed by").unwrap(); - writeln!(f, " TBR branch-swapping, TNT's own default replication count),").unwrap(); - writeln!(f, " a reasonable first-pass strategy on this data's memory").unwrap(); - writeln!(f, " footprint, xmult's ratchet/drift/tree-fusion buffers ran").unwrap(); - writeln!(f, " this out of RAM at TNT's default mxram on this dataset.").unwrap(); - writeln!(f, " export - F : write the trees held in the buffer to file F, in").unwrap(); - writeln!(f, " TNT/Hennig86 format ('-' means trees, as opposed to data).").unwrap(); - writeln!(f, ";").unwrap(); - writeln!(f, "hold 20;").unwrap(); - writeln!(f, "mult;").unwrap(); - writeln!(f, "export - {tre_name};").unwrap(); - - info!( - "TNT script → {path} (costs scaled x{cost_scale:.0}, runs a default `hold 20; mult;` \ - search and writes trees to {tre_name} in TNT's working directory — edit the trailing \ - comment block in the script to change this)\n\ - Run it with:\n \ - printf 'proc {path};\\nquit;\\n' | tnt\n\ - (or start `tnt` interactively and type `proc {path};`)" - ); -} - -/// Scale `matrix` by `cost_scale` and round to integers (TNT's smatrix/cost -/// commands reject decimals), then take the *metric 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 `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 -/// TNT-version-dependent. -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 { - for j in 0..16 { - m[i][j] = (matrix[i][j] * cost_scale).round() as i64; - } - } - for k in 0..16 { - for i in 0..16 { - for j in 0..16 { - let via = m[i][k] + m[k][j]; - if via < m[i][j] { - m[i][j] = via; - } - } - } - } - m -} - -// ── Sankoff cost matrix → PhyG custom-alphabet TCM + ready-to-run script ──── -// -// PhyG's `tcm:STRING` format needs no alphabet recoding, unlike `--tnt`: -// its parser reads the alphabet straight from the tcm file's own first -// line, so `--sankoff`'s own `_sankoff.fasta` (already IUPAC+`0`) is reused -// as-is via `prefasta:`. PhyG auto-adds its own indel/gap state as an -// (n+1)-th row/column of the tcm — inert here since the alignment already -// encodes absence as an ordinary state (`0`), never as `-` (see -// `write_sankoff_alignment_fasta`'s own comment on why, and the RAxML-era -// bug that motivated it). The gap row/column below reuses `matrix[i][0]`/ -// `matrix[0][j]` (cost to/from `∅`) as the closest principled value for a -// state that, in practice, is never actually triggered. - -fn write_sankoff_phyg(matrix: &[[f64; 16]; 16], output: &Option, cost_scale: f64) { - let scaled_matrix = scaled_metric_matrix(matrix, cost_scale); - - let basename = |suffix: &str| -> String { - output.as_ref() - .and_then(|p| p.file_name()) - .map(|n| format!("{}{suffix}", n.to_string_lossy())) - .unwrap_or_else(|| format!("sankoff{suffix}")) - }; - let full_path = |suffix: &str| -> String { - output.as_ref() - .map(|p| format!("{}{suffix}", p.display())) - .unwrap_or_else(|| format!("sankoff{suffix}")) - }; - - let tcm_path = full_path("_sankoff.tcm"); - let mut f = BufWriter::new(std::fs::File::create(&tcm_path).unwrap_or_else(|e| { - eprintln!("error creating {tcm_path}: {e}"); - std::process::exit(1); - })); - let alphabet_line = STATE_SYMBOL.iter().map(|c| c.to_string()).collect::>().join(" "); - writeln!(f, "{alphabet_line}").unwrap(); - for i in 0..16 { - let mut row: Vec = (0..16).map(|j| scaled_matrix[i][j]).collect(); - row.push(scaled_matrix[i][0]); // gap column: same cost as to/from ∅ - writeln!(f, "{}", row.iter().map(|v| v.to_string()).collect::>().join(" ")).unwrap(); - } - let mut gap_row: Vec = (0..16).map(|j| scaled_matrix[0][j]).collect(); - gap_row.push(0); - writeln!(f, "{}", gap_row.iter().map(|v| v.to_string()).collect::>().join(" ")).unwrap(); - info!("PhyG TCM → {tcm_path}"); - - let pg_path = full_path("_sankoff.pg"); - let mut f = BufWriter::new(std::fs::File::create(&pg_path).unwrap_or_else(|e| { - eprintln!("error creating {pg_path}: {e}"); - std::process::exit(1); - })); - let fasta_name = basename("_sankoff.fasta"); - let tcm_name = basename("_sankoff.tcm"); - let tre_name = basename("_sankoff.tre"); - writeln!(f, "read(prefasta:\"{fasta_name}\", tcm:\"{tcm_name}\")").unwrap(); - writeln!(f, "search(seconds:300, instances:4)").unwrap(); - writeln!(f, "report(\"{tre_name}\", graphs, newick, overwrite)").unwrap(); - - let pg_dir = std::path::Path::new(&pg_path).parent() - .filter(|d| !d.as_os_str().is_empty()) - .map(|d| d.display().to_string()) - .unwrap_or_else(|| ".".into()); - let pg_name = std::path::Path::new(&pg_path).file_name() - .map(|n| n.to_string_lossy().into_owned()) - .unwrap_or_else(|| pg_path.clone()); - info!( - "PhyG script → {pg_path} (costs scaled x{cost_scale:.0}, runs a default 300s/4-instance \ - search and writes trees to {tre_name})\n\ - Run it with:\n \ - cd {pg_dir} && phyg {pg_name}\n\ - (`phyg` must run from that directory — `read()`/`report()` in the script use relative \ - file names)" - ); -} - // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── fn upgma_to_newick(dendro: &kodama::Dendrogram, names: &[String]) -> String { diff --git a/src/obikmer/src/cmd/distance/phyg.rs b/src/obikmer/src/cmd/distance/phyg.rs new file mode 100644 index 00000000..f0d0167d --- /dev/null +++ b/src/obikmer/src/cmd/distance/phyg.rs @@ -0,0 +1,81 @@ +use std::io::{BufWriter, Write}; +use std::path::PathBuf; + +use tracing::info; + +use super::sankoff::{STATE_SYMBOL, scaled_metric_matrix}; + +// ── Sankoff cost matrix → PhyG custom-alphabet TCM + ready-to-run script ──── +// +// PhyG's `tcm:STRING` format needs no alphabet recoding, unlike `--tnt`: +// its parser reads the alphabet straight from the tcm file's own first +// line, so `--sankoff`'s own `_sankoff.fasta` (already IUPAC+`0`) is reused +// as-is via `prefasta:`. PhyG auto-adds its own indel/gap state as an +// (n+1)-th row/column of the tcm — inert here since the alignment already +// encodes absence as an ordinary state (`0`), never as `-` (see +// `sankoff::write_sankoff_alignment_fasta`'s own comment on why, and the +// RAxML-era bug that motivated it). The gap row/column below reuses +// `matrix[i][0]`/`matrix[0][j]` (cost to/from `∅`) as the closest +// principled value for a state that, in practice, is never actually +// triggered. + +pub(super) fn write_sankoff_phyg(matrix: &[[f64; 16]; 16], output: &Option, cost_scale: f64) { + let scaled_matrix = scaled_metric_matrix(matrix, cost_scale); + + let basename = |suffix: &str| -> String { + output.as_ref() + .and_then(|p| p.file_name()) + .map(|n| format!("{}{suffix}", n.to_string_lossy())) + .unwrap_or_else(|| format!("sankoff{suffix}")) + }; + let full_path = |suffix: &str| -> String { + output.as_ref() + .map(|p| format!("{}{suffix}", p.display())) + .unwrap_or_else(|| format!("sankoff{suffix}")) + }; + + let tcm_path = full_path("_sankoff.tcm"); + let mut f = BufWriter::new(std::fs::File::create(&tcm_path).unwrap_or_else(|e| { + eprintln!("error creating {tcm_path}: {e}"); + std::process::exit(1); + })); + let alphabet_line = STATE_SYMBOL.iter().map(|c| c.to_string()).collect::>().join(" "); + writeln!(f, "{alphabet_line}").unwrap(); + for i in 0..16 { + let mut row: Vec = (0..16).map(|j| scaled_matrix[i][j]).collect(); + row.push(scaled_matrix[i][0]); // gap column: same cost as to/from ∅ + writeln!(f, "{}", row.iter().map(|v| v.to_string()).collect::>().join(" ")).unwrap(); + } + let mut gap_row: Vec = (0..16).map(|j| scaled_matrix[0][j]).collect(); + gap_row.push(0); + writeln!(f, "{}", gap_row.iter().map(|v| v.to_string()).collect::>().join(" ")).unwrap(); + info!("PhyG TCM → {tcm_path}"); + + let pg_path = full_path("_sankoff.pg"); + let mut f = BufWriter::new(std::fs::File::create(&pg_path).unwrap_or_else(|e| { + eprintln!("error creating {pg_path}: {e}"); + std::process::exit(1); + })); + let fasta_name = basename("_sankoff.fasta"); + let tcm_name = basename("_sankoff.tcm"); + let tre_name = basename("_sankoff.tre"); + writeln!(f, "read(prefasta:\"{fasta_name}\", tcm:\"{tcm_name}\")").unwrap(); + writeln!(f, "search(seconds:300, instances:4)").unwrap(); + writeln!(f, "report(\"{tre_name}\", graphs, newick, overwrite)").unwrap(); + + let pg_dir = std::path::Path::new(&pg_path).parent() + .filter(|d| !d.as_os_str().is_empty()) + .map(|d| d.display().to_string()) + .unwrap_or_else(|| ".".into()); + let pg_name = std::path::Path::new(&pg_path).file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_else(|| pg_path.clone()); + info!( + "PhyG script → {pg_path} (costs scaled x{cost_scale:.0}, runs a default 300s/4-instance \ + search and writes trees to {tre_name})\n\ + Run it with:\n \ + cd {pg_dir} && phyg {pg_name}\n\ + (`phyg` must run from that directory — `read()`/`report()` in the script use relative \ + file names)" + ); +} diff --git a/src/obikmer/src/cmd/distance/sankoff.rs b/src/obikmer/src/cmd/distance/sankoff.rs new file mode 100644 index 00000000..76676237 --- /dev/null +++ b/src/obikmer/src/cmd/distance/sankoff.rs @@ -0,0 +1,202 @@ +use std::io::{BufWriter, Write}; +use std::path::PathBuf; + +use obifastwrite::{JsonVal, write_record}; +use obikindex::{BasePairTally, PHatEstimate, SankoffWeights, SnpAlignment}; +use tracing::info; + +// ── Sankoff pseudo-alignment → FASTA ──────────────────────────────────────── +// +// Same data as `--snp`'s pseudo-alignment (`SnpAlignment`/ +// `snp_pseudo_alignment`), re-coded so its symbols match the accompanying +// `--sankoff-matrix` output exactly: `0` for the empty/absent state instead +// of `-`, which TNT/PhyG would otherwise read as their own gap character +// rather than our "family absent" state. + +pub(super) fn write_sankoff_alignment_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_sankoff.fasta", p.display())) + .unwrap_or_else(|| "sankoff.fasta".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); + for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { + let recoded: Vec = seq.iter().map(|&b| if b == b'-' { b'0' } else { b }).collect(); + write_record(&recoded, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| { + eprintln!("error writing {path}: {e}"); + std::process::exit(1); + }); + } + info!("Sankoff pseudo-alignment → {path} ({n_sites} site{})", + if n_sites == 1 { "" } else { "s" }); +} + +// ── Sankoff cost matrix → CSV ──────────────────────────────────────────────── +// +// 16 states indexed by bitmask (bit 0=A, 1=C, 2=G, 3=T; state 0 is `∅`), +// matching the convention already used for `--snp`'s IUPAC-coded output and +// for the external TNT/PhyG scripts this feeds. Calibration report (p_hat, +// its variance, how many pairs/loci went into it, the resulting c_ctx) goes +// to the log, not the CSV, since it's a run-level fact, not per-cell data. + +// IUPAC ambiguity code per state (same mapping as `siblings::iupac_code`, +// already used for `--snp`'s pseudo-alignment — a biologist reads "R" as +// "A or G" without needing this file's convention explained), with `0` +// standing in for the empty state (`-` would collide with TNT/PhyG's own +// gap/range syntax). Bit order: 0=A, 1=C, 2=G, 3=T. This project's +// canonical alphabet — `tnt::write_sankoff_tnt` recodes it to TNT's own +// default alphabet at the adapter boundary, rather than using it here. +pub(super) const STATE_SYMBOL: [char; 16] = [ + '0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N', +]; + +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()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + write!(f, "state").unwrap(); + for sym in STATE_SYMBOL { write!(f, ",{sym}").unwrap(); } + writeln!(f).unwrap(); + for (s, row) in matrix.iter().enumerate() { + write!(f, "{}", STATE_SYMBOL[s]).unwrap(); + for cost in row { write!(f, ",{cost:.4}").unwrap(); } + writeln!(f).unwrap(); + } + info!("Sankoff cost matrix → {path}"); +} + +// ── Sankoff calibration parameters → YAML report ──────────────────────────── +// +// Everything `--sankoff` estimates from real data, in one durable, +// machine-readable file: `p_hat` and its variance (with how many pairs/loci +// went into it), the derived `c_ctx`, and the base-pair substitution tally +// (raw counts, not just the derived costs) — kept for the same reason raw +// counts are kept anywhere else in this project: costs are a modelling +// choice built *from* the counts, and reproducing/re-deriving them later +// needs the counts, not just their current derived value. Structured (YAML, +// not an ad hoc key=value text file) so R/Python/etc. can load it directly +// rather than re-parsing free text. + +#[derive(serde::Serialize)] +struct SankoffSubstitution { + pair: String, + count: u64, + cost: 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, +} + +pub(super) fn write_sankoff_params( + estimate: &PHatEstimate, + tally: &BasePairTally, + weights: &SankoffWeights, + 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 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 path = output.as_ref() + .map(|p| format!("{}_sankoff_params.yaml", p.display())) + .unwrap_or_else(|| "sankoff_params.yaml".into()); + let f = std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + }); + serde_yaml::to_writer(f, &report).unwrap_or_else(|e| { + eprintln!("error writing {path}: {e}"); + std::process::exit(1); + }); + info!("Sankoff calibration parameters → {path}"); +} + +/// Scale `matrix` by `cost_scale` and round to integers (TNT's smatrix/cost +/// and PhyG's `tcm:` commands both reject decimals), then take the *metric +/// 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. +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 { + for j in 0..16 { + m[i][j] = (matrix[i][j] * cost_scale).round() as i64; + } + } + for k in 0..16 { + for i in 0..16 { + for j in 0..16 { + let via = m[i][k] + m[k][j]; + if via < m[i][j] { + m[i][j] = via; + } + } + } + } + m +} diff --git a/src/obikmer/src/cmd/distance/tnt.rs b/src/obikmer/src/cmd/distance/tnt.rs new file mode 100644 index 00000000..977994c4 --- /dev/null +++ b/src/obikmer/src/cmd/distance/tnt.rs @@ -0,0 +1,119 @@ +use std::io::{BufWriter, Write}; +use std::path::PathBuf; + +use obikindex::SnpAlignment; +use tracing::info; + +use super::sankoff::{STATE_SYMBOL, scaled_metric_matrix}; + +// ── Sankoff cost matrix + alignment → ready-to-run TNT script ────────────── +// +// TNT's *default* xread reader only accepts its own 0-9A-F alphabet (see +// its manual: "up to 16 states are allowed by xread, using symbols 0-9 ... +// and A-F") — the wider IUPAC set `STATE_SYMBOL` uses is rejected as an +// "alien symbol" unless `nstates dna` is set, which imposes TNT's own fixed +// DNA encoding instead, incompatible with a custom smatrix. And TNT's +// `smatrix`/`cost` commands reject decimal costs ("found symbol . when +// reading transformation costs"). So: recode to TNT's alphabet and +// integer-scale the costs here, at this adapter's boundary, rather than +// degrading the project's own canonical (IUPAC, real-valued) output. + +const TNT_STATE_SYMBOL: [char; 16] = [ + '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', +]; + +pub(super) fn write_sankoff_tnt( + matrix: &[[f64; 16]; 16], + alignment: &SnpAlignment, + labels: &[String], + output: &Option, + cost_scale: f64, +) { + let path = output.as_ref() + .map(|p| format!("{}_sankoff.tnt", p.display())) + .unwrap_or_else(|| "sankoff.tnt".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + + // IUPAC-ish symbol -> bitmask, to translate the alignment (which uses + // `STATE_SYMBOL`, `-` already normalised to `0` by `snp_pseudo_alignment` + // callers) into TNT's alphabet without re-deriving state indices. + let mut iupac_to_state = [0u8; 128]; + for (state, &sym) in STATE_SYMBOL.iter().enumerate() { + iupac_to_state[sym as usize] = state as u8; + } + + let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); + writeln!(f, "xread").unwrap(); + writeln!(f, "'obikmer central-position SNP families, calibrated Sankoff 16-state encoding'").unwrap(); + writeln!(f, "{n_sites} {}", labels.len()).unwrap(); + for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { + write!(f, "{label} ").unwrap(); + for &b in seq { + let b = if b == b'-' { b'0' } else { b }; + let state = iupac_to_state[b as usize]; + write!(f, "{}", TNT_STATE_SYMBOL[state as usize]).unwrap(); + } + writeln!(f).unwrap(); + } + writeln!(f, ";\n").unwrap(); + + let scaled_matrix = scaled_metric_matrix(matrix, cost_scale); + + writeln!(f, "smatrix =0 (family16)").unwrap(); + for i in 0..16 { + for j in (i + 1)..16 { + writeln!(f, "{}/{} {}", TNT_STATE_SYMBOL[i], TNT_STATE_SYMBOL[j], scaled_matrix[i][j]).unwrap(); + } + } + writeln!(f, ";\n").unwrap(); + + writeln!(f, "ccode ( 0.{} ;", n_sites - 1).unwrap(); + writeln!(f, "smatrix +0 0.{} ;", n_sites - 1).unwrap(); + writeln!(f).unwrap(); + + // Basename only (not the full `path`/`output` prefix): TNT's natural + // workflow is to `cd` into the output directory before `proc`-ing the + // script, and an absolute path here would break if that directory is + // later moved or copied elsewhere. + let tre_name = output.as_ref() + .and_then(|p| p.file_name()) + .map(|n| format!("{}_sankoff.tre", n.to_string_lossy())) + .unwrap_or_else(|| "sankoff.tre".into()); + + // TNT's plain command parser has no comment syntax of its own — `/* */` + // and `[ ]` are only recognised inside the (separately-enabled) macro + // scripting language, and fail with "No command!" here otherwise + // (verified against this file with the local TNT binary). `quote` is + // the closest working equivalent: it prints free text and does not + // otherwise affect parsing, so it doubles as an explanation of the + // defaults below when the script is run. `;` ends a `quote` block like + // any other TNT command, so the text itself must avoid semicolons. + writeln!(f, "quote").unwrap(); + writeln!(f, "Default search below (edit or delete this block to run your own strategy):").unwrap(); + writeln!(f, " hold N : size of TNT's tree buffer (how many equally-parsimonious").unwrap(); + writeln!(f, " trees it keeps in memory at once), 20 is a small, fast").unwrap(); + writeln!(f, " default, raise it if mult reports it had to drop trees.").unwrap(); + writeln!(f, " mult : traditional search (random addition sequences followed by").unwrap(); + writeln!(f, " TBR branch-swapping, TNT's own default replication count),").unwrap(); + writeln!(f, " a reasonable first-pass strategy on this data's memory").unwrap(); + writeln!(f, " footprint, xmult's ratchet/drift/tree-fusion buffers ran").unwrap(); + writeln!(f, " this out of RAM at TNT's default mxram on this dataset.").unwrap(); + writeln!(f, " export - F : write the trees held in the buffer to file F, in").unwrap(); + writeln!(f, " TNT/Hennig86 format ('-' means trees, as opposed to data).").unwrap(); + writeln!(f, ";").unwrap(); + writeln!(f, "hold 20;").unwrap(); + writeln!(f, "mult;").unwrap(); + writeln!(f, "export - {tre_name};").unwrap(); + + info!( + "TNT script → {path} (costs scaled x{cost_scale:.0}, runs a default `hold 20; mult;` \ + search and writes trees to {tre_name} in TNT's working directory — edit the trailing \ + comment block in the script to change this)\n\ + Run it with:\n \ + printf 'proc {path};\\nquit;\\n' | tnt\n\ + (or start `tnt` interactively and type `proc {path};`)" + ); +} -- 2.54.0 From 28d841c7beda8b7fe740680c925724c037c1f30e Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 12 Aug 2026 18:23:51 +0200 Subject: [PATCH 03/46] feat(distance): add IQ-TREE output and optimize state index mapping Introduce `--iqtree` and `--raw-snp-counts` flags to generate IQ-TREE model files, recoded FASTA alignments, and per-pair diagnostic counts. Centralize alphabet conversion by extracting a precomputed state index lookup table into the Sankoff module, eliminating redundant iterations across downstream adapters. --- .gitignore | 1 + src/obikmer/src/cmd/distance/iqtree.rs | 183 ++++++++++++++++++++++++ src/obikmer/src/cmd/distance/mod.rs | 91 +++++++++++- src/obikmer/src/cmd/distance/sankoff.rs | 11 ++ src/obikmer/src/cmd/distance/tnt.rs | 7 +- 5 files changed, 285 insertions(+), 8 deletions(-) create mode 100644 src/obikmer/src/cmd/distance/iqtree.rs diff --git a/.gitignore b/.gitignore index 3cd5920d..32189eb2 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,4 @@ benchmark/specific_index_count benchmark/specific_index_presence TNT phyg +*.tnt diff --git a/src/obikmer/src/cmd/distance/iqtree.rs b/src/obikmer/src/cmd/distance/iqtree.rs new file mode 100644 index 00000000..0a1ac655 --- /dev/null +++ b/src/obikmer/src/cmd/distance/iqtree.rs @@ -0,0 +1,183 @@ +use std::io::{BufWriter, Write}; +use std::path::PathBuf; + +use obifastwrite::{JsonVal, write_record}; +use obikindex::SnpAlignment; +use tracing::info; + +use super::sankoff::state_index_table; + +// ── Sankoff-calibrated data → IQ-TREE custom ML model + recoded alignment ── +// +// Not itself a Sankoff computation — IQ-TREE does maximum likelihood, not +// parsimony. Only the *source data* is shared with `--tnt`/`--phyg` (the +// calibrated 16-state cost matrix, the pseudo-alignment); the operation +// performed on it here is different, hence no "sankoff" in these names, +// unlike `tnt::write_sankoff_tnt`/`phyg::write_sankoff_phyg`. +// +// Verified against the locally installed `iqtree3` binary/source (not just +// its docs — see `docmd/theory/evolutionary_distances.md`, "Next +// direction: genuine ML branch lengths"), because the web documentation's +// `-mdef` NEXUS route turned out not to apply to a plain (non-mixture) +// custom morphology model. The real mechanism: pass a **file path** +// directly as `-m`, containing (as whitespace/newline-separated numbers) +// the lower-triangular exchangeability matrix `R` (`k(k-1)/2` values, PAML +// row-major order) immediately followed by the `k` state frequencies `π` +// on the same stream — `ModelMarkov::readRates`/`readStateFreq` read them +// in that exact order, no header, no separator required. +// +// `R` is recovered from the calibrated Sankoff cost matrix via +// `R(a,b) = exp(-cost(a,b))` (the cost is `-ln(rate)`, see "A concrete +// Sankoff cost matrix"), symmetric by construction (the underlying tally +// never captured direction). `π` is the real, empirical, non-uniform +// marginal frequency of each state across the whole alignment — precise +// enough at this sample size (908k+ sites) without spending IQ-TREE's own +// `+FO` ML degrees of freedom re-estimating it (checked: IQ-TREE's `+F` +// doesn't work as a shortcut here either, a custom-file model always +// requires the frequency line in the file itself). IQ-TREE reconstructs +// the (generally asymmetric) rate matrix internally as +// `Q(i,j) = R(i,j)·π_j` — reversible for *any* `π`, not just uniform, +// because `R` is symmetric. +// +// IQ-TREE infers its state count from the highest-ordinal symbol actually +// present in the alignment, not from a declared count (`--seqtype +// MORPH{N}` was tested and does not override this for real ML analysis, +// only for the `--alisim` simulator). So states that never occur anywhere +// in this particular alignment are dropped, and the survivors are +// renumbered compactly (`0..k-1`, order preserved) rather than leaving +// gaps that would silently misalign every value IQ-TREE reads. Both the +// model and the alignment must agree on this same renumbering, so it's +// computed once (`CompactAlphabet`) and shared between them. + +const IQTREE_STATE_SYMBOL: [char; 16] = [ + '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', +]; + +struct CompactAlphabet { + /// Canonical (0..16) state index -> compact index, for states that occur. + old_to_compact: [Option; 16], + /// Compact index -> canonical state index, order-preserving. + compact_to_old: Vec, + /// Empirical frequency of each compact-indexed state (sums to 1). + freq: Vec, +} + +impl CompactAlphabet { + fn k(&self) -> usize { + self.compact_to_old.len() + } +} + +fn compact_alphabet(alignment: &SnpAlignment) -> CompactAlphabet { + let iupac_to_state = state_index_table(); + + let mut occurs = [false; 16]; + let mut counts = [0u64; 16]; + for seq in &alignment.sequences { + for &b in seq { + let b = if b == b'-' { b'0' } else { b }; + let state = iupac_to_state[b as usize] as usize; + occurs[state] = true; + counts[state] += 1; + } + } + + let mut old_to_compact: [Option; 16] = [None; 16]; + let mut compact_to_old: Vec = Vec::new(); + for old in 0..16 { + if occurs[old] { + old_to_compact[old] = Some(compact_to_old.len() as u8); + compact_to_old.push(old as u8); + } + } + + let total: u64 = compact_to_old.iter().map(|&old| counts[old as usize]).sum(); + let freq: Vec = compact_to_old.iter() + .map(|&old| counts[old as usize] as f64 / total as f64) + .collect(); + + CompactAlphabet { old_to_compact, compact_to_old, freq } +} + +/// Write the `R` (exchangeability) + `π` (frequencies) model file IQ-TREE's +/// `-m +ASC` reads. Returns the path, so the caller can print a +/// single combined "how to run this" message once the alignment is also +/// written. +fn write_iqtree_model(matrix: &[[f64; 16]; 16], alphabet: &CompactAlphabet, output: &Option) -> String { + let rate = |old_i: u8, old_j: u8| (-matrix[old_i as usize][old_j as usize]).exp(); + + let model_path = output.as_ref() + .map(|p| format!("{}_iqtree.model", p.display())) + .unwrap_or_else(|| "iqtree.model".into()); + let mut f = BufWriter::new(std::fs::File::create(&model_path).unwrap_or_else(|e| { + eprintln!("error creating {model_path}: {e}"); + std::process::exit(1); + })); + for i in 1..alphabet.k() { + let row: Vec = (0..i) + .map(|j| format!("{:.6}", rate(alphabet.compact_to_old[i], alphabet.compact_to_old[j]))) + .collect(); + writeln!(f, "{}", row.join(" ")).unwrap(); + } + writeln!(f, "{}", alphabet.freq.iter().map(|p| format!("{p:.6}")).collect::>().join(" ")).unwrap(); + info!("IQ-TREE model file → {model_path} ({} of 16 states present in the alignment)", alphabet.k()); + model_path +} + +/// Write the pseudo-alignment recoded to the same compact `0..k-1` alphabet +/// as `write_iqtree_model`'s matrix — not `--sankoff`'s own IUPAC alphabet, +/// since IQ-TREE needs the symbol ordinal itself to match the surviving +/// state count (see this module's own doc comment on `MORPH{N}`). +fn write_iqtree_alignment( + alignment: &SnpAlignment, + labels: &[String], + alphabet: &CompactAlphabet, + output: &Option, +) -> (String, usize) { + let iupac_to_state = state_index_table(); + + let fasta_path = output.as_ref() + .map(|p| format!("{}_iqtree.fasta", p.display())) + .unwrap_or_else(|| "iqtree.fasta".into()); + let mut f = BufWriter::new(std::fs::File::create(&fasta_path).unwrap_or_else(|e| { + eprintln!("error creating {fasta_path}: {e}"); + std::process::exit(1); + })); + let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); + for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { + let recoded: Vec = seq.iter().map(|&b| { + let b = if b == b'-' { b'0' } else { b }; + let old = iupac_to_state[b as usize] as usize; + let compact = alphabet.old_to_compact[old] + .expect("state occurs in the alignment, so it must have a compact index"); + IQTREE_STATE_SYMBOL[compact as usize] as u8 + }).collect(); + write_record(&recoded, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| { + eprintln!("error writing {fasta_path}: {e}"); + std::process::exit(1); + }); + } + (fasta_path, n_sites) +} + +pub(super) fn write_iqtree( + matrix: &[[f64; 16]; 16], + alignment: &SnpAlignment, + labels: &[String], + output: &Option, +) { + let alphabet = compact_alphabet(alignment); + let model_path = write_iqtree_model(matrix, &alphabet, output); + let (fasta_path, n_sites) = write_iqtree_alignment(alignment, labels, &alphabet, output); + + let prefix_name = output.as_ref() + .and_then(|p| p.file_name()) + .map(|n| format!("{}_iqtree", n.to_string_lossy())) + .unwrap_or_else(|| "iqtree".into()); + info!( + "IQ-TREE alignment → {fasta_path} ({n_sites} sites, {} states)\n\ + Run with:\n \ + iqtree3 -s {fasta_path} --seqtype MORPH -m {model_path}+ASC --prefix {prefix_name} -T AUTO", + alphabet.k() + ); +} diff --git a/src/obikmer/src/cmd/distance/mod.rs b/src/obikmer/src/cmd/distance/mod.rs index bd698caa..3dfe08d3 100644 --- a/src/obikmer/src/cmd/distance/mod.rs +++ b/src/obikmer/src/cmd/distance/mod.rs @@ -1,3 +1,4 @@ +mod iqtree; mod phyg; mod sankoff; mod tnt; @@ -16,6 +17,7 @@ use obikindex::{ use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; +use iqtree::write_iqtree; use phyg::write_sankoff_phyg; use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params}; use tnt::write_sankoff_tnt; @@ -98,6 +100,19 @@ pub struct DistanceArgs { #[arg(long)] pub raw_snp_distance: bool, + /// Write the raw per-pair counts (`n_snp`, `n_shared`, `n_eligible`) + /// behind `--raw-snp-distance`'s ratio, one row per genome pair — a + /// diagnostic table, not a matrix. The ratio alone can't distinguish + /// "identical at every eligible locus" from "almost no eligible loci + /// at all" (e.g. `0.0` from 0/2 looks the same as `0.0` from 0/2000), + /// and that distinction matters a lot for genome pairs near the edge + /// of what central-position families can resolve (see + /// `docmd/theory/evolutionary_distances.md`, "Run 3" and the + /// IQ-TREE/Mash comparison). Same annex requirement as + /// `--raw-snp-distance`. + #[arg(long)] + pub raw_snp_counts: bool, + /// Write a SNP-only pseudo-alignment (FASTA, IUPAC-coded) from an /// already-built sibling annex — one row per genome, one column per /// variable family (monomorphic families skipped), no flanking @@ -141,6 +156,24 @@ pub struct DistanceArgs { #[arg(long)] pub phyg: bool, + /// Also write _iqtree.model and _iqtree.fasta, a + /// custom-model file and a matching + /// recoded alignment for genuine maximum-likelihood inference with + /// IQ-TREE (`iqtree3 -s ... --seqtype MORPH -m ...+ASC`) — real branch + /// lengths, unlike `--tnt`/`--phyg`'s parsimony step counts. The model + /// is the reversible `Q(i,j) = R(i,j)·π_j` construction: `R` + /// (exchangeability, symmetric) recovered from the same calibrated + /// cost matrix `--sankoff` computes, `π` the real empirical state + /// frequencies counted from the alignment (not IQ-TREE's `+FO`/`+F` — + /// neither applies to a custom-file model, see + /// `docmd/theory/evolutionary_distances.md`). Only the states that + /// actually occur in this alignment are kept, compactly renumbered + /// (IQ-TREE infers its state count from the alignment itself, and a + /// gap in the numbering would silently misalign the model file). + /// Implies `--sankoff`. + #[arg(long)] + pub iqtree: bool, + /// Scale factor applied before rounding real-valued costs to the /// integers both `--tnt`'s smatrix/cost commands and `--phyg`'s `tcm:` /// matrix require. Keep this small: the total tree score is this scale @@ -157,10 +190,12 @@ pub struct DistanceArgs { pub sankoff_cost_scale: f64, /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _rawsnp.csv, _snp.fasta, + /// _siblings.csv, _rawsnp.csv, _rawsnp_counts.csv, + /// _snp.fasta, /// _sankoff_matrix.csv, _sankoff_params.yaml, /// _sankoff.fasta, _sankoff.tnt, _sankoff.tcm, - /// _sankoff.pg, _nj.nwk, _upgma.nwk. + /// _sankoff.pg, _iqtree.model, _iqtree.fasta, + /// _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -208,6 +243,13 @@ pub fn run(args: DistanceArgs) { }); write_raw_snp_distance_csv(&result, &labels, &args.output); } + if args.raw_snp_counts { + let result = idx.raw_snp_distance().unwrap_or_else(|e| { + eprintln!("error computing raw SNP distance: {e}"); + std::process::exit(1); + }); + write_raw_snp_counts_csv(&result, &labels, &args.output); + } if args.snp { let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { eprintln!("error computing SNP pseudo-alignment: {e}"); @@ -215,7 +257,7 @@ pub fn run(args: DistanceArgs) { }); write_snp_fasta(&alignment, &labels, &args.output); } - if args.sankoff || args.tnt || args.phyg { + if args.sankoff || args.tnt || args.phyg || args.iqtree { let raw = idx.raw_snp_distance().unwrap_or_else(|e| { eprintln!("error computing raw SNP distance: {e}"); std::process::exit(1); @@ -248,6 +290,9 @@ pub fn run(args: DistanceArgs) { if args.phyg { write_sankoff_phyg(&matrix, &args.output, args.sankoff_cost_scale); } + if args.iqtree { + write_iqtree(&matrix, &alignment, &labels, &args.output); + } } // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp`/ @@ -260,10 +305,12 @@ pub fn run(args: DistanceArgs) { if args.sibling_annex || args.sibling_stats || args.raw_snp_distance + || args.raw_snp_counts || args.snp || args.sankoff || args.tnt || args.phyg + || args.iqtree { return; } @@ -450,6 +497,44 @@ fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], info!("raw single-copy SNP distance matrix → {path}"); } +// ── Raw single-copy SNP distance → per-pair diagnostic counts ────────────── +// +// A pair table (one row per unordered genome pair), not a matrix: the ratio +// alone can't distinguish "identical across every eligible locus" from +// "almost no eligible locus at all" — both can read `0.0`/`NA` in +// `--raw-snp-distance`'s output. Distinguishing them matters most exactly +// where it's easy to miss: genome pairs near the edge of what +// central-position families can resolve at all (deep cross-lineage splits, +// see `docmd/theory/evolutionary_distances.md`, "Run 3" and the later +// IQ-TREE/Mash comparison — a `ratio=0.0` backed by 2 eligible loci is not +// the same claim as one backed by 2000). + +fn write_raw_snp_counts_csv(result: &RawSnpDistanceOutput, labels: &[String], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_rawsnp_counts.csv", p.display())) + .unwrap_or_else(|| "rawsnp_counts.csv".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + let n = labels.len(); + writeln!(f, "genome_a,genome_b,n_snp,n_shared,n_eligible,ratio").unwrap(); + for i in 0..n { + for j in (i + 1)..n { + let snp = result.snp[[i, j]]; + let shared = result.shared[[i, j]]; + let eligible = snp + shared; + write!(f, "{},{},{snp},{shared},{eligible}", labels[i], labels[j]).unwrap(); + if eligible == 0 { + writeln!(f, ",NA").unwrap(); + } else { + writeln!(f, ",{:.6}", snp as f64 / eligible as f64).unwrap(); + } + } + } + info!("raw single-copy SNP distance counts (diagnostic) → {path}"); +} + // ── SNP-only pseudo-alignment → FASTA ─────────────────────────────────────── // // One record per genome, IUPAC-coded, no flanking sequence — see diff --git a/src/obikmer/src/cmd/distance/sankoff.rs b/src/obikmer/src/cmd/distance/sankoff.rs index 76676237..5e02fda6 100644 --- a/src/obikmer/src/cmd/distance/sankoff.rs +++ b/src/obikmer/src/cmd/distance/sankoff.rs @@ -52,6 +52,17 @@ pub(super) const STATE_SYMBOL: [char; 16] = [ '0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N', ]; +/// `STATE_SYMBOL` byte -> state index (0..16), for adapters that need to +/// translate an alignment written in this alphabet into their own. Shared +/// rather than rebuilt per adapter (`tnt`, `iqtree`). +pub(super) fn state_index_table() -> [u8; 128] { + let mut table = [0u8; 128]; + for (state, &sym) in STATE_SYMBOL.iter().enumerate() { + table[sym as usize] = state as u8; + } + table +} + pub(super) fn write_sankoff_matrix_csv( matrix: &[[f64; 16]; 16], estimate: &PHatEstimate, diff --git a/src/obikmer/src/cmd/distance/tnt.rs b/src/obikmer/src/cmd/distance/tnt.rs index 977994c4..48f5ad53 100644 --- a/src/obikmer/src/cmd/distance/tnt.rs +++ b/src/obikmer/src/cmd/distance/tnt.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use obikindex::SnpAlignment; use tracing::info; -use super::sankoff::{STATE_SYMBOL, scaled_metric_matrix}; +use super::sankoff::{scaled_metric_matrix, state_index_table}; // ── Sankoff cost matrix + alignment → ready-to-run TNT script ────────────── // @@ -40,10 +40,7 @@ pub(super) fn write_sankoff_tnt( // IUPAC-ish symbol -> bitmask, to translate the alignment (which uses // `STATE_SYMBOL`, `-` already normalised to `0` by `snp_pseudo_alignment` // callers) into TNT's alphabet without re-deriving state indices. - let mut iupac_to_state = [0u8; 128]; - for (state, &sym) in STATE_SYMBOL.iter().enumerate() { - iupac_to_state[sym as usize] = state as u8; - } + let iupac_to_state = state_index_table(); let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); writeln!(f, "xread").unwrap(); -- 2.54.0 From c26623fa00a990884883d3ecd5ddd374cbc2f27d Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 12 Aug 2026 19:51:35 +0200 Subject: [PATCH 04/46] Add repeatable --exclude-genome flag to obikmer distance command Integrate in-memory row/column zeroing and alignment filtering across SNP, Sankoff, TNT, and IQ-TREE output paths. Add strict validation for missing labels, implement `write_iqtree` with empirical stationary frequencies, and introduce `--raw-snp-counts` diagnostic CSV output. Update documentation to reflect experimental validation of backbone resolution limits and theoretical considerations for CTMC rate matrices. --- docmd/theory/evolutionary_distances.md | 207 ++++++++++++++++++++++++- src/obikmer/src/cmd/distance/mod.rs | 83 +++++++++- 2 files changed, 282 insertions(+), 8 deletions(-) diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index 6fd21d8b..c7e2ebfe 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -897,7 +897,212 @@ Models](https://iqtree.github.io/doc/Complex-Models). (see "A concrete Sankoff cost matrix" above) is already a log-rate — a genuine CTMC rate matrix `Q` could plausibly be recovered as `rate(a,b) = exp(-cost(a,b))`, renormalised so each row sums to zero, once -the stationary-frequency gap above is closed. Not yet attempted. +the stationary-frequency gap above is closed. Superseded by the two +sections below, which implement and then substantially revise this. + +### `R`/`π` implemented; `--exclude-genome` added; rogue-taxon test negative (2026-08-12) + +`write_iqtree` (`obikmer/src/cmd/distance/iqtree.rs`) implements exactly +the `R = exp(-cost)` / empirical-`π` design above: writes +`_iqtree.model` (lower-triangular `R`, PAML order, then `π`) and +`_iqtree.fasta` (alignment recoded to the compact `0..k-1` +alphabet of states actually present), states/frequencies restricted to +whichever of the 16 canonical states actually occur, compactly +renumbered — see the `MORPH{N}` risk above. Verified end to end on the +20-genome benchmark: real, non-zero, varied branch lengths (`Total tree +length: 6.773`), converged log-likelihood, `State frequencies:` in +IQ-TREE's own output matching the computed `π` exactly. + +**Rogue-taxon hypothesis, tested and refuted.** The backbone (*Yersinia*, +*Proteus*, *Opitutus*, *Shouchella*, *Wolbachia*...) resolves as a +near-linear comb with ~8 near-zero branch lengths — visible on the real +tree, and independently on a Mash+NJ tree built from an entirely +different signal (whole-genome k-mer distance, no relation to the Sankoff +pipeline), which shows the same comb shape. Hypothesis: `Saccharolobus`/ +`Candidozyma` (near-zero real signal — see below) destabilise the +heuristic tree search enough to also blur resolution elsewhere ("rogue +taxa", a documented phenomenon). Tested directly: reran `iqtree3` after +removing both taxa (`--exclude-genome`, `+ASC` recomputed on the +resulting variable-sites-only alignment since removing taxa turns some +columns invariant) — **still exactly 8 near-zero backbone branches**, +identical comb shape. Refuted for this dataset: the backbone's weak +resolution is a property of the character system's signal at that +divergence depth (matches "Run 3" above), not rogue-taxon interference. + +**Diagnosed why `Saccharolobus`/`Candidozyma` place so poorly**, using a +new diagnostic (`--raw-snp-counts`, `_rawsnp_counts.csv`: one row +per genome pair, `n_snp,n_shared,n_eligible,ratio` — the counts +`--raw-snp-distance`'s ratio-only matrix discards, needed because +`ratio=0.0` from 2 eligible loci and from 2000 look identical in the +ratio alone). `Saccharolobus` has 21,020 non-`∅` sites in the real +alignment (comparable to other taxa) — but **21,017 of them (100.0%) are +private**: no other genome has a non-`∅` state at the same site. Only 3 +sites are shared with any other genome at all (1, 2, and 14 +co-occurring genomes respectively). Real informativeness for placement +tracks shared sites, not raw non-`∅` count — with no other archaeon (or +even archaea-adjacent bacterium) in the dataset, there is nothing to +anchor `Saccharolobus`'s position to, regardless of how much of its own +data exists. `Candidozyma` shows the same pattern, more extreme (almost +all `NA`, the 3 non-`NA` pairs all exactly `0.0`, never `1.0` — itself a +tell: with `n_eligible=1`, the ratio can only be exactly `0` or `1`, so +3-for-3 landing on `0.0` is more than sampling noise alone would predict; +possibly ascertainment bias — see below — or possibly a few genuinely +ultra-conserved loci; not resolved). + +**`--exclude-genome LABEL`** (repeatable, `obikmer distance`) added for +exactly this kind of test: zeroes the excluded genome's row/column in +`RawSnpDistanceOutput` after `raw_snp_distance` runs (a pair with zero +counts is already skipped by `calibrate_p_hat`/`base_pair_tally` — no +`obikindex` traversal change needed) and drops its row from +`SnpAlignment` before any output is written. Deliberately *not* index +surgery (a new, smaller on-disk index) — genome sets to exclude are +expected to change between quick tests, so an in-memory filter is the +right tool, not a new index-rewriting subsystem. Scoped to the sibling-annex +family of computations (`--raw-snp-distance`/`--raw-snp-counts`/`--snp`/ +`--sankoff` and everything it implies) — does not affect the plain +`--metric` distance matrix/NJ/UPGMA path (a different, unrelated +computation on `idx.distance()`, not touched). + +**Caveat surfaced while reusing a stale `π`**: rerunning IQ-TREE on a +genome-reduced alignment while keeping the *original* (20-genome) +model file is inconsistent — `π`'s composition shifts once low-cardinality-0 +columns that were only variable because of the removed taxa drop out. +Measured directly: `π(∅)` `0.9048` (20 genomes) → `0.8869` (18 genomes, +variable sites only) — real (~378k affected cells, from `Saccharolobus`'s +~21k private sites × 18 remaining genomes) but modest in *proportion* +(~2 points) because the alignment was already monomorphic-filtered +before `Saccharolobus` was ever added, so removing it only drops the +subset of columns that were variable *because of* it specifically, not +every column it appears in. + +### `R` via `exp(-cost)` is wrong for a CTMC; cardinality/composition decomposition (open, 2026-08-12) + +**The flaw in `R = exp(-cost)`, precisely.** `cost` (`build_cost_matrix`'s +output) is a *shortest-path closure* over an elementary-edit graph +(Floyd-Warshall) — correct and required for Sankoff parsimony, where +`cost(a,b)` must be a metric. But a CTMC's own matrix exponential +(`exp(Qt)`, computed internally by IQ-TREE) *already* sums over +paths of every length through the elementary rates — that's the whole +mechanism by which a CTMC generates indirect transitions. Feeding it a +pre-summed, multi-hop shortest-path cost and exponentiating that *again* +as if each entry were a direct edge double-applies the "compose multiple +steps" logic once in log-space (Floyd-Warshall, additive) and once more +inside IQ-TREE's own exponential — systematically over-penalising +non-adjacent state pairs (e.g. `∅→{A,C}` priced as two chained edges, +`2×c_ctx`, when it should be one direct lookup). + +**Resolution, in two parts — both estimated directly from the real +alignment, not smoothed through a small parametric formula:** + +1. **Cardinality model**: a 5-state (`0,1,2,3,4`) first-order Markov + chain, estimated from the empirical cardinality co-occurrence table + (pooled across all *included* genome pairs — same + saturated/`NA`-pair exclusion discipline as `calibrate_p_hat`, 14 + saturated + 18 `NA` pairs excluded of 190 in the benchmark; the + result is materially different from the unfiltered version and more + internally consistent, not just "cleaner"). Diagonal included (the + probability of a family *staying* at the same cardinality is part of + the model, not assumed away). Measured, cost `= -ln(observed/expected + under independence)`, on the 20-genome benchmark (`π` here from the + marginal cardinality distribution: `90.481%, 5.261%, 4.183%, 0.067%, + 0.008%` for `c=0..4` respectively — note `c=2` is *not* rare, almost + as common as `c=1`): + + ``` + c=0/c=0: 0.166 c=0/c=1: 0.160 c=0/c=2: 0.157 c=0/c=3: 0.059 c=0/c=4: 0.006 + c=1/c=1: -0.803 (enriched) c=1/c=2: 0.381 c=1/c=3: 1.027 c=1/c=4: 1.794 + c=2/c=2: 4.032 (sharply suppressed) c=2/c=3: 2.620 c=2/c=4: 2.416 + c=3/c=3: -0.254 c=3/c=4: -0.411 (too few observations to trust) + ``` + Two robust findings, both stable under the saturation filter: (a) all + `∅`-involving costs are low and close to each other (`0.006`–`0.166`) + regardless of how many members are gained/lost at once — sharply at + odds with the current model's implicit `×2`/`×3` multi-hop scaling; + (b) `c=2/c=2` (two genomes both showing an ambiguous 2-member state at + the same site) is dramatically under-represented (~1.5–1.8% of the + independence expectation) — a real, robust anomaly, not explained. + +2. **Composition model**: unchanged — `sub_cost` (the existing 6-category + Ts/Tv-biased calibration), estimated **only from unambiguous sites** + (cardinality-1 ↔ cardinality-1 pairs), because composition bias is a + substitution phenomenon and only means something when cardinality is + conserved. Checked whether composition bias also appears in pure + gain/loss events (no substitution involved, so no bias expected a + priori): single-base "gain" (`∅→`single base) is close to uniform + (`24.35/25.92/26.22/23.51%` for A/C/G/T) — consistent with "no + mutational mechanism, so no bias" as expected. Two-base "gain" + (`∅→`2-member state) is *not* uniform even after correcting for the + real (non-25/25/25/25) marginal base frequencies: `{A,G}`/`{C,T}` + (the transition-linked pairs) mildly enriched (`obs/exp` `1.12`/ + `1.16`), `{C,G}` sharply suppressed (`obs/exp 0.563`, expected to be + the *most* common pair under independence since C and G are + individually the two most frequent bases, observed the least) — an + unexplained anomaly, deliberately **not** built into the model (no + mechanistic story for why gain/loss would carry a `{C,G}`-specific + bias), left as an open puzzle rather than fit. + +**Composing the two into a full pairwise cost — the actual replacement +for both the elementary-edge graph and its Floyd-Warshall closure.** For +any two states `A`, `B` (not just the "clean" same-cardinality or +pure-subset cases the current graph handles directly): let `shared = A∩B` +(free), `lost = A\B`, `gained = B\A`. Pair off `k = min(|lost|,|gained|)` +elements between `lost` and `gained` as substitution events, choosing the +pairing that minimises total `sub_cost` (a trivial assignment problem — +at most 4 elements per side). The `|lost|-k` (or `|gained|-k`) leftover, +unpaired elements are a *pure* cardinality change, `|A|→|B|`, costed by +**one direct lookup** in the cardinality model above — no chaining, no +Floyd-Warshall. Example: `{A,C}→{G}` (cardinality 2→1, no shared base): +pair 1 substitution (cheaper of `A→G`, `C→G`), 1 base left over unpaired +→ cost `= sub_cost(chosen pair) + cardinality_cost(2→1)`. This is a +direct, closed-form cost for *every* pair of the 16 states, replacing +`build_cost_matrix`'s graph-plus-shortest-path construction outright — +and specifically fixes the CTMC double-counting problem, since every +entry is now a single decomposed lookup, never a sum of chained edges. + +**Reframed as a likelihood (product of probabilities), not a cost (sum of +`-ln`s) — same content, but forces the diagonals (no-change cases) to be +kept rather than implicitly dropped.** Both sub-models are proper +transition *probability* matrices, diagonal included — `P_cardinality` +includes "stay at the same cardinality", `P_composition` includes "stay +the same base" (e.g. `P(G→G)`, not just `sub_cost`'s off-diagonal +entries). Composing: + +``` +P(A→B) = P_cardinality(|A|→|B|) + × ∏_{x ∈ A∩B} P_composition(x→x) (shared bases: "stayed") + × ∏_{(x,y) paired} P_composition(x→y) (parsimony-paired substitutions) +``` + +with the unpaired leftover `lost`/`gained` elements (if `|lost|≠|gained|`) +contributing *nothing further* beyond the `P_cardinality` term already +counted — consistent with the finding above that pure gain/loss carries +no separate composition bias worth modelling. Worked example, +`{A,C}→{A,G}` (`shared={A}`, one paired substitution `C→G`, nothing left +over): `P = P_cardinality(2→2) · P_composition(A→A) · P_composition(C→G)`. + +The row-wise product of these two independently-calibrated models isn't +guaranteed to already sum to exactly `1` across all `B` for a fixed `A` +(the two aren't perfectly independent in reality) — so each row of the +resulting 16×16 matrix is renormalised (divided by its own sum) after +composition, **not** the matrix as a whole (which would produce a joint +distribution over `(A,B)` pairs, the wrong object — a transition matrix +needs each row, "given I start in `A`", to be a valid distribution over +where I end up). + +This now covers every pair of the 16 states with no unhandled case +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). ## Heterozygosity, ploidy, and consensus-assembly inputs diff --git a/src/obikmer/src/cmd/distance/mod.rs b/src/obikmer/src/cmd/distance/mod.rs index 3dfe08d3..c3b1035e 100644 --- a/src/obikmer/src/cmd/distance/mod.rs +++ b/src/obikmer/src/cmd/distance/mod.rs @@ -85,6 +85,24 @@ pub struct DistanceArgs { #[arg(long)] pub sibling_annex: bool, + /// 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`. + #[arg(long = "exclude-genome", value_name = "LABEL")] + pub exclude_genome: Vec, + /// Tally the sibling-count distribution (CSV) of an already-built annex /// (run with `--sibling-annex` first, in this invocation or an earlier /// one). A separate, occasional diagnostic pass — not run every time the @@ -211,6 +229,52 @@ pub fn run(args: DistanceArgs) { let labels: Vec = idx.meta().genomes.iter().map(|g| g.label.clone()).collect(); let n = labels.len(); + // ── 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 + // reach TNT/PhyG/IQ-TREE as a real (empty) taxon. + let exclude_mask: Vec = { + let mut mask = vec![false; n]; + for label in &args.exclude_genome { + match labels.iter().position(|l| l == label) { + Some(i) => mask[i] = true, + None => { + eprintln!("error: --exclude-genome {label:?} does not match any genome in this index"); + std::process::exit(1); + } + } + } + mask + }; + let zero_excluded_pairs = |result: &mut RawSnpDistanceOutput| { + for i in 0..n { + if !exclude_mask[i] { + continue; + } + for j in 0..n { + result.snp[[i, j]] = 0; + result.snp[[j, i]] = 0; + result.shared[[i, j]] = 0; + result.shared[[j, i]] = 0; + } + } + }; + let drop_excluded = |alignment: SnpAlignment| -> (SnpAlignment, Vec) { + let sequences = alignment.sequences.into_iter().enumerate() + .filter(|(i, _)| !exclude_mask[*i]) + .map(|(_, seq)| seq) + .collect(); + let kept_labels = labels.iter().enumerate() + .filter(|(i, _)| !exclude_mask[*i]) + .map(|(_, l)| l.clone()) + .collect(); + (SnpAlignment { sequences }, kept_labels) + }; + // ── Sibling-count/minorant annex (independent of the distance metric) ── // Construction (`--sibling-annex`) and stats (`--sibling-stats`) are // deliberately decoupled: the annex is meant to be (re)built routinely, @@ -237,17 +301,19 @@ pub fn run(args: DistanceArgs) { write_sibling_stats_csv(&stats, &labels, &args.output); } if args.raw_snp_distance { - let result = idx.raw_snp_distance().unwrap_or_else(|e| { + let mut result = idx.raw_snp_distance().unwrap_or_else(|e| { eprintln!("error computing raw SNP distance: {e}"); std::process::exit(1); }); + zero_excluded_pairs(&mut result); write_raw_snp_distance_csv(&result, &labels, &args.output); } if args.raw_snp_counts { - let result = idx.raw_snp_distance().unwrap_or_else(|e| { + let mut result = idx.raw_snp_distance().unwrap_or_else(|e| { eprintln!("error computing raw SNP distance: {e}"); std::process::exit(1); }); + zero_excluded_pairs(&mut result); write_raw_snp_counts_csv(&result, &labels, &args.output); } if args.snp { @@ -255,13 +321,15 @@ pub fn run(args: DistanceArgs) { eprintln!("error computing SNP pseudo-alignment: {e}"); std::process::exit(1); }); - write_snp_fasta(&alignment, &labels, &args.output); + let (alignment, kept_labels) = drop_excluded(alignment); + write_snp_fasta(&alignment, &kept_labels, &args.output); } if args.sankoff || args.tnt || args.phyg || args.iqtree { - let raw = idx.raw_snp_distance().unwrap_or_else(|e| { + let mut raw = idx.raw_snp_distance().unwrap_or_else(|e| { eprintln!("error computing raw SNP distance: {e}"); 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; @@ -282,16 +350,17 @@ pub fn run(args: DistanceArgs) { eprintln!("error computing SNP pseudo-alignment: {e}"); std::process::exit(1); }); - write_sankoff_alignment_fasta(&alignment, &labels, &args.output); + let (alignment, kept_labels) = drop_excluded(alignment); + write_sankoff_alignment_fasta(&alignment, &kept_labels, &args.output); if args.tnt { - write_sankoff_tnt(&matrix, &alignment, &labels, &args.output, args.sankoff_cost_scale); + write_sankoff_tnt(&matrix, &alignment, &kept_labels, &args.output, args.sankoff_cost_scale); } if args.phyg { write_sankoff_phyg(&matrix, &args.output, args.sankoff_cost_scale); } if args.iqtree { - write_iqtree(&matrix, &alignment, &labels, &args.output); + write_iqtree(&matrix, &alignment, &kept_labels, &args.output); } } -- 2.54.0 From 0c86ea03850d6ddc621350a310a3d5199c4d5d8a Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 13 Aug 2026 09:29:44 +0200 Subject: [PATCH 05/46] 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 { -- 2.54.0 From 6acafa7f2c5e6fd5a33a0cc4b1dc2f36979f44c9 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 13 Aug 2026 10:01:29 +0200 Subject: [PATCH 06/46] docs: add obikmer user guide and MkDocs build configuration Introduces a comprehensive documentation set covering theoretical foundations, CLI usage, installation, and system architecture. Adds MkDocs configuration and Makefile targets to generate, serve with live reload, and clean the documentation site. Includes citation styles and bibliography files for academic references. --- Makefile | 21 +- UserDocMD/architecture.md | 33 +++ UserDocMD/ecology-letters.csl | 230 +++++++++++++++++++ UserDocMD/formats/index_layout.md | 56 +++++ UserDocMD/index.md | 68 ++++++ UserDocMD/installation.md | 84 +++++++ UserDocMD/references.bib | 261 ++++++++++++++++++++++ UserDocMD/theory/encoding.md | 28 +++ UserDocMD/theory/entropy_filter.md | 36 +++ UserDocMD/theory/indexing_architecture.md | 32 +++ UserDocMD/theory/kmers_and_superkmers.md | 24 ++ UserDocMD/theory/minimizer_selection.md | 42 ++++ UserDocMD/usage/annotate.md | 25 +++ UserDocMD/usage/distance.md | 97 ++++++++ UserDocMD/usage/dump.md | 25 +++ UserDocMD/usage/estimate.md | 18 ++ UserDocMD/usage/filter.md | 47 ++++ UserDocMD/usage/index_command.md | 50 +++++ UserDocMD/usage/merge.md | 29 +++ UserDocMD/usage/pack.md | 15 ++ UserDocMD/usage/predicates.md | 46 ++++ UserDocMD/usage/query.md | 38 ++++ UserDocMD/usage/reindex.md | 25 +++ UserDocMD/usage/select.md | 35 +++ UserDocMD/usage/superkmer.md | 27 +++ UserDocMD/usage/unitig.md | 19 ++ UserDocMD/usage/utils.md | 26 +++ mkdocs-user.yml | 61 +++++ 28 files changed, 1497 insertions(+), 1 deletion(-) create mode 100644 UserDocMD/architecture.md create mode 100644 UserDocMD/ecology-letters.csl create mode 100644 UserDocMD/formats/index_layout.md create mode 100644 UserDocMD/index.md create mode 100644 UserDocMD/installation.md create mode 100644 UserDocMD/references.bib create mode 100644 UserDocMD/theory/encoding.md create mode 100644 UserDocMD/theory/entropy_filter.md create mode 100644 UserDocMD/theory/indexing_architecture.md create mode 100644 UserDocMD/theory/kmers_and_superkmers.md create mode 100644 UserDocMD/theory/minimizer_selection.md create mode 100644 UserDocMD/usage/annotate.md create mode 100644 UserDocMD/usage/distance.md create mode 100644 UserDocMD/usage/dump.md create mode 100644 UserDocMD/usage/estimate.md create mode 100644 UserDocMD/usage/filter.md create mode 100644 UserDocMD/usage/index_command.md create mode 100644 UserDocMD/usage/merge.md create mode 100644 UserDocMD/usage/pack.md create mode 100644 UserDocMD/usage/predicates.md create mode 100644 UserDocMD/usage/query.md create mode 100644 UserDocMD/usage/reindex.md create mode 100644 UserDocMD/usage/select.md create mode 100644 UserDocMD/usage/superkmer.md create mode 100644 UserDocMD/usage/unitig.md create mode 100644 UserDocMD/usage/utils.md create mode 100644 mkdocs-user.yml diff --git a/Makefile b/Makefile index dedc3bed..5889db68 100644 --- a/Makefile +++ b/Makefile @@ -10,6 +10,11 @@ DOC_FILE := mkdocs.yml DOC_SITE := doc DOC_PORT := 8001 +DOC_USER_DIR := UserDocMD +DOC_USER_FILE := mkdocs-user.yml +DOC_USER_SITE := doc-user +DOC_USER_PORT := 8002 + # ── virtualenv ──────────────────────────────────────────────────────────────── $(VENV)/bin/activate: @@ -60,8 +65,22 @@ doc-serve: $(MKDOCS) clean-doc: rm -rf $(DOC_SITE)/ +.PHONY: doc-user +doc-user: $(MKDOCS) + $(MKDOCS) build -f $(DOC_USER_FILE) + +.PHONY: doc-user-serve +doc-user-serve: $(MKDOCS) + $(MKDOCS) serve -f $(DOC_USER_FILE) \ + --dev-addr=127.0.0.1:$(DOC_USER_PORT) \ + --livereload + +.PHONY: clean-doc-user +clean-doc-user: + rm -rf $(DOC_USER_SITE)/ + .PHONY: clean -clean: clean-doc +clean: clean-doc clean-doc-user rm -rf $(VENV) # ── release ─────────────────────────────────────────────────────────────────── diff --git a/UserDocMD/architecture.md b/UserDocMD/architecture.md new file mode 100644 index 00000000..87bf9b91 --- /dev/null +++ b/UserDocMD/architecture.md @@ -0,0 +1,33 @@ +# Architecture notes for advanced use + +This page describes execution-level behavior relevant to sizing and running `obikmer` on large datasets or multi-socket machines. It complements the [index format](formats/index_layout.md) and [theory](theory/indexing_architecture.md) pages. + +## Sequence invariant + +Every input sequence is treated purely as a compact representation of a set of overlapping kmers: + +- Only the `A`/`C`/`G`/`T` alphabet (case-insensitive) is recognized; a sequence is cut at any other character (including IUPAC ambiguity codes), so runs containing them are not represented in the index. +- Sequences are internally processed in chunks of at most 256 nucleotides; a chunk shorter than k is dropped. This is invisible to the user beyond the ACGT-only, minimum-length-k constraints above. +- Kmers are always handled in canonical form (see [DNA encoding](theory/encoding.md)), so the tool is strand-agnostic throughout: a kmer and its reverse complement are always the same entry. + +## Index dimensioning + +An index directory is organized as `KmerIndex → partitions → layers`, with a canonical kmer belonging to exactly one (partition, layer) pair. This is what makes set operations (merge, filter, distance) parallel and coordination-free across partitions. + +- **Partition count** (`-p`/`--partitions`, rounded up to a power of 2) is the main dimensioning knob: more partitions means more independent parallel units and a smaller working set per partition, at the cost of more open files during construction. +- **Layers** accumulate as an index grows through successive merges; per-partition query cost grows with the number of layers (worst case linear, expected constant since most kmer lookups resolve in the first layer they could plausibly be in). +- Genome columns (count or presence data) are kept at a consistent width across every layer and partition after a merge, which is what allows whole-index aggregate distances (Jaccard, Bray-Curtis, Euclidean, Hellinger, …) to be computed as a two-pass cascade (local partial sums per partition, then a global combination) with no double counting. + +## Parallel execution and NUMA awareness + +Partition-level work (index construction, `merge`, `filter`, `reindex`, `select`, `distance`'s sibling-annex/Sankoff computations) is dispatched by a partition runner that adapts to the machine's memory topology, detected automatically at startup via hwloc: + +- On a multi-socket / multi-NUMA-node machine, one thread pool is pinned per NUMA node, and each partition is processed entirely by threads pinned to one node — keeping the memory a partition touches local to that node's DRAM. This matters because touching kmer data across NUMA nodes without pinning can degrade throughput by an order of magnitude or more on large multi-socket machines. +- On a single-socket machine, Apple Silicon, or if hwloc cannot report NUMA topology, all cores are treated as one node with no pinning and negligible overhead — this is the default behavior on macOS. +- Within a node, the number of active worker threads ramps up progressively rather than being fixed up front: it starts conservatively and grows in steps, but only as long as measured CPU efficiency or disk I/O throughput keeps improving. If neither improves after a growth step, the runner stops adding workers — avoiding oversubscription on stages that are memory-bandwidth-bound rather than CPU- or I/O-bound. Ramp speed scales with the number of cores per node, so a single-node machine ramps just as fast as a large multi-node one. + +No CLI flag controls this directly; it is fully automatic at runtime. NUMA-aware pinning can be compiled out (Cargo feature `numa`, on by default), in which case a plain global thread pool is used instead. + +## Kmer filtering (`filter`) + +[`filter`](usage/filter.md) evaluates predicates against the genome metadata matrix directly whenever every active filter can be expressed as a column-level test (e.g. "any outgroup column non-zero"), producing a per-slot keep/drop decision without touching kmer sequence data at all. If any active filter cannot be expressed this way, evaluation falls back to a per-kmer, row-level check. Either way, the result is always written as a single, freshly compacted layer (`unitigs.bin` and the MPHF are rebuilt from the surviving kmers), never as an additional layer on top of the source index. diff --git a/UserDocMD/ecology-letters.csl b/UserDocMD/ecology-letters.csl new file mode 100644 index 00000000..6a96c856 --- /dev/null +++ b/UserDocMD/ecology-letters.csl @@ -0,0 +1,230 @@ + + diff --git a/UserDocMD/formats/index_layout.md b/UserDocMD/formats/index_layout.md new file mode 100644 index 00000000..09633b09 --- /dev/null +++ b/UserDocMD/formats/index_layout.md @@ -0,0 +1,56 @@ +# Index construction and on-disk layout + +## Construction pipeline + +Building an index ([`index`](../usage/index_command.md)) proceeds through a fixed sequence of phases, each operating independently per partition (see [Partitioning and indexing architecture](../theory/indexing_architecture.md)): + +1. **Scatter.** A single streaming pass over the input. Each sequence fragment is cut at non-ACGT bases, passed through the low-complexity entropy filter (see [Low-complexity kmer filter](../theory/entropy_filter.md)), and any resulting segment shorter than k is dropped. Surviving segments are decomposed into super-kmers, canonicalized, and routed by `hash(minimizer) mod n_partitions` into one file per partition. +2. **Dereplication.** Within each partition, identical super-kmer sequences are merged and their occurrence counts summed. This count is per super-kmer, not per kmer — a kmer's true abundance is the sum of the counts of every super-kmer containing it. +3. **Exact counting.** Every kmer in every dereplicated super-kmer is enumerated and its exact total count computed. A per-genome kmer frequency spectrum is produced at this stage. +4. **Quorum filtering.** Kmers outside the `--min-abundance`/`--max-abundance` range are dropped, and super-kmers are recompacted around the surviving kmer set. +5. **Local assembly.** The surviving kmers of each partition are assembled into unitigs — maximal non-branching runs of a local de Bruijn graph — such that every kmer appears exactly once, at one (unitig, offset) location. +6. **MPHF and evidence construction.** A minimal perfect hash function is built over the canonical kmers of each partition, together with the evidence structure needed to verify that a queried kmer was genuinely indexed (see below). Per-genome counts or presence bits are recorded alongside if requested. + +Phases 1–5 are independent per partition and run in parallel; phase 6 finalizes each partition once its kmer set is fixed. + +## Minimal perfect hash function (MPHF) + +Each partition's surviving kmers are mapped to a dense range of integer slots by a minimal perfect hash function: no collisions, near-optimal space (a few bits per key), O(1) lookup. Because an MPHF maps *any* input to some slot — including kmers that were never indexed — a lookup alone cannot distinguish a genuinely indexed kmer from an arbitrary one; every lookup is followed by an evidence check. + +## Evidence: exact vs. approximate + +Two verification modes are available, selected at build time (`index --approx`) and convertible afterwards ([`reindex`](../usage/reindex.md)): + +- **Exact** (default): the hashed slot stores a pointer back into the partition's unitig data. At query time the kmer is reconstructed from that location and compared directly to the query. Zero false positives, at the cost of one extra random read per lookup. +- **Approximate** (`--approx`): the slot stores a short fingerprint (`--evidence-bits` bits) instead of a pointer; verification is a single fingerprint comparison. This trades a small, bounded false-positive rate ($1/2^b$ per kmer, reduced further to about $1/2^{b \cdot z}$ for a read requiring $z$ consecutive matching kmers via the `-z`/`--findere-z` parameter) for lower memory and disk usage, since no reconstruction index is needed. See [`estimate`](../usage/estimate.md) to explore this trade-off before building. + +## On-disk layout + +``` +/ + index.meta global configuration (k, minimizer size, partition count, + evidence mode, whether counts are stored) and genome list/metadata + scatter.done / count.done / index.done build-progress sentinels + spectrums/