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.
This commit is contained in:
Eric Coissac
2026-08-12 20:05:10 +02:00
parent 14aa82521d
commit 55d7fa2067
17 changed files with 1494 additions and 180 deletions
+123 -49
View File
@@ -379,84 +379,156 @@ per pair loses. Not yet run.
### A concrete Sankoff cost matrix for the 16-state alphabet ### 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 **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`). `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 Elements present in both cost nothing. Pair up to `min(a,b)` of the
remaining elements as **substitutions** (cheaper than treating them as an 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 choice satisfies); whatever is left over after pairing is a pure
**gain/loss**: **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 `{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 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 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 same-size sets share nothing); `{A,C,G}->{A,C,T}` = 1 (`G`/`T` is the only
mismatched pair, `A,C` shared). mismatched pair, `A,C` shared).
**`∅` is a flat-cost special case, not a instance of the general formula.** **`∅` is *not* a flat-cost special case — a corrected position, reversing
Applying the formula naively to e.g. `{A,C,G} -> ∅` would charge `3*c_gl` an earlier draft of this design.** An earlier version charged
(three independent losses). Reject that: total context disappearance `cost(X, ∅) = c_ctx` flat, independent of `|X|`, on the grounds that every
(flank-breaking, or true structural loss) is plausibly **one** event, not member of a family shares identical flanking sequence, so one mutation
`|X|` of them, so: 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 cost(∅, ∅) = 0
``` ```
`c_ctx` should not be guessed — it is exactly the value already derived for **Computing `c_ctx`.** Not guessed — built from the value already derived
the `>1` bucket in "Context, detectability, and a 3-way ordinal distance per for the `>1` bucket in "Context, detectability, and a 3-way ordinal
pair" above (`(2m*p_hat)/(1-(1-p_hat)^(2m)) + p_hat`), reused rather than distance per pair" above:
invented.
```
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 **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 closed-form formula directly.** Build a graph on all 16 states, `∅`
types — substitution edges between same-cardinality sets differing by one included, with two edge types — substitution edges between same-cardinality
element (weight `c_sub`, or `c_ts`/`c_tv` if split further below), and sets differing by one element (weight `c_sub`, or `c_ts`/`c_tv` if split
gain/loss edges between sets whose cardinality differs by one (weight further below), and gain/loss edges between sets whose cardinality differs
`c_gl`) — then define `cost(X,Y)` as shortest-path distance in that graph, by one (weight `c_ctx`, per the unification above; `∅` connects to each
precomputed once (16 nodes, trivial) into a dense 16x16 matrix before singleton state this way, being a strict subset of it) — then define
feeding it to Sankoff/TNT. Verified equivalent to the closed-form formula `cost(X,Y)` as shortest-path distance in that graph, uniformly for every
above in the uniform-cost case (checked by hand on `{A}->{C,G,T}`: both give `X,Y` including `∅`, precomputed once (16 nodes, trivial) into a dense
`c_sub + 2*c_gl`). The graph construction is not just a reformulation for 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 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)` costs stop being uniform (next point) — the closed-form's `min(a,b)`
counting silently assumes *any* pairing costs the same, which breaks the counting silently assumes *any* pairing costs the same, which breaks the
moment `c_sub` depends on which two bases are involved. moment `c_sub` depends on which two bases are involved.
**Transition/transversion refinement.** Split `c_sub` into `c_ts` (A<->G or **Substitution refinement — implemented as the full 6-category symmetric
C<->T) and `c_tv` (the other four pairs) — already well-defined in canonical matrix, not just Ts/Tv.** `c_sub` was originally going to split into just
space (see "Canonical invariance" above: a transition maps to a transition, `c_ts` (A<->G or C<->T) and `c_tv` (the other four pairs) — already
a transversion to a transversion, regardless of orientation). This turns 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 "pick `min(a,b)` substitution pairs" into a genuine (tiny, <=4 elements per
side, trivially enumerable) minimum-cost bipartite matching problem instead side, trivially enumerable) minimum-cost bipartite matching problem instead
of a plain count — the state-graph shortest-path construction handles this of a plain count — the state-graph shortest-path construction handles this
automatically, no separate logic needed. Calibration is not a guess either: automatically, no separate logic needed. Calibration is not a guess
the "Sufficient statistic: 4x4 base-pair tally" section below already plans either: `KmerIndex::base_pair_tally` collects the pooled `(centre_i,
to collect the joint `(centre_i, centre_j)` distribution over resolved (`0` centre_j)` distribution over resolved (SNP) sites — restricted to the same
or `1`) sites — that tally directly gives the empirical Ts/Tv ratio, from ratio-ceiling-included genome pairs as `p_hat`'s own calibration, for the
which `c_ts`/`c_tv` follow (e.g. `cost ∝ -log(observed rate)`, the standard same saturation-exclusion reason — and `substitution_costs_from_tally`
generalised-parsimony step-weighting heuristic), reusing a statistic already turns that into `cost(a,b) = -ln(rate(a,b))`, normalised so the most
planned rather than adding a new one. frequent category costs `1.0` (the standard generalised-parsimony
step-weighting heuristic, generalised from 2 categories to 6). A second
**`gamma`/`mu` (i.e. `c_gl`/`c_sub`) sensitivity sweep before calibration.** full pass over the annex is required for this (see `base_pair_tally`'s own
No strong prior on whether gain/loss events should cost more or less than a docs for why it can't share `raw_snp_distance`'s single pass) — both now
point substitution — duplication/deletion rates are not a priori equal to share one traversal helper, `scan_family_pairs`, rather than duplicating
point-mutation rates, but the direction isn't obvious, and setting it too the per-family reconciliation logic.
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.
**Caveat carried over from the `D_F = min(a,b)` rejection earlier:** this **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 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 **Feasibility confirmed:** TNT's `costs` command accepts custom step
matrices for multistate characters, so this whole construction (16x16 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 tunable parameters) is directly usable there — no new tooling required
before testing it. 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 `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 the unweighted mean of the pairwise `snp/(snp+shared)` ratios already
present in `rawsnp.csv`, restricted to the relevant taxon subset. Used 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/ with the observed pairwise ratios for genomes within Enterobacteriaceae/
Eubacteria, which cluster around 1.5-2 and barely move as the taxon set 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 widens — the context signal is stable, not sensitive to which subset is
+1
View File
@@ -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);
+27
View File
@@ -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
+136 -18
View File
@@ -48,6 +48,15 @@ dependencies = [
"as-slice", "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]] [[package]]
name = "allocator-api2" name = "allocator-api2"
version = "0.2.21" version = "0.2.21"
@@ -497,6 +506,18 @@ dependencies = [
"windows-sys 0.59.0", "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]] [[package]]
name = "const-oid" name = "const-oid"
version = "0.10.2" version = "0.10.2"
@@ -820,6 +841,26 @@ dependencies = [
"syn", "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]] [[package]]
name = "equivalent" name = "equivalent"
version = "1.0.2" version = "1.0.2"
@@ -907,6 +948,16 @@ dependencies = [
"percent-encoding", "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]] [[package]]
name = "fs_at" name = "fs_at"
version = "0.2.1" version = "0.2.1"
@@ -1228,13 +1279,26 @@ version = "0.17.11"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
dependencies = [ dependencies = [
"console", "console 0.15.11",
"number_prefix", "number_prefix",
"portable-atomic", "portable-atomic",
"unicode-width", "unicode-width",
"web-time", "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]] [[package]]
name = "infer" name = "infer"
version = "0.19.0" version = "0.19.0"
@@ -1326,12 +1390,9 @@ dependencies = [
[[package]] [[package]]
name = "kodama" name = "kodama"
version = "0.2.3" version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7a44f3a71a44fbf49ce38152db7dc9adf959d4fe5c29344cd1858bdbda8d9091" checksum = "f868feceed4703c842925c14232667c5f29c8059d0354154b2d06ec3bf9daf82"
dependencies = [
"num-traits",
]
[[package]] [[package]]
name = "lazy_static" name = "lazy_static"
@@ -1365,6 +1426,12 @@ dependencies = [
"pkg-config", "pkg-config",
] ]
[[package]]
name = "linux-raw-sys"
version = "0.4.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab"
[[package]] [[package]]
name = "linux-raw-sys" name = "linux-raw-sys"
version = "0.12.1" version = "0.12.1"
@@ -1695,7 +1762,7 @@ version = "0.1.0"
dependencies = [ dependencies = [
"crossbeam-channel", "crossbeam-channel",
"hwlocality", "hwlocality",
"indicatif", "indicatif 0.17.11",
"ndarray", "ndarray",
"obicompactvec", "obicompactvec",
"obikpartitionner", "obikpartitionner",
@@ -1720,7 +1787,7 @@ version = "1.1.44"
dependencies = [ dependencies = [
"clap", "clap",
"csv", "csv",
"indicatif", "indicatif 0.18.6",
"kodama", "kodama",
"obidebruinj", "obidebruinj",
"obifastwrite", "obifastwrite",
@@ -1737,7 +1804,9 @@ dependencies = [
"obitaxonomy", "obitaxonomy",
"pprof", "pprof",
"rayon", "rayon",
"serde",
"serde_json", "serde_json",
"serde_yaml",
"speedytree", "speedytree",
"tracing", "tracing",
"tracing-subscriber", "tracing-subscriber",
@@ -1749,7 +1818,7 @@ version = "0.1.0"
dependencies = [ dependencies = [
"cacheline-ef", "cacheline-ef",
"epserde", "epserde",
"indicatif", "indicatif 0.17.11",
"memmap2", "memmap2",
"niffler 3.0.0", "niffler 3.0.0",
"obicompactvec", "obicompactvec",
@@ -1853,7 +1922,7 @@ dependencies = [
"memmap2", "memmap2",
"niffler 3.0.0", "niffler 3.0.0",
"obikseq", "obikseq",
"rustix", "rustix 1.1.4",
"serde", "serde",
"serde_json", "serde_json",
"tempfile", "tempfile",
@@ -1863,7 +1932,8 @@ dependencies = [
name = "obisys" name = "obisys"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"indicatif", "fs4",
"indicatif 0.17.11",
"libc", "libc",
"sysinfo", "sysinfo",
"tracing", "tracing",
@@ -2017,10 +2087,11 @@ dependencies = [
[[package]] [[package]]
name = "pprof" name = "pprof"
version = "0.13.0" version = "0.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ef5c97c51bd34c7e742402e216abdeb44d415fbe6ae41d56b114723e953711cb" checksum = "38a01da47675efa7673b032bf8efd8214f1917d89685e07e395ab125ea42b187"
dependencies = [ dependencies = [
"aligned-vec",
"backtrace", "backtrace",
"cfg-if", "cfg-if",
"findshlibs", "findshlibs",
@@ -2028,15 +2099,15 @@ dependencies = [
"log", "log",
"nix 0.26.4", "nix 0.26.4",
"once_cell", "once_cell",
"parking_lot",
"prost", "prost",
"prost-build", "prost-build",
"prost-derive", "prost-derive",
"sha2", "sha2",
"smallvec", "smallvec",
"spin",
"symbolic-demangle", "symbolic-demangle",
"tempfile", "tempfile",
"thiserror 1.0.69", "thiserror 2.0.18",
] ]
[[package]] [[package]]
@@ -2386,6 +2457,19 @@ dependencies = [
"semver", "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]] [[package]]
name = "rustix" name = "rustix"
version = "1.1.4" version = "1.1.4"
@@ -2395,7 +2479,7 @@ dependencies = [
"bitflags 2.11.1", "bitflags 2.11.1",
"errno", "errno",
"libc", "libc",
"linux-raw-sys", "linux-raw-sys 0.12.1",
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
@@ -2521,6 +2605,19 @@ dependencies = [
"zmij", "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]] [[package]]
name = "sha2" name = "sha2"
version = "0.10.9" version = "0.10.9"
@@ -2588,6 +2685,15 @@ dependencies = [
"rb_tree", "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]] [[package]]
name = "stable_deref_trait" name = "stable_deref_trait"
version = "1.2.1" version = "1.2.1"
@@ -2736,7 +2842,7 @@ dependencies = [
"fastrand", "fastrand",
"getrandom 0.3.4", "getrandom 0.3.4",
"once_cell", "once_cell",
"rustix", "rustix 1.1.4",
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
@@ -2920,6 +3026,18 @@ version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" 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]] [[package]]
name = "untrusted" name = "untrusted"
version = "0.9.0" version = "0.9.0"
@@ -3421,7 +3539,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156" checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156"
dependencies = [ dependencies = [
"libc", "libc",
"rustix", "rustix 1.1.4",
] ]
[[package]] [[package]]
+6 -1
View File
@@ -8,6 +8,7 @@ mod merge;
mod numa; mod numa;
mod rebuild; mod rebuild;
mod reindex; mod reindex;
mod sankoff;
mod select; mod select;
mod siblings; mod siblings;
mod stats; mod stats;
@@ -19,4 +20,8 @@ pub use merge::MergeMode;
pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME}; pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME};
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED}; pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
pub use stats::IndexBitsPerKmer; 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,
};
+402
View File
@@ -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::<u64>::zeros((3, 3));
let mut shared = Array2::<u64>::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);
}
}
+207 -99
View File
@@ -14,18 +14,25 @@
//! mask by callers, not stored (see `FamilyMask` and //! mask by callers, not stored (see `FamilyMask` and
//! [`sibling_annex_stats`](KmerIndex::sibling_annex_stats) below). //! [`sibling_annex_stats`](KmerIndex::sibling_annex_stats) below).
//! //!
//! Per layer, an `obipipeline` `Flat` stage (throttled — see //! Per layer, an `obipipeline` batch transform (throttled — see
//! `obipipeline::throttle`) generates each k-mer's 3 central variants, //! `obipipeline::throttle`) generates a whole batch's central variants at
//! interleaved across many in-flight k-mers by the scheduler's shared //! once (`BATCH_SIZE` source k-mers in, that batch's variants out as one
//! worker pool rather than processed on a single thread. The actual //! pipeline message), interleaved across many in-flight batches by the
//! cross-partition lookup, though, reuses //! scheduler's shared worker pool rather than processed on a single
//! `KmerPartition::query_partition_with` — the same partition-batching //! thread. The actual cross-partition lookup reuses a `PartitionCache` of
//! mechanism `obikmer query` already uses (open a partition's files once, //! every partition's already-open MPHF layers, built once for the whole
//! answer a whole batch of queries against it) — rather than a per-item //! `build_sibling_annex` run, rather than reopening files per lookup or
//! pipeline stage: an earlier per-item design reopened/re-mmap'd every //! per source layer. Two earlier, coarser-grained designs were tried and
//! target partition's files on every single lookup, which was fine at //! measured (not guessed) to be worse, in order: (1) reopening/re-mmap'ing
//! toy scale but manifested as ~90% system time against a real index. //! every target partition's files on every single lookup — fine at toy
//! See `docmd/theory/evolutionary_distances.md`, Step 2b, "Mechanism". //! 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::path::Path;
use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::atomic::{AtomicU8, Ordering};
@@ -107,34 +114,30 @@ fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize {
// ── obipipeline data types ───────────────────────────────────────────────── // ── obipipeline data types ─────────────────────────────────────────────────
/// One distinct k-mer of the layer currently being processed, at its local /// A batch of this layer's distinct k-mers (local MPHF slot + k-mer), the
/// MPHF slot — the pipeline's source item. Carries a throttle slot (shared, /// pipeline's source item — batched, not one k-mer per item, so that
/// `Arc`-wrapped so it can be cloned into each of the (up to 3) items this /// pipeline messages and their synchronisation cost stay amortised over
/// one fans out into via the `Flat` stage) that is only released once every /// thousands of lookups (see `build_layer_sibling_annex`'s comment on
/// one of those descendants has been fully processed — see the module docs' /// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved
/// "Throttling" note for why this is required, not optional, once a `Flat` /// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform,
/// stage is in the pipeline. /// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
struct SourceItem { struct SourceBatch {
slot: usize, items: Vec<(usize, CanonicalKmer)>,
kmer: CanonicalKmer, _permit: ThrottleGuard,
_permit: Arc<ThrottleGuard>,
} }
/// One of a source k-mer's (up to 3) central-substitution variants, already /// One batch's worth of central-substitution variants (up to 3 per source
/// routed to its destination partition and carrying its own central base /// k-mer), each already routed to its destination partition and carrying
/// (0=A/1=C/2=G/3=T) — the mask bit it will set on a hit. Carries a clone of /// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a
/// the source item's throttle permit. /// hit. `(dest_partition, variant, source_slot, base)` per entry.
struct VariantQuery { struct VariantBatch {
source_slot: usize, items: Vec<(usize, CanonicalKmer, usize, u8)>,
dest_partition: usize, _permit: ThrottleGuard,
variant: CanonicalKmer,
base: u8,
_permit: Arc<ThrottleGuard>,
} }
enum SibData { enum SibData {
Item(SourceItem), Batch(SourceBatch),
Query(VariantQuery), Variants(VariantBatch),
} }
/// Every partition's already-open MPHF layers, built **once** for the whole /// 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); mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
} }
// ── obipipeline: Flat stage generates variants only — the actual // ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one —
// cross-partition lookup reuses `KmerPartition::query_partition_with` // the actual cross-partition lookup reuses
// (the same batching mechanism `obikmer query` already uses: open a // `KmerPartition::query_partition_with` (the same batching mechanism
// partition's files once, answer a whole batch of queries against // `obikmer query` already uses: open a partition's files once,
// it) instead of one lookup per pipeline item. A per-item lookup // answer a whole batch of queries against it) instead of one lookup
// (tried first) reopened/re-mmap'd every target partition's files on // per pipeline item. A per-item lookup (tried first) reopened/
// every single variant — fine at the scale of a handful of test // re-mmap'd every target partition's files on every single variant
// k-mers, but with billions of lookups against a real index this // — fine at toy scale, but ~90% system time against a real index,
// manifested as ~90% system time, observed in practice. ─────────── // 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 n_workers = obisys::effective_parallelism();
let capacity = 256; let capacity = 256;
// Throttling is not optional once a `Flat` stage is in the pipeline // Throttling limits how many *batches* are in flight at once — the
// (see `obipipeline::throttle`'s docs): without it, every worker can // permit is acquired per batch (not per k-mer) in the source
// become a simultaneous `Flat` producer, saturate the shared output // thread, and released once its `VariantBatch` has been read out of
// channel, and deadlock against the scheduler's own dispatch loop — // the pipeline by the accumulation loop below. See
// also observed in practice. The permit acquired here for a source // `obipipeline::throttle`'s docs for why this is required, not
// k-mer is held (via the `Arc`-shared guard carried through // optional, once a `Flat`-style stage sits in the pipeline.
// `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.
let sources: Vec<(usize, CanonicalKmer)> = slot_kmer let sources: Vec<(usize, CanonicalKmer)> = slot_kmer
.iter() .iter()
.enumerate() .enumerate()
.filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer))) .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer)))
.collect(); .collect();
let throttled = obipipeline::throttle(sources.into_iter(), n_workers).map(|t| SourceItem { let batches: Vec<Vec<(usize, CanonicalKmer)>> = sources
slot: t.item.0, .chunks(BATCH_SIZE)
kmer: t.item.1, .map(|chunk| chunk.to_vec())
_permit: Arc::new(t.guard), .collect();
let throttled = obipipeline::throttle(batches.into_iter(), n_workers).map(|t| SourceBatch {
items: t.item,
_permit: t.guard,
}); });
let pipe = obipipeline::make_pipe! { let pipe = obipipeline::make_pipe! {
SibData : SourceItem => VariantQuery, SibData : SourceBatch => VariantBatch,
|| { | {
move |item: SourceItem| -> Vec<VariantQuery> { move |batch: SourceBatch| -> VariantBatch {
let kmer = item.kmer; let mut items = Vec::with_capacity(batch.items.len() * 3);
let permit = item._permit; for (slot, kmer) in batch.items {
kmer.central_canonical_neighbors() for variant in kmer.central_canonical_neighbors() {
.into_iter() if variant == kmer {
.filter(|variant| *variant != kmer) continue;
.map(|variant| VariantQuery { }
source_slot: item.slot, items.push((
dest_partition: partition_of(variant, n_parts), partition_of(variant, n_parts),
variant, variant,
base: central_base(variant, k), slot,
_permit: Arc::clone(&permit), central_base(variant, k),
}) ));
.collect::<Vec<_>>() }
}
VariantBatch { items, _permit: batch._permit }
} }
} : Item => Query, } : Batch => Variants,
}; };
// ── Group generated variants by destination partition. `cache` // ── 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 // lookup is not free just because the file isn't reopened. Grouping
// keeps one partition's pages hot while its whole batch is resolved, // keeps one partition's pages hot while its whole batch is resolved,
// instead of faulting pages in and out as lookups jump between // instead of faulting pages in and out as lookups jump between
// partitions in whatever order the `Flat` stage happens to produce // partitions in whatever order the pipeline happens to produce
// them. The throttle permit drops here, once accumulated. ───────── // them. Each batch's throttle permit drops here, once accumulated.
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect(); let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
for vq in pipe.apply(throttled, n_workers, capacity) { for vb in pipe.apply(throttled, n_workers, capacity) {
outgoing[vq.dest_partition].push((vq.variant, vq.source_slot, vq.base)); 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 // ── Resolve each partition's batch against the cache in one
@@ -635,9 +655,27 @@ pub struct RawSnpDistanceOutput {
} }
impl KmerIndex { impl KmerIndex {
/// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex /// Shared traversal behind [`raw_snp_distance`](Self::raw_snp_distance)
/// (run [`build_sibling_annex`](Self::build_sibling_annex) first). /// and [`base_pair_tally`](Self::base_pair_tally): for every family
pub fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> { /// (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<Acc, F, C>(
&self,
label: &str,
zero: impl Fn() -> Acc + Sync,
on_pair: F,
combine: C,
) -> OKIResult<Acc>
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_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len(); let n_genomes = self.meta.genomes.len();
let with_counts = self.meta.config.with_counts; 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 pb = progress_bar(label, layer_dirs.len() as u64, "layers");
let partials: Vec<(Array2<u64>, Array2<u64>)> = layer_dirs let partials: Vec<Acc> = layer_dirs
.par_iter() .par_iter()
.map(|layer_dir| -> OKIResult<(Array2<u64>, Array2<u64>)> { .map(|layer_dir| -> OKIResult<Acc> {
let mut snp = Array2::<u64>::zeros((n_genomes, n_genomes)); let mut acc = zero();
let mut shared = Array2::<u64>::zeros((n_genomes, n_genomes));
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); 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 meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
@@ -754,31 +791,102 @@ impl KmerIndex {
continue; continue;
} }
let Some(bj) = single_form[j] else { continue }; let Some(bj) = single_form[j] else { continue };
if bi == bj { on_pair(&mut acc, i, j, bi, bj);
shared[[i, j]] += 1;
shared[[j, i]] += 1;
} else {
snp[[i, j]] += 1;
snp[[j, i]] += 1;
}
} }
} }
} }
pb.inc(1); pb.inc(1);
Ok((snp, shared)) Ok(acc)
}) })
.collect::<OKIResult<Vec<_>>>()?; .collect::<OKIResult<Vec<_>>>()?;
pb.finish_and_clear(); pb.finish_and_clear();
let mut snp = Array2::<u64>::zeros((n_genomes, n_genomes)); let mut total = zero();
let mut shared = Array2::<u64>::zeros((n_genomes, n_genomes)); for partial in partials {
for (s, sh) in partials { total = combine(total, partial);
snp += &s;
shared += &sh;
} }
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<RawSnpDistanceOutput> {
let n_genomes = self.meta.genomes.len();
let (snp, shared) = self.scan_family_pairs(
"raw_snp_distance",
|| (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::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 }) 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<BasePairTally> {
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 /// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set
+5 -3
View File
@@ -22,15 +22,17 @@ obikindex = { path = "../obikindex", default-features = false }
obitaxonomy = { path = "../obitaxonomy" } obitaxonomy = { path = "../obitaxonomy" }
obilayeredmap = { path = "../obilayeredmap" } obilayeredmap = { path = "../obilayeredmap" }
clap = { version = "4", features = ["derive"] } clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1" serde_json = "1"
serde_yaml = "0.9.33"
csv = "1" csv = "1"
kodama = "0.2" kodama = "0.3.0"
speedytree = "0.1" speedytree = "0.1"
rayon = "1" rayon = "1"
indicatif = "0.17" indicatif = "0.18"
tracing = "0.1.44" tracing = "0.1.44"
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } 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] [features]
default = ["numa"] default = ["numa"]
+495 -10
View File
@@ -4,7 +4,11 @@ use std::path::PathBuf;
use clap::Args; use clap::Args;
use kodama::{Method, linkage}; use kodama::{Method, linkage};
use obifastwrite::{JsonVal, write_record}; 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 speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
use tracing::info; use tracing::info;
@@ -94,9 +98,61 @@ pub struct DistanceArgs {
#[arg(long)] #[arg(long)]
pub snp: bool, 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 <prefix>_sankoff.tnt, a ready-to-run TNT script (`proc
/// <file>;`) 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 <prefix>_sankoff.tcm and <prefix>_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: <prefix>_dist.csv, <prefix>_shared.csv, /// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
/// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_snp.fasta, /// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_snp.fasta,
/// <prefix>_nj.nwk, <prefix>_upgma.nwk. /// <prefix>_sankoff_matrix.csv, <prefix>_sankoff_params.yaml,
/// <prefix>_sankoff.fasta, <prefix>_sankoff.tnt, <prefix>_sankoff.tcm,
/// <prefix>_sankoff.pg, <prefix>_nj.nwk, <prefix>_upgma.nwk.
/// If omitted, the distance matrix is written to stdout. /// If omitted, the distance matrix is written to stdout.
#[arg(short, long)] #[arg(short, long)]
pub output: Option<PathBuf>, pub output: Option<PathBuf>,
@@ -117,6 +173,13 @@ pub fn run(args: DistanceArgs) {
// deliberately decoupled: the annex is meant to be (re)built routinely, // deliberately decoupled: the annex is meant to be (re)built routinely,
// the distribution only occasionally, on demand. // the distribution only occasionally, on demand.
if args.sibling_annex { 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"); info!("building sibling-count/minorant annex");
idx.build_sibling_annex().unwrap_or_else(|e| { idx.build_sibling_annex().unwrap_or_else(|e| {
eprintln!("error building sibling annex: {e}"); eprintln!("error building sibling annex: {e}");
@@ -144,15 +207,56 @@ pub fn run(args: DistanceArgs) {
}); });
write_snp_fasta(&alignment, &labels, &args.output); 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 let tally = idx.base_pair_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| {
// their own operation, not a modifier on top of a distance-metric eprintln!("error computing base-pair tally: {e}");
// computation — a metric was never requested by asking for any of them, std::process::exit(1);
// so there is nothing for the rest of this function to compute. Not a });
// historical accident to keep: stop here rather than always also let sub_cost = substitution_costs_from_tally(&tally);
// running a Jaccard (or whichever `--metric` defaults to) pass and let mean_sub_cost = mean_substitution_cost(&tally, &sub_cost);
// printing an unrequested matrix. let c_ctx = c_ctx_from_p_hat(estimate.p_hat, m, mean_sub_cost);
if args.sibling_annex || args.sibling_stats || args.raw_snp_distance || args.snp {
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; return;
} }
@@ -364,6 +468,387 @@ fn write_snp_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option<
if n_sites == 1 { "" } else { "s" }); 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<PathBuf>) {
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<u8> = 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<PathBuf>,
) {
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<SankoffSubstitution>,
}
fn write_sankoff_params(
estimate: &PHatEstimate,
tally: &BasePairTally,
weights: &SankoffWeights,
ratio_ceiling: f64,
m: usize,
mean_sub_cost: f64,
output: &Option<PathBuf>,
) {
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<PathBuf>,
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<PathBuf>, 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::<Vec<_>>().join(" ");
writeln!(f, "{alphabet_line}").unwrap();
for i in 0..16 {
let mut row: Vec<i64> = (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::<Vec<_>>().join(" ")).unwrap();
}
let mut gap_row: Vec<i64> = (0..16).map(|j| scaled_matrix[0][j]).collect();
gap_row.push(0);
writeln!(f, "{}", gap_row.iter().map(|v| v.to_string()).collect::<Vec<_>>().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 ─────────────────────────────────────── // ── UPGMA Newick from kodama dendrogram ───────────────────────────────────────
fn upgma_to_newick(dendro: &kodama::Dendrogram<f64>, names: &[String]) -> String { fn upgma_to_newick(dendro: &kodama::Dendrogram<f64>, names: &[String]) -> String {
+7
View File
@@ -78,6 +78,13 @@ pub fn run(args: FilterCmdArgs) {
filters.push(Box::new(MinComplexity { level_max: args.complexity_level_max, theta })); 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(); let mut rep = Reporter::new();
KmerIndex::rebuild(&args.output, &src, &filters, mode, args.force, &mut rep) KmerIndex::rebuild(&args.output, &src, &filters, mode, args.force, &mut rep)
.unwrap_or_else(|e| { .unwrap_or_else(|e| {
+10
View File
@@ -156,6 +156,16 @@ pub fn run(args: IndexArgs) {
let output = args.output.clone(); let output = args.output.clone();
let mut rep = Reporter::new(); 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 ──────────────────────────────────────────────── // ── Resolve evidence kind ────────────────────────────────────────────────
let (evidence, effective_kmer_size) = if args.approx { let (evidence, effective_kmer_size) = if args.approx {
let (z, b, fp) = resolve_approx_params(args.findere_z, args.evidence_bits, args.fp); let (z, b, fp) = resolve_approx_params(args.findere_z, args.evidence_bits, args.fp);
+7
View File
@@ -64,6 +64,13 @@ pub fn run(args: MergeArgs) {
sources.len(), n_genomes, args.output.display() 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(); 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| { 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}"); eprintln!("error merging: {e}");
+7
View File
@@ -12,6 +12,13 @@ pub struct PackArgs {
} }
pub fn run(args: 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| { let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
eprintln!("error opening index: {e}"); eprintln!("error opening index: {e}");
std::process::exit(1); std::process::exit(1);
+7
View File
@@ -47,6 +47,13 @@ pub fn run(args: ReindexArgs) {
IndexMode::Exact 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| { let mut idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
eprintln!("error opening index: {e}"); eprintln!("error opening index: {e}");
std::process::exit(1); std::process::exit(1);
+10
View File
@@ -215,6 +215,16 @@ pub fn run(args: SelectArgs) {
std::process::exit(1); 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| { let mut src = KmerIndex::open(&args.source).unwrap_or_else(|e| {
eprintln!("error opening source index: {e}"); eprintln!("error opening source index: {e}");
std::process::exit(1); std::process::exit(1);
+1
View File
@@ -8,3 +8,4 @@ libc = "0.2"
sysinfo = "0.33" sysinfo = "0.33"
indicatif = "0.17" indicatif = "0.17"
tracing = "0.1" tracing = "0.1"
fs4 = "0.9"
+43
View File
@@ -8,6 +8,49 @@ use tracing::{debug, info, warn};
const BRAILLE: &[&str] = &["", "", "", "", "", "", "", "", "", ""]; 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<Self> {
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 ────────────────────────────────────────────────────────────────── // ── TracedBar ──────────────────────────────────────────────────────────────────
/// Wrapper around `ProgressBar` that emits `tracing` events when stderr is not /// Wrapper around `ProgressBar` that emits `tracing` events when stderr is not