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