Merge pull request 'Push twnyxtspltux' (#76) from push-twnyxtspltux into main
Reviewed-on: #76
This commit was merged in pull request #76.
This commit is contained in:
@@ -2249,9 +2249,85 @@ F84/TN93 and is included below. `snp-` prefix on every CLI value.
|
||||
**`+Γ` rate-heterogeneity modifier, applicable to `snp-jc`, `snp-k2p`,
|
||||
`snp-k81`, `snp-t92`, `snp-f84`, `snp-hky85`, `snp-tn93`** (not `snp-raw`,
|
||||
nothing to correct; not `snp-logdet`, no standard gamma formulation) — same
|
||||
formula as the base correction, weighted by a shape parameter `α` supplied
|
||||
by the user (`--gamma-shape <alpha>`), not estimated by ML. A modifier on
|
||||
existing values, not a separate enum arm per distance.
|
||||
formula as the base correction, weighted by a shape parameter `α` either
|
||||
supplied by the user (`--gamma-shape <alpha>`) or estimated from the data
|
||||
(`--gamma-shape auto`/`estimate`, method-of-moments — not ML; see
|
||||
"Automatic α estimation" below). A modifier on existing values, not a
|
||||
separate enum arm per distance.
|
||||
|
||||
### Automatic α estimation (`--gamma-shape auto`)
|
||||
|
||||
**Correction (verified against the primary source, 2026-09-11):** Jin &
|
||||
Nei (1990) itself (*"Limitations of the Evolutionary Parsimony Method of
|
||||
Phylogenetic Analysis"*, Mol. Biol. Evol. 7(2):82–102 — the paper this
|
||||
whole `+Γ` correction is cited from, confirmed algebraically to match this
|
||||
codebase's `corrected_log`/`k2p` exactly against their eq. A4, general, and
|
||||
A8, the `a = 1` case) contains **no** data-driven α-estimation procedure.
|
||||
Their own recommendation (p. 98) is a fixed default: *"we suggest that the
|
||||
gamma distance with a = 1 [eq. A8] be used. However, one may choose a
|
||||
different gamma distance, estimating a from data. Wilson et al. (1989)
|
||||
recently used a distance with a = 1/2 for restriction-site data of
|
||||
mitochondrial DNA in hominoids."* — i.e. Jin & Nei explicitly punt
|
||||
data-driven estimation to a *different* paper (Wilson et al. 1989), not
|
||||
read/verified here. The estimator below is therefore **not** "Jin & Nei's
|
||||
method" under any framing — that attribution (present in an earlier
|
||||
revision of this section) was wrong, not just under-cited.
|
||||
|
||||
**Implemented** (`PartitionDispersion`, `obikphylo/src/siblings/algorithms/pairwise.rs`)
|
||||
as an independent method-of-moments estimator, unrelated to any specific
|
||||
published procedure: pools substitution counts by **partition** rather
|
||||
than by genome pair, during the same `reduce_pairwise` pass that builds
|
||||
`PairwiseTally` (no second scan).
|
||||
|
||||
For partition `i`: `n_i` = substitutions pooled over every genome pair,
|
||||
`L_i` = eligible loci pooled over every genome pair, `R_i = n_i / L_i`.
|
||||
Modeling among-site rate heterogeneity the same way as the `+Γ` correction
|
||||
itself (a `Gamma(α, α)`-distributed, mean-1, multiplicative rate shared by
|
||||
every locus in a partition — the classical Poisson-Gamma/negative-binomial
|
||||
mixture, the general identity behind gamma-rate-heterogeneity corrections,
|
||||
also behind Uzzell & Corbin's (1971) original observation that substitution
|
||||
counts across sites/regions are over-dispersed relative to Poisson):
|
||||
|
||||
\[
|
||||
\mathbb{E}[R_i] = \mu \qquad \mathrm{Var}[R_i] = \frac{\mu}{L_i} + \frac{\mu^2}{\alpha}
|
||||
\]
|
||||
|
||||
Weighting each partition's squared deviation by its own `L_i` cancels the
|
||||
Poisson term before attributing what's left to `α`:
|
||||
|
||||
\[
|
||||
\hat\mu = \frac{\sum_i n_i}{\sum_i L_i} \qquad
|
||||
V = \frac{\sum_i L_i (R_i-\hat\mu)^2}{\sum_i L_i} \qquad
|
||||
\hat\alpha = \frac{\hat\mu^2}{V - \hat\mu/\bar L}
|
||||
\]
|
||||
|
||||
where `\bar L` is the mean partition size. Returns "no estimate" (falls
|
||||
back to the uncorrected formula, warns) when fewer than 2 partitions have
|
||||
data, `\hat\mu \le 0`, or `V` doesn't exceed the Poisson floor
|
||||
`\hat\mu/\bar L` — no detectable over-dispersion, `α` would be unbounded.
|
||||
|
||||
**Caveat, stated explicitly rather than left implicit**: unlike every
|
||||
closed-form correction in `snp_distance.rs` (each verified line-by-line
|
||||
against `ape`'s `dist_dna.c`, and now also against Jin & Nei 1990 directly
|
||||
for the base `+Γ` formula), this estimator is derived from first
|
||||
principles (the general Poisson-Gamma/negative-binomial identity) with no
|
||||
primary-source procedure behind it at all — not Jin & Nei's (confirmed
|
||||
above), and Wilson et al. (1989), the paper they point to instead, hasn't
|
||||
been read/verified either. Mathematically self-consistent (re-derived and
|
||||
checked, not guessed), but a from-scratch method, not a literature
|
||||
implementation. If `--gamma-shape` needs a value with a literature
|
||||
pedigree rather than an estimated one, Jin & Nei's own stated default,
|
||||
`α = 1` (`--gamma-shape 1`), is the better-supported choice today.
|
||||
|
||||
Deliberately **not** gated by `--sankoff-ratio-ceiling` the way
|
||||
`base_pair_tally` is (same precedent as `cardinality_tally` — see its own
|
||||
doc comment): that filter excludes individual saturated *pairs* from a
|
||||
composition estimate computed once at the very end, from the complete
|
||||
`PairwiseTally`; the partition axis needed here only exists transiently,
|
||||
one partition at a time, while `PairwiseTally` is still being built — long
|
||||
before any pair's final SNP ratio (and thus its ratio_ceiling eligibility)
|
||||
is known. `--exclude-genome` isn't applied either, matching
|
||||
`reduce_pairwise`'s own raw per-pair fold.
|
||||
|
||||
**Implemented now: `snp-raw`, `snp-jc`, `snp-k2p`, `snp-k81`, `snp-f81`,
|
||||
`snp-t92`, `snp-tn93`, `snp-tv`, all with `+Γ` except `raw`/`tv`** — see
|
||||
@@ -2391,6 +2467,397 @@ mechanical substitution but weren't independently checked against an
|
||||
`ape`-equivalent reference for those three specifically — flagged here, not
|
||||
silently assumed correct.
|
||||
|
||||
### α-estimation literature survey and design discussion (2026-09-11)
|
||||
|
||||
Follow-up discussion after implementing `PartitionDispersion::estimate_alpha`
|
||||
(previous section), triggered by the user pointing out that the literature
|
||||
conflates two distinct problems: *estimating α from an alignment* versus
|
||||
*using a given α in a Jin-Nei-style corrected distance*. Jin & Nei (1990)
|
||||
itself is squarely in the second camp — α is a user-supplied parameter
|
||||
there, not something their paper estimates (confirmed by reading the paper
|
||||
directly, see previous section). This section catalogs the actual
|
||||
α-estimation literature and records which approaches were considered and
|
||||
why each was or wasn't adopted.
|
||||
|
||||
**Bibliography** (methodological papers, not exhaustive users of `+Γ`):
|
||||
|
||||
| Method | Reference | Needs a tree? | Estimates "true" α or a task-specific one? | Cost |
|
||||
|---|---|---|---|---|
|
||||
| ML, continuous Γ (foundational) | **Gu, Fu & Li (1995)**, MBE 12:546–557 — read in full | yes, topology fixed a priori | true α (+ θ jointly, invariant+Γ) | very high, ≤5–6 taxa in practice |
|
||||
| ML, discrete Γ | Yang (1993/94), MBE 10:1396; JME 39:306 | yes | true α | high |
|
||||
| Parsimony-based ML on inferred changes | Yang & Kumar (1996), MBE 13:650 | yes, equal branch lengths assumed | true α, biased when branches unequal | low |
|
||||
| Corrected substitutions/site → NB fit | **Gu & Zhang (1997)**, MBE 14:1106 — read in full | yes, + ancestral reconstruction | true α, close to ML accuracy | low (given a tree) |
|
||||
| Topology-optimizing grid search | **Guindon & Gascuel (2002)**, MBE 19:534 — read in full | yes, but built *from* the candidate distances themselves (BIONJ) | **not** true α — deliberately biased *upward* for topological accuracy | low–moderate (grid × tree build) |
|
||||
| Free (non-parametric) rate distribution | Susko et al. (2003), Syst. Biol. 52:594 | yes | tests the Γ assumption itself | high |
|
||||
| Parsimony counts/site → NB fit | **Wakeley (1993)**, JME 37:613 — read in full | yes, NJ tree + Fitch parsimony ancestral reconstruction | true α, biased *upward* by parsimony (conservative) | low (given a tree) |
|
||||
| Poisson-Gamma/negative-binomial, foundational | Uzzell & Corbin (1971), Science 172:1089 — **not read**, only seen via citations | unknown — likely also count-per-site/tree-based, unverified | true α (indirect, via NB overdispersion) | unknown |
|
||||
| Bayesian posterior over α | BPP/MrBayes/BEAST2/RevBayes | yes | posterior, not point estimate | very high |
|
||||
|
||||
**Correction (2026-09-11, after reading Wakeley in full):** an earlier
|
||||
revision of this table filed Wakeley (1993) under "no tree needed" —
|
||||
wrong. Wakeley's method: build a tree (neighbor-joining, Saitou & Nei
|
||||
1987, on the real mtDNA data; coalescent-simulated for the simulation
|
||||
study), reconstruct ancestral states at internal nodes by Fitch (1971)
|
||||
parsimony, count the *minimum* number of changes per site this implies,
|
||||
then fit either a two-rate Poisson mixture or a gamma-distributed-rates
|
||||
model (⟹ negative binomial, parameters by the methods of Cohen 1965 and
|
||||
Johnson & Kotz 1969) to the resulting per-site count distribution — the
|
||||
same general "tree + ancestral reconstruction + per-site counts + NB fit"
|
||||
shape as Gu & Zhang (1997), just using raw parsimony counts directly
|
||||
rather than Gu & Zhang's multiple-hit-corrected `k̂`. Confirmed by
|
||||
simulation (Wakeley's own Fig. 1) that parsimony underestimates both the
|
||||
mean *and* (more severely) the variance of per-site change counts, which
|
||||
biases `α` **upward** (toward apparent uniformity) — consistent with Gu &
|
||||
Zhang's own citation of Wakeley for exactly this bias, and with the
|
||||
general "parsimony overestimates α" pattern noted throughout this
|
||||
literature. Real hypervariable-region-1 mtDNA data gave `α̂ ≈ 0.44–0.60`
|
||||
per data set, `0.47` combined — despite the conservative bias, still
|
||||
solidly in the "strong heterogeneity" range, which is the paper's own
|
||||
point: the method is biased but usefully conservative, not useless.
|
||||
|
||||
**Where `PartitionDispersion::estimate_alpha` actually sits**: nowhere in
|
||||
this table's tree-dependent rows. It shares only the *abstract
|
||||
mathematical identity* (Poisson-Gamma mixture ⟹ negative binomial
|
||||
over-dispersion) with Wakeley/Uzzell-Corbin/Gu&Zhang, applied via a
|
||||
genuinely different, tree-free and ancestral-reconstruction-free route:
|
||||
pairwise genome-to-genome substitution counts pooled by *partition*
|
||||
instead of per-site counts pooled across an alignment's columns after
|
||||
ancestral state reconstruction. As far as this survey has established
|
||||
(Uzzell & Corbin not read, so not fully ruled out), **no tree and no
|
||||
ancestral reconstruction of any kind** appears to be a genuinely
|
||||
distinguishing property of `PartitionDispersion`, not something it
|
||||
inherits from prior art — consistent with it being an independent
|
||||
derivation rather than a literature implementation (see previous
|
||||
section's caveat, now on firmer footing).
|
||||
|
||||
**Why a tree-free estimate is even valid — the identifiability question.**
|
||||
Raised directly by the user after reading Gu, Fu & Li (1995), which states
|
||||
explicitly: *"one cannot estimate the rate heterogeneity when the number
|
||||
of sequences is <3."* This is a real non-identifiability result, not a
|
||||
practical inconvenience: with only 2 sequences, a site's history is
|
||||
summarized by one binary observation (differ / don't differ), collapsed
|
||||
across all sites into two aggregate numbers (transition and transversion
|
||||
proportions) — divergence time and among-site rate variance are
|
||||
confounded in that pair of numbers, with no way to separate them. Every
|
||||
method surveyed above (Gu,Fu&Li, Yang, Wakeley, Gu&Zhang) needs ≥3
|
||||
sequences *sharing one evolutionary history* (a tree) specifically to get
|
||||
multiple independent looks at the *same* site's rate across different
|
||||
lineages, which is what identifies its variance.
|
||||
|
||||
`PartitionDispersion` doesn't violate this, because it isn't estimating
|
||||
the same quantity. The literature above estimates rate variation **across
|
||||
sites, within one shared tree**. `PartitionDispersion` estimates rate
|
||||
variation **across partitions, pooled over every genome pair in the
|
||||
index** — each genome pair stands in for one independent draw of a
|
||||
partition's relative rate, the same identifying role multiple lineages
|
||||
play in the tree-based methods, just substituting "genome pair" for
|
||||
"lineage" and "partition" for "site." With `P` genome pairs (potentially
|
||||
in the hundreds here) all contributing to every partition's pooled count,
|
||||
the ≥3-samples identifiability requirement is met by the pair count, not
|
||||
by tree depth — closer in spirit to Jin & Nei's own original motivating
|
||||
scenario (comparing several genomic *regions* across one fixed panel of
|
||||
taxa) than to single-tree site-rate estimation.
|
||||
|
||||
This substitution carries its own assumption, which must be named rather
|
||||
than left implicit: that a partition's *relative* rate (fast/slow) is
|
||||
reasonably stable **across genome pairs**, not just across sites within
|
||||
one pair — the direct analogue of the "no lineage-specific rate variation,
|
||||
only site-specific" assumption every tree-based `+Γ` method already makes,
|
||||
just moved from lineage→pair and site→partition. If this breaks — e.g. an
|
||||
index mixing very closely related and very divergent genome pairs, where a
|
||||
partition's saturation behavior differs qualitatively between the two —
|
||||
the estimate could be misled in a way a tree-based method would at least
|
||||
have the topology to detect and a tree-free, pool-everything method
|
||||
cannot. Not yet tested against real data with strongly heterogeneous
|
||||
pairwise divergence; worth keeping in mind as the main open validity
|
||||
question for `--gamma-shape auto`, not the citation question (now settled)
|
||||
this whole discussion started from.
|
||||
|
||||
**Gu, Fu & Li (1995), read in full — the foundational ML method, not a
|
||||
candidate to implement.** Models rate variation as invariant+Γ (`θ` =
|
||||
proportion of invariant sites, `α` = gamma shape among the variable sites),
|
||||
and derives the exact joint likelihood over nucleotide *configurations*
|
||||
across all `n` sequences on a **fixed, a-priori-known tree topology** —
|
||||
for 3 sequences this is closed-form (Jukes-Cantor), for `n` sequences it's
|
||||
a sum over `4^n` configurations weighted by coefficients tied to the
|
||||
specific topology's branching structure (their eq. 32–34), with branch
|
||||
lengths, `α`, and `θ` jointly optimized by Newton-Raphson (Hessian over all
|
||||
free parameters). The site-rate integral has the same closed form
|
||||
`E = θ + (1-θ)(1+D/α)^{-α}` that underlies `+Γ` distance corrections
|
||||
generally (the same mechanism as `corrected_log`, not a coincidence — Jin &
|
||||
Nei's own formula is the `θ=0` special case applied to a pairwise `D`
|
||||
rather than a whole-tree one). Explicitly stated by the authors as
|
||||
infeasible beyond about 5–6 taxa (their own simulations stop at 5); not a
|
||||
candidate for obikmer's typical genome counts. Notable contribution worth
|
||||
keeping in mind regardless: `ρ = (1+θα)/(1+α)` is shown to be a more robust
|
||||
summary of rate heterogeneity than `α` alone whenever `θ` isn't negligible
|
||||
— `α` and `θ` are confounded (nearly-invariant sites can come from either a
|
||||
small `α`'s left tail or from a nonzero `θ`), so `α` alone can be very
|
||||
unstable while `ρ` stays well-behaved. `PartitionDispersion` has no `θ`
|
||||
term at all (no invariant-site component), so this confound doesn't arise
|
||||
for it the same way, but it's worth remembering if an invariant+Γ variant
|
||||
is ever considered.
|
||||
|
||||
**Gu & Zhang (1997), read in full — rejected for now.** Their procedure:
|
||||
(1) a tree with least-squares branch lengths must already exist; (2)
|
||||
ancestral states are reconstructed at every internal node (they use Zhang &
|
||||
Nei 1997's likelihood method, but note plain parsimony would work, just
|
||||
less accurately); (3) per site, branches are split into "changed"/"unchanged"
|
||||
given the ancestral reconstruction, and the *expected* substitution count
|
||||
`k̂` (corrected for multiple hits) is obtained by solving their eq. (7)
|
||||
(or eq. 12 for the generalized model) numerically — a root-find per site,
|
||||
not a closed form; (4) the site-level `k̂` values (real-valued, not
|
||||
integers) are fit to a negative binomial (Uzzell & Corbin 1971) by ML to
|
||||
get α — again no closed form, numerical optimization. This is a real
|
||||
per-site pipeline requiring a tree *before* distances/α can be computed —
|
||||
the reverse of obikmer's current dependency order (`snp_distance` computes
|
||||
distances, which `--nj`/`--upgma` only turn into a tree *afterward*). Their
|
||||
own answer to this chicken-and-egg problem is an iterative
|
||||
distance→tree→α→distance loop, which they themselves flag as expensive.
|
||||
Adopting this would mean a new ancestral-reconstruction subsystem (even a
|
||||
parsimony/Fitch-only version) plus a per-family nonlinear solve plus an NB
|
||||
ML fit — a much bigger addition than `PartitionDispersion`, not something
|
||||
to build without a concrete need beyond "closer to Gu&Zhang's own
|
||||
simulated accuracy than our moment estimator."
|
||||
|
||||
**Guindon & Gascuel (2002), read in full — plausible future addition, not implemented.**
|
||||
Their **Q criterion**: for a candidate α, build a tree (they use BIONJ) from
|
||||
the α-corrected distance matrix; for every internal branch, group taxa into
|
||||
the 4 subtrees it separates (A, B, C, D), compute mean inter-subtree
|
||||
distances \(\bar\delta_{AB}\), \(\bar\delta_{CD}\), etc., and let `S ≤ M ≤ L`
|
||||
be the three pairings' sums (`{AB,CD}`, `{AC,BD}`, `{AD,BC}`); the branch's
|
||||
reliability score is `Q_branch = L − M` (zero when the four points are
|
||||
perfectly tree-additive). The whole-tree `Q` is the mean over internal
|
||||
branches (negative/zero branches excluded). `α* = argmin_α Q(α)` over a
|
||||
grid (their run: ~60–100 points, 0.1 to 5000, finer spacing where
|
||||
sensitivity is highest). Complexity: `O(n²l)` once for the base
|
||||
frequencies, then `O(n²r)` for the `r` candidate distance matrices, then
|
||||
`O(n³r)` worst case for building `r` trees and evaluating `Q` on each
|
||||
(same order as the tree-building step itself, so "free" relative to it).
|
||||
|
||||
Central finding, **the opposite direction from an initial misreading in
|
||||
this discussion**: `α_opt` (Guindon-Gascuel's topology-optimal value) is
|
||||
**always ≥ the true α**, not the reverse — underestimating rate
|
||||
heterogeneity (i.e., picking a larger α than reality) reduces the variance
|
||||
of distance estimates and thereby improves NJ/BIONJ topological accuracy,
|
||||
especially when the molecular clock roughly holds (where `α_opt → ∞`, i.e.
|
||||
no correction at all is topologically best). This is an *empirical
|
||||
regularity* observed across their tested conditions (η ∈ {0.5, 2.0}, true
|
||||
α ∈ {0.1, 0.7}, 20-taxon trees) — not a theorem — so treating it as a hard
|
||||
guarantee would be another unverified-claim mistake of the kind this whole
|
||||
discussion has been correcting.
|
||||
|
||||
**Proposed integration (not yet implemented)**: since `α_true ≤ α_opt`
|
||||
empirically, `PartitionDispersion::estimate_alpha()`'s output could seed
|
||||
the *lower bound* of Guindon-Gascuel's grid (search `[α̂, ∞)` instead of
|
||||
`[0.1, 5000]`), narrowing the grid substantially and reusing the existing
|
||||
estimator rather than needing an independent "true α" source — this also
|
||||
sidesteps Gu & Zhang's tree-first dependency problem entirely, since
|
||||
Guindon-Gascuel's own grid already builds a tree per candidate. Still
|
||||
requires: (a) a **BIONJ** implementation (see below), (b) the Q-criterion
|
||||
subtree-grouping logic, (c) a grid-search driver. Not started.
|
||||
|
||||
**BIONJ availability in the Rust ecosystem — checked 2026-09-11, none found.**
|
||||
Searched crates.io directly (API query for "bionj": 0 results) and read the
|
||||
docs.rs pages of the most plausible candidates: `phylotree` (builds/reads/
|
||||
manipulates trees but has no distance-matrix reconstruction — no NJ, no
|
||||
BIONJ, no UPGMA despite web-search summaries claiming otherwise, corrected
|
||||
after checking the actual docs), `speedytree` (NJ only — Canonical/RapidNJ/
|
||||
Hybrid variants, explicitly no BIONJ), `nj` (plain NJ only), `phylo` (no
|
||||
distance-based reconstruction at all). `obikphylo::neighbor_joining`
|
||||
(`tree.rs`) is plain NJ, not BIONJ — BIONJ (Gascuel 1997) differs only in
|
||||
how branch lengths are weighted during agglomeration (accounts for
|
||||
distance-estimate variance, not just the sum), a modest delta over an
|
||||
existing NJ implementation rather than a rewrite, but would need to be
|
||||
written in-house if the Guindon-Gascuel integration above is ever pursued
|
||||
(their own paper notes other tree-building methods gave similar results in
|
||||
their simulations, so plain NJ may be an acceptable substitute if BIONJ
|
||||
proper is judged not worth the effort).
|
||||
|
||||
### Sharing one sample across algorithms (`--sankoff`/`--tnt`/`--phyg`/`--iqtree` + `snp-*` `--distance`) — implemented 2026-09-11
|
||||
|
||||
**Correctness bug, not just a performance one.** `--sankoff`/`--tnt`/`--phyg`/
|
||||
`--iqtree` (via `SiblingExt::sankoff_bundle`) and a `snp-*` `--distance`
|
||||
(via `SiblingExt::snp_distance`) always consume the *exact same* selection
|
||||
parameters when requested together in one `obikmer phylo` invocation
|
||||
(`n`/`free_loss`/`no_ambiguity`/excluded-set/`entropy_bias` — the CLI has
|
||||
no way to give them different values in one run, see `cmd/phylo/args.rs`).
|
||||
Before this fix, each independently called `sample_index`, and
|
||||
`sample_layer`'s `rand::rng()` (`subsample.rs:235`) is a **thread-local
|
||||
generator that advances across calls, not reseeded each time** — so the
|
||||
second call silently drew a *different* random sample of sites than the
|
||||
first, even with identical parameters. Verified empirically (throwaway
|
||||
8-genome index): running `snp-jc` alone vs. combined with `--sankoff`
|
||||
produced visibly different distance matrices from the same index/params.
|
||||
This defeated the actual point of combining these flags — running several
|
||||
algorithms (Sankoff/TNT/PhyG/IQ-TREE calibration, a `snp-*` distance) on
|
||||
one *identical* site selection for direct comparison.
|
||||
|
||||
**Fix**: `SankoffBundle` (`sankoff.rs`) now retains its `PairwiseTally`/
|
||||
`PartitionDispersion` internally (previously consumed into `raw`/
|
||||
`base_pair_tally`/`cardinality_tally` and dropped) and exposes
|
||||
`SankoffBundle::snp_distance(kind, gamma_shape)`, which computes the
|
||||
matrix from that *same* already-sampled tally — no second `sample_index`
|
||||
call. The shared post-sampling logic (the `--gamma-shape` support check,
|
||||
`alpha` resolution, final matrix build — previously all inline in the
|
||||
`snp_distance()` free function) was factored into
|
||||
`pub(crate) fn distance_matrix(tally, dispersion, kind, gamma_shape)` in
|
||||
`snp_distance.rs`, called by both the standalone `snp_distance()` (after
|
||||
its own fresh sampling) and `SankoffBundle::snp_distance` (reusing the
|
||||
bundle's). `cmd/phylo/mod.rs` keeps the `Option<SankoffBundle>` alive past
|
||||
the Sankoff-family `if` block and, when a `snp-*` `--distance` is also
|
||||
requested, calls `bundle.snp_distance(...)` instead of
|
||||
`cache.snp_distance(...)` whenever a bundle was built — since both branches
|
||||
are driven by the same `args.*` fields, the parameters trivially always
|
||||
match when both fire; no runtime "do the params match" check needed.
|
||||
Verified end-to-end: the reused path logs `(reusing the Sankoff bundle's
|
||||
sample)` and the `snp_distance` stage timer reads `0ms` (`formula(kind)`
|
||||
is `O(n²)` post-processing, no I/O), vs. ~150-230ms for a fresh sample on
|
||||
the same tiny test index.
|
||||
|
||||
#### Explicit `--session` (implemented 2026-09-12)
|
||||
|
||||
The in-process fix above only covered reuse *within one command*.
|
||||
`--session DIR` (`cmd/phylo/args.rs`) names a directory (outside the
|
||||
index) that persists the `snp-*` `--distance` sample — the
|
||||
`PairwiseTally`/`PartitionDispersion` pair `sample_index` would otherwise
|
||||
rebuild from scratch — across *separate* `obikmer phylo` invocations.
|
||||
|
||||
**Crate placement, as planned**: a new **`obiksession`** crate
|
||||
(`src/obiksession/`), domain-agnostic — no `PairwiseTally`/"site" concept
|
||||
anywhere in it. `Session::open(dir, params: &[u8], force: bool)` handles
|
||||
directory lifecycle and the params-conflict check against an opaque byte
|
||||
blob; `Session::store(name, bytes)`/`Session::restore(name) -> Option<Mmap>`
|
||||
handle a checksummed artifact cache. `obikphylo`
|
||||
(`siblings/algorithms/session_cache.rs`, shared by `snp_distance.rs` and
|
||||
`sankoff.rs` — see below) is the consumer: it serializes its own
|
||||
`PairwiseTally`/`PartitionDispersion`/`SnpAlignment` to bytes (via `rkyv`)
|
||||
and hands them to `obiksession`, which never sees their type.
|
||||
|
||||
**Serialization**: `rkyv` 0.8.18, added as planned. One deviation from the
|
||||
original zero-copy pitch, scoped down deliberately: artifacts are restored
|
||||
via `rkyv::from_bytes` (full owned deserialize) rather than `rkyv::access`
|
||||
(zero-copy over the `mmap`). True zero-copy would require every
|
||||
`PairwiseTally` query method (`categories`, `base_freq`, `pair`, ...) to
|
||||
work generically over `Archived<PairwiseTally>` as well as the owned type
|
||||
— a separate, larger change not needed to get the actual win (skipping
|
||||
`sample_index`'s expensive re-scan of the sibling annex; deserializing an
|
||||
already-in-memory-sized count array is comparatively cheap CPU, not I/O).
|
||||
`#[derive(Archive, Serialize, Deserialize)]` was added to `PairStats`,
|
||||
`PairwiseTally`, and `PartitionDispersion` (`pairwise.rs`) for this.
|
||||
|
||||
**Startup behavior, as designed**: `Session::open` creates `DIR` if
|
||||
missing (fresh, params recorded); if `DIR` already holds different saved
|
||||
params, `cmd/phylo` reports a hard error and exits rather than silently
|
||||
preferring one side, unless `--session-force` is given (which discards the
|
||||
directory's cached artifacts and starts over under the new params) —
|
||||
implements exactly the "explicit escape hatch, not silent override"
|
||||
decision from the design discussion. `obiksession` itself doesn't decode
|
||||
either side's bytes to produce a field-by-field diff (true to "opaque
|
||||
blob," left to whoever needs it); `cmd/phylo`'s conflict message is
|
||||
currently a single generic sentence listing which flags could be the
|
||||
cause, not a computed diff — an honest v1 simplification, not a limitation
|
||||
of `obiksession` itself.
|
||||
|
||||
**Per-artifact caching, as designed**: only `PairwiseTally`/
|
||||
`PartitionDispersion` are cached (`snp_distance.rs`'s `restore_tally`/
|
||||
`store_tally`) — `ratio_ceiling`/`gamma_shape` remain pure post-processing
|
||||
over whatever tally is in hand, cached or fresh, exactly as planned.
|
||||
|
||||
**Locking, reconsidered and kept as the whole-session `DirLock`.** First
|
||||
written up as a corner cut ("per-entry locking was sketched, not
|
||||
implemented"). On actually working through the concrete race, whole-session
|
||||
locking turned out to be the *correct* choice, not a simplification:
|
||||
`--session-force` resets a session by deleting its whole `artifacts/`
|
||||
directory. If the directory lock were only held briefly (during
|
||||
`Session::open`, then released so a long-running `Session` doesn't block
|
||||
other processes from touching *different* artifacts), a second process
|
||||
calling `open(..., force: true)` could acquire the lock, see mismatched
|
||||
params, and wipe `artifacts/` **while a first process's already-open
|
||||
`Session` is mid-checkpoint** — genuine corruption risk, not just wasted
|
||||
work. Holding the lock for the `Session`'s entire lifetime makes that
|
||||
sequence impossible: a `force` reset must wait for every currently-open
|
||||
`Session` on that directory to be dropped first. The concurrency this
|
||||
gives up (two *different* processes reading/writing *different* artifacts
|
||||
of the same session at once) isn't exercised by anything today — only one
|
||||
code path (`snp_distance`/`sankoff_bundle`, sequentially per invocation)
|
||||
uses sessions at all.
|
||||
|
||||
**Per-layer chunked/resumable dumps — implemented.** `subsample::sample_index`
|
||||
gained a `skip_layers: usize` parameter (skip the first N (partition,
|
||||
layer) pairs, in its own deterministic iteration order, entirely — no
|
||||
I/O, no `on_layer` call) and its `on_layer` callback gained a `raw_index`
|
||||
argument (that pair's 0-based position in the same order), so a caller can
|
||||
tell exactly how far a run has gotten. `snp_distance`/`sankoff_bundle`
|
||||
checkpoint their tally/dispersion (and, for `sankoff_bundle`, the
|
||||
in-progress `SnpAlignment`) to the session every
|
||||
`session_cache::CHECKPOINT_INTERVAL_LAYERS` (8) layers or
|
||||
`CHECKPOINT_INTERVAL` (30s) of wall time, whichever comes first — bounding
|
||||
how much completed scanning a crash can lose, without paying a full
|
||||
tally-rewrite on every one of what can be thousands of layers. A `progress`
|
||||
artifact records how many layers are reflected in the checkpoint, with a
|
||||
`PROGRESS_COMPLETE` (`u64::MAX`) sentinel once every layer has been
|
||||
processed (distinct from any real layer count, so "finished" is never
|
||||
confused with "stopped after N layers where N happens to equal the total").
|
||||
|
||||
Verified with a real kill: built a 10-genome/64-partition index, ran
|
||||
`--distance snp-jc --session DIR` under `timeout 0.4`, confirmed a
|
||||
checkpoint at `progress=8` (exactly `CHECKPOINT_INTERVAL_LAYERS`) was on
|
||||
disk, then re-ran the same command and confirmed it resumed from layer 8
|
||||
and completed successfully — 0-effort verification that the mechanism
|
||||
itself is real, not just plausible-sounding.
|
||||
|
||||
**Important correction from that same test — resume is *not* bit-reproducible,
|
||||
and shouldn't be.** The output matrix from the resumed run differed from
|
||||
an uninterrupted control run against the same index/params. Root cause:
|
||||
`sample_layer`'s random draws come from `rand::rng()`, a per-process,
|
||||
OS-seeded generator with no continuity across a process boundary — the
|
||||
layers processed before a kill and the layers processed after a resume
|
||||
come from two independent random streams. An earlier revision of this doc
|
||||
(and of the code comments) claimed resume "reproduces the exact same
|
||||
sample an uninterrupted run would have" — that's wrong, and was corrected
|
||||
after the user pointed out the actual reasoning error: true randomness has
|
||||
no notion of continuity in the first place, so there was never a
|
||||
correctness requirement for a resumed run to match an uninterrupted one
|
||||
bit-for-bit. The sample produced after a resume is exactly as legitimate a
|
||||
random draw as one from an uninterrupted run — just a *different* one,
|
||||
same as running the command twice without `--session` already gives two
|
||||
different samples today. That variance is explicitly wanted (see below),
|
||||
not a defect to fix. Comments in `subsample.rs`/`session_cache.rs` were
|
||||
corrected to state this plainly instead of the false reproducibility claim.
|
||||
|
||||
**Explicitly declined: deterministic/seeded sampling.** Raised as a
|
||||
possible fix for the above (seed `sample_layer`'s RNG from a per-session
|
||||
value plus `(partition, layer)`, making a resume reproduce bit-identical
|
||||
results, and, if applied everywhere, making *any* two runs with identical
|
||||
parameters produce identical samples). The user explicitly declined this
|
||||
for now: repeated runs are relied on to measure sampling variance, and a
|
||||
`--seed` flag would need its own explicit design if wanted later — this is
|
||||
not an accidental gap, it's a stated preference to revisit only if asked.
|
||||
|
||||
**`SankoffBundle` integration — implemented.** `sankoff_bundle` accepts
|
||||
the same `session: Option<&Session>` and uses the *same* artifact names
|
||||
(`session_cache::{TALLY_ARTIFACT, DISPERSION_ARTIFACT, PROGRESS_ARTIFACT}`)
|
||||
as `snp_distance`, plus its own `ALIGNMENT_ARTIFACT` (`sankoff.rs`) folded
|
||||
in lockstep with the tally at every checkpoint (same `on_layer` call
|
||||
produces both, so they can never drift out of sync). Consequence verified
|
||||
directly: a tally checkpointed by `--sankoff --session DIR` is restored by
|
||||
a *later, separate* `--distance snp-k2p --session DIR` invocation with no
|
||||
`--sankoff` at all, and vice versa — cross-consumer reuse, not just
|
||||
same-command reuse, exactly the synergy this design aimed for.
|
||||
|
||||
**Per-artifact caching, as designed**: only the tally/dispersion (plus,
|
||||
for `sankoff_bundle`, the alignment) are cached — `ratio_ceiling`/
|
||||
`gamma_shape` remain pure post-processing over whatever tally is in hand,
|
||||
cached or fresh, exactly as planned.
|
||||
|
||||
**Still not done — genuine follow-up, not a justified tradeoff:**
|
||||
**exhaustive-sample caching for `--subsample`-varying exploration.** The
|
||||
design discussion's suggestion — cache the exhaustive (`n = None`) sample
|
||||
once and subsample *from* it on each run instead of one session per exact
|
||||
`n` — wasn't implemented; a session today is still scoped to one exact
|
||||
parameter tuple, `--subsample` included.
|
||||
|
||||
### Output format: PHYLIP-relaxed by default for the distance matrix
|
||||
|
||||
**Implemented.** The primary distance-matrix output
|
||||
|
||||
@@ -17,7 +17,7 @@ obikmer phylo INDEX [OPTIONS]
|
||||
| Option | Default | Description |
|
||||
|---|---|---|
|
||||
| `--distance` | `jaccard` | See the two tables below for the full list of accepted values |
|
||||
| `--gamma-shape ALPHA` | none | Rate-heterogeneity correction, for `snp-*` values that support it (see below). No effect on the other values; rejected if given together with a value that doesn't support it |
|
||||
| `--gamma-shape ALPHA\|auto` | none | Rate-heterogeneity correction, for `snp-*` values that support it (see below). Either a fixed $\alpha$ or `auto`/`estimate` to fit it from the data (see "Automatic $\alpha$ estimation" below). No effect on the other values; rejected if given together with a value that doesn't support it |
|
||||
| `--presence-threshold` | `1` | Minimum count for a kmer to be considered present, for `jaccard`/`mash` on a count index |
|
||||
| `--csv` | off | Write the matrix as plain CSV instead of the default relaxed-PHYLIP format |
|
||||
| `--shared-kmers` | off | Also write the shared-kmer count matrix. Only valid with a whole-index metric, not a `snp-*` value |
|
||||
@@ -121,6 +121,28 @@ $$d = Q$$
|
||||
|
||||
`--gamma-shape ALPHA` applies to every value above except `snp-raw` and `snp-tv`: each $-\ln(x)$ term in the formulas above is replaced by $\alpha\left(x^{-1/\alpha}-1\right)$ (the same weight, same $x$).
|
||||
|
||||
### Automatic $\alpha$ estimation (`--gamma-shape auto`)
|
||||
|
||||
`--gamma-shape auto` (or the equivalent `--gamma-shape estimate`) fits $\alpha$ from the index itself instead of requiring a user-supplied value, using a method-of-moments estimator computed once, from the same sampling pass that builds the pairwise substitution tally — no extra scan of the index.
|
||||
|
||||
The estimator pools substitution counts by **partition** rather than by genome pair: for partition $i$, let $n_i$ be the total number of substitutions observed across every genome pair, and $L_i$ the total number of eligible loci across every genome pair, in that partition. Define the partition's observed substitution rate:
|
||||
|
||||
$$R_i = \frac{n_i}{L_i}$$
|
||||
|
||||
Under a single shared substitution rate with no among-site heterogeneity, each $R_i$ would vary only by Poisson sampling noise. Rate heterogeneity is modeled, as elsewhere in this correction, by a $\mathrm{Gamma}(\alpha,\alpha)$-distributed multiplicative rate (mean 1) shared by every locus in a partition — the classical Poisson–Gamma (negative-binomial) mixture. Under that model:
|
||||
|
||||
$$\mathbb{E}[R_i] = \mu \qquad \mathrm{Var}[R_i] = \frac{\mu}{L_i} + \frac{\mu^2}{\alpha}$$
|
||||
|
||||
where $\mu$ is the pooled substitution rate across every partition. Weighting each partition's squared deviation by its own $L_i$ removes the first (Poisson) term before attributing what's left to genuine rate heterogeneity:
|
||||
|
||||
$$\hat\mu = \frac{\sum_i n_i}{\sum_i L_i} \qquad V = \frac{\sum_i L_i\,(R_i-\hat\mu)^2}{\sum_i L_i} \qquad \bar L = \frac{\sum_i L_i}{\text{number of partitions}}$$
|
||||
|
||||
$$\hat\alpha = \frac{\hat\mu^2}{V - \hat\mu/\bar L}$$
|
||||
|
||||
If the measured variance $V$ doesn't exceed the Poisson floor $\hat\mu/\bar L$ (no detectable over-dispersion across partitions — the data are consistent with a single shared rate), $\alpha$ is left undefined: the correction is silently disabled for that run rather than applying a fabricated value, and a warning is logged. When an estimate is produced, it's logged at the `info` level before the distance matrix is computed.
|
||||
|
||||
Note: this is a method-of-moments estimator derived from the standard Poisson–Gamma relationship between substitution counts and gamma-distributed rate variation, applied per-partition — it is not part of Jin & Nei's (1990) original publication, which only defines the `+Γ` distance formula itself and, absent an estimate, recommends the fixed default $\alpha = 1$ (`--gamma-shape 1`) rather than proposing a way to estimate it from data. `alpha < 1` indicates strong among-site rate heterogeneity (many near-invariant loci, a few fast ones); `alpha` growing large makes the correction converge to the uncorrected formula.
|
||||
|
||||
### Output
|
||||
|
||||
Without `-o`, the matrix goes to stdout in relaxed-PHYLIP format (`n` on the first line, then one `label<TAB>value...` row per genome). With `--csv`, the format is instead a header row `genome,<label1>,<label2>,...` followed by one `<label>,<value1>,<value2>,...` row per genome, 6 decimals. Both formats are symmetric with a zero diagonal, except where noted below.
|
||||
@@ -180,6 +202,23 @@ A family is eligible for a genome pair $(i,j)$ only if both genomes carry exactl
|
||||
|
||||
`--subsample`, `--free-loss`, `--no-ambiguity`, `--entropy`/`--entropy-sd` are shared by `--pseudo-alignment`, `--sankoff` (and everything it implies: `--tnt`/`--phyg`/`--iqtree`), and a `snp-*` `--distance` value — one draw feeds all of them in a single invocation. `--subsample` is mandatory for `--pseudo-alignment`/`--sankoff`; for a `snp-*` `--distance` value it is optional (omitted means every non-monomorphic family in the index, not an approximation).
|
||||
|
||||
Combining `--sankoff` (or `--tnt`/`--phyg`/`--iqtree`) with a `snp-*` `--distance` value in the same command reuses that one draw for both — the distance and the Sankoff calibration/alignment are guaranteed to be computed from the *identical* set of sampled sites, never two independent samples, so the two outputs are directly comparable. This only holds within a single command; running them as two separate `obikmer phylo` invocations draws two independent samples even with the same flags — unless `--session` is used (below).
|
||||
|
||||
Every invocation, with or without `--session`, draws its own fresh random sample by default — running the same command twice gives two different (but equally valid) samples, which is useful for measuring sampling variance and is kept that way deliberately. `--session` does not change this: it makes a *specific* sample reusable on request, it does not make sampling itself reproducible from one independent run to the next.
|
||||
|
||||
### `--session`: reusing a sample across separate commands
|
||||
|
||||
| Option | Default | Description |
|
||||
|---|---|---|
|
||||
| `--session DIR` | none | Persist the sample (and, for `--sankoff`, its calibration/alignment) in `DIR` so a later, separate `obikmer phylo` invocation with the exact same selection parameters restores it instead of resampling |
|
||||
| `--session-force` | off | With `--session DIR`: overwrite its saved parameters and cached sample instead of erroring out when this run's parameters don't match. No effect without `--session` |
|
||||
|
||||
`DIR` is created if it doesn't exist. If it already holds a sample built with different `--subsample`/`--free-loss`/`--no-ambiguity`/`--exclude-genome`/`--min-shared-family`/`--entropy`/`--entropy-sd` values than this run, the command exits with an error rather than silently using either the old or the new values — pass `--session-force` to discard the old sample and rebuild under the new parameters, or point `--session` at a different directory to keep both.
|
||||
|
||||
A `--sankoff`-family run and a `snp-*` `--distance` run share the same cached sample when pointed at the same `--session DIR` — build it once with either, reuse it from the other, in either order, across separate commands.
|
||||
|
||||
If a run using `--session` is interrupted (crash, kill, `Ctrl-C`), the next run against the same `--session DIR` resumes from the last automatic checkpoint (roughly every 8 partitions'-worth of sampling progress) instead of starting over. The resumed run's sample is **not** guaranteed to be identical to what an uninterrupted run would have produced past that checkpoint — each process draws its own independent random sequence, same as any two separate invocations do — but nothing already checkpointed is lost, and no work needs redoing beyond that point.
|
||||
|
||||
Without `--subsample`, every variable family (family size ≥ 2) is used. With `--subsample N`, roughly `N` families are kept instead, drawn in proportion to how many candidate families each part of the index actually holds, so the sample stays representative of the whole index. If the index has fewer than `N` candidate families, `--subsample` has no effect.
|
||||
|
||||
### `--shannon`: measuring how informative a family is
|
||||
|
||||
Generated
+147
-1
@@ -213,6 +213,29 @@ version = "3.20.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb"
|
||||
|
||||
[[package]]
|
||||
name = "bytecheck"
|
||||
version = "0.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "26333eeac754f0ad8a6bcd0eb0ac012156302e4e16b852b72ee399aea4f12c29"
|
||||
dependencies = [
|
||||
"bytecheck_derive",
|
||||
"ptr_meta",
|
||||
"rancor",
|
||||
"simdutf8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bytecheck_derive"
|
||||
version = "0.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "46d07918caa9eeaaf06b7873925c53a61daac173539b4f7715090745e44e4e69"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "byteorder"
|
||||
version = "1.5.0"
|
||||
@@ -1236,6 +1259,26 @@ dependencies = [
|
||||
"windows 0.48.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "munge"
|
||||
version = "0.4.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5e17401f259eba956ca16491461b6e8f72913a0a114e39736ce404410f915a0c"
|
||||
dependencies = [
|
||||
"munge_macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "munge_macro"
|
||||
version = "0.4.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4568f25ccbd45ab5d5603dc34318c1ec56b117531781260002151b8530a9f931"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nanorand"
|
||||
version = "0.6.1"
|
||||
@@ -1480,7 +1523,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "obikmer"
|
||||
version = "1.3.4"
|
||||
version = "1.3.5"
|
||||
dependencies = [
|
||||
"clap",
|
||||
"csv",
|
||||
@@ -1499,6 +1542,7 @@ dependencies = [
|
||||
"obikrope",
|
||||
"obikselect",
|
||||
"obikseq",
|
||||
"obiksession",
|
||||
"obikstats",
|
||||
"obipipeline",
|
||||
"obiread",
|
||||
@@ -1541,6 +1585,7 @@ dependencies = [
|
||||
"obikindex",
|
||||
"obikindexer",
|
||||
"obikseq",
|
||||
"obiksession",
|
||||
"obipipeline",
|
||||
"obiread",
|
||||
"obiskbuilder",
|
||||
@@ -1549,6 +1594,7 @@ dependencies = [
|
||||
"petgraph",
|
||||
"rand 0.10.2",
|
||||
"rayon",
|
||||
"rkyv",
|
||||
"speedytree",
|
||||
"tempfile",
|
||||
"tracing",
|
||||
@@ -1615,6 +1661,17 @@ dependencies = [
|
||||
"xxhash-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "obiksession"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"memmap2",
|
||||
"obisys",
|
||||
"tempfile",
|
||||
"tracing",
|
||||
"xxhash-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "obikstats"
|
||||
version = "0.1.0"
|
||||
@@ -1963,6 +2020,26 @@ dependencies = [
|
||||
"xxhash-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ptr_meta"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "743da816b98c921cdbe8628ef7381b76f25ecf4da599fc80aca90eae7ef70cc0"
|
||||
dependencies = [
|
||||
"ptr_meta_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ptr_meta_derive"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1c8d9ca532f185d5d4db7a7c9d51420b452168ea1c2b913953281bd6fe1fcbd0"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.45"
|
||||
@@ -1990,6 +2067,15 @@ version = "0.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09"
|
||||
|
||||
[[package]]
|
||||
name = "rancor"
|
||||
version = "0.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9b534442d0fcdb55d66f373d9cac6d33b6293a2335bc2136dbd06ce0e87d2572"
|
||||
dependencies = [
|
||||
"ptr_meta",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.8.6"
|
||||
@@ -2151,6 +2237,15 @@ version = "0.8.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
|
||||
|
||||
[[package]]
|
||||
name = "rend"
|
||||
version = "0.5.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "663ba70707f96e871406fe10d68128412e619b06d1d47cb91c3a4c6501176240"
|
||||
dependencies = [
|
||||
"bytecheck",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
version = "0.17.14"
|
||||
@@ -2165,6 +2260,36 @@ dependencies = [
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rkyv"
|
||||
version = "0.8.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d9776093b7ca170454ab1406954f7b7d97a57c51dc6c0642957fb2ef25c2d399"
|
||||
dependencies = [
|
||||
"bytecheck",
|
||||
"bytes",
|
||||
"hashbrown",
|
||||
"indexmap",
|
||||
"munge",
|
||||
"ptr_meta",
|
||||
"rancor",
|
||||
"rend",
|
||||
"rkyv_derive",
|
||||
"tinyvec",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rkyv_derive"
|
||||
version = "0.8.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1c25ef604ac7dd839d44d64648952ea23c97866f124ff671b0ed2cf3ad9bb06e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-hash"
|
||||
version = "2.1.2"
|
||||
@@ -2359,6 +2484,12 @@ version = "0.3.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214"
|
||||
|
||||
[[package]]
|
||||
name = "simdutf8"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e"
|
||||
|
||||
[[package]]
|
||||
name = "smallvec"
|
||||
version = "1.15.1"
|
||||
@@ -2587,6 +2718,21 @@ dependencies = [
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tinyvec"
|
||||
version = "1.13.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4cf0ded5c4e56918d8f8a339e1bb67d038d3bc6d144ac407904015ba2e4cde9b"
|
||||
dependencies = [
|
||||
"tinyvec_macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tinyvec_macros"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
|
||||
|
||||
[[package]]
|
||||
name = "tracing"
|
||||
version = "0.1.44"
|
||||
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
[workspace]
|
||||
resolver = "3"
|
||||
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obiskio","obidebruinj", "obicompactvec", "obisys", "obikindex", "obikindexer", "obikquery", "obikdump", "obikfilter", "obikselect", "obikrebuild", "obikmerge", "obikstats", "obikidxcache", "obitaxonomy", "obikentropy", "obikphylo", "obikalgorithm"]
|
||||
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obiskio","obidebruinj", "obicompactvec", "obisys", "obikindex", "obikindexer", "obikquery", "obikdump", "obikfilter", "obikselect", "obikrebuild", "obikmerge", "obikstats", "obikidxcache", "obitaxonomy", "obikentropy", "obikphylo", "obikalgorithm", "obiksession"]
|
||||
[profile.release]
|
||||
debug = 1
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "obikmer"
|
||||
version = "1.3.4"
|
||||
version = "1.3.5"
|
||||
edition = "2024"
|
||||
|
||||
[[bin]]
|
||||
@@ -24,6 +24,7 @@ obikstats = { path = "../obikstats" }
|
||||
obikquery = { path = "../obikquery" }
|
||||
obikidxcache = { path = "../obikidxcache" }
|
||||
obikphylo = { path = "../obikphylo" }
|
||||
obiksession = { path = "../obiksession" }
|
||||
obikrope = { path = "../obikrope" }
|
||||
obifastwrite = { path = "../obifastwrite" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
|
||||
@@ -1,7 +1,36 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
use clap::Args;
|
||||
use obikphylo::{DistanceMetric, SnpDistanceKind};
|
||||
use obikphylo::{DistanceMetric, GammaShape, SnpDistanceKind};
|
||||
|
||||
/// `--gamma-shape` CLI value: a fixed `alpha`, or `auto`/`estimate` to fit
|
||||
/// it from the data (see [`GammaShape::Auto`]). Kept as its own tiny type
|
||||
/// (rather than parsing straight to `GammaShape`) so the CLI layer stays
|
||||
/// free of `GammaShape::Disabled`, which the `Option<GammaShapeArg>` around
|
||||
/// this already expresses on its own (`None` = flag not given at all).
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum GammaShapeArg {
|
||||
Fixed(f64),
|
||||
Auto,
|
||||
}
|
||||
|
||||
impl From<GammaShapeArg> for GammaShape {
|
||||
fn from(arg: GammaShapeArg) -> Self {
|
||||
match arg {
|
||||
GammaShapeArg::Fixed(a) => GammaShape::Fixed(a),
|
||||
GammaShapeArg::Auto => GammaShape::Auto,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_gamma_shape(s: &str) -> Result<GammaShapeArg, String> {
|
||||
if s.eq_ignore_ascii_case("auto") || s.eq_ignore_ascii_case("estimate") {
|
||||
return Ok(GammaShapeArg::Auto);
|
||||
}
|
||||
s.parse::<f64>()
|
||||
.map(GammaShapeArg::Fixed)
|
||||
.map_err(|_| format!("invalid --gamma-shape value {s:?}: expected a positive number or \"auto\"/\"estimate\""))
|
||||
}
|
||||
|
||||
/// `--distance` value — either one of `obikphylo::DistanceMetric`'s
|
||||
/// whole-index metrics (routed to `IndexCache::distance`) or one of
|
||||
@@ -195,6 +224,27 @@ pub struct PhyloArgs {
|
||||
#[arg(long)]
|
||||
pub entropy_sd: Option<f64>,
|
||||
|
||||
/// Persist the `snp-*` `--distance` sample (the pairwise substitution
|
||||
/// tally) in `DIR` across separate `obikmer phylo` invocations: a later
|
||||
/// run against the same index with the exact same
|
||||
/// `--subsample`/`--free-loss`/`--no-ambiguity`/`--exclude-genome`/
|
||||
/// `--min-shared-family`/`--entropy`/`--entropy-sd` restores it instead
|
||||
/// of resampling. `DIR` is created if it doesn't exist. If it already
|
||||
/// holds a sample built under *different* selection parameters, this
|
||||
/// run is rejected with an error listing what changed, unless
|
||||
/// `--session-force` is also given (which discards the old sample and
|
||||
/// starts fresh under this run's parameters). Only covers a `snp-*`
|
||||
/// `--distance` today — `--sankoff`/`--tnt`/`--phyg`/`--iqtree` don't
|
||||
/// read or write a `--session` yet.
|
||||
#[arg(long, value_name = "DIR")]
|
||||
pub session: Option<PathBuf>,
|
||||
|
||||
/// With `--session DIR`: overwrite its saved selection parameters (and
|
||||
/// discard its cached sample) instead of erroring out when this run's
|
||||
/// parameters don't match. No effect without `--session`.
|
||||
#[arg(long, requires = "session")]
|
||||
pub session_force: bool,
|
||||
|
||||
/// Calibrate a 16-state Sankoff cost matrix (and its matching
|
||||
/// pseudo-alignment) from an already-built sibling annex — requires
|
||||
/// `--subsample <N>`, and shares `--free-loss`/`--no-ambiguity`/
|
||||
@@ -292,8 +342,11 @@ pub struct PhyloArgs {
|
||||
/// `snp-raw`/`snp-tv`, which have nothing to correct/are deliberately
|
||||
/// uncorrected). Has no effect on the whole-index metrics. Rejected at
|
||||
/// runtime if given alongside an unsupported `--distance` value.
|
||||
#[arg(long, value_name = "ALPHA")]
|
||||
pub gamma_shape: Option<f64>,
|
||||
/// Either a positive number (fixed `alpha`) or `auto`/`estimate` to fit
|
||||
/// `alpha` from the data itself — method-of-moments over the
|
||||
/// over-dispersion of substitution rate across partitions.
|
||||
#[arg(long, value_name = "ALPHA|auto", value_parser = parse_gamma_shape)]
|
||||
pub gamma_shape: Option<GammaShapeArg>,
|
||||
|
||||
/// Minimum count to consider a kmer present when computing Jaccard on count indexes
|
||||
#[arg(long, default_value = "1")]
|
||||
|
||||
@@ -11,8 +11,8 @@ use std::sync::Arc;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::KmerIndex;
|
||||
use obikphylo::siblings::{
|
||||
EntropyBias, SiblingExt, cardinality_transition_probs, composition_transition_probs,
|
||||
pairwise_cost_matrix,
|
||||
EntropyBias, GammaShape, SiblingExt, cardinality_transition_probs,
|
||||
composition_transition_probs, pairwise_cost_matrix,
|
||||
};
|
||||
use obikphylo::{Metrics, neighbor_joining, upgma};
|
||||
use obisys::{Reporter, Stage};
|
||||
@@ -24,6 +24,43 @@ use phylip::write_phylip_relaxed;
|
||||
use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params};
|
||||
use tnt::write_sankoff_tnt;
|
||||
|
||||
/// Deterministic byte encoding of the `snp-*` `--distance` selection
|
||||
/// parameters, for `obiksession::Session::open`'s params-conflict check.
|
||||
/// `obiksession` treats this as an opaque blob (see its own crate docs) —
|
||||
/// only this module needs to know the layout, and only well enough to
|
||||
/// produce/compare it, never to pretty-print a per-field diff (v1 keeps
|
||||
/// the conflict message generic rather than decoding both sides field by
|
||||
/// field — see its use in `run` below).
|
||||
fn session_params_blob(
|
||||
n: Option<usize>,
|
||||
free_loss: bool,
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
) -> Vec<u8> {
|
||||
let mut buf = Vec::new();
|
||||
match n {
|
||||
Some(v) => {
|
||||
buf.push(1);
|
||||
buf.extend_from_slice(&(v as u64).to_le_bytes());
|
||||
}
|
||||
None => buf.push(0),
|
||||
}
|
||||
buf.push(free_loss as u8);
|
||||
buf.push(no_ambiguity as u8);
|
||||
buf.extend_from_slice(&(excluded.len() as u64).to_le_bytes());
|
||||
buf.extend(excluded.iter().map(|&b| b as u8));
|
||||
match entropy_bias {
|
||||
Some(EntropyBias { mu, sigma }) => {
|
||||
buf.push(1);
|
||||
buf.extend_from_slice(&mu.to_le_bytes());
|
||||
buf.extend_from_slice(&sigma.to_le_bytes());
|
||||
}
|
||||
None => buf.push(0),
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
pub use args::PhyloArgs;
|
||||
|
||||
pub fn run(args: PhyloArgs) {
|
||||
@@ -236,6 +273,47 @@ pub fn run(args: PhyloArgs) {
|
||||
None
|
||||
};
|
||||
|
||||
// ── `--session`: persist the snp-* sample across separate invocations ──
|
||||
// Only covers `snp-*` `--distance` today (see `--session`'s own docs);
|
||||
// `None` when `--session` isn't given, matching today's always-resample
|
||||
// default exactly.
|
||||
let session = args.session.as_ref().map(|dir| {
|
||||
let params = session_params_blob(
|
||||
args.subsample,
|
||||
args.free_loss,
|
||||
args.no_ambiguity,
|
||||
&snp_exclude_mask,
|
||||
entropy_bias,
|
||||
);
|
||||
match obiksession::Session::open(dir, ¶ms, args.session_force) {
|
||||
Ok(Ok(session)) => {
|
||||
match session.outcome() {
|
||||
obiksession::SessionOutcome::Created => {
|
||||
info!("--session {}: starting fresh", dir.display());
|
||||
}
|
||||
obiksession::SessionOutcome::Reused => {
|
||||
info!("--session {}: found, will reuse its cached sample if present", dir.display());
|
||||
}
|
||||
}
|
||||
session
|
||||
}
|
||||
Ok(Err(_conflict)) => {
|
||||
eprintln!(
|
||||
"error: --session {} was built with different selection parameters \
|
||||
(--subsample/--free-loss/--no-ambiguity/--exclude-genome/--min-shared-family/\
|
||||
--entropy/--entropy-sd) than this run — pass --session-force to discard its \
|
||||
cached sample and start fresh, or point --session at a different directory",
|
||||
dir.display()
|
||||
);
|
||||
std::process::exit(1);
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("error opening --session {}: {e}", dir.display());
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// ── SNP pseudo-alignment (`--pseudo-alignment`) ─────────────────────────────
|
||||
if args.pseudo_alignment {
|
||||
let Some(subsample_n) = args.subsample else {
|
||||
@@ -268,6 +346,14 @@ pub fn run(args: PhyloArgs) {
|
||||
}
|
||||
|
||||
// ── Sankoff cost-matrix calibration (`--sankoff`, `--tnt`, `--phyg`, `--iqtree`) ──
|
||||
// `bundle` is kept around (not dropped at the end of this block) so the
|
||||
// `snp-*` `--distance` branch below can reuse its already-sampled
|
||||
// tally instead of resampling — see `SankoffBundle::snp_distance`'s own
|
||||
// docs for why that's required for correctness, not just speed, when
|
||||
// both are requested in the same invocation (they always share the
|
||||
// same `n`/`free_loss`/`no_ambiguity`/excluded-set/`entropy_bias`: the
|
||||
// CLI has no way to give them different values in one run).
|
||||
let mut bundle = None;
|
||||
if args.sankoff || args.tnt || args.phyg || args.iqtree {
|
||||
let Some(subsample_n) = args.subsample else {
|
||||
eprintln!("error: --sankoff requires --subsample <N>");
|
||||
@@ -276,7 +362,7 @@ pub fn run(args: PhyloArgs) {
|
||||
|
||||
info!("sampling Sankoff calibration bundle (target {subsample_n} site(s))");
|
||||
let t = Stage::start("sankoff_bundle");
|
||||
let bundle = cache
|
||||
let b = cache
|
||||
.sankoff_bundle(
|
||||
subsample_n,
|
||||
args.free_loss,
|
||||
@@ -284,6 +370,7 @@ pub fn run(args: PhyloArgs) {
|
||||
&snp_exclude_mask,
|
||||
entropy_bias,
|
||||
args.sankoff_ratio_ceiling,
|
||||
session.as_ref(),
|
||||
)
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!("error computing Sankoff calibration bundle: {e}");
|
||||
@@ -291,25 +378,25 @@ pub fn run(args: PhyloArgs) {
|
||||
});
|
||||
rep.push(t.stop());
|
||||
|
||||
let p_card = cardinality_transition_probs(&bundle.cardinality_tally);
|
||||
let p_comp = composition_transition_probs(&bundle.base_pair_tally);
|
||||
let p_card = cardinality_transition_probs(&b.cardinality_tally);
|
||||
let p_comp = composition_transition_probs(&b.base_pair_tally);
|
||||
let matrix = pairwise_cost_matrix(&p_card, &p_comp, args.free_loss);
|
||||
|
||||
write_sankoff_matrix_csv(&matrix, &args.output);
|
||||
write_sankoff_params(
|
||||
&bundle.cardinality_tally,
|
||||
&b.cardinality_tally,
|
||||
&p_card,
|
||||
&bundle.base_pair_tally,
|
||||
&b.base_pair_tally,
|
||||
&p_comp,
|
||||
args.sankoff_ratio_ceiling,
|
||||
&args.output,
|
||||
);
|
||||
write_sankoff_alignment_fasta(&bundle.alignment, &labels, &args.output, args.free_loss);
|
||||
write_sankoff_alignment_fasta(&b.alignment, &labels, &args.output, args.free_loss);
|
||||
|
||||
if args.tnt {
|
||||
write_sankoff_tnt(
|
||||
&matrix,
|
||||
&bundle.alignment,
|
||||
&b.alignment,
|
||||
&labels,
|
||||
&args.output,
|
||||
args.sankoff_cost_scale,
|
||||
@@ -322,13 +409,14 @@ pub fn run(args: PhyloArgs) {
|
||||
if args.iqtree {
|
||||
write_iqtree(
|
||||
&matrix,
|
||||
&bundle.alignment,
|
||||
&b.alignment,
|
||||
&labels,
|
||||
&args.output,
|
||||
args.free_loss,
|
||||
args.iqtree_min_freq,
|
||||
);
|
||||
}
|
||||
bundle = Some(b);
|
||||
}
|
||||
|
||||
// ── Distance computation: classic whole-index metric vs. `snp-*` ───────────
|
||||
@@ -354,28 +442,40 @@ pub fn run(args: PhyloArgs) {
|
||||
std::process::exit(1);
|
||||
}
|
||||
let kind = args.distance.as_snp().expect("DistanceArg is always classic or snp");
|
||||
info!(
|
||||
"computing {kind:?} SNP distance for {n} genome(s){}",
|
||||
match args.subsample {
|
||||
Some(n) => format!(" (subsampled, target {n} site(s))"),
|
||||
None => " (exhaustive)".into(),
|
||||
}
|
||||
);
|
||||
let gamma_shape = args.gamma_shape.map_or(GammaShape::Disabled, Into::into);
|
||||
let t = Stage::start("snp_distance");
|
||||
let matrix = cache
|
||||
.snp_distance(
|
||||
kind,
|
||||
args.subsample,
|
||||
args.free_loss,
|
||||
args.no_ambiguity,
|
||||
&snp_exclude_mask,
|
||||
entropy_bias,
|
||||
args.gamma_shape,
|
||||
)
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!("error computing SNP distance: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
let matrix = match &bundle {
|
||||
Some(b) => {
|
||||
info!(
|
||||
"computing {kind:?} SNP distance for {n} genome(s) \
|
||||
(reusing the Sankoff bundle's sample)"
|
||||
);
|
||||
b.snp_distance(kind, gamma_shape)
|
||||
}
|
||||
None => {
|
||||
info!(
|
||||
"computing {kind:?} SNP distance for {n} genome(s){}",
|
||||
match args.subsample {
|
||||
Some(n) => format!(" (subsampled, target {n} site(s))"),
|
||||
None => " (exhaustive)".into(),
|
||||
}
|
||||
);
|
||||
cache.snp_distance(
|
||||
kind,
|
||||
args.subsample,
|
||||
args.free_loss,
|
||||
args.no_ambiguity,
|
||||
&snp_exclude_mask,
|
||||
entropy_bias,
|
||||
gamma_shape,
|
||||
session.as_ref(),
|
||||
)
|
||||
}
|
||||
}
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!("error computing SNP distance: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
(matrix, None)
|
||||
}
|
||||
|
||||
@@ -12,12 +12,14 @@ obicompactvec = { path = "../obicompactvec" }
|
||||
obikidxcache = { path = "../obikidxcache" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
obipipeline = { path = "../obipipeline" }
|
||||
obiksession = { path = "../obiksession" }
|
||||
kodama = "0.3.0"
|
||||
memmap2 = "0.9"
|
||||
ndarray = "0.17"
|
||||
petgraph = "0.6.4"
|
||||
rand = "0.10"
|
||||
rayon = "1"
|
||||
rkyv = "0.8.18"
|
||||
speedytree = "0.1"
|
||||
tracing = "0.1.44"
|
||||
|
||||
|
||||
@@ -19,5 +19,5 @@ mod tree;
|
||||
pub mod siblings;
|
||||
|
||||
pub use distance::{DistanceMetric, DistanceOutput, Metrics};
|
||||
pub use siblings::SnpDistanceKind;
|
||||
pub use siblings::{GammaShape, SnpDistanceKind};
|
||||
pub use tree::{Tree, neighbor_joining, upgma};
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
use rkyv::{Archive, Deserialize, Serialize};
|
||||
|
||||
use super::masking::{iupac_code, masked_state};
|
||||
use super::subsample::{EntropyBias, SurvivingFamily, sample_index};
|
||||
@@ -17,6 +18,10 @@ use super::subsample::{EntropyBias, SurvivingFamily, sample_index};
|
||||
/// index simply doesn't appear in `genome_indices`. `pub`, not
|
||||
/// `pub(crate)`: part of the public signature of
|
||||
/// [`crate::siblings::extensions::SiblingExt::snp_pseudo_alignment`].
|
||||
///
|
||||
/// `Archive`/`Serialize`/`Deserialize` (`rkyv`): `--session` persistence
|
||||
/// for `SankoffBundle` — same reasoning as `PairwiseTally`'s own derive.
|
||||
#[derive(Archive, Serialize, Deserialize)]
|
||||
pub struct SnpAlignment {
|
||||
pub sequences: Vec<Vec<u8>>,
|
||||
/// This index's own genome numbering (matching `IndexCache::meta().genomes()`'s
|
||||
@@ -54,7 +59,8 @@ pub(crate) fn snp_pseudo_alignment(
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|_partition, _layer, survivors| {
|
||||
0,
|
||||
|_partition, _layer, _raw_index, survivors| {
|
||||
reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences);
|
||||
},
|
||||
)?;
|
||||
|
||||
@@ -22,6 +22,7 @@ mod masking;
|
||||
mod minorant_selection;
|
||||
mod pairwise;
|
||||
mod sankoff;
|
||||
mod session_cache;
|
||||
mod snp_distance;
|
||||
mod stats;
|
||||
mod subsample;
|
||||
@@ -34,7 +35,7 @@ pub use cardcomp::{
|
||||
};
|
||||
pub use pairwise::{BasePairTally, CardinalityTally, RawSnpDistanceOutput};
|
||||
pub use sankoff::SankoffBundle;
|
||||
pub use snp_distance::SnpDistanceKind;
|
||||
pub use snp_distance::{GammaShape, SnpDistanceKind};
|
||||
pub use stats::SiblingAnnexStats;
|
||||
pub use subsample::{EntropyBias, SurvivingFamily};
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@
|
||||
//! machines), not something to optimise away.
|
||||
|
||||
use ndarray::Array2;
|
||||
use rkyv::{Archive, Deserialize, Serialize};
|
||||
|
||||
use super::subsample::SurvivingFamily;
|
||||
|
||||
@@ -30,7 +31,12 @@ use super::subsample::SurvivingFamily;
|
||||
/// [`BasePairTally`] the latter (they are *not* the same population — see
|
||||
/// `BasePairTally::same`'s own docs). `cardinality` is restricted to
|
||||
/// variable families only, matching [`CardinalityTally`]'s own scope.
|
||||
#[derive(Default, Clone, Copy)]
|
||||
///
|
||||
/// `Archive`/`Serialize`/`Deserialize`: part of [`PairwiseTally`]'s own
|
||||
/// `--session` persistence (see that struct's docs) — `rkyv` needs every
|
||||
/// field type to derive these too, hence deriving here on an otherwise
|
||||
/// purely internal struct.
|
||||
#[derive(Default, Clone, Copy, Archive, Serialize, Deserialize)]
|
||||
struct PairStats {
|
||||
subst: [[u64; 4]; 4],
|
||||
same_all: [u64; 4],
|
||||
@@ -40,6 +46,13 @@ struct PairStats {
|
||||
|
||||
/// See the module docs. `pub(crate)`: consumed by
|
||||
/// `algorithms::sankoff`'s orchestration, not yet exposed past this crate.
|
||||
///
|
||||
/// `Archive`/`Serialize`/`Deserialize` (`rkyv`): lets `--session` dump/
|
||||
/// restore a tally across separate CLI invocations as opaque bytes handed
|
||||
/// to `obiksession::Session::store`/`restore` — `obiksession` itself knows
|
||||
/// nothing about this type, it only stores/checksums/mmaps the bytes this
|
||||
/// crate produces (see `obiksession`'s own crate docs on that split).
|
||||
#[derive(Archive, Serialize, Deserialize)]
|
||||
pub(crate) struct PairwiseTally {
|
||||
n_genomes: usize,
|
||||
/// Upper triangle only (`i < j`), flat-indexed via [`Self::flat_index`].
|
||||
@@ -52,6 +65,10 @@ impl PairwiseTally {
|
||||
Self { n_genomes, pairs: vec![PairStats::default(); n_pairs] }
|
||||
}
|
||||
|
||||
pub(crate) fn n_genomes(&self) -> usize {
|
||||
self.n_genomes
|
||||
}
|
||||
|
||||
/// See `crate::siblings::helpers::triangle_index`.
|
||||
fn flat_index(&self, i: usize, j: usize) -> usize {
|
||||
crate::siblings::helpers::triangle_index(self.n_genomes, i, j)
|
||||
@@ -273,14 +290,137 @@ pub struct CardinalityTally {
|
||||
pub counts: [[u64; 5]; 5],
|
||||
}
|
||||
|
||||
/// One layer's worth of [`SurvivingFamily`] folded into `tally` — the
|
||||
/// "reduce" half of the map/reduce [`super::subsample::sample_index`]
|
||||
/// drives, self-contained (no sampling/masking logic of its own) so it can
|
||||
/// run alongside any other reduce over the same batch (see
|
||||
/// `algorithms::alignment::reduce_alignment`, called from the same
|
||||
/// `on_layer` closure in `algorithms::sankoff::sankoff_bundle`).
|
||||
pub(crate) fn reduce_pairwise(survivors: &[SurvivingFamily], tally: &mut PairwiseTally) {
|
||||
/// Per-partition, pooled-over-every-genome-pair `(substitutions, eligible
|
||||
/// loci)` counts — feeds [`Self::estimate_alpha`], the method-of-moments
|
||||
/// `--gamma-shape auto` estimator (an independent derivation — see that
|
||||
/// method's own docs for why this is *not* Jin & Nei's (1990) method,
|
||||
/// verified against their paper directly, despite using the same `+Γ`
|
||||
/// distance formula). Pooling across pairs (rather than keeping one series
|
||||
/// per pair) treats every included genome pair as one more observation of
|
||||
/// the same underlying rate-heterogeneity signal — comparable in spirit to
|
||||
/// combining information across several genomic regions/genes for a fixed
|
||||
/// set of taxa, generalised here to every pair at once.
|
||||
///
|
||||
/// Deliberately **not** gated by `--sankoff-ratio-ceiling` the way
|
||||
/// [`PairwiseTally::base_pair_tally`] is (see
|
||||
/// [`PairwiseTally::cardinality_tally`]'s own docs for the same
|
||||
/// precedent): that filter excludes individual saturated *pairs* from a
|
||||
/// composition estimate computed once at the very end, from the complete
|
||||
/// `PairwiseTally`; the partition axis needed here only exists transiently,
|
||||
/// one partition at a time, while `PairwiseTally` is still being built —
|
||||
/// long before any pair's final SNP ratio (and thus its ratio_ceiling
|
||||
/// eligibility) is known. Genome exclusion (`--exclude-genome`) isn't
|
||||
/// applied here either, matching `reduce_pairwise`'s own raw per-pair fold
|
||||
/// (`PairStats` accumulates over every genome pair unconditionally;
|
||||
/// exclusion is a filter applied only in derived, post-hoc views).
|
||||
#[derive(Default, Clone, Archive, Serialize, Deserialize)]
|
||||
pub(crate) struct PartitionDispersion {
|
||||
/// One entry per partition id seen so far (grows lazily — partitions
|
||||
/// are visited in increasing order in practice, but nothing here
|
||||
/// assumes that).
|
||||
counts: Vec<(u64, u64)>,
|
||||
}
|
||||
|
||||
impl PartitionDispersion {
|
||||
/// Add one partition/layer's worth of pooled counts (from one
|
||||
/// [`reduce_pairwise`] call) to that partition's running total —
|
||||
/// additive so a partition with several layers accumulates correctly
|
||||
/// across multiple calls.
|
||||
fn record(&mut self, partition: usize, subst: u64, eligible: u64) {
|
||||
if partition >= self.counts.len() {
|
||||
self.counts.resize(partition + 1, (0, 0));
|
||||
}
|
||||
let entry = &mut self.counts[partition];
|
||||
entry.0 += subst;
|
||||
entry.1 += eligible;
|
||||
}
|
||||
|
||||
/// Method-of-moments estimate of the gamma shape parameter `alpha`
|
||||
/// used by the `+Γ` distance correction, from the over-dispersion of
|
||||
/// per-partition substitution *rates* relative to what a single shared
|
||||
/// rate would produce under pure Poisson sampling noise — the general
|
||||
/// negative-binomial (Poisson-Gamma mixture) identity, the same kind of
|
||||
/// reasoning behind Uzzell & Corbin's (1971) original observation that
|
||||
/// substitution counts across sites/regions are over-dispersed
|
||||
/// relative to Poisson.
|
||||
///
|
||||
/// **Not Jin & Nei's (1990) method** — checked directly against their
|
||||
/// paper (*"Limitations of the Evolutionary Parsimony Method of
|
||||
/// Phylogenetic Analysis"*, Mol. Biol. Evol. 7(2):82-102, the source of
|
||||
/// the `+Γ` distance formula this alpha feeds into, confirmed to match
|
||||
/// `corrected_log`/`k2p` exactly against their eq. A4/A8): that paper
|
||||
/// contains no data-driven alpha estimator. Their own recommendation
|
||||
/// (p. 98) is the fixed default `alpha = 1`; for estimating alpha from
|
||||
/// data they point to a different paper (Wilson et al. 1989, not
|
||||
/// read/verified here) rather than proposing their own procedure. This
|
||||
/// estimator is therefore a from-scratch derivation (Poisson count
|
||||
/// within a partition, with a Gamma(alpha, alpha)-distributed (mean 1)
|
||||
/// multiplicative rate shared by every locus in that partition), not a
|
||||
/// literature implementation — unlike every closed-form correction in
|
||||
/// `snp_distance.rs`, each verified line-by-line against `ape`'s
|
||||
/// `dist_dna.c`.
|
||||
///
|
||||
/// Partitions have unequal eligible-locus counts (`l_i`), so a plain
|
||||
/// unweighted variance across raw per-partition rates would conflate
|
||||
/// genuine rate heterogeneity with the extra Poisson noise smaller
|
||||
/// partitions carry. Weighting each partition's squared deviation by
|
||||
/// its own `l_i` cancels that: for `R_i = subst_i / l_i`, `E[R_i] =
|
||||
/// mu`, `Var[R_i] = mu/l_i + mu²/alpha`, so `E[Σ l_i (R_i-mu)² / Σ l_i]
|
||||
/// = mu/mean(l) + mu²/alpha` — the Poisson floor
|
||||
/// (`mu/mean(l)`) is subtracted below before solving for `alpha`.
|
||||
///
|
||||
/// Returns `None` when there's nothing to estimate from: fewer than 2
|
||||
/// partitions with data, a pooled mean rate of exactly 0, or a
|
||||
/// measured variance at or below the Poisson floor (no detectable
|
||||
/// over-dispersion — `alpha` would be unbounded, not just large).
|
||||
pub(crate) fn estimate_alpha(&self) -> Option<f64> {
|
||||
let rows: Vec<(f64, f64)> = self
|
||||
.counts
|
||||
.iter()
|
||||
.filter(|&&(_, eligible)| eligible > 0)
|
||||
.map(|&(subst, eligible)| (subst as f64 / eligible as f64, eligible as f64))
|
||||
.collect();
|
||||
if rows.len() < 2 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let weight_total: f64 = rows.iter().map(|&(_, w)| w).sum();
|
||||
let mean: f64 = rows.iter().map(|&(rate, w)| rate * w).sum::<f64>() / weight_total;
|
||||
if mean <= 0.0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let weighted_variance: f64 =
|
||||
rows.iter().map(|&(rate, w)| w * (rate - mean).powi(2)).sum::<f64>() / weight_total;
|
||||
let mean_eligible = weight_total / rows.len() as f64;
|
||||
let poisson_floor = mean / mean_eligible;
|
||||
let excess = weighted_variance - poisson_floor;
|
||||
if excess <= 0.0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(mean * mean / excess)
|
||||
}
|
||||
}
|
||||
|
||||
/// One layer's worth of [`SurvivingFamily`] folded into `tally` (and
|
||||
/// `dispersion`, pooled under `partition` — see [`PartitionDispersion`]'s
|
||||
/// own docs for why the partition axis must be captured here rather than
|
||||
/// after the fact) — the "reduce" half of the map/reduce
|
||||
/// [`super::subsample::sample_index`] drives, self-contained (no
|
||||
/// sampling/masking logic of its own) so it can run alongside any other
|
||||
/// reduce over the same batch (see `algorithms::alignment::reduce_alignment`,
|
||||
/// called from the same `on_layer` closure in
|
||||
/// `algorithms::sankoff::sankoff_bundle`).
|
||||
pub(crate) fn reduce_pairwise(
|
||||
survivors: &[SurvivingFamily],
|
||||
partition: usize,
|
||||
tally: &mut PairwiseTally,
|
||||
dispersion: &mut PartitionDispersion,
|
||||
) {
|
||||
let n_genomes = tally.n_genomes;
|
||||
let mut partition_subst: u64 = 0;
|
||||
let mut partition_eligible: u64 = 0;
|
||||
for family in survivors {
|
||||
let variable = family.mask.family_size() >= 2;
|
||||
let genome_mask = &family.genome_mask;
|
||||
@@ -296,7 +436,9 @@ pub(crate) fn reduce_pairwise(survivors: &[SurvivingFamily], tally: &mut Pairwis
|
||||
let stats = tally.pair_mut(i, j);
|
||||
|
||||
if let (Some(bi), Some(bj)) = (bi, single_form(j)) {
|
||||
partition_eligible += 1;
|
||||
if bi != bj {
|
||||
partition_subst += 1;
|
||||
stats.subst[bi as usize][bj as usize] += 1;
|
||||
stats.subst[bj as usize][bi as usize] += 1;
|
||||
} else {
|
||||
@@ -317,4 +459,59 @@ pub(crate) fn reduce_pairwise(survivors: &[SurvivingFamily], tally: &mut Pairwis
|
||||
}
|
||||
}
|
||||
}
|
||||
dispersion.record(partition, partition_subst, partition_eligible);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod dispersion_tests {
|
||||
use super::PartitionDispersion;
|
||||
|
||||
#[test]
|
||||
fn fewer_than_two_partitions_gives_no_estimate() {
|
||||
let mut d = PartitionDispersion::default();
|
||||
assert!(d.estimate_alpha().is_none());
|
||||
d.record(0, 5, 100);
|
||||
assert!(d.estimate_alpha().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_mean_rate_gives_no_estimate() {
|
||||
let mut d = PartitionDispersion::default();
|
||||
d.record(0, 0, 100);
|
||||
d.record(1, 0, 100);
|
||||
assert!(d.estimate_alpha().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identical_per_partition_rate_gives_no_estimate() {
|
||||
// Every partition sees exactly the same rate — variance across
|
||||
// partitions is *below* the Poisson floor (it's exactly 0), so
|
||||
// there's no over-dispersion to attribute to rate heterogeneity.
|
||||
let mut d = PartitionDispersion::default();
|
||||
for _ in 0..5 {
|
||||
d.record(d.counts.len(), 10, 100);
|
||||
}
|
||||
assert!(d.estimate_alpha().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clearly_over_dispersed_rates_give_a_finite_positive_alpha() {
|
||||
// Two partitions with no substitutions, two with twice the pooled
|
||||
// mean rate — far more spread than Poisson noise alone would
|
||||
// produce at these counts, so a finite alpha should come out.
|
||||
let mut d = PartitionDispersion::default();
|
||||
d.record(0, 0, 100);
|
||||
d.record(1, 20, 100);
|
||||
d.record(2, 0, 100);
|
||||
d.record(3, 20, 100);
|
||||
|
||||
// Hand-computed from the weighted-moments formula in
|
||||
// `estimate_alpha`'s own docs: mean = 0.1, weighted variance =
|
||||
// 0.01, Poisson floor = mean/mean_eligible = 0.001, alpha =
|
||||
// mean² / (variance - floor) = 0.01 / 0.009.
|
||||
let expected = 0.01 / 0.009;
|
||||
let alpha = d.estimate_alpha().expect("clear over-dispersion should yield an estimate");
|
||||
assert!(alpha > 0.0, "alpha should be positive, got {alpha}");
|
||||
assert!((alpha - expected).abs() < 1e-9, "got {alpha}, expected {expected}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,22 +10,80 @@
|
||||
//! that removes the old two-pass `raw_snp_distance`-before-`base_pair_tally`
|
||||
//! dependency entirely).
|
||||
|
||||
use ndarray::Array2;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obiksession::Session;
|
||||
|
||||
use super::alignment::{SnpAlignment, reduce_alignment};
|
||||
use super::pairwise::{BasePairTally, CardinalityTally, PairwiseTally, RawSnpDistanceOutput, reduce_pairwise};
|
||||
use super::pairwise::{
|
||||
BasePairTally, CardinalityTally, PairwiseTally, PartitionDispersion, RawSnpDistanceOutput,
|
||||
reduce_pairwise,
|
||||
};
|
||||
use super::session_cache::{
|
||||
CHECKPOINT_INTERVAL, CHECKPOINT_INTERVAL_LAYERS, PROGRESS_COMPLETE, restore_tally_checkpoint,
|
||||
store_tally_checkpoint,
|
||||
};
|
||||
use super::snp_distance::{GammaShape, SnpDistanceKind, distance_matrix};
|
||||
use super::subsample::{EntropyBias, sample_index};
|
||||
|
||||
/// The alignment's own `--session` artifact — kept in lockstep with the
|
||||
/// shared `pairwise_tally`/`partition_dispersion`/`progress` artifacts
|
||||
/// (`session_cache`): both are folded from the *same* `on_layer` call, so
|
||||
/// checkpointing them together, under the same `progress` marker, is what
|
||||
/// keeps a resumed alignment's column count consistent with how many
|
||||
/// layers the resumed tally actually reflects.
|
||||
const ALIGNMENT_ARTIFACT: &str = "snp_alignment";
|
||||
|
||||
fn restore_alignment(session: &Session) -> Option<SnpAlignment> {
|
||||
let bytes = session.restore(ALIGNMENT_ARTIFACT).ok().flatten()?;
|
||||
rkyv::from_bytes::<SnpAlignment, rkyv::rancor::Error>(&bytes[..]).ok()
|
||||
}
|
||||
|
||||
fn store_alignment(session: &Session, alignment: &SnpAlignment) -> OKIResult<()> {
|
||||
let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(alignment)
|
||||
.map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?;
|
||||
session.store(ALIGNMENT_ARTIFACT, &bytes).map_err(OKIError::Io)
|
||||
}
|
||||
|
||||
/// Every output the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` pipeline needs,
|
||||
/// computed together from one shared, possibly-subsampled/entropy-biased
|
||||
/// selection — see the module docs. `pub`: part of the public signature of
|
||||
/// [`crate::siblings::extensions::SiblingExt::sankoff_bundle`].
|
||||
///
|
||||
/// `tally`/`dispersion` are kept (not just consumed into `raw`/
|
||||
/// `base_pair_tally`/`cardinality_tally`) so [`Self::snp_distance`] can
|
||||
/// compute a `snp-*` `--distance` matrix from this *exact* sample — see
|
||||
/// that method's own docs for why reusing it, rather than a fresh
|
||||
/// `SiblingExt::snp_distance` call, is a correctness requirement here, not
|
||||
/// just a performance nicety.
|
||||
pub struct SankoffBundle {
|
||||
pub alignment: SnpAlignment,
|
||||
pub raw: RawSnpDistanceOutput,
|
||||
pub base_pair_tally: BasePairTally,
|
||||
pub cardinality_tally: CardinalityTally,
|
||||
tally: PairwiseTally,
|
||||
dispersion: PartitionDispersion,
|
||||
}
|
||||
|
||||
impl SankoffBundle {
|
||||
/// A `snp-*` `--distance` matrix computed from this bundle's own
|
||||
/// already-sampled tally — no new [`sample_index`] call.
|
||||
///
|
||||
/// Exists specifically for the case where `--sankoff`/`--tnt`/`--phyg`/
|
||||
/// `--iqtree` and a `snp-*` `--distance` are requested in the same CLI
|
||||
/// invocation: both consume exactly the same selection parameters
|
||||
/// (`n`/`free_loss`/`no_ambiguity`/excluded set/`entropy_bias` — the
|
||||
/// CLI has no way to give them different values in one run), so a
|
||||
/// second independent `SiblingExt::snp_distance` call wouldn't just
|
||||
/// waste the sampling work, it would silently sample a **different**
|
||||
/// set of sites (`sample_layer`'s `rand::rng()` is a thread-local
|
||||
/// generator that advances across calls, not reseeded each time) —
|
||||
/// defeating the actual point of running several algorithms together,
|
||||
/// which is comparing them on one identical site selection.
|
||||
pub fn snp_distance(&self, kind: SnpDistanceKind, gamma_shape: GammaShape) -> OKIResult<Array2<f64>> {
|
||||
distance_matrix(&self.tally, &self.dispersion, kind, gamma_shape)
|
||||
}
|
||||
}
|
||||
|
||||
/// See [`crate::siblings::extensions::SiblingExt::sankoff_bundle`] for the
|
||||
@@ -33,6 +91,14 @@ pub struct SankoffBundle {
|
||||
/// [`BasePairTally`] only (see [`PairwiseTally::cardinality_tally`]'s own
|
||||
/// docs for why [`CardinalityTally`] uses a different — `excluded`-only —
|
||||
/// inclusion rule).
|
||||
///
|
||||
/// `session`: `--session`, `None` means no persistence (always resample).
|
||||
/// `Some(session)` must already be open under the exact selection
|
||||
/// parameters this call is about to use — the caller's job, same contract
|
||||
/// as [`super::snp_distance::snp_distance`]'s own `session` parameter,
|
||||
/// which this shares an artifact namespace with (see `session_cache`'s
|
||||
/// module docs): a tally checkpointed here is restorable by a later
|
||||
/// `--distance snp-*` call against the same session, and vice versa.
|
||||
pub(crate) fn sankoff_bundle(
|
||||
cache: &IndexCache,
|
||||
n: usize,
|
||||
@@ -41,26 +107,72 @@ pub(crate) fn sankoff_bundle(
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
ratio_ceiling: f64,
|
||||
session: Option<&Session>,
|
||||
) -> OKIResult<SankoffBundle> {
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
let genome_indices: Vec<usize> = (0..n_genomes)
|
||||
.filter(|&g| !excluded.get(g).copied().unwrap_or(false))
|
||||
.collect();
|
||||
let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); genome_indices.len()];
|
||||
let mut tally = PairwiseTally::new(n_genomes);
|
||||
|
||||
sample_index(
|
||||
cache,
|
||||
n,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|_partition, _layer, survivors| {
|
||||
reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences);
|
||||
reduce_pairwise(&survivors, &mut tally);
|
||||
},
|
||||
)?;
|
||||
let restored = session.and_then(|s| Some((restore_tally_checkpoint(s)?, restore_alignment(s)?)));
|
||||
let (mut tally, mut dispersion, mut sequences, skip_layers) = match restored {
|
||||
Some(((tally, dispersion, skip_layers), alignment)) => {
|
||||
(tally, dispersion, alignment.sequences, skip_layers)
|
||||
}
|
||||
None => {
|
||||
let sequences = vec![Vec::new(); genome_indices.len()];
|
||||
(PairwiseTally::new(n_genomes), PartitionDispersion::default(), sequences, 0)
|
||||
}
|
||||
};
|
||||
|
||||
if skip_layers != usize::MAX {
|
||||
// Same checkpoint cadence and rationale as `snp_distance`'s own
|
||||
// resumable loop (see `session_cache::CHECKPOINT_INTERVAL*`'s own
|
||||
// docs) — the alignment's `sequences` are folded in lockstep with
|
||||
// the tally so a resumed run's column count always matches how
|
||||
// many layers the restored tally actually reflects.
|
||||
let mut last_checkpoint_layer = skip_layers;
|
||||
let mut last_checkpoint_at = std::time::Instant::now();
|
||||
sample_index(
|
||||
cache,
|
||||
n,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
skip_layers,
|
||||
|partition, _layer, raw_index, survivors| {
|
||||
reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences);
|
||||
reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
|
||||
let Some(session) = session else { return };
|
||||
let layers_done = raw_index + 1;
|
||||
let due_by_count =
|
||||
layers_done.saturating_sub(last_checkpoint_layer) >= CHECKPOINT_INTERVAL_LAYERS;
|
||||
let due_by_time = last_checkpoint_at.elapsed() >= CHECKPOINT_INTERVAL;
|
||||
if !(due_by_count || due_by_time) {
|
||||
return;
|
||||
}
|
||||
let alignment = SnpAlignment { sequences: sequences.clone(), genome_indices: genome_indices.clone() };
|
||||
let stored = store_tally_checkpoint(session, &tally, &dispersion, layers_done as u64)
|
||||
.and_then(|()| store_alignment(session, &alignment));
|
||||
match stored {
|
||||
Ok(()) => {
|
||||
last_checkpoint_layer = layers_done;
|
||||
last_checkpoint_at = std::time::Instant::now();
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("--session: failed to checkpoint at layer {layers_done}: {e}");
|
||||
}
|
||||
}
|
||||
},
|
||||
)?;
|
||||
|
||||
if let Some(session) = session {
|
||||
let alignment = SnpAlignment { sequences: sequences.clone(), genome_indices: genome_indices.clone() };
|
||||
store_tally_checkpoint(session, &tally, &dispersion, PROGRESS_COMPLETE)?;
|
||||
store_alignment(session, &alignment)?;
|
||||
}
|
||||
}
|
||||
|
||||
let raw = tally.raw_snp_distance();
|
||||
let included = tally.included(ratio_ceiling, excluded);
|
||||
@@ -72,5 +184,7 @@ pub(crate) fn sankoff_bundle(
|
||||
raw,
|
||||
base_pair_tally,
|
||||
cardinality_tally,
|
||||
tally,
|
||||
dispersion,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
//! Shared `--session` checkpoint plumbing for [`super::snp_distance`] and
|
||||
//! [`super::sankoff`] — both fold [`super::subsample::sample_index`]'s
|
||||
//! layer-by-layer output into a [`PairwiseTally`]/[`PartitionDispersion`]
|
||||
//! pair, and both persist it under the *same* artifact names on purpose:
|
||||
//! a tally checkpointed by one is restorable by the other. Concretely,
|
||||
//! `--sankoff --session DIR` today followed by `--distance snp-k2p
|
||||
//! --session DIR` tomorrow restores the first run's tally instead of
|
||||
//! resampling, even though only `--sankoff` originally built it —
|
||||
//! `snp_distance`'s own selection-parameter conflict check (via
|
||||
//! `obiksession::Session::open`, the caller's responsibility) is what
|
||||
//! guarantees that's actually safe to do.
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obiksession::Session;
|
||||
|
||||
use super::pairwise::{PairwiseTally, PartitionDispersion};
|
||||
|
||||
pub(crate) const TALLY_ARTIFACT: &str = "pairwise_tally";
|
||||
pub(crate) const DISPERSION_ARTIFACT: &str = "partition_dispersion";
|
||||
pub(crate) const PROGRESS_ARTIFACT: &str = "progress";
|
||||
/// Sentinel `progress` value meaning "every layer was processed — this
|
||||
/// tally is complete," distinct from any real layer count so a resumed run
|
||||
/// can tell "finished" apart from "stopped after N layers" without relying
|
||||
/// on N happening to equal the index's actual layer count.
|
||||
pub(crate) const PROGRESS_COMPLETE: u64 = u64::MAX;
|
||||
|
||||
/// How much in-progress sampling work a checkpoint can lose: at most
|
||||
/// `CHECKPOINT_INTERVAL_LAYERS` layers *or* `CHECKPOINT_INTERVAL` of wall
|
||||
/// time, whichever comes first — not "checkpoint every layer" (rewriting
|
||||
/// the whole, potentially large, tally/dispersion after every single one of
|
||||
/// what can be thousands of layers would make the checkpointing overhead
|
||||
/// itself dominate), and not "only at the end" either (which would defeat
|
||||
/// the point: on a long run against a real production index, losing 90% of
|
||||
/// an already-completed scan to a crash is exactly the expensive case
|
||||
/// resumability exists to avoid — the actual justification, not "the tally
|
||||
/// is cheap to rebuild," which was an earlier, wrong framing of this
|
||||
/// tradeoff caught on review).
|
||||
pub(crate) const CHECKPOINT_INTERVAL_LAYERS: usize = 8;
|
||||
pub(crate) const CHECKPOINT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30);
|
||||
|
||||
/// Restore a checkpointed `(PairwiseTally, PartitionDispersion, resume-from
|
||||
/// layer index)` triple from `session` — `None` on any miss (an artifact
|
||||
/// not cached, an I/O error, or a deserialization failure), all treated
|
||||
/// identically: something's missing or unusable, fall back to starting
|
||||
/// from scratch (`skip_layers = 0`) rather than surfacing a hard error over
|
||||
/// what's meant to be a transparent cache. The resume index is
|
||||
/// [`PROGRESS_COMPLETE`] when a prior run finished every layer. Layers
|
||||
/// processed after a resume draw a fresh, independent random sample (see
|
||||
/// `subsample::sample_index`'s own `skip_layers` docs) — deliberately not
|
||||
/// bit-identical to what an uninterrupted run would have produced,
|
||||
/// consistent with wanting run-to-run sampling variance rather than fixed
|
||||
/// reproducibility.
|
||||
///
|
||||
/// Full owned deserialization (`rkyv::from_bytes`) for the tally/dispersion,
|
||||
/// not zero-copy (`rkyv::access`) — simpler to integrate with
|
||||
/// `PairwiseTally`'s existing query methods, which all take `&self` by
|
||||
/// value semantics, not an `Archived<PairwiseTally>`. Still skips the
|
||||
/// expensive part (`sample_index` re-scanning the sibling annex for every
|
||||
/// already-completed layer); the deserialize of an already-in-memory-sized
|
||||
/// byte buffer is comparatively cheap. Making this genuinely zero-copy
|
||||
/// later would mean every `PairwiseTally` accessor (`categories`,
|
||||
/// `base_freq`, ...) working generically over `Archived<PairwiseTally>`
|
||||
/// too — a bigger, separate change.
|
||||
pub(crate) fn restore_tally_checkpoint(session: &Session) -> Option<(PairwiseTally, PartitionDispersion, usize)> {
|
||||
let tally_bytes = session.restore(TALLY_ARTIFACT).ok().flatten()?;
|
||||
let dispersion_bytes = session.restore(DISPERSION_ARTIFACT).ok().flatten()?;
|
||||
let progress_bytes = session.restore(PROGRESS_ARTIFACT).ok().flatten()?;
|
||||
let tally = rkyv::from_bytes::<PairwiseTally, rkyv::rancor::Error>(&tally_bytes[..]).ok()?;
|
||||
let dispersion =
|
||||
rkyv::from_bytes::<PartitionDispersion, rkyv::rancor::Error>(&dispersion_bytes[..]).ok()?;
|
||||
let progress_raw: [u8; 8] = progress_bytes.get(..8)?.try_into().ok()?;
|
||||
let progress = u64::from_le_bytes(progress_raw);
|
||||
let skip_layers = if progress == PROGRESS_COMPLETE { usize::MAX } else { progress as usize };
|
||||
if skip_layers == usize::MAX {
|
||||
tracing::info!("--session: restored a complete cached SNP tally — skipping resampling");
|
||||
} else {
|
||||
tracing::info!("--session: resuming a cached SNP tally from layer {skip_layers} onward");
|
||||
}
|
||||
Some((tally, dispersion, skip_layers))
|
||||
}
|
||||
|
||||
/// Store `tally`/`dispersion` plus how far sampling has gotten
|
||||
/// (`layers_done`, or [`PROGRESS_COMPLETE`] once every layer has been
|
||||
/// processed) — a checkpoint [`restore_tally_checkpoint`] can resume from
|
||||
/// later, whether that "later" is after a crash or a clean, separate
|
||||
/// invocation.
|
||||
pub(crate) fn store_tally_checkpoint(
|
||||
session: &Session,
|
||||
tally: &PairwiseTally,
|
||||
dispersion: &PartitionDispersion,
|
||||
progress: u64,
|
||||
) -> OKIResult<()> {
|
||||
let tally_bytes = rkyv::to_bytes::<rkyv::rancor::Error>(tally)
|
||||
.map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?;
|
||||
session.store(TALLY_ARTIFACT, &tally_bytes).map_err(OKIError::Io)?;
|
||||
let dispersion_bytes = rkyv::to_bytes::<rkyv::rancor::Error>(dispersion)
|
||||
.map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?;
|
||||
session.store(DISPERSION_ARTIFACT, &dispersion_bytes).map_err(OKIError::Io)?;
|
||||
session.store(PROGRESS_ARTIFACT, &progress.to_le_bytes()).map_err(OKIError::Io)?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -24,8 +24,13 @@
|
||||
use ndarray::Array2;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
use obiksession::Session;
|
||||
|
||||
use super::pairwise::PairwiseTally;
|
||||
use super::pairwise::{PairwiseTally, PartitionDispersion};
|
||||
use super::session_cache::{
|
||||
CHECKPOINT_INTERVAL, CHECKPOINT_INTERVAL_LAYERS, PROGRESS_COMPLETE, restore_tally_checkpoint,
|
||||
store_tally_checkpoint,
|
||||
};
|
||||
use super::sibling_family_size_histogram;
|
||||
use super::subsample::{EntropyBias, sample_index};
|
||||
|
||||
@@ -69,6 +74,19 @@ impl SnpDistanceKind {
|
||||
}
|
||||
}
|
||||
|
||||
/// `--gamma-shape` value. `Disabled` (no flag given) leaves every
|
||||
/// `corrected_log` call at its plain `-ln(x)` term; `Fixed` threads the
|
||||
/// given `alpha` through unchanged; `Auto` estimates `alpha` once, up
|
||||
/// front, from [`PartitionDispersion::estimate_alpha`] over the same
|
||||
/// sampling pass that builds the [`PairwiseTally`] — see that method's own
|
||||
/// docs for the estimator itself.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum GammaShape {
|
||||
Disabled,
|
||||
Fixed(f64),
|
||||
Auto,
|
||||
}
|
||||
|
||||
/// `-ln(x)`, gamma-mixture corrected when `alpha` is given: replaces the
|
||||
/// single-rate `-ln(x)` term with the standard Jin-Nei (1990) gamma
|
||||
/// substitution `alpha * (x^(-1/alpha) - 1)` — the mechanical term-wise
|
||||
@@ -209,10 +227,26 @@ fn formula(kind: SnpDistanceKind) -> fn(&PairwiseTally, usize, usize, Option<f64
|
||||
/// ignored (forced to `None`) when `n` is `None`, regardless of what the
|
||||
/// caller passes.
|
||||
///
|
||||
/// `gamma_shape`: `--gamma-shape`, `None` disables the correction. Rejected
|
||||
/// with [`obikindex::OKIError::InvalidInput`] if given alongside a `kind`
|
||||
/// that doesn't support it ([`SnpDistanceKind::supports_gamma`]) — checked
|
||||
/// here rather than left to silently no-op.
|
||||
/// `gamma_shape`: `--gamma-shape`, [`GammaShape::Disabled`] leaves the
|
||||
/// correction off. Rejected with [`obikindex::OKIError::InvalidInput`] if
|
||||
/// anything but `Disabled` is given alongside a `kind` that doesn't support
|
||||
/// it ([`SnpDistanceKind::supports_gamma`]) — checked here rather than left
|
||||
/// to silently no-op. [`GammaShape::Auto`] estimates `alpha` once, from the
|
||||
/// same sampling pass, via [`PartitionDispersion::estimate_alpha`]; if that
|
||||
/// estimation finds no measurable signal (see its own docs for when), the
|
||||
/// correction is silently left off rather than applied with a fabricated
|
||||
/// value — logged either way.
|
||||
///
|
||||
/// `session`: `--session`, `None` means no persistence (today's default
|
||||
/// behavior — always resample). `Some(session)` is expected to already be
|
||||
/// open under the *exact* selection parameters this call is about to use
|
||||
/// (`n`/`free_loss`/`no_ambiguity`/`excluded`/`entropy_bias`) — checking
|
||||
/// that is the caller's job (`obiksession::Session::open`'s own
|
||||
/// params-conflict mechanism), not re-validated here. If a cached tally is
|
||||
/// found it's restored and `sample_index` is skipped entirely; otherwise
|
||||
/// the fresh tally is stored back into the session before this returns, so
|
||||
/// a later call under the same session/params restores instead of
|
||||
/// resampling.
|
||||
pub(crate) fn snp_distance(
|
||||
cache: &IndexCache,
|
||||
kind: SnpDistanceKind,
|
||||
@@ -221,43 +255,134 @@ pub(crate) fn snp_distance(
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
gamma_shape: Option<f64>,
|
||||
gamma_shape: GammaShape,
|
||||
session: Option<&Session>,
|
||||
) -> OKIResult<Array2<f64>> {
|
||||
if gamma_shape.is_some() && !kind.supports_gamma() {
|
||||
// Fail fast, before paying for `sample_index`, not just inside
|
||||
// `distance_matrix` (which runs after sampling either way — the right
|
||||
// place for the check when called from `SankoffBundle::snp_distance`,
|
||||
// where sampling already happened as part of `--sankoff` itself).
|
||||
if gamma_shape != GammaShape::Disabled && !kind.supports_gamma() {
|
||||
return Err(obikindex::OKIError::InvalidInput(
|
||||
"--gamma-shape has no effect on this --distance value".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
let mut tally = PairwiseTally::new(n_genomes);
|
||||
|
||||
let (target, entropy_bias) = match n {
|
||||
Some(target) => (target, entropy_bias),
|
||||
let (mut tally, mut dispersion, skip_layers) = match session.and_then(restore_tally_checkpoint) {
|
||||
Some(triple) => triple,
|
||||
None => {
|
||||
let counts = sibling_family_size_histogram(cache)?;
|
||||
let total_eligible = (counts[1] + counts[2] + counts[3]) as usize;
|
||||
(total_eligible, None)
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
(PairwiseTally::new(n_genomes), PartitionDispersion::default(), 0)
|
||||
}
|
||||
};
|
||||
|
||||
if target > 0 {
|
||||
sample_index(
|
||||
cache,
|
||||
target,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|_partition, _layer, survivors| {
|
||||
super::pairwise::reduce_pairwise(&survivors, &mut tally);
|
||||
},
|
||||
)?;
|
||||
if skip_layers != usize::MAX {
|
||||
let (target, entropy_bias) = match n {
|
||||
Some(target) => (target, entropy_bias),
|
||||
None => {
|
||||
let counts = sibling_family_size_histogram(cache)?;
|
||||
let total_eligible = (counts[1] + counts[2] + counts[3]) as usize;
|
||||
(total_eligible, None)
|
||||
}
|
||||
};
|
||||
|
||||
if target > 0 {
|
||||
// Checkpoint every `CHECKPOINT_INTERVAL_LAYERS` layers or
|
||||
// `CHECKPOINT_INTERVAL` of wall time, whichever comes first —
|
||||
// see those constants' own docs for why "just recompute
|
||||
// everything" isn't the right call to make on a long run
|
||||
// against a real production index.
|
||||
let mut last_checkpoint_layer = skip_layers;
|
||||
let mut last_checkpoint_at = std::time::Instant::now();
|
||||
sample_index(
|
||||
cache,
|
||||
target,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
skip_layers,
|
||||
|partition, _layer, raw_index, survivors| {
|
||||
super::pairwise::reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
|
||||
let Some(session) = session else { return };
|
||||
let layers_done = raw_index + 1;
|
||||
let due_by_count =
|
||||
layers_done.saturating_sub(last_checkpoint_layer) >= CHECKPOINT_INTERVAL_LAYERS;
|
||||
let due_by_time = last_checkpoint_at.elapsed() >= CHECKPOINT_INTERVAL;
|
||||
if !(due_by_count || due_by_time) {
|
||||
return;
|
||||
}
|
||||
match store_tally_checkpoint(session, &tally, &dispersion, layers_done as u64) {
|
||||
Ok(()) => {
|
||||
last_checkpoint_layer = layers_done;
|
||||
last_checkpoint_at = std::time::Instant::now();
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("--session: failed to checkpoint at layer {layers_done}: {e}");
|
||||
}
|
||||
}
|
||||
},
|
||||
)?;
|
||||
}
|
||||
|
||||
if let Some(session) = session {
|
||||
store_tally_checkpoint(session, &tally, &dispersion, PROGRESS_COMPLETE)?;
|
||||
}
|
||||
}
|
||||
|
||||
distance_matrix(&tally, &dispersion, kind, gamma_shape)
|
||||
}
|
||||
|
||||
/// The post-sampling half of [`snp_distance`] — everything downstream of
|
||||
/// an already-built [`PairwiseTally`]/[`PartitionDispersion`]: the
|
||||
/// `--gamma-shape` support check, `alpha` resolution (fixed/auto/disabled),
|
||||
/// and the final `n×n` matrix build. Factored out so
|
||||
/// [`super::sankoff::SankoffBundle::snp_distance`] can reuse a tally it
|
||||
/// already has (from `--sankoff`/`--tnt`/`--phyg`/`--iqtree` sharing the
|
||||
/// exact same selection parameters as a `snp-*` `--distance` in the same
|
||||
/// CLI invocation) instead of re-running [`sample_index`] — which, beyond
|
||||
/// the wasted work, would also draw a **different** random sample the
|
||||
/// second time (`rand::rng()` is a thread-local generator that advances
|
||||
/// across calls, never reseeded per call), silently breaking the
|
||||
/// "same site selection for every algorithm in this run" guarantee that's
|
||||
/// the actual point of combining these flags in one invocation.
|
||||
pub(crate) fn distance_matrix(
|
||||
tally: &PairwiseTally,
|
||||
dispersion: &PartitionDispersion,
|
||||
kind: SnpDistanceKind,
|
||||
gamma_shape: GammaShape,
|
||||
) -> OKIResult<Array2<f64>> {
|
||||
if gamma_shape != GammaShape::Disabled && !kind.supports_gamma() {
|
||||
return Err(obikindex::OKIError::InvalidInput(
|
||||
"--gamma-shape has no effect on this --distance value".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let alpha = match gamma_shape {
|
||||
GammaShape::Disabled => None,
|
||||
GammaShape::Fixed(a) => Some(a),
|
||||
GammaShape::Auto => match dispersion.estimate_alpha() {
|
||||
Some(a) => {
|
||||
tracing::info!(
|
||||
"--gamma-shape auto: estimated alpha = {a:.4} (method-of-moments, \
|
||||
over-dispersion of substitution rate across partitions)"
|
||||
);
|
||||
Some(a)
|
||||
}
|
||||
None => {
|
||||
tracing::warn!(
|
||||
"--gamma-shape auto: no measurable among-partition rate over-dispersion \
|
||||
detected — proceeding without the gamma correction"
|
||||
);
|
||||
None
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
let n_genomes = tally.n_genomes();
|
||||
let f = formula(kind);
|
||||
Ok(Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
|
||||
if i == j { 0.0 } else { f(&tally, i, j, gamma_shape) }
|
||||
if i == j { 0.0 } else { f(tally, i, j, alpha) }
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -302,7 +427,8 @@ mod tests {
|
||||
});
|
||||
}
|
||||
let mut tally = PairwiseTally::new(2);
|
||||
reduce_pairwise(&families, &mut tally);
|
||||
let mut dispersion = PartitionDispersion::default();
|
||||
reduce_pairwise(&families, 0, &mut tally, &mut dispersion);
|
||||
tally
|
||||
}
|
||||
|
||||
|
||||
@@ -124,8 +124,33 @@ fn circular_entropy_values(
|
||||
/// its own `rayon::join`/`scope` worker, since `Arc` (unlike a plain
|
||||
/// borrowed slice) satisfies the `Send + 'static` those need.
|
||||
///
|
||||
/// Returns the actual number of sites kept, which may be less than `n` (see
|
||||
/// the module docs).
|
||||
/// `skip_layers`: the number of (partition, layer) pairs, in this
|
||||
/// function's own iteration order (`cache.partitions()` then `0..n_layer`,
|
||||
/// both deterministic for an unchanged index), to skip *entirely* —
|
||||
/// neither their (cheap) eligibility-bitset lookup's `sample_layer` call
|
||||
/// nor `on_layer` runs for them. Exists for `--session` resume: quotas are
|
||||
/// computed from `total_eligible` (summed over *every* layer regardless of
|
||||
/// `skip_layers`), so a resumed run reproduces the exact same *per-layer
|
||||
/// quotas* an uninterrupted one would have — i.e. how many sites each
|
||||
/// remaining layer contributes stays correct. It does **not** reproduce
|
||||
/// the exact same *sample*: `sample_layer`'s draws come from `rand::rng()`
|
||||
/// (`Self::sample_layer`'s own docs), a per-process, OS-seeded generator
|
||||
/// with no continuity across a process boundary, so the layers processed
|
||||
/// after a resume draw different (but equally valid) random outcomes than
|
||||
/// an uninterrupted run would have. Confirmed by direct test (kill mid-run,
|
||||
/// resume, compare to an uninterrupted run against the same index/params:
|
||||
/// different distance matrices, both legitimate samples). This is
|
||||
/// deliberate, not a defect: run-to-run sampling variance is wanted (see
|
||||
/// `DevDocMD/theory/evolutionary_distances.md`, "`--session`" section) —
|
||||
/// making resume bit-reproducible would require seeding deterministically,
|
||||
/// which would also remove that variance for ordinary repeated runs
|
||||
/// unless scoped carefully, and was explicitly declined for now. `0` for
|
||||
/// every caller not doing checkpointed resume.
|
||||
///
|
||||
/// Returns the actual number of *newly processed* sites kept (not
|
||||
/// counting whatever `skip_layers` layers already contributed on an
|
||||
/// earlier, interrupted call) — may be less than `n` (see the module
|
||||
/// docs).
|
||||
pub(crate) fn sample_index(
|
||||
cache: &IndexCache,
|
||||
n: usize,
|
||||
@@ -133,7 +158,8 @@ pub(crate) fn sample_index(
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
mut on_layer: impl FnMut(usize, usize, Arc<Vec<SurvivingFamily>>),
|
||||
skip_layers: usize,
|
||||
mut on_layer: impl FnMut(usize, usize, usize, Arc<Vec<SurvivingFamily>>),
|
||||
) -> OKIResult<usize> {
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
let fast_mode = is_fast_mode(cache);
|
||||
@@ -171,7 +197,10 @@ pub(crate) fn sample_index(
|
||||
}
|
||||
|
||||
let mut total_kept = 0usize;
|
||||
for (part, l, mut eligible, n_minorants) in layers {
|
||||
for (raw_index, (part, l, mut eligible, n_minorants)) in layers.into_iter().enumerate() {
|
||||
if raw_index < skip_layers {
|
||||
continue;
|
||||
}
|
||||
let count_layer = eligible.view().count_ones();
|
||||
if count_layer == 0 {
|
||||
continue;
|
||||
@@ -184,6 +213,7 @@ pub(crate) fn sample_index(
|
||||
cache,
|
||||
part,
|
||||
l,
|
||||
raw_index,
|
||||
&mut eligible,
|
||||
n_minorants,
|
||||
quota,
|
||||
@@ -207,6 +237,7 @@ fn sample_layer(
|
||||
cache: &IndexCache,
|
||||
partition: usize,
|
||||
layer_idx: usize,
|
||||
raw_index: usize,
|
||||
eligible: &mut TempBitVecBuilder,
|
||||
n_minorants: usize,
|
||||
quota: usize,
|
||||
@@ -216,7 +247,7 @@ fn sample_layer(
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
on_layer: &mut impl FnMut(usize, usize, Arc<Vec<SurvivingFamily>>),
|
||||
on_layer: &mut impl FnMut(usize, usize, usize, Arc<Vec<SurvivingFamily>>),
|
||||
) -> OKIResult<usize> {
|
||||
if n_minorants == 0 {
|
||||
return Ok(0);
|
||||
@@ -330,7 +361,7 @@ fn sample_layer(
|
||||
}
|
||||
|
||||
if !survivors.is_empty() {
|
||||
on_layer(partition, layer_idx, Arc::new(survivors));
|
||||
on_layer(partition, layer_idx, raw_index, Arc::new(survivors));
|
||||
}
|
||||
|
||||
Ok(kept)
|
||||
|
||||
@@ -15,8 +15,8 @@ use obikindex::{OKIError, OKIResult};
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::siblings::algorithms::{
|
||||
EntropyBias, SankoffBundle, Selection, SiblingAnnexStats, SnpAlignment, SnpDistanceKind,
|
||||
build_layer_sibling_annex, family_entropy, family_entropy_4, is_fast_mode,
|
||||
EntropyBias, GammaShape, SankoffBundle, Selection, SiblingAnnexStats, SnpAlignment,
|
||||
SnpDistanceKind, build_layer_sibling_annex, family_entropy, family_entropy_4, is_fast_mode,
|
||||
sankoff_bundle, scan_layer_families, sibling_annex_stats, sibling_family_size_histogram,
|
||||
snp_distance, snp_pseudo_alignment,
|
||||
};
|
||||
@@ -113,6 +113,14 @@ pub trait SiblingExt {
|
||||
/// `crate::siblings::CardinalityTally`'s own docs for why that
|
||||
/// wouldn't make sense: cardinality reflects each genome's own
|
||||
/// coverage/duplication structure, not the pair's mutual divergence).
|
||||
/// `session`: `--session`, `None` disables persistence (always
|
||||
/// resample). `Some` must already be open under these exact selection
|
||||
/// parameters (the caller's responsibility) — restores a cached tally/
|
||||
/// alignment when present instead of resampling, and checkpoints a
|
||||
/// fresh one back as sampling progresses otherwise. Shares its
|
||||
/// artifact namespace with [`Self::snp_distance`]'s own `session`
|
||||
/// parameter internally: a tally checkpointed by one is restorable by
|
||||
/// the other, given the same session and selection parameters.
|
||||
fn sankoff_bundle(
|
||||
&self,
|
||||
n: usize,
|
||||
@@ -121,6 +129,7 @@ pub trait SiblingExt {
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
ratio_ceiling: f64,
|
||||
session: Option<&obiksession::Session>,
|
||||
) -> OKIResult<SankoffBundle>;
|
||||
|
||||
/// A `snp-*` `--distance` matrix — one of the closed-form corrections
|
||||
@@ -141,9 +150,19 @@ pub trait SiblingExt {
|
||||
/// subsample).
|
||||
///
|
||||
/// `gamma_shape`: `--gamma-shape`, the Jin-Nei rate-heterogeneity
|
||||
/// correction — `None` disables it, `Some(alpha)` is rejected with
|
||||
/// correction — [`GammaShape::Disabled`] turns it off,
|
||||
/// [`GammaShape::Fixed`]/[`GammaShape::Auto`] are rejected with
|
||||
/// [`OKIError::InvalidInput`] for a `kind` that doesn't support it
|
||||
/// ([`SnpDistanceKind::supports_gamma`]).
|
||||
/// ([`SnpDistanceKind::supports_gamma`]). `Auto` estimates `alpha` from
|
||||
/// the data itself (method-of-moments over per-partition substitution
|
||||
/// rate dispersion) rather than requiring a user-supplied value.
|
||||
///
|
||||
/// `session`: `--session`, `None` disables persistence (always
|
||||
/// resample). `Some` must already be open under these exact selection
|
||||
/// parameters (the caller's responsibility, via
|
||||
/// `obiksession::Session::open`'s params-conflict check) — restores a
|
||||
/// cached tally when present instead of resampling, and stores a
|
||||
/// freshly computed one back for next time otherwise.
|
||||
fn snp_distance(
|
||||
&self,
|
||||
kind: SnpDistanceKind,
|
||||
@@ -152,7 +171,8 @@ pub trait SiblingExt {
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
gamma_shape: Option<f64>,
|
||||
gamma_shape: GammaShape,
|
||||
session: Option<&obiksession::Session>,
|
||||
) -> OKIResult<Array2<f64>>;
|
||||
|
||||
/// The Family Overlap annex — number of shared *variable* families per
|
||||
@@ -319,8 +339,9 @@ impl SiblingExt for IndexCache {
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
ratio_ceiling: f64,
|
||||
session: Option<&obiksession::Session>,
|
||||
) -> OKIResult<SankoffBundle> {
|
||||
sankoff_bundle(self, n, free_loss, no_ambiguity, excluded, entropy_bias, ratio_ceiling)
|
||||
sankoff_bundle(self, n, free_loss, no_ambiguity, excluded, entropy_bias, ratio_ceiling, session)
|
||||
}
|
||||
|
||||
fn snp_distance(
|
||||
@@ -331,9 +352,10 @@ impl SiblingExt for IndexCache {
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
gamma_shape: Option<f64>,
|
||||
gamma_shape: GammaShape,
|
||||
session: Option<&obiksession::Session>,
|
||||
) -> OKIResult<Array2<f64>> {
|
||||
snp_distance(self, kind, n, free_loss, no_ambiguity, excluded, entropy_bias, gamma_shape)
|
||||
snp_distance(self, kind, n, free_loss, no_ambiguity, excluded, entropy_bias, gamma_shape, session)
|
||||
}
|
||||
|
||||
fn family_overlap(&self) -> OKIResult<FamilyOverlap> {
|
||||
|
||||
@@ -37,7 +37,7 @@ pub use iter::{
|
||||
pub use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
|
||||
|
||||
pub use algorithms::{
|
||||
BasePairTally, CardinalityTally, EntropyBias, RawSnpDistanceOutput, SankoffBundle,
|
||||
BasePairTally, CardinalityTally, EntropyBias, GammaShape, RawSnpDistanceOutput, SankoffBundle,
|
||||
SiblingAnnexStats, SnpAlignment, SnpDistanceKind, cardinality_transition_probs,
|
||||
composition_transition_probs, pairwise_cost_matrix,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
[package]
|
||||
name = "obiksession"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
obisys = { path = "../obisys", default-features = false }
|
||||
memmap2 = "0.9.11"
|
||||
xxhash-rust = { version = "0.8.18", features = ["xxh3"] }
|
||||
tracing = "0.1"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
@@ -0,0 +1,19 @@
|
||||
//! Domain-agnostic session persistence: a named directory that survives
|
||||
//! across separate CLI invocations, holding an opaque parameter blob (used
|
||||
//! to detect stale/mismatched reuse) plus a checksummed cache of named
|
||||
//! binary artifacts.
|
||||
//!
|
||||
//! Deliberately knows nothing about what the parameters or artifacts
|
||||
//! *mean* — no k-mer/genome/tally concept anywhere in this crate. A
|
||||
//! caller (e.g. `obikphylo`, for its `--session` support around
|
||||
//! `PairwiseTally`/`PartitionDispersion`/`SnpAlignment`) serializes its own
|
||||
//! parameter struct and artifacts to bytes however it likes (`rkyv`,
|
||||
//! `bincode`, ...) and hands raw `&[u8]` to this crate; [`Session`] only
|
||||
//! deals in byte blobs, directory lifecycle, and locking. Same split
|
||||
//! already used elsewhere in this project between a generic mechanism
|
||||
//! crate and the domain crate that plugs into it (`obisys`/`obikindex`,
|
||||
//! `obicompactvec`/`obikindexer`).
|
||||
|
||||
mod session;
|
||||
|
||||
pub use session::{ParamsConflict, Session, SessionOutcome};
|
||||
@@ -0,0 +1,249 @@
|
||||
use std::fs;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use memmap2::Mmap;
|
||||
use obisys::DirLock;
|
||||
|
||||
const PARAMS_FILE: &str = "params.bin";
|
||||
const ARTIFACTS_DIR: &str = "artifacts";
|
||||
const CHECKSUM_EXT: &str = "xxh3";
|
||||
|
||||
/// Whether [`Session::open`] found a fresh (just-created, or just-reset by
|
||||
/// `force`) directory or reused one whose saved parameters matched the
|
||||
/// requested ones exactly.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum SessionOutcome {
|
||||
/// New directory, or an existing one reset because its saved
|
||||
/// parameters didn't match and `force: true` was passed. No cached
|
||||
/// artifact in it can be trusted — there either isn't one yet, or the
|
||||
/// ones that were there have just been deleted.
|
||||
Created,
|
||||
/// Existing directory whose saved parameters matched byte-for-byte.
|
||||
/// Every artifact already cached under it was produced under these
|
||||
/// same parameters and is safe to [`Session::restore`].
|
||||
Reused,
|
||||
}
|
||||
|
||||
/// Returned by [`Session::open`] when the directory already holds a
|
||||
/// parameter blob that doesn't match the one just requested, and `force`
|
||||
/// wasn't set. This crate has no idea what the bytes mean, so it can't
|
||||
/// produce a useful diagnostic on its own — the caller, who serialized
|
||||
/// both, is expected to deserialize `saved`/`requested` itself and report
|
||||
/// exactly which field(s) differ before deciding whether to error out,
|
||||
/// point at a different directory, or retry [`Session::open`] with
|
||||
/// `force: true`.
|
||||
#[derive(Debug)]
|
||||
pub struct ParamsConflict {
|
||||
pub saved: Vec<u8>,
|
||||
pub requested: Vec<u8>,
|
||||
}
|
||||
|
||||
/// An open session directory: exclusively locked (via [`obisys::DirLock`])
|
||||
/// for as long as this value lives, so no other `obiksession`-using
|
||||
/// process can read a half-written artifact or reset the directory out
|
||||
/// from under this one.
|
||||
#[derive(Debug)]
|
||||
pub struct Session {
|
||||
dir: PathBuf,
|
||||
outcome: SessionOutcome,
|
||||
_lock: DirLock,
|
||||
}
|
||||
|
||||
impl Session {
|
||||
/// Open (or create, or reset) a session directory — see
|
||||
/// [`SessionOutcome`]/[`ParamsConflict`] for the three possible
|
||||
/// outcomes. Blocks until the directory's lock is free (see
|
||||
/// [`obisys::DirLock::acquire`]) — a second process pointed at the
|
||||
/// same session directory waits rather than racing it.
|
||||
pub fn open(dir: &Path, params: &[u8], force: bool) -> io::Result<Result<Session, ParamsConflict>> {
|
||||
let lock = DirLock::acquire(dir)?;
|
||||
let params_path = dir.join(PARAMS_FILE);
|
||||
|
||||
let saved = read_checked(¶ms_path)?;
|
||||
match saved {
|
||||
Some(saved) if saved == params => Ok(Ok(Session {
|
||||
dir: dir.to_path_buf(),
|
||||
outcome: SessionOutcome::Reused,
|
||||
_lock: lock,
|
||||
})),
|
||||
Some(saved) if !force => Ok(Err(ParamsConflict { saved, requested: params.to_vec() })),
|
||||
_ => {
|
||||
// No saved params yet, or a mismatch with `force: true`:
|
||||
// start clean. Any artifact cached under different params
|
||||
// is stale by construction — deleting the whole
|
||||
// `artifacts/` directory is simpler and safer than trying
|
||||
// to figure out which entries are still valid.
|
||||
let artifacts_dir = dir.join(ARTIFACTS_DIR);
|
||||
if artifacts_dir.exists() {
|
||||
fs::remove_dir_all(&artifacts_dir)?;
|
||||
}
|
||||
fs::create_dir_all(&artifacts_dir)?;
|
||||
write_checked(¶ms_path, params)?;
|
||||
Ok(Ok(Session {
|
||||
dir: dir.to_path_buf(),
|
||||
outcome: SessionOutcome::Created,
|
||||
_lock: lock,
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dir(&self) -> &Path {
|
||||
&self.dir
|
||||
}
|
||||
|
||||
pub fn outcome(&self) -> SessionOutcome {
|
||||
self.outcome
|
||||
}
|
||||
|
||||
fn artifact_path(&self, name: &str) -> PathBuf {
|
||||
self.dir.join(ARTIFACTS_DIR).join(name)
|
||||
}
|
||||
|
||||
/// Restore a previously [`store`](Self::store)d artifact, `mmap`ed
|
||||
/// read-only for zero-copy access by the caller (e.g. `rkyv::access`
|
||||
/// straight over the returned bytes).
|
||||
///
|
||||
/// `Ok(None)` if `name` isn't cached yet, *or* if its checksum doesn't
|
||||
/// match what was recorded at [`store`](Self::store) time — treated
|
||||
/// the same as "not cached" rather than an error, so a partial dump
|
||||
/// left behind by a killed process just triggers a recompute on the
|
||||
/// next run instead of a hard failure.
|
||||
pub fn restore(&self, name: &str) -> io::Result<Option<Mmap>> {
|
||||
verify_mmap(&self.artifact_path(name))
|
||||
}
|
||||
|
||||
/// Store `bytes` under `name`, plus a checksum sidecar, atomically:
|
||||
/// both are written to a temporary file and renamed into place, so a
|
||||
/// crash mid-write can never leave a corrupt file that a later
|
||||
/// [`restore`](Self::restore) mistakes for valid (the checksum file is
|
||||
/// only ever renamed into place *after* the data file's rename
|
||||
/// succeeds, so the two can't disagree about which write completed).
|
||||
pub fn store(&self, name: &str, bytes: &[u8]) -> io::Result<()> {
|
||||
write_checked(&self.artifact_path(name), bytes)
|
||||
}
|
||||
}
|
||||
|
||||
fn checksum_path_for(path: &Path) -> PathBuf {
|
||||
let mut os_str = path.as_os_str().to_owned();
|
||||
os_str.push(".");
|
||||
os_str.push(CHECKSUM_EXT);
|
||||
PathBuf::from(os_str)
|
||||
}
|
||||
|
||||
/// `mmap` `path` and verify it against its checksum sidecar. `Ok(None)`
|
||||
/// when `path` doesn't exist, its checksum sidecar doesn't exist/parse, or
|
||||
/// the checksum doesn't match — every case here means "nothing usable is
|
||||
/// cached," never an error the caller has to handle specially.
|
||||
fn verify_mmap(path: &Path) -> io::Result<Option<Mmap>> {
|
||||
if !path.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
let Ok(expected_str) = fs::read_to_string(checksum_path_for(path)) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Ok(expected) = expected_str.trim().parse::<u64>() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let file = fs::File::open(path)?;
|
||||
// SAFETY: the file is only ever written by `write_checked` (atomic
|
||||
// temp-then-rename) and this whole directory is held under this
|
||||
// process's exclusive `DirLock` for the `Session`'s entire lifetime —
|
||||
// no other process can be concurrently modifying it underneath this
|
||||
// mapping.
|
||||
let mmap = unsafe { Mmap::map(&file)? };
|
||||
let actual = xxhash_rust::xxh3::xxh3_64(&mmap[..]);
|
||||
if actual != expected {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
"session artifact failed checksum verification — treating as absent"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(mmap))
|
||||
}
|
||||
|
||||
fn read_checked(path: &Path) -> io::Result<Option<Vec<u8>>> {
|
||||
Ok(verify_mmap(path)?.map(|mmap| mmap[..].to_vec()))
|
||||
}
|
||||
|
||||
fn write_checked(path: &Path, bytes: &[u8]) -> io::Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent)?;
|
||||
}
|
||||
let tmp = path.with_extension("tmp");
|
||||
fs::write(&tmp, bytes)?;
|
||||
fs::rename(&tmp, path)?;
|
||||
|
||||
let checksum = xxhash_rust::xxh3::xxh3_64(bytes);
|
||||
let checksum_path = checksum_path_for(path);
|
||||
let checksum_tmp = checksum_path.with_extension("tmp");
|
||||
fs::write(&checksum_tmp, checksum.to_string())?;
|
||||
fs::rename(&checksum_tmp, checksum_path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn fresh_directory_is_created() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path().join("session");
|
||||
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
|
||||
assert_eq!(session.outcome(), SessionOutcome::Created);
|
||||
assert!(session.restore("missing").unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matching_params_reuse_and_restore_artifact() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path().join("session");
|
||||
{
|
||||
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
|
||||
session.store("tally", b"some serialized bytes").unwrap();
|
||||
}
|
||||
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
|
||||
assert_eq!(session.outcome(), SessionOutcome::Reused);
|
||||
let restored = session.restore("tally").unwrap().unwrap();
|
||||
assert_eq!(&restored[..], b"some serialized bytes");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mismatched_params_without_force_conflicts() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path().join("session");
|
||||
Session::open(&dir, b"params-v1", false).unwrap().unwrap();
|
||||
let conflict = Session::open(&dir, b"params-v2", false).unwrap().unwrap_err();
|
||||
assert_eq!(conflict.saved, b"params-v1");
|
||||
assert_eq!(conflict.requested, b"params-v2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mismatched_params_with_force_resets_artifacts() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path().join("session");
|
||||
{
|
||||
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
|
||||
session.store("tally", b"old data").unwrap();
|
||||
}
|
||||
let session = Session::open(&dir, b"params-v2", true).unwrap().unwrap();
|
||||
assert_eq!(session.outcome(), SessionOutcome::Created);
|
||||
assert!(session.restore("tally").unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corrupted_artifact_is_treated_as_absent() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path().join("session");
|
||||
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
|
||||
session.store("tally", b"good data").unwrap();
|
||||
// Corrupt the data file after the fact without touching the
|
||||
// checksum sidecar — simulates a partial/garbled write.
|
||||
std::fs::write(dir.join("artifacts").join("tally"), b"corrupted!").unwrap();
|
||||
assert!(session.restore("tally").unwrap().is_none());
|
||||
}
|
||||
}
|
||||
@@ -14,6 +14,7 @@ use tracing::info;
|
||||
/// Windows) via `std::fs::File::lock`/`try_lock`, not a hand-rolled PID
|
||||
/// file: the OS releases it automatically on process exit, including a
|
||||
/// crash — no stale-lock cleanup logic needed.
|
||||
#[derive(Debug)]
|
||||
pub struct DirLock {
|
||||
_file: std::fs::File,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user