diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md new file mode 100644 index 0000000..3d6be69 --- /dev/null +++ b/docmd/theory/evolutionary_distances.md @@ -0,0 +1,329 @@ +# Central-position SNP distance (discussion) + +Not implemented. Design discussion for a substitution-rate estimator that +observes SNPs directly from paired-genome k-mer comparison, as an alternative +to Mash's Poisson-Jaccard inference (see [obicompactvec](../implementation/obicompactvec.md) +for the implemented Jaccard/Mash distances). + +## Motivation + +Mash infers a mutation rate from a single scalar (Jaccard) via a model that +assumes independence across the k positions. For k odd, a substitution at the +**central base** of a k-mer, with the `2m` flanking bases (`m = (k-1)/2`) +otherwise conserved, is directly observable: it is a SNP. This gives access not +just to a rate but to the substitution's nature (transition/transversion), +enabling classical corrected distances. + +## Statistic and correspondence with `shared` + +A genomic position `p` is covered by `k` overlapping k-mer windows. Requiring +the substitution to sit at the window's **center** makes exactly one window +per SNP eligible — a 1:1 correspondence between SNP and center-neighbor k-mer +pair, avoiding the ~k-fold overcount of an any-position neighbor search. + +A locus with a fully conserved `k`-window (flanks **and** center) is precisely +an exact-shared k-mer — already produced by the existing `shared_kmers` matrix +(`--shared-kmers`, `BitPartials::partial_jaccard` / `CountPartials::partial_threshold_jaccard`). +A locus with conserved flanks but a substituted center is a "central SNP". Both +count each locus exactly once, in matching units: + +``` +p_hat[i,j] = SNP[i,j] / (SNP[i,j] + shared[i,j]) +``` + +`p_hat` is `P(center substituted | 2m flanks conserved)`. Only the numerator +`SNP[i,j]` needs computing; the denominator term is already available. + +**Canonical invariance**: for odd k, the central position maps to itself under +reverse-complement (`m -> k-1-m = m`, base complemented). A transition maps to +a transition, a transversion to a transversion — the transition/transversion +split is well-defined in canonical space. + +## Sufficient statistic: 4x4 base-pair tally + +Tabulating the joint distribution of `(center_i, center_j)` over conserved-flank +loci, per genome pair, is sufficient for every downstream correction: + +| Estimator | Input | Formula | +|---|---|---| +| Raw p-distance | total off-diagonal / total | `p = SNP / (SNP + shared)` | +| Jukes-Cantor | p | `d = -3/4 * ln(1 - 4p/3)` | +| Kimura 2-parameter | transition rate P, transversion rate Q | `d = 1/2 ln(1/(1-2P-Q)) + 1/4 ln(1/(1-2Q))` | +| LogDet/paralinear | full 4x4 + base-composition margins | `d ~= -1/4 ln det(F)`, robust to non-stationary base composition | + +JC/K2P need only the total and the transition/transversion split. LogDet needs +the full 4x4 (adds a diagonal tally of shared k-mers' central base — cheap). + +Memory for the 4x4 tally: `n^2 * 16` counters. Trivial for the project's +genome-scale use case (tens to hundreds of genomes); ~13 GB at n=10^4 — outside +scope but worth flagging if n grows. + +## Biases (properties of the estimator, not defects) + +1. **Conserved-flank ascertainment bias.** Only SNPs with intact `2m`-base + flanks are visible; window-intact probability decays as `(1-p)^{2m}`. For + k=31 (2m=30): 0.74 at p=1%, 0.21 at p=5%, 0.04 at p=10%. This estimator + targets **closely related genomes**. Under rate heterogeneity across sites + (universal in practice), conserved flanks correlate with slow centers, so + `p_hat` underestimates the genome-wide average rate — it specifically + estimates the substitution rate of **conserved regions**. +2. **Bias toward isolated SNPs.** Two SNPs within k of each other disqualify + each other's flanks. Hypervariable regions are invisible by construction. +3. **Indels are invisible.** A frameshift destroys k-mer matches in a block; + this channel captures substitutions only. Indel divergence shows up as lost + shared k-mers (lower Jaccard/Mash), not as SNP signal. +4. **k-dependent specificity.** "A k-mer match implies common ancestry" is + quantitative. For a 3 Gbp genome, expected random flank-30 collisions + (k=31): `(3e9)^2 / 4^30 ~= 8` — negligible. At k=21: `(3e9)^2 / 4^20 ~= 2e6` + — no longer negligible. k=31 is safe; k<=21 is marginal to unreliable for + large genomes. The large k that guarantees homology is the same k that + shrinks the detectable-divergence window — an inherent tension. + +## Implementation: avoid materializing a de Bruijn graph + +A central-SNP pair is topologically a simple bubble in the colored de Bruijn +graph (source/sink k-mer shared, two length-k branches differing only at the +midpoint). Classical bubble-calling (Cortex/discoSNP-style) finds these, but +requires the graph — nodes plus adjacency for ~10^9 colored k-mers — resident +in memory. **Rejected**: prohibitive RAM for this project's scale. + +A naive per-pair generalisation of variant lookup across n genomes (query each +non-shared k-mer's 3 central variants against every counterpart genome's +index) costs `O(n^2 . N . 3)` random lookups, with the same k-mer's 3 variants +regenerated and requeried once per counterpart genome — pure redundant work. +**Rejected** as the basis for an n-genome design. + +## Implementation: sequential per-partition sweep (no scratch, no graph) + +`KmerIndex::distance()` already opens every partition's `presence_store`/ +`count_store` simultaneously, memory-mapped, into one `LayeredStore` +(`distance.rs:73-77`). "Querying another partition" is therefore not a new +I/O pattern to design — it is the same O(1) MPHF+evidence lookup the `query` +command already performs at scale. This lets the SNP tally be computed with +**no scratch files and no auxiliary graph**, by sweeping partitions once each +as a source: + +1. For source partition `p`, enumerate its **distinct** k-mers (one per MPHF + slot; each already carries its full multi-genome presence/count vector — + no need to explode per (k-mer, genome) occurrence). +2. For each, generate the 3 central-substitution variants and **canonicalise + each independently** (`min(kmer, revcomp)`, exactly as any normal query) — + this avoids the orientation edge case a masked-flank grouping would have + (a substitution that flips canonical orientation is handled correctly + because each variant is canonicalised on its own, not inferred from a + fixed-orientation flank key). +3. Compute each variant's target partition `q` via its minimizer; batch/sort + the partition's outgoing variant queries by `q` for locality. +4. Look up each variant in `q`'s already-mmap'd MPHF+evidence; on a hit, + combine the source's presence vector (base `a`) with the variant's + presence vector (base `b`): for every `i` carrying `a` and `j` carrying + `b`, `tally[i,j][a,b] += 1`. + +**Deduplication needs no persisted state.** Sweeping partitions in a fixed +order `p = 0, 1, ..., P-1` and only acting on a variant when its target +partition `q >= p` guarantees each unordered SNP pair is counted exactly +once: a pair with `q < p` was already resolved earlier, when `q` was itself +the source partition and `p` (being `>= q`) was a valid forward target. No +cross-partition flag array is needed — the sweep order *is* the +deduplication rule. Within the same partition (`q == p`), a lightweight +transient tie-break suffices: either a `#slots(p)`-bit scratch flag reset per +partition, or simply comparing the two k-mers' raw `u64` encodings and only +counting when `kmer_source < kmer_variant` — no storage at all. + +This is a distinct computation stage, not a `partial_*` in the existing +additive-by-partition sense: step 3-4 read across partition boundaries by +construction, unlike the row-local `partial_jaccard`/`partial_threshold_jaccard` +primitives. But it requires no new index files, permanent or scratch: +`unitigs.bin`, `mphf.bin`, `evidence.bin`, and the presence/count columns are +read as-is, and the only extra memory is the current partition's small +outgoing-query batch (`#kmers(p) * 3`, released once `p` is done) plus the +persistent `tally` accumulator (`n^2 * 16` counters, see above). + +**Outer loop (over source partitions `p`) must stay sequential.** Two +independent reasons, not just one: (a) memory — the bounded-footprint claim +above only holds with one partition's outgoing-query batch in flight; running +`T` source partitions concurrently multiplies that batch by `T`, exactly the +blowup the design avoids; (b) correctness — the `q >= p` deduplication rule +requires partitions to be claimed as sources in a fixed order; running `p1 < +p2` concurrently gives no guarantee `p1` has finished claiming its `q >= p1` +targets before `p2` starts claiming its own, breaking the "counted exactly +once" property. + +**Inner loop (target-partition lookups for a fixed `p`) parallelises safely.** +Each lookup is an O(1) read against an already-mmap'd structure, independent +of the others, with no growing allocation — no memory blowup, no ordering +dependency between different `q`. The only shared mutable state is `tally`; +give each worker thread a **thread-local partial tally** (fixed `n^2 * 16` +size, independent of partition size) and merge into the global `tally` once +`p`'s inner loop completes — the same reduce-then-merge pattern Rayon already +uses elsewhere in this codebase to open partitions in parallel. Extra memory: +`#threads * n^2 * 16`, negligible (~512 MB at n~1000, 32 threads) and +unrelated to partition size. + +**Cost**: `3 * N_distinct` MPHF lookups total across the whole index (each +partition swept once as source) — the same order of magnitude and the same +operation as running `query` over the index's entire k-mer content against +itself, three times. This is the tool's already-optimized regime, not a new +I/O profile to validate. + +## Cheaper: subsampling + +Since the target is a ratio, restricting the source-partition sweep to a +bottom-`s` hash sketch (only enumerate k-mers with `hash < threshold` as +sources) divides the lookup count by the sampling factor without biasing +`p_hat`. Mash-like tradeoff: rate estimated from a sample, not the full +k-mer set. + +## Recommendation + +Sequential per-partition sweep (Route D): reuse the already-mmap'd +per-partition MPHF/evidence/presence structures for O(1) variant lookups, +dedup via the fixed sweep-order rule (`q >= p`, plus an in-partition +tie-break), no scratch files, no graph materialisation. Denominator reused +from the existing `shared_kmers` matrix. Distances (p, JC, K2P, LogDet) as +finalisations of the resulting 4x4 tally, mirroring the `partial_* -> *_dist_matrix` +pattern used for Jaccard/Mash/Bray-Curtis/etc. + +## Detailed implementation plan + +Grounded in the current codebase. File/type references are anchors, not +prescriptions; adjust to reality when implementing. + +### Step 0 — new low-level primitives (`obikseq`) + +Two helpers do not yet exist and are prerequisites: + +1. **Central neighbours.** `CanonicalKmerOf` already exposes + `left_canonical_neighbors()` / `right_canonical_neighbors()` + (`obikseq/src/kmer.rs`), each returning the 4 canonicalised neighbours at + an end position. Add `central_canonical_neighbors()` returning the 4 + variants at position `m = (k-1)/2` (each independently canonicalised via + `.canonical()`). The 3 that differ from the source are the query variants; + skip the identity. Building on `nucleotide(i)` / the raw 2-bit layout keeps + it O(1). +2. **Lone-k-mer minimiser.** Routing a *synthetic* variant to its partition + needs its minimiser, but `RollingStat` (`obiskbuilder/src/rolling_stat.rs`) + only computes minimisers incrementally along a sequence. Add a standalone + `minimizer(kmer) -> Minimizer` that scans the `k-m+1` m-mer windows + (`PackedSeq::mmer`, `obikseq/src/packed_seq.rs`), canonicalises each, and + takes the min by `seq_hash()` — the same selection `RollingStat` performs, + evaluated once. Partition index is then + `(minimizer.seq_hash() & (n_partitions - 1)) as usize`, exactly as + `QueryBatch::from_records` (`obikmer/src/cmd/query.rs:142`); `n_partitions` + is a power of two so the mask is valid. + +### Step 1 — the tally accumulator (`obikindex`) + +A `SnpTally` holding, per genome pair, the 4x4 joint count of central bases: +`n * n * 4 * 4` `u64` (or a packed lower-triangular form since it is +symmetric). Provide `merge(&mut self, other: &SnpTally)` for the thread-local +reduce, and accessors yielding, per pair `(i,j)`: total off-diagonal (SNP), +transition count `P`, transversion count `Q`. The diagonal need not be stored +for JC/K2P (the denominator comes from `shared_kmers`); store it only if +LogDet is wanted, by also tallying shared k-mers' central base. + +### Step 2 — the sweep (`obikindex`, new `snp.rs`) + +Mirror `distance.rs`: open the presence or count store per partition. But +instead of a per-partition `partial_*`, run the sequential source sweep: + +```text +for p in 0..n_partitions: # OUTER — sequential + open source partition p's layers (QueryLayer-style, obikpartitionner) + enumerate distinct canonical k-mers of p (one per MPHF slot) with their + presence/count vectors # column-major, as query stage 2 + par_iter over these source k-mers: # INNER — rayon, thread-local tally + for each of the 3 central variants: + q = partition_of(variant) + if q < p: continue # dedup: forward targets only + if q == p and variant <= source.raw(): continue # in-partition tie-break + slot = layers[q].find_slot(variant) # MphfLayer::find, mmap'd + if hit: + vb = variant presence/count vector + for i in genomes with source base a: + for j in genomes with variant base b: + thread_tally[i,j][a,b] += 1 + merge thread-local tallies into global SnpTally +``` + +The inner lookup is precisely `QueryLayer::find_slot` + +`col_value(g, slot)` (`obikpartitionner/src/query_layer.rs`) — reuse or factor +out that path rather than reimplementing MPHF access. Enumerating "all distinct +k-mers of a partition with their vectors" is the `dump`/`query` stage-2 +column-major scan already implemented in `dump_layer.rs` / +`query_partition_with`; factor a reusable iterator if none fits. + +`presence_threshold` applies exactly as elsewhere: a genome "carries base b" +iff its count at that slot is `>= presence_threshold` (trivially `>= 1` for +presence indexes). + +### Step 3 — finalisation (`obikindex`) + +From the global `SnpTally` + the existing `shared_kmers` matrix, derive n x n +distance matrices, each a pure function of the accumulated counts (same shape +as `jaccard_to_mash`): + +- `p_hat[i,j] = SNP / (SNP + shared)` +- Jukes-Cantor, Kimura-2P (from `P`, `Q`), optionally LogDet (needs the + diagonal + base-composition margins). + +Guard the singularities (`p >= 3/4` for JC, `1-2P-Q <= 0` or `1-2Q <= 0` for +K2P) by clamping to a max distance, as `jaccard_to_mash` clamps `J <= 0`. + +### Step 4 — surfacing (`obikindex` + `obikmer` CLI) + +These metrics do not fit `DistanceMetric`'s current `LayeredStore`-partial +dispatch (they need the cross-partition sweep and produce a different +intermediate). Two options, to decide: + +- **(a)** New `DistanceMetric` variants (`Pdistance`, `JukesCantor`, + `Kimura2P`, `LogDet`) whose `KmerIndex::distance` arm calls the sweep + (`snp.rs`) instead of the partial path, still returning `DistanceOutput`. + Keeps one CLI surface (`--metric jukes-cantor`), at the cost of a branch in + `distance()` that ignores the `LayeredStore` it built. +- **(b)** A dedicated pathway (`KmerIndex::snp_distance`) and a distinct CLI + entry, if mixing a cross-partition sweep into the partition-local `distance` + command is judged architecturally muddy. + +Recommendation: (a) for user ergonomics (all pairwise distances under +`distance`, all feeding NJ/UPGMA/`--shared-kmers` unchanged), but compute the +sweep lazily only when an SNP-family metric is requested, so the existing +metrics keep their partition-local fast path untouched. + +### Step 5 — subsampling flag + +Add `--snp-sample ` (or a bottom-`s` hash threshold): restrict the +source-k-mer enumeration in Step 2 to `seq_hash(kmer) < threshold`. Divides +lookups proportionally; `p_hat` is unbiased. Off by default (exact). + +### Testing + +- **Primitive unit tests**: `central_canonical_neighbors` on hand-checked + k-mers incl. palindrome-boundary cases; lone-k-mer `minimizer` against + `RollingStat`'s incremental result on the same k-mer. +- **End-to-end tiny index**: two 1-genome indexes differing by a handful of + known isolated SNPs (transitions and transversions placed by hand), assert + exact `SNP`, `P`, `Q` counts and the resulting JC/K2P values. +- **Dedup invariant**: assert the tally is identical regardless of genome/ + partition order and that no pair is double-counted (compare against a + brute-force all-pairs reference on a small index). +- **Subsampling**: `p_hat` within sampling error of the exact run. + +### Suggested phasing + +1. Step 0 primitives + their unit tests (self-contained, no distance wiring). + This also unblocks the long-declared-but-unimplemented `query --mismatch` + (`obikmer/src/cmd/query.rs:676`, currently a warning), which needs the same + neighbour + routing machinery. +2. `SnpTally` + finalisation math with a brute-force (non-swept) reference + backend, validated on a tiny index. +3. The real per-partition sweep (Step 2) behind the same finalisation; assert + it matches the brute-force backend. +4. CLI surfacing (Step 4a) and NJ/UPGMA integration (already generic over the + matrix). +5. Subsampling (Step 5). + +## References + +The Mash mutation-rate model this discussion contrasts with: +[@Mash-distances-doc; @Fan2015-mash-formula]. diff --git a/mkdocs.yml b/mkdocs.yml index 7973e78..f0a6718 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -36,6 +36,7 @@ nav: - Entropy filter: theory/entropy.md - Minimizer selection: theory/minimizer.md - Partitioning architecture: theory/indexing.md + - Central-position SNP distance (discussion): theory/evolutionary_distances.md - Implementation: - SuperKmer: implementation/superkmer.md - Kmer: implementation/kmer.md