From 2610a4af79838702c33dd54f9d01eb9b636c29bf Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 9 Jul 2026 11:37:52 +0200 Subject: [PATCH 1/8] feat: add Mash distance metric and rolling entropy support Implement the Mash distance metric across the CLI, index, and compact vector traits. This includes adding a `Mash` variant to the `DistanceMetric` enum and `MetricArg` CLI argument, implementing the conversion from Jaccard distances using the standard mutation-rate estimator formula, and updating documentation with supported metrics and algorithmic references. Additionally, add an `entropy` method to rolling statistics for computing order-specific entropy. --- docmd/implementation/obicompactvec.md | 13 +++++++++++++ docmd/index.md | 2 +- docmd/references.bib | 18 ++++++++++++++++++ src/obicompactvec/src/traits.rs | 23 +++++++++++++++++++++++ src/obikindex/src/distance.rs | 4 ++++ src/obikmer/src/cmd/distance.rs | 2 ++ src/obiskbuilder/src/rolling_stat.rs | 7 ------- 7 files changed, 61 insertions(+), 8 deletions(-) diff --git a/docmd/implementation/obicompactvec.md b/docmd/implementation/obicompactvec.md index 301b021..34c2024 100644 --- a/docmd/implementation/obicompactvec.md +++ b/docmd/implementation/obicompactvec.md @@ -347,11 +347,24 @@ Provided finalisations: | `relfreq_euclidean_dist_matrix()` | `√partial_relfreq_euclidean[i,j]` | | `hellinger_dist_matrix()` | `√partial_hellinger[i,j] / √2` | | `hellinger_euclidean_dist_matrix()` | `√partial_hellinger[i,j]` | +| `threshold_mash_dist_matrix(k, t)` | Mash distance, derived from `threshold_jaccard_dist_matrix(t)` — no separate partial | ### BitPartials Required: `partial_jaccard() -> (Array2, Array2)`, `partial_hamming() -> Array2`. Both additive across layers and partitions. +Provided finalisations also include `jaccard_dist_matrix()`, `hamming_dist_matrix()`, and `mash_dist_matrix(k)`. + +### Mash distance + +`mash_dist_matrix`/`threshold_mash_dist_matrix` add no new additive primitive: both are a pointwise transform of the existing Jaccard distance matrix, per the Mash mutation-rate estimator [@Mash-distances-doc; @Fan2015-mash-formula]: + +``` +D = -1/k · ln(2J / (1+J)), J = 1 - d_jaccard +``` + +`J ≤ 0` (i.e. `d_jaccard ≥ 1`, no shared k-mers) maps to `D = 1` (maximal distance) rather than the `ln` singularity at `J = 0`. + --- ## Temp-file-backed types diff --git a/docmd/index.md b/docmd/index.md index 9f6d650..812cc86 100644 --- a/docmd/index.md +++ b/docmd/index.md @@ -13,7 +13,7 @@ | `query` | Query an index with sequences and annotate matches | | `dump` | Dump all indexed k-mers as CSV (kmer + per-genome counts or presence); supports the shared [kmer filtering](implementation/filtering.md) system; `--head N` limits output to the first N k-mers | | `annotate` | Add or update genome metadata from a CSV file; or dump metadata as CSV | -| `distance` | Compute pairwise distance matrix between genomes; optionally build NJ/UPGMA trees; `--presence-threshold N` sets the minimum count to consider a k-mer present when computing Jaccard on count indexes (default 1) | +| `distance` | Compute pairwise distance matrix between genomes (`--metric jaccard\|mash\|hamming\|bray-curtis\|relfreq-bray-curtis\|euclidean\|relfreq-euclidean\|hellinger\|hellinger-euclidean`); optionally build NJ/UPGMA trees; `--presence-threshold N` sets the minimum count to consider a k-mer present when computing Jaccard/Mash on count indexes (default 1) | | `unitig` | Build a global de Bruijn graph across all partitions and enumerate its unitigs as FASTA; supports the shared [kmer filtering](implementation/filtering.md) system | | `select` | Project and/or aggregate genome columns into a new or in-place index; the column-axis counterpart of `filter` (see [select](implementation/select.md)) | | `estimate` | Estimate approximate-index parameters (z, evidence bits, FP rates) before indexing | diff --git a/docmd/references.bib b/docmd/references.bib index a195b12..fca4c8b 100644 --- a/docmd/references.bib +++ b/docmd/references.bib @@ -241,3 +241,21 @@ volume = 33, year = 2017, bdsk-url-1 = {http://dx.doi.org/10.1093/bioinformatics/btw832}} + +@misc{Mash-distances-doc, + author = {{Marbl Lab}}, + howpublished = {Mash documentation}, + title = {Mash Distance}, + url = {https://mash.readthedocs.io/en/latest/distances.html}, + urldate = {2026-07-09}, + year = 2026} + +@article{Fan2015-mash-formula, + author = {Fan, Huan and Ives, Anthony R and Surget-Groba, Yann and Cannon, Charles H}, + doi = {10.1186/s12864-015-1647-5}, + journal = {BMC Genomics}, + number = 1, + title = {An assembly and alignment-free method of phylogeny reconstruction from next-generation sequencing data}, + url = {https://doi.org/10.1186/s12864-015-1647-5}, + volume = 16, + year = 2015} diff --git a/src/obicompactvec/src/traits.rs b/src/obicompactvec/src/traits.rs index cc52bc1..f3c0440 100644 --- a/src/obicompactvec/src/traits.rs +++ b/src/obicompactvec/src/traits.rs @@ -1,5 +1,16 @@ use ndarray::{Array1, Array2}; +/// Convert a Jaccard distance matrix (`1 - J`) into a Mash distance matrix, per +/// https://mash.readthedocs.io/en/latest/distances.html: +/// `D = -1/k * ln(2J / (1+J))`. +fn jaccard_to_mash(d_jaccard: &Array2, k: usize) -> Array2 { + d_jaccard.mapv(|d| { + let j = 1.0 - d; + if j <= 0.0 { 1.0 } + else { -1.0 / k as f64 * (2.0 * j / (1.0 + j)).ln() } + }) +} + // ── Column-level weight statistic — total count or presence count per column. /// Additive across layers and partitions; used as denominator in normalised distances. /// @@ -74,6 +85,12 @@ pub trait CountPartials: ColumnWeights { m } + /// Mash distance (https://mash.readthedocs.io/en/latest/distances.html), derived + /// from the presence-threshold Jaccard distance. + fn threshold_mash_dist_matrix(&self, k: usize, threshold: u32) -> Array2 { + jaccard_to_mash(&self.threshold_jaccard_dist_matrix(threshold), k) + } + fn relfreq_bray_dist_matrix(&self) -> Array2 { let global = self.col_weights(); let mut m = self.partial_relfreq_bray(&global).mapv(|v| 1.0 - v); @@ -126,6 +143,12 @@ pub trait BitPartials: ColumnWeights { m } + /// Mash distance (https://mash.readthedocs.io/en/latest/distances.html), derived + /// from the Jaccard distance. + fn mash_dist_matrix(&self, k: usize) -> Array2 { + jaccard_to_mash(&self.jaccard_dist_matrix(), k) + } + fn hamming_dist_matrix(&self) -> Array2 { self.partial_hamming() } diff --git a/src/obikindex/src/distance.rs b/src/obikindex/src/distance.rs index 867cff0..f5d63be 100644 --- a/src/obikindex/src/distance.rs +++ b/src/obikindex/src/distance.rs @@ -14,6 +14,8 @@ pub enum DistanceMetric { Jaccard, /// Hamming distance (number of differing kmer positions) on presence/absence data. Hamming, + /// Mash distance on presence/absence data (Jaccard-derived mutation-rate estimate). + Mash, /// Bray-Curtis dissimilarity on raw counts. BrayCurtis, /// Bray-Curtis dissimilarity normalised by per-genome total counts. @@ -84,6 +86,7 @@ impl KmerIndex { DistanceMetric::Hellinger => CountPartials::hellinger_dist_matrix(&global), DistanceMetric::HellingerEuclidean => CountPartials::hellinger_euclidean_dist_matrix(&global), DistanceMetric::Jaccard => CountPartials::threshold_jaccard_dist_matrix(&global, presence_threshold), + DistanceMetric::Mash => CountPartials::threshold_mash_dist_matrix(&global, self.kmer_size(), presence_threshold), DistanceMetric::Hamming => { return Err(OKIError::InvalidInput( "Hamming is only available for presence/absence indexes".into(), @@ -108,6 +111,7 @@ impl KmerIndex { let matrix = match metric { DistanceMetric::Jaccard => BitPartials::jaccard_dist_matrix(&global), + DistanceMetric::Mash => BitPartials::mash_dist_matrix(&global, self.kmer_size()), DistanceMetric::Hamming => { BitPartials::hamming_dist_matrix(&global).mapv(|v| v as f64) } diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index f66d78f..999beba 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -10,6 +10,7 @@ use tracing::info; #[derive(clap::ValueEnum, Clone, Copy, Debug)] pub enum MetricArg { Jaccard, + Mash, Hamming, BrayCurtis, #[value(name = "relfreq-bray-curtis")] @@ -26,6 +27,7 @@ impl From for DistanceMetric { fn from(m: MetricArg) -> Self { match m { MetricArg::Jaccard => DistanceMetric::Jaccard, + MetricArg::Mash => DistanceMetric::Mash, MetricArg::Hamming => DistanceMetric::Hamming, MetricArg::BrayCurtis => DistanceMetric::BrayCurtis, MetricArg::RelfreqBrayCurtis => DistanceMetric::RelfreqBrayCurtis, diff --git a/src/obiskbuilder/src/rolling_stat.rs b/src/obiskbuilder/src/rolling_stat.rs index 09b3f8d..e8f1bc1 100644 --- a/src/obiskbuilder/src/rolling_stat.rs +++ b/src/obiskbuilder/src/rolling_stat.rs @@ -196,13 +196,6 @@ impl RollingStat { .map(|raw| Minimizer::from_raw_unchecked(raw << (64 - self.m * 2))) } - pub fn entropy(&self, order: usize) -> Option { - if !self.ready() { - return None; - } - Some(self.entropy.entropy(order)) - } - pub fn normalized_entropy(&self) -> Option { if !self.ready() { return None; -- 2.54.0 From 45df9919e55c3b58382725b12fb944936906aa16 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Fri, 10 Jul 2026 09:47:19 +0200 Subject: [PATCH 2/8] docs: add central-position SNP distance estimator spec Introduces a design specification for inferring substitution rates directly from k-mers with conserved flanks. The document details a memory-efficient implementation that computes 4x4 base-pair tallies using existing MPHF structures, enabling classical corrections without de Bruijn graph materialization. Updates MkDocs navigation to include the new theory page. --- docmd/theory/evolutionary_distances.md | 329 +++++++++++++++++++++++++ mkdocs.yml | 1 + 2 files changed, 330 insertions(+) create mode 100644 docmd/theory/evolutionary_distances.md 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 -- 2.54.0 From 8bc6d533e5b8446551d6345d7cf1aa6a425f6507 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 12:35:27 +0200 Subject: [PATCH 3/8] feat: support negative count filters as group size offsets Updates CLI parsing to accept negative integers for count filters, interpreting them as offsets from the group size (e.g., `-1` means all but one). A resolution closure enforces a floor of 1 to prevent unconstrained filtering on small groups. Additionally, refines evolutionary distance documentation to condition comparisons on local homology, replacing union-based Jaccard with a self-contained `SnpTally`. This unified approach streamlines SNP and shared count computation, incorporates paralogy and heterozygosity handling, and enables direct derivation of corrected distance matrices without external dependencies. --- docmd/implementation/filtering.md | 49 ++- docmd/theory/evolutionary_distances.md | 417 +++++++++++++++++++++++-- src/obikmer/src/cmd/predicate.rs | 46 ++- 3 files changed, 464 insertions(+), 48 deletions(-) diff --git a/docmd/implementation/filtering.md b/docmd/implementation/filtering.md index ea6d4a2..b45938e 100644 --- a/docmd/implementation/filtering.md +++ b/docmd/implementation/filtering.md @@ -92,18 +92,48 @@ For each genome: | Flag | Applies to | Meaning | |------|-----------|---------| -| `--min-count N` | ingroup | k-mer present in at least N ingroup genomes | -| `--max-count N` | ingroup | k-mer present in at most N ingroup genomes | +| `--min-count N` | ingroup | k-mer present in at least N ingroup genomes (N may be negative, see below) | +| `--max-count N` | ingroup | k-mer present in at most N ingroup genomes (N may be negative, see below) | | `--min-frac F` | ingroup | k-mer present in at least fraction F of ingroup genomes | | `--max-frac F` | ingroup | k-mer present in at most fraction F of ingroup genomes | -| `--min-outgroup-count N` | outgroup | k-mer present in at least N outgroup genomes | -| `--max-outgroup-count N` | outgroup | k-mer present in at most N outgroup genomes | +| `--min-outgroup-count N` | outgroup | k-mer present in at least N outgroup genomes (N may be negative, see below) | +| `--max-outgroup-count N` | outgroup | k-mer present in at most N outgroup genomes (N may be negative, see below) | | `--min-outgroup-frac F` | outgroup | k-mer present in at least fraction F of outgroup genomes | | `--max-outgroup-frac F` | outgroup | k-mer present in at most fraction F of outgroup genomes | | `--min-total-count N` | all genomes | sum of per-genome counts ≥ N (`filter` only) | | `--max-total-count N` | all genomes | sum of per-genome counts ≤ N (`filter` only) | | `--presence-threshold N` | all | per-genome count > N to be considered "present" (default 0) | +### Negative counts — offset from group size + +The four integer count flags (`--min-count`, `--max-count`, `--min-outgroup-count`, +`--max-outgroup-count`) accept **negative** values, interpreted as an offset counted +down from the group size `n`, resolved at run time once `n` is known: + +| Value | Effective threshold | +|-------|---------------------| +| `N ≥ 0` | literal absolute count `N` | +| `-x` (x > 0) | `max(1, n − x)` — "all but x" | + +`-1` literally means *all but one*, `-2` *all but two*, and so on. This expresses +a quorum relative to the group size that a plain fraction cannot state exactly +(e.g. "present in every genome except at most one" is `n−1`, which is `0.9` for +`n = 10` but `0.857…` for `n = 7`). + +The threshold is **floored at 1**, never 0: the negative form always keeps +constraining the group. Without the floor, `--min-count -1` on a singleton +ingroup (`n = 1`) would resolve to `0` ("at least 0") and silently drop the +constraint; the floor makes it `1` ("present in that one genome") instead. + +To express a count of `0` (e.g. "absent from the ingroup"), use the literal `0`, +not a negative — `0` and `-0` are indistinguishable, so the offset form starts at +`-1`. + +> **Edge case** — on an *empty* group (`n = 0`, e.g. a predicate matching no +> genome), a negative count still resolves to `1`, an impossible constraint that +> rejects every k-mer. This is consistent with an empty group letting nothing +> through, but differs from the "no constraint" behaviour of the fraction flags. + **Conditional defaults** — the defaults for `--min-frac` and `--max-outgroup-count` depend on two conditions: whether the corresponding group was declared, **and** whether any quorum flag for that group was explicitly set. @@ -215,6 +245,17 @@ obikmer filter src --output dst \ --max-outgroup-count 0 ``` +Noise-tolerant core — keep k-mers present in *all but one* ingroup genome +(`-1` = `n−1`) and absent from *all but one* of the outgroup: + +```sh +obikmer filter src --output dst \ + --ingroup "genus=Betula" \ + --outgroup "*" \ + --min-count -1 \ + --max-outgroup-count -1 +``` + To dump only k-mers specific to *Betula nana*: ```sh diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index 3d6be69..6ff266b 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -7,12 +7,30 @@ 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. +**Primary intent: restrict the comparison to what is actually comparable.** +Mash's Jaccard is computed over the **union** of both genomes' k-mer content: +anything not identically shared is folded into a single undifferentiated +mass, whether the cause is a point substitution, a genuinely absent +homologous region (lineage-specific content, gene-family expansion, HGT, +genome-size asymmetry), or a diverged paralogous copy. The model then +back-infers a single mutation rate from that mass, silently attributing +non-homology to mutation. The central-SNP approach instead conditions every +comparison on local, positive evidence of homology: a locus only enters the +statistic if its `2m` flanking bases (`m = (k-1)/2`) are found intact in +*both* genomes — genuinely absent or non-homologous content is excluded from +the comparison entirely (neither numerator nor denominator), rather than +silently counted as divergence. This is a conditioning on comparability, not +just a richer summary statistic — see "Statistic and correspondence with +`shared`" below for how it plays out against genome-size asymmetry and +diverged gene families, and "Heterozygosity, ploidy, and consensus-assembly +inputs" for the corresponding paralogy/heterozygosity filter. + +**Secondary benefit: access to the substitution's nature.** Because the +central base of an odd-k window is directly observable once the flanks are +confirmed conserved, this also yields more than a rate — the +transition/transversion split — enabling classical corrected distances +(Jukes-Cantor, Kimura 2-parameter, LogDet) that a single Jaccard scalar +cannot support. ## Statistic and correspondence with `shared` @@ -21,24 +39,125 @@ 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: +A locus with a fully conserved `k`-window (flanks **and** center) is an +exact-shared k-mer at that locus; 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. +`p_hat` is `P(center substituted | 2m flanks conserved)`. `shared[i,j]` here +is **not** the general-purpose `shared_kmers` matrix used by Jaccard/Mash +(`--shared-kmers`, `BitPartials::partial_jaccard` / +`CountPartials::partial_threshold_jaccard`) — that matrix counts raw k-mer +identity with no per-genome copy-number constraint, whereas `p_hat`'s +denominator applies the eligibility rule defined below (raw or +paralogy-filtered). Both `SNP` and `shared` are accumulated by the same +sweep, from the same per-locus candidate set (source k-mer + 3 variants), +under the same eligibility rule — see "Locus eligibility" below, and +"Heterozygosity, ploidy, and consensus-assembly inputs" for why the +copy-number constraint matters and what it costs. **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. +## Locus eligibility: raw definition vs. paralogy filter + +For each k-mer `x` observed in genome A (source, one MPHF slot; the 3 +central-position variants generated as in the sweep below): check whether +A's locus (flanks fixed) is resolvable in genome B under one of the 4 +central forms. + +**Raw / no model.** The locus counts in the denominator iff at least one of +the 4 forms is present in B; it counts in the numerator iff the form found in +B differs from A's own. No constraint on A's or B's own copy number at this +locus. Open question, not resolved: what if **more than one** of the 4 forms +is present in B simultaneously (ambiguous target — count once arbitrarily, +count all, or drop)? The stringent filter below sidesteps the question by +construction rather than answering it. + +**Stringent / paralogy-aware.** The locus counts only if exactly one of the +4 forms is present in A **and** exactly one is present in B (`count == 1` at +that slot too, when a count index is available, to also exclude same-allele +duplicates that presence alone cannot see). This drops the raw definition's +ambiguous-B case automatically, at the cost of also dropping heterozygous +sites indiscriminately alongside true duplications (see "Heterozygosity, +ploidy, and consensus-assembly inputs" below). + +**Rejected: parsimony-based multiset pairing for multiplicity > 1.** Rather +than dropping ambiguous loci, pair identical alleles between A and B first +(0-mutation explanation preferred), then take `min(unmatched_A, unmatched_B)` +as inferred SNP pairs. Rejected on two grounds: (1) circularity — selecting +pairs by minimal apparent divergence, then measuring divergence on those same +pairs, deflates the estimate by construction, not a neutral heuristic; (2) +the discriminating signal is a single base among 4 possible values, and the +flanks are *already* guaranteed identical for every candidate by +construction (that is how the locus was selected) — no information remains +in a k-mer window to tell which copy in B truly corresponds to which copy in +A once multiplicity > 1 on either side. Any pairing rule invents a +correspondence the data cannot support. Multiplicity > 1 is treated as +non-identifiable, not as a puzzle to solve with a heuristic. + +## Heterozygosity, ploidy, and consensus-assembly inputs + +A within-genome multiplicity signal (more than one of the 4 central forms +present at a locus) is produced identically by two distinct causes: +paralogous duplication and diploid/polyploid heterozygosity. K-mer data alone +cannot distinguish them. The one real discriminator is sequencing depth +(heterozygous site: total depth of the present forms ~= the genome's +single-copy average; duplication: ~2x or more) — but that signal only exists +if genome "counts" are raw-read depth (FASTQ input), not occurrence counts in +an assembled FASTA, where per-locus depth is not preserved. + +**Magnitude is taxon- and mating-system-dependent, not universal.** +Heterozygosity density: mammals ~1 site / 1-1.5 kb (~0.1%); highly +outcrossing plants (maize, poplar) reported an order of magnitude higher +(~1%); self-fertilising plants (*Arabidopsis thaliana*) near zero — but with +a documented failure mode where segmental duplication masquerades as +"pseudo-heterozygosity"; fungi split between haploid vegetative stages +(non-issue) and dikaryotic Basidiomycetes, where two long-diverged haploid +nuclei coexist without fusing. The estimator's target use case (closely +related genomes, k=31) is exactly where the stringent filter above costs the +least for low-heterozygosity taxa and the most for outcrossing/dikaryotic +ones — no universal threshold; this is a scope caveat to document, not a +problem to solve generically. + +**Why assembled-consensus inputs don't make measured distances wrong.** +Phylogenetic inputs are near-universally assemblies, not raw reads, and +assemblers collapse heterozygous sites to one consensus allele per +position — effectively an arbitrary, largely uncorrelated-between-assemblies +choice at each het site. This does not inject unbounded noise: standard +population genetics gives `d_xy = d_a + (pi_A + pi_B)/2` — the expected +pairwise difference between a random allele of population A and a random +allele of population B equals the net (fixed) divergence `d_a` plus the +average of the two populations' own within-population diversity `pi`. +Consensus flattening realises exactly this random-allele draw, so the +measured genome-to-genome distance is a `d_xy`-like quantity, not `d_a` — +inflated by heterozygosity by a well-characterised additive term, not +distorted unpredictably. The term is negligible when `pi << d_xy` (the common +case for cross-species comparisons), and becomes material precisely in the +two cases already flagged above: very closely related genomes (this +estimator's explicit target) and highly heterozygous outcrossing organisms, +where `pi` and `d_xy` are the same order of magnitude. + +Caveat: this assumes the flattening is uncorrelated with the phylogenetic +signal — plausible for de novo assembly, not guaranteed for reference-guided +assembly biased toward one allele (e.g. the reference's) at each het site, +which would turn the noise term into a systematic bias toward the reference +lineage. Not evaluated here. + +**Forward-looking implication, not part of the current design.** The +multiplicity > 1 signal discarded by the stringent filter is a crude +per-genome proxy for `pi` (under low background paralogy). If a `pi_hat` per +genome were tallied alongside `SnpTally`, a `d_a` correction +(`p_hat - mean(pi_hat_i, pi_hat_j)/2`, roughly) could recover an estimate +closer to net divergence instead of `d_xy` — a possible extension, not +scoped here. + ## Sufficient statistic: 4x4 base-pair tally Tabulating the joint distribution of `(center_i, center_j)` over conserved-flank @@ -51,8 +170,9 @@ loci, per genome pair, is sufficient for every downstream correction: | 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). +JC/K2P need only the total and the transition/transversion split (the +diagonal collapses to a single "shared" total). LogDet needs the full 4x4, +already populated at no extra cost (see Step 1/2 below). 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 @@ -67,6 +187,16 @@ scope but worth flagging if n grows. (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**. + Two distinct factors are at play here, not one: `P(centre of a given + window is a SNP) = p` exactly, **independent of k** — a direct restatement + of the raw per-site rate via the bijective window<->centre-position + correspondence (Statistic section above), not a k-dependent quantity. + `(1-p)^{2m}` is the *separate*, genuinely k-dependent ascertainment factor + (are the flanks also intact). The two multiply: + `P(usable window showing a central SNP) = p * (1-p)^{2m}` — e.g. at + p=1/31 (~3.2%), k=31: `p * (1-p)^30 ~= 0.0323 * 0.374 ~= 1.2%`, i.e. about + 1 window in 83, not 1 in 31 (which is only the centre-mutated fraction, + before requiring intact flanks). 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; @@ -179,10 +309,13 @@ k-mer set. 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. +tie-break), no scratch files, no graph materialisation. Both the SNP +(off-diagonal) and shared (diagonal) counts are accumulated by this same +sweep, under the locus-eligibility rule chosen (raw or paralogy-filtered) — +not reused from the general-purpose `shared_kmers` matrix, whose raw-identity +definition does not apply the same copy-number constraint. 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 @@ -218,9 +351,13 @@ 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. +diagonal (shared, i.e. `p_hat`'s denominator minus SNP), transition count +`P`, transversion count `Q`. The diagonal is always populated — it is not an +optional LogDet-only extra, since `p_hat`'s denominator is no longer sourced +from the external `shared_kmers` matrix (see "Locus eligibility" and +"Statistic and correspondence with `shared`" above): the source k-mer's own +presence/count vector, already in hand when it is enumerated, supplies the +diagonal entry directly, at no extra lookup cost. ### Step 2 — the sweep (`obikindex`, new `snp.rs`) @@ -233,6 +370,11 @@ for p in 0..n_partitions: # OUTER — sequential 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 + apply eligibility rule to the source's own vector (raw: none; # diagonal + stringent: exactly one of the 4 forms present in each genome) # gate + for i in genomes eligible with source base a: + for j in genomes eligible with source base a: + thread_tally[i,j][a,a] += 1 # diagonal — no extra lookup for each of the 3 central variants: q = partition_of(variant) if q < p: continue # dedup: forward targets only @@ -240,8 +382,9 @@ for p in 0..n_partitions: # OUTER — sequential 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: + apply eligibility rule to vb (as above) + for i in eligible genomes with source base a: + for j in eligible genomes with variant base b: thread_tally[i,j][a,b] += 1 merge thread-local tallies into global SnpTally ``` @@ -257,11 +400,231 @@ column-major scan already implemented in `dump_layer.rs` / iff its count at that slot is `>= presence_threshold` (trivially `>= 1` for presence indexes). +### Open problem (unresolved, session end — not yet fully convinced) + +The `q >= p` / tie-break dedup rule in Step 2's pseudocode above is **flawed** +for the stringent (paralogy-filtered) eligibility rule: it only ever brings +two family members into view at once (the source and one looked-up variant), +never all four simultaneously, and which subset gets compared depends on +partition sweep order. "Exactly one of the 4 forms present in genome A" is a +whole-family property and cannot be decided correctly from a sequence of +pairwise, order-dependent glimpses — the pseudocode above needs revision, not +just the eligibility gate bolted onto it as written. + +Direction discussed, **not yet settled**: + +1. **Every distinct source k-mer looks up all 3 variants unconditionally** + (drop the `q < p` skip entirely) so that every observed family member + independently gathers all 4 vectors (its own + whichever of the 3 + variants exist) at once — a whole-family, order-independent view, computed + redundantly once per observed member. Same total lookup order of + magnitude as already budgeted (`3 * N_distinct`), just organised + differently (no lookup actually skipped, versus the original rule which + skipped roughly half). +2. **Tie-break after gathering, not before**: only the member whose own + canonical encoding is the smallest *among the members actually observed* + (now known, since all were just looked up) writes to `SnpTally`; the + others silently discard their redundant computation. Deterministic, + order-independent — as a side effect this also removes the "outer loop + must stay sequential" constraint from the cost/parallelism discussion + above, since no step depends on partition processing order any more. +3. **Proposed optimisation**: precompute, once at index build time, a + compact global (not per-genome) annex per MPHF slot — the count of + *other* family members observed anywhere in the dataset (0-3). Slots with + count 0 (majority under low divergence and few genomes, but see the + scaling caveat below) need no cross-lookup at all: eligibility reduces to + a local `count == 1` check at that single slot, and only slots with count + >= 1 enter the 3-lookup sweep machinery above. Revised (see Step 2b + below): minorant status *is* stored alongside the count after all, on 3 + bits rather than 2 — it comes for free from the same lookups needed to + count siblings, and storing it lets the sweep discard non-minorant slots + without re-fetching anything. + +**Minorant/sibling-count relationship, worked out precisely.** "Minorant" is +a one-way implication from sibling count, not an equivalence: `0 siblings +=> minorant` (trivially — with no other observed member, the k-mer is by +definition the smallest of the observed set, itself alone), and its +contrapositive `not minorant => >= 1 sibling`. The converse does not hold: +being the minorant says nothing about sibling count — a minorant can have 0, +1, 2 or 3 siblings, all with larger encodings than itself. Consequence: this +confirms, as a logical necessity rather than a heuristic, that a 0-sibling +slot can always write its diagonal contribution with zero ambiguity and no +lookup (it is unconditionally its own minorant) — but it gives no shortcut +for the >= 1-sibling case, where minorant status still requires the actual +comparison of gathered encodings; sibling count alone never determines it. + +**When to compute the annex, and cache invalidation.** Sibling count is a +property of the whole set of columns (genomes/groups) currently in the +index, not of any single genome — it cannot be computed correctly at +mono-genome build time (a family may gain siblings, or its minorant may +change, once more genomes are merged in later). Computing it eagerly at +every `merge` would also waste work on intermediate merged states nobody +ever queries. Instead: compute it lazily, on first `distance` call against a +given index, and persist the result alongside that index for subsequent +calls — the same lazy-derived-cache pattern `PersistentBitMatrix` already +uses for `Columnar` -> `Packed`. This requires no explicit invalidation for +`merge` or `filter` (`obikindex/src/merge.rs`, `obikmer/src/cmd/filter.rs`): +both only ever write to a fresh `--output` directory, never mutate an input +index in place, so a re-merged/re-filtered index is simply a new state with +no annex yet. `select --in-place` (`select_layer.rs:139-235`) is the +exception: it aggregates genome columns into groups (Any/All/None/Sum/Min/ +Max) by mutating the existing index's files without changing its location. +It does not remove k-mer rows, but it can still change eligibility and +sibling counts derived from those rows (e.g. a `Sum` over several +single-copy genomes can read as multi-copy at the group level). Because it +mutates in place, **`select --in-place` must explicitly invalidate (delete +or mark stale) any cached sibling-count annex for that index** — the one +operation in the current pipeline where this doesn't happen for free. + +Not yet convinced this is the right shape, and Step 2's pseudocode above has +not been rewritten to match — flagged for the next pass rather than resolved +here. + +### Step 2b — sibling-count / minorant annex (consolidated plan) + +Scope: only the precursor annex (sibling count 0-3 per slot, minorant +decided on demand) — not the SNP tally itself, whose Step 2 sweep remains +unresolved above. This piece is simpler than the sweep, because it writes to +an independent per-slot value, not a shared cross-k-mer accumulator, so it +needs no dedup/ownership logic at all at this stage. + +1. **Primitive.** Reuse `central_canonical_neighbors()` from Step 0 + unchanged — the 3 canonicalised central-substitution variants of a k-mer. +2. **New annex type** (`obicompactvec`, alongside `bitmatrix.rs`): a 3-bit- + per-slot packed array, one per partition — same on-disk shape family as + `PersistentBitMatrix`'s `Packed` variant, but simpler (no per-genome + columns, a single derived read-only value per slot). 3 bits, not 2: + revised to also store minorant status alongside sibling count, since it + comes for free from the same lookups (point 3 below) — 5 real states + (not-minorant; minorant with 0/1/2/3 siblings) fit in 3 bits (8 states, + 3 unused). This lets the future SNP sweep discard a non-minorant slot + **instantly**, with no lookup at all, instead of having to regenerate and + look up its siblings just to rediscover it isn't the designated writer — + moving that cost into this one-time, cached pass instead of repeating it + on every future sweep. The otherwise-unreachable combination + "not-minorant + 0 siblings" (impossible: 0 siblings always implies + minorant, see below) doubles as a free **"not yet computed" sentinel** — + annex files for all partitions/layers can be pre-initialised to this + value before the computation pass runs, distinguishing genuinely-computed + 0-sibling slots from not-yet-processed ones with no extra storage. +3. **Computation pass** (`obikindex`, new `siblings.rs`): **one + `obipipeline` run per layer, iterated sequentially over the index's + layers** — settled after two false starts, worth recording both. + - *False start 1*: "fully parallel over every partition/slot at once, + no ordering at all". Correctness is fine with this (sibling count and + minorant are order-independent, unlike the old `q >= p` dedup they + replace), but it reintroduces, at a larger scale, exactly the + memory-blowup the original Step 2 sweep's sequential-outer-loop + constraint existed to prevent: scattering every source partition at + once multiplies the in-flight outgoing-query volume by the number of + partitions. + - *False start 2*: push the layer loop itself into the pipeline (source + = the index's layers, a first `Flat` stage expands each layer into + its k-mers). `obipipeline`'s scheduler already bounds memory on its + own — it dispatches every item through a **shared** worker pool at + each stage boundary (`scheduler.rs:217-372`, `dispatch()` into a + common `worker_tx` queue, any free worker picks up any pending item; + not "one worker owns a chunk end to end"), with a biased `Select` + that prioritises draining items already advanced in the chain over + admitting new source items (`scheduler.rs:271-282`: stage results + outrank the source, "vider le pipeline en priorité" / "dernier + recours" for new data) — so bounded channel `capacity` plus this + drain-first bias already caps in-flight work without any external + sequential discipline. Correct, but it means k-mers from several + layers can be completing concurrently, so the sink would need to + track several open per-layer annex-file writers at once — real, + avoidable complexity. + - **Settled design**: keep the layer loop external and sequential — + not for memory (the pipeline's own `capacity`/priority mechanism + already provides that, for free, regardless), but so each pipeline + run's sink targets exactly one layer's annex file, no concurrent + multi-writer bookkeeping. Per layer: source = that layer's distinct + k-mers; a `Flat` (1->N) stage generates the 3 central variants of a + k-mer, each tagged with its origin (local slot); a transform stage + routes each variant to its target partition (unchanged per-k-mer + minimiser); a transform stage performs the lookup (existence-only — + `find_slot` hit/miss, cheaper than the SNP sweep's full column + fetch); a final stage/sink folds each answer into its origin's + running state (below) and, once a layer's k-mers are all resolved, + flushes the completed array to that layer's annex file. Many small, + single-purpose stages on purpose, to let the scheduler interleave + them finely across many in-flight items — this deliberately does + **not** mirror how `obipipeline` is used elsewhere today: `query.rs`'s + `process_chunk` lumps parse+route+query+serialise into one closure + (`query.rs:325,743-758`), and `scatter.rs` only pipelines file- + reading/superkmer construction, routing partitions afterwards in a + plain sequential loop (`KmerPartition::write_batch`, + `partition.rs:140`) — both under-use the fine-grained scheduling the + mechanism offers, so they are not precedents to copy, only existing + (and arguably improvable, out of scope here) usages. Cross-partition + lookups (querying another layer's MPHF for a variant) remain + necessary as before — only the *output* side is kept single-layer. + - **Reconciliation**: processed at the granularity of one *answer batch + per destination partition*, not one source k-mer at a time — this is a + proper shuffle, not a per-k-mer wait. Each source partition `p` holds a + small array of running states `(minorant = true, siblings = 0)`, one + per local slot, initialised at scatter time and **persisting across + however many destination-partition batches answer it** (up to 3, one + per variant, not necessarily all from the same `q`). Every scattered + query carries an origin tag (source partition + local slot) so its + answer can be routed back. When target partition `q` returns its batch + (all answers for every query that named `q`, regardless of which source + k-mer or which source partition they came from), that batch is walked + once, locally, and each answer updates — via its origin tag — the + matching entry in *its* source partition's array: a miss changes + nothing; a hit does `siblings += 1`, and if the found sibling's own + encoding is smaller than the source's, `minorant = false`. A given + source k-mer's state is final only once every destination batch + concerning it has been folded in; its partition's array is flushed to + the persistent annex once complete. Commutative per entry, so the order + in which destination batches arrive and get folded in doesn't matter. + + **Open optimisation, not adopted yet — real tradeoff, not a free win.** + Since looking up sibling `y` from `x`'s visit already yields everything + needed to fill `y`'s own annex entry too, one visit per *family* could in + principle replace one visit per *observed family member* — cutting this + pass's cost roughly by the average family size instead of paying + `3 * N_distinct` regardless. But it means threads processing different + source k-mers can end up writing the *same* sibling's slot concurrently — + the fully independent, ownership-free parallelism of the plan above is + deliberately traded away for this gain. It stays safe only because the + computed value for a given slot is deterministic regardless of who + computes it, so redundant concurrent writes converge to the same + value — correct as long as each write is atomic, no locking needed — but + it is a real design complexity increase over "every member redoes its + own 3 lookups independently," not a strict improvement to adopt by + default. +4. **Trigger and caching** (`obikindex::KmerIndex`/`distance.rs`): compute + lazily on first `distance` call for an SNP-family metric against a given + index; check for an existing annex file first (mirrors + `PersistentBitMatrix::open()`'s auto-detect-and-fall-back, + `bitmatrix.rs:264-287`); if absent, run step 3 and persist; if present, + mmap and reuse. +5. **Invalidation.** `merge` and `filter` always write to a fresh `--output` + directory (`obikindex/src/merge.rs`, `obikmer/src/cmd/filter.rs`) so a + re-merged/re-filtered index simply has no annex yet — nothing to + invalidate. `select --in-place` (`select_layer.rs:139-235`) mutates + columns of an existing index without changing its location, which can + change sibling counts without removing rows — it must explicitly delete + any cached annex for that index as part of its in-place rewrite. +6. **Testing**: hand-built tiny indexes with known sibling counts (0-3); + order-independence (recompute twice on a static index, identical + result, given the fully-parallel no-ownership design); invalidation + (annex absent/correctly recomputed after `select --in-place`); once + Step 2's sweep is fixed, a regression check that sibling_count == 0 + slots are never looked up cross-partition during the sweep. + +Cost: `3 * N_distinct` existence-only lookups, computed once per index +state and amortised over every subsequent `distance` call that reuses the +cached annex — cheaper per-lookup than the sweep itself (hit/miss only, no +column fetch). + ### 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`): +From the global `SnpTally` alone (diagonal and off-diagonal both populated by +the sweep, see Step 1/2 — no dependency on the external `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 diff --git a/src/obikmer/src/cmd/predicate.rs b/src/obikmer/src/cmd/predicate.rs index 47baab9..b9ae40f 100644 --- a/src/obikmer/src/cmd/predicate.rs +++ b/src/obikmer/src/cmd/predicate.rs @@ -151,12 +151,14 @@ pub struct FilterArgs { pub outgroup: Vec, /// Minimum number of ingroup genomes containing the k-mer - #[arg(long)] - pub min_count: Option, + /// (negative: offset from group size, e.g. -1 = all but one) + #[arg(long, allow_hyphen_values = true)] + pub min_count: Option, /// Maximum number of ingroup genomes containing the k-mer - #[arg(long)] - pub max_count: Option, + /// (negative: offset from group size, e.g. -1 = all but one) + #[arg(long, allow_hyphen_values = true)] + pub max_count: Option, /// Minimum fraction of ingroup genomes containing the k-mer [0.0–1.0] /// (default 1.0 when --ingroup is set, 0.0 otherwise) @@ -168,13 +170,15 @@ pub struct FilterArgs { pub max_frac: Option, /// Minimum number of outgroup genomes containing the k-mer - #[arg(long)] - pub min_outgroup_count: Option, + /// (negative: offset from outgroup size, e.g. -1 = all but one) + #[arg(long, allow_hyphen_values = true)] + pub min_outgroup_count: Option, /// Maximum number of outgroup genomes containing the k-mer - /// (default 0 when --outgroup is set, no constraint otherwise) - #[arg(long)] - pub max_outgroup_count: Option, + /// (default 0 when --outgroup is set, no constraint otherwise; + /// negative: offset from outgroup size, e.g. -1 = all but one) + #[arg(long, allow_hyphen_values = true)] + pub max_outgroup_count: Option, /// Minimum fraction of outgroup genomes containing the k-mer [0.0–1.0] #[arg(long)] @@ -239,12 +243,12 @@ pub fn matching_genome_indices(pred_str: &str, genomes: &[GenomeInfo]) -> Result pub struct GroupFilterParams { pub threshold: u32, - pub min_count: Option, - pub max_count: Option, + pub min_count: Option, + pub max_count: Option, pub min_frac: Option, pub max_frac: Option, - pub min_outgroup_count: Option, - pub max_outgroup_count: Option, + pub min_outgroup_count: Option, + pub max_outgroup_count: Option, pub min_outgroup_frac: Option, pub max_outgroup_frac: Option, } @@ -279,12 +283,20 @@ pub fn build_group_filter( let default_min_frac = if !ingroup_preds.is_empty() && !ingroup_quorum_explicit { 1.0 } else { 0.0 }; let default_max_outgroup_count = if !outgroup_preds.is_empty() && !outgroup_quorum_explicit { 0 } else { out_size }; - let min_count = p.min_count.unwrap_or(0); - let max_count = p.max_count.unwrap_or(in_size); + // Resolve a signed count: negative means an offset from the group size + // (e.g. -1 = all but one), floored at 1 so the negative form always keeps + // constraining the group — even a singleton group, where n-1 would be 0 + // and would otherwise drop the constraint entirely. + let resolve = |v: isize, size: usize| -> usize { + if v < 0 { (size as isize + v).max(1) as usize } else { v as usize } + }; + + let min_count = p.min_count.map(|v| resolve(v, in_size)).unwrap_or(0); + let max_count = p.max_count.map(|v| resolve(v, in_size)).unwrap_or(in_size); let min_frac = p.min_frac.unwrap_or(default_min_frac); let max_frac = p.max_frac.unwrap_or(1.0); - let min_outgroup_count = p.min_outgroup_count.unwrap_or(0); - let max_outgroup_count = p.max_outgroup_count.unwrap_or(default_max_outgroup_count); + let min_outgroup_count = p.min_outgroup_count.map(|v| resolve(v, out_size)).unwrap_or(0); + let max_outgroup_count = p.max_outgroup_count.map(|v| resolve(v, out_size)).unwrap_or(default_max_outgroup_count); let min_outgroup_frac = p.min_outgroup_frac.unwrap_or(0.0); let max_outgroup_frac = p.max_outgroup_frac.unwrap_or(1.0); -- 2.54.0 From ea914bb536dda84530958404c8dbbb5fae6bdbc7 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 14:56:44 +0200 Subject: [PATCH 4/8] feat: implement per-k-mer sibling counts and central neighbor generation Introduce the siblingannex module in obicompactvec to store per-slot minorant flags and sibling counts in a memory-mapped annex file. Add a scatter-gather pipeline in obikindex to compute these values across index layers and write them to .psib files. Implement central_canonical_neighbors in obikseq for generating strand-aware k-mer variants around the middle base. Expose rolling statistics in obiskbuilder and update dependency graphs accordingly. --- src/Cargo.lock | 3 + src/obicompactvec/src/lib.rs | 2 + src/obicompactvec/src/siblingannex.rs | 198 +++++++++++++++ src/obikindex/Cargo.toml | 5 + src/obikindex/src/lib.rs | 1 + src/obikindex/src/siblings.rs | 335 ++++++++++++++++++++++++++ src/obikseq/src/kmer.rs | 21 ++ src/obikseq/src/tests/kmer.rs | 42 ++++ src/obiskbuilder/src/lib.rs | 3 +- 9 files changed, 609 insertions(+), 1 deletion(-) create mode 100644 src/obicompactvec/src/siblingannex.rs create mode 100644 src/obikindex/src/siblings.rs diff --git a/src/Cargo.lock b/src/Cargo.lock index 0dc087d..87ae4a4 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1701,11 +1701,14 @@ dependencies = [ "obikpartitionner", "obikseq", "obilayeredmap", + "obiread", + "obiskbuilder", "obiskio", "obisys", "rayon", "serde", "serde_json", + "tempfile", "tracing", ] diff --git a/src/obicompactvec/src/lib.rs b/src/obicompactvec/src/lib.rs index 9041ab7..9a136bb 100644 --- a/src/obicompactvec/src/lib.rs +++ b/src/obicompactvec/src/lib.rs @@ -7,6 +7,7 @@ mod intmatrix; mod layer_meta; mod meta; mod reader; +mod siblingannex; mod tempbitvec; mod tempintvec; mod views; @@ -18,6 +19,7 @@ pub use builder::PersistentCompactIntVecBuilder; pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask}; pub use intmatrix::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, pack_compact_int_matrix}; pub use layer_meta::LayerMeta; +pub use siblingannex::{SiblingAnnex, SiblingAnnexBuilder, SiblingInfo}; pub use reader::{PersistentCompactIntVec, Iter as CompactIntVecIter}; pub use tempbitvec::{TempBitVec, TempBitVecBuilder}; pub use tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder}; diff --git a/src/obicompactvec/src/siblingannex.rs b/src/obicompactvec/src/siblingannex.rs new file mode 100644 index 0000000..94cb5cc --- /dev/null +++ b/src/obicompactvec/src/siblingannex.rs @@ -0,0 +1,198 @@ +//! Sibling-count / minorant annex: a compact, read-only-after-build, per-slot +//! derived value used by the central-position SNP distance estimator (see +//! `docmd/theory/evolutionary_distances.md`, "Step 2b"). +//! +//! One byte is stored per MPHF slot of a partition/layer, encoding two +//! independent facts about the slot's k-mer's "family" (the up-to-4 k-mers +//! sharing the same flanks, differing only at the central base), both +//! properties of the whole current multi-genome index, not of any one +//! genome: +//! +//! - bit 0: `minorant` — is this k-mer's canonical encoding the smallest +//! among the family members actually observed in the index? +//! - bits 1-2: `siblings` — how many *other* family members (0-3) are +//! observed anywhere in the index. +//! +//! Byte value 0 (`minorant = false`, `siblings = 0`) is logically +//! unreachable as a real result (0 siblings always implies minorant — see +//! the design doc) and is reused as the "not yet computed" sentinel: annex +//! files are pre-initialised to all-zero, and a real value is only ever +//! written once, by the computation pass. +//! +//! Deliberately simpler than a true 3-bit pack (1 byte/slot instead of 3 +//! bits/slot): correctness and simplicity first: for a first implementation. +//! Packing to 3 bits/slot is a pure storage-density follow-up, not a +//! behavioural change, left for later. + +use std::fs::{File, OpenOptions}; +use std::io; +use std::path::{Path, PathBuf}; + +use memmap2::{Mmap, MmapMut}; + +const MAGIC: [u8; 4] = *b"PSIB"; + +// Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (1 byte/slot) follows. +const HEADER_SIZE: usize = 16; + +/// Decoded value of one slot's annex entry. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SiblingInfo { + pub minorant: bool, + pub siblings: u8, // 0..=3 +} + +impl SiblingInfo { + #[inline] + fn encode(self) -> u8 { + (self.siblings << 1) | (self.minorant as u8) + } + + #[inline] + fn decode(byte: u8) -> Option { + if byte == 0 { + // The unreachable "not minorant + 0 siblings" combination — + // reserved as the "not yet computed" sentinel. + return None; + } + Some(SiblingInfo { + minorant: byte & 1 != 0, + siblings: (byte >> 1) & 0b11, + }) + } +} + +// ── SiblingAnnex (reader) ─────────────────────────────────────────────────── + +pub struct SiblingAnnex { + mmap: Mmap, + n: usize, + path: PathBuf, +} + +impl SiblingAnnex { + pub fn open(path: &Path) -> io::Result { + let mmap = unsafe { Mmap::map(&File::open(path)?)? }; + if mmap.len() < HEADER_SIZE { + return Err(io::Error::new(io::ErrorKind::InvalidData, "PSIB file too short")); + } + if mmap[0..4] != MAGIC { + return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PSIB magic")); + } + let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize; + if mmap.len() < HEADER_SIZE + n { + return Err(io::Error::new(io::ErrorKind::InvalidData, "PSIB file truncated")); + } + Ok(Self { mmap, n, path: path.to_path_buf() }) + } + + pub fn path(&self) -> &Path { &self.path } + pub fn len(&self) -> usize { self.n } + pub fn is_empty(&self) -> bool { self.n == 0 } + + /// `None` means the slot has not (yet) been computed — see module docs. + pub fn get(&self, slot: usize) -> Option { + SiblingInfo::decode(self.mmap[HEADER_SIZE + slot]) + } +} + +// ── SiblingAnnexBuilder (writer) ──────────────────────────────────────────── + +pub struct SiblingAnnexBuilder { + mmap: MmapMut, + n: usize, + path: PathBuf, +} + +impl SiblingAnnexBuilder { + /// Create a new annex of `n` slots at `path`, pre-initialised to the + /// "not yet computed" sentinel (all-zero). + pub fn new(n: usize, path: &Path) -> io::Result { + let file_size = HEADER_SIZE + n; + let file = OpenOptions::new() + .read(true).write(true).create(true).truncate(true) + .open(path)?; + file.set_len(file_size as u64)?; + let mut mmap = unsafe { MmapMut::map_mut(&file)? }; + mmap[0..4].copy_from_slice(&MAGIC); + mmap[4..8].copy_from_slice(&[0u8; 4]); + mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes()); + // Data region left at 0 by `set_len`/mmap — the sentinel value. + Ok(Self { mmap, n, path: path.to_path_buf() }) + } + + pub fn len(&self) -> usize { self.n } + pub fn is_empty(&self) -> bool { self.n == 0 } + + pub fn get(&self, slot: usize) -> Option { + SiblingInfo::decode(self.mmap[HEADER_SIZE + slot]) + } + + pub fn set(&mut self, slot: usize, info: SiblingInfo) { + // Redundant concurrent writes from independent recomputation paths + // converge to the same encoded byte for a given slot, so a plain + // store here is safe even without external synchronisation, as long + // as the byte write itself is atomic (true for a single aligned + // byte on every platform this project targets). + self.mmap[HEADER_SIZE + slot] = info.encode(); + } + + pub fn close(self) -> io::Result<()> { self.mmap.flush() } + + pub fn finish(self) -> io::Result { + let path = self.path.clone(); + self.close()?; + SiblingAnnex::open(&path) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[test] + fn sentinel_is_zero_and_unset_slots_read_as_uncomputed() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test.psib"); + let builder = SiblingAnnexBuilder::new(4, &path).unwrap(); + for slot in 0..4 { + assert_eq!(builder.get(slot), None); + } + builder.close().unwrap(); + } + + #[test] + fn roundtrip_all_valid_states() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test.psib"); + let mut builder = SiblingAnnexBuilder::new(5, &path).unwrap(); + + let cases = [ + SiblingInfo { minorant: true, siblings: 0 }, + SiblingInfo { minorant: true, siblings: 1 }, + SiblingInfo { minorant: true, siblings: 2 }, + SiblingInfo { minorant: true, siblings: 3 }, + SiblingInfo { minorant: false, siblings: 2 }, + ]; + for (slot, info) in cases.iter().enumerate() { + builder.set(slot, *info); + } + let annex = builder.finish().unwrap(); + for (slot, info) in cases.iter().enumerate() { + assert_eq!(annex.get(slot), Some(*info)); + } + } + + #[test] + fn not_minorant_zero_siblings_is_unreachable_via_set_and_decodes_as_sentinel() { + // Documented invariant, not enforced by the type: callers must never + // construct this combination. If they do, it is indistinguishable + // from "not computed" — exercised here to pin the behaviour down. + let dir = tempdir().unwrap(); + let path = dir.path().join("test.psib"); + let mut builder = SiblingAnnexBuilder::new(1, &path).unwrap(); + builder.set(0, SiblingInfo { minorant: false, siblings: 0 }); + assert_eq!(builder.get(0), None); + } +} diff --git a/src/obikindex/Cargo.toml b/src/obikindex/Cargo.toml index c8fef8e..05d3d77 100644 --- a/src/obikindex/Cargo.toml +++ b/src/obikindex/Cargo.toml @@ -10,6 +10,7 @@ obiskio = { path = "../obiskio" } obisys = { path = "../obisys" } obicompactvec = { path = "../obicompactvec" } obilayeredmap = { path = "../obilayeredmap" } +obiskbuilder = { path = "../obiskbuilder" } ndarray = "0.16" rayon = "1" crossbeam-channel = "0.5" @@ -19,6 +20,10 @@ indicatif = "0.17" tracing = "0.1.44" hwlocality = { version = "1.0.0-alpha.11", features = ["vendored"], optional = true } +[dev-dependencies] +obiread = { path = "../obiread" } +tempfile = "3" + [features] default = ["numa"] numa = ["hwlocality"] diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 371a05c..1f769c6 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -9,6 +9,7 @@ mod numa; mod rebuild; mod reindex; mod select; +mod siblings; mod stats; pub use error::{OKIError, OKIResult}; diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs new file mode 100644 index 0000000..73b09b9 --- /dev/null +++ b/src/obikindex/src/siblings.rs @@ -0,0 +1,335 @@ +//! Sibling-count / minorant annex construction. +//! +//! See `docmd/theory/evolutionary_distances.md`, "Step 2b — sibling-count / +//! minorant annex", for the full design discussion this implements. +//! +//! For each distinct k-mer of each layer of the (already built/merged) +//! index, computes two facts about its "family" (the up to 4 k-mers sharing +//! its flanks, differing only at the central base — well-defined for odd +//! k), both properties of the whole current multi-genome index rather than +//! of any one genome: +//! - how many *other* family members (0-3) are observed anywhere in the +//! index; +//! - whether this k-mer is the "minorant" of its family — the smallest +//! canonical encoding among the members actually observed. +//! +//! Implementation note (deviation from the fully staged `obipipeline` design +//! discussed at length in the doc): this first implementation processes each +//! layer with a straightforward sequential scatter (batch the layer's +//! outgoing variant queries by destination partition) then gather (one +//! `query_partition_with` call per destination partition) — not the +//! multi-stage elementary `obipipeline` pipeline the design settled on. The +//! external semantics (one annex file per layer, sequential outer loop over +//! layers, order-independent reconciliation) match the design exactly; only +//! the internal execution mechanism is simplified, as a scope trade-off. +//! Revisiting this to use `obipipeline` with elementary stages, as designed, +//! is a follow-up, not a behavioural change. + +use std::collections::HashMap; +use std::path::Path; + +use obicompactvec::{SiblingAnnexBuilder, SiblingInfo}; +use obikpartitionner::{KmerDesc, QueryHit}; +use obikseq::{CanonicalKmer, Minimizer}; +use obilayeredmap::{MphfLayer, OLMError}; +use obilayeredmap::meta::PartitionMeta; +use obiskbuilder::rolling_stat::RollingStat; +use obiskio::UnitigFileReader; + +use crate::error::{OKIError, OKIResult}; +use crate::index::KmerIndex; + +const INDEX_SUBDIR: &str = "index"; +const ANNEX_FILE_NAME: &str = "siblings.psib"; + +fn olm_to_ok(e: OLMError) -> OKIError { + match e { + OLMError::Io(e) => OKIError::Io(e), + other => OKIError::InvalidInput(format!("layered-map error: {other}")), + } +} + +/// Minimiser of a single, isolated canonical k-mer (not part of a streamed +/// sequence). `RollingStat` computes minimisers incrementally along a +/// sequence; this feeds one k-mer's bases through a fresh instance to get +/// the same selection for a single, disconnected k-mer. Not the leanest +/// possible primitive (an O(1)-amortised dedicated scan, as originally +/// sketched in the design doc's Step 0, would avoid the ASCII round-trip and +/// `RollingStat` allocation) but correct and reuses already-tested logic; +/// left as a follow-up optimisation. +fn lone_kmer_minimizer(kmer: CanonicalKmer) -> Minimizer { + let ascii = kmer.to_ascii(); + let mut rs = RollingStat::new(0); + for b in ascii { + rs.push(b); + } + rs.canonical_minimizer() + .expect("RollingStat must be ready after k bases of a valid k-mer") +} + +/// Destination partition for a (possibly synthetic) canonical k-mer, using +/// the same routing rule as the rest of the index (`minimiser.seq_hash() & +/// mask`, `n_partitions` is a power of two). +fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize { + let mask = (n_partitions as u64) - 1; + (lone_kmer_minimizer(kmer).seq_hash() & mask) as usize +} + +/// Running reconciliation state for one source k-mer, initialised to the +/// trivial "no siblings observed yet" state and folded incrementally (in any +/// order — commutative) as query answers come back. +#[derive(Clone, Copy)] +struct RunningState { + minorant: bool, + siblings: u8, +} + +impl Default for RunningState { + fn default() -> Self { + RunningState { minorant: true, siblings: 0 } + } +} + +impl KmerIndex { + /// Build the sibling-count/minorant annex for every layer of every + /// partition of this (already built) index, writing one annex file per + /// layer alongside its existing index files. Safe to call again later + /// (e.g. after a fresh `merge`) — each run simply overwrites the annex + /// files of the index it is called on. + /// + /// Cross-partition/cross-layer lookups are required (a k-mer's siblings + /// can live in any partition), but the layer loop itself — and thus the + /// annex file this produces — stays local to one layer at a time. + pub fn build_sibling_annex(&self) -> OKIResult<()> { + let n_parts = self.n_partitions(); + + for part in 0..n_parts { + let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + self.build_layer_sibling_annex(&layer_dir, &meta, n_parts)?; + } + } + + Ok(()) + } + + fn build_layer_sibling_annex( + &self, + layer_dir: &Path, + meta: &PartitionMeta, + n_parts: usize, + ) -> OKIResult<()> { + let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + let n_slots = mphf.n(); + + // ── Enumerate this layer's distinct k-mers, one per slot ──────────── + let mut slot_kmer: Vec> = vec![None; n_slots]; + let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) + .map_err(OKIError::Partition)?; + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if let Some(slot) = mphf.find(kmer) { + slot_kmer[slot] = Some(kmer); + } + } + + // ── Scatter: bucket outgoing variant queries by destination partition ── + // `KmerDesc.seq_idx` carries the origin (this layer's local slot); + // `KmerDesc.pos` is repurposed as a 0/1 flag: 1 iff this specific + // variant's own encoding is smaller than the source's — decided here, + // at scatter time, since both encodings are already in hand; the + // query response only needs to confirm existence (hit/miss). + let mut outgoing: Vec>> = + (0..n_parts).map(|_| HashMap::new()).collect(); + + for (slot, maybe_kmer) in slot_kmer.iter().enumerate() { + let Some(kmer) = maybe_kmer else { continue }; + for variant in kmer.central_canonical_neighbors() { + if variant == *kmer { + continue; // identity substitution — not a real variant + } + let smaller = variant.raw() < kmer.raw(); + let dest = partition_of(variant, n_parts); + outgoing[dest].entry(variant).or_default().push(KmerDesc { + seq_idx: slot as u32, + pos: smaller as u32, + }); + } + } + + // ── Gather + reconcile ─────────────────────────────────────────────── + let mut state = vec![RunningState::default(); n_slots]; + let n_genomes = self.meta.genomes.len(); + let with_counts = self.meta.config.with_counts; + + for (dest, kmers) in outgoing.iter().enumerate() { + if kmers.is_empty() { + continue; + } + self.partition() + .query_partition_with(dest, kmers, n_genomes, with_counts, |hit| { + if let QueryHit::Found(descs) = hit { + for d in descs { + let slot = d.seq_idx as usize; + let st = &mut state[slot]; + st.siblings = (st.siblings + 1).min(3); + if d.pos == 1 { + st.minorant = false; + } + } + } + }) + .map_err(OKIError::Partition)?; + } + + // ── Write the layer's annex file ───────────────────────────────────── + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?; + for (slot, st) in state.iter().enumerate() { + if slot_kmer[slot].is_none() { + continue; // unused MPHF slot, if any — leave at the sentinel + } + builder.set(slot, SiblingInfo { minorant: st.minorant, siblings: st.siblings }); + } + builder.close()?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::io::Write; + use std::path::Path; + + use obicompactvec::SiblingAnnex; + use obikseq::{Kmer, Sequence}; + use obisys::Reporter; + use tempfile::tempdir; + + use crate::meta::{GenomeInfo, IndexConfig}; + use crate::merge::MergeMode; + + use super::*; + + // k must be >= 11 (project constraint, "k ∈ [11,31]"); k=11, level_max=1, + // theta=0.0 mirror `obiskbuilder`'s own tests (smaller k/level_max + // combinations trip an unrelated pre-existing bug in `obikentropy`'s + // sliding-window ring buffer — not this feature's concern). + const K: usize = 11; + const M: usize = 5; + + /// Build a single-genome index from one in-memory FASTA sequence, driving + /// the same primitives `obikmer`'s `scatter` step uses (minus the + /// multi-file `obipipeline` wrapper — a single sequence needs none of + /// that): normalise -> build superkmers -> route -> write. + fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex { + let fasta_path = dir.join(format!("{label}.fasta")); + let mut f = std::fs::File::create(&fasta_path).unwrap(); + writeln!(f, ">{label}").unwrap(); + f.write_all(seq).unwrap(); + writeln!(f).unwrap(); + drop(f); + + let index_path = dir.join(format!("{label}.idx")); + let config = IndexConfig { + kmer_size: K, + minimizer_size: M, + n_bits: 0, // 1 partition — keeps the test deterministic and simple + with_counts: false, + evidence: obilayeredmap::IndexMode::Exact, + block_bits: 0, + }; + let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)), false) + .expect("create"); + + let mut rep = Reporter::new(); + let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta"); + for page in stream { + let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0); + idx.partition_mut().write_batch(batch).expect("write_batch"); + } + idx.partition_mut().close().expect("close partition writers"); + idx.mark_scattered().expect("mark_scattered"); + idx.dereplicate_and_count(false, &mut rep).expect("dereplicate_and_count"); + idx.build_layers(1, None, false, &mut rep).expect("build_layers"); + idx + } + + fn canonical(ascii: &[u8]) -> CanonicalKmer { + Kmer::from_ascii(ascii).unwrap().canonical() + } + + /// Read back the annex entry for a given canonical k-mer from the merged + /// index's (single) partition/layer, asserting it was found at all. + fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> SiblingInfo { + let index_dir = idx.partition().part_dir(0).join(INDEX_SUBDIR); + let meta = PartitionMeta::load(&index_dir).unwrap(); + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let mphf = MphfLayer::open(&layer_dir, &meta.mode).unwrap(); + if let Some(slot) = mphf.find(kmer) { + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).unwrap(); + return annex.get(slot).expect("slot must have a computed annex entry"); + } + } + panic!("kmer not found in any layer of partition 0"); + } + + fn merge_two(dir: &Path, g1: &KmerIndex, g2: &KmerIndex) -> KmerIndex { + let mut rep = Reporter::new(); + KmerIndex::merge( + &dir.join("merged.idx"), + &[g1, g2], + MergeMode::Presence, + false, + false, + 1.0, + &mut rep, + ) + .expect("merge") + } + + #[test] + fn sibling_annex_one_sibling_each() { + // k=11, centre = index 5 (0-based). Two genomes, each exactly one + // k-mer, sharing every base except the centre: + // g1 = "AACCGCTTAAG" (centre 'C') + // g2 = "AACCGGTTAAG" (centre 'G') + // Hand-verified: both stay forward-oriented under canonicalisation + // (each is lexicographically smaller than its own reverse + // complement, since both start with "AA"), and raw(g1) < raw(g2) + // (only differing base: C=0b01 < G=0b10 at the centre) — so g1 is + // the minorant, g2 is not, and each is the other's one sibling. + let dir = tempdir().unwrap(); + let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG"); + let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG"); + let merged = merge_two(dir.path(), &g1, &g2); + merged.build_sibling_annex().expect("build_sibling_annex"); + + let a = annex_info_for(&merged, canonical(b"AACCGCTTAAG")); + assert_eq!(a, SiblingInfo { minorant: true, siblings: 1 }, "AACCGCTTAAG"); + + let b = annex_info_for(&merged, canonical(b"AACCGGTTAAG")); + assert_eq!(b, SiblingInfo { minorant: false, siblings: 1 }, "AACCGGTTAAG"); + } + + #[test] + fn sibling_annex_zero_siblings_when_identical_and_no_variant_exists() { + // Same k-mer in both genomes, no other genome around to carry a + // variant -> 0 siblings, trivially its own minorant. + let dir = tempdir().unwrap(); + let g1 = build_single_genome_index(dir.path(), "g1", b"GATTACAGATC"); + let g2 = build_single_genome_index(dir.path(), "g2", b"GATTACAGATC"); + let merged = merge_two(dir.path(), &g1, &g2); + merged.build_sibling_annex().expect("build_sibling_annex"); + + let info = annex_info_for(&merged, canonical(b"GATTACAGATC")); + assert_eq!(info, SiblingInfo { minorant: true, siblings: 0 }, "GATTACAGATC"); + } +} diff --git a/src/obikseq/src/kmer.rs b/src/obikseq/src/kmer.rs index 7dba41c..365a785 100644 --- a/src/obikseq/src/kmer.rs +++ b/src/obikseq/src/kmer.rs @@ -341,6 +341,27 @@ impl CanonicalKmerOf { ] } + /// Return the four central canonical neighbours (each already canonical), + /// substituting the base at the middle position `m = (L::len()-1)/2` + /// (well-defined for odd `L::len()`). Each of the 4 substitutions is + /// canonicalised independently — this correctly handles the case where a + /// substitution flips the canonical orientation, unlike inferring the + /// variant from a fixed-orientation flank key. One of the 4 equals + /// `self`'s own canonical form (the identity substitution); callers that + /// only want the 3 genuine variants should skip it. + pub fn central_canonical_neighbors(&self) -> [CanonicalKmerOf; 4] { + let k = L::len(); + let m = (k - 1) / 2; + let shift = KMER_BITS - 2 - 2 * m; + let cleared = self.0 & !((0b11 as RawKmer) << shift); + [ + KmerOf::(cleared | ((0 as RawKmer) << shift), PhantomData).canonical(), + KmerOf::(cleared | ((1 as RawKmer) << shift), PhantomData).canonical(), + KmerOf::(cleared | ((2 as RawKmer) << shift), PhantomData).canonical(), + KmerOf::(cleared | ((3 as RawKmer) << shift), PhantomData).canonical(), + ] + } + /// Return the inner value as a raw [`KmerOf`]. #[inline] pub fn into_kmer(self) -> KmerOf { diff --git a/src/obikseq/src/tests/kmer.rs b/src/obikseq/src/tests/kmer.rs index 2aafdd6..5d9d29c 100644 --- a/src/obikseq/src/tests/kmer.rs +++ b/src/obikseq/src/tests/kmer.rs @@ -210,4 +210,46 @@ mod tests { check!(31); check!(32); } + + // ── central_canonical_neighbors ───────────────────────────────────────── + + #[test] + fn central_canonical_neighbors_hand_checked_k3() { + // k=3, centre = index 1. For "ACG", every one of the 4 central + // substitutions ("AAG","ACG","AGG","ATG") happens to stay in forward + // orientation when canonicalised (verified by hand: each is already + // lexicographically <= its own reverse complement), so this case + // exercises the substitution logic without the RC-flip edge case. + let ck = KmerOf::>::from_ascii(b"ACG").unwrap().canonical(); + let neighbours = ck.central_canonical_neighbors(); + let ascii: Vec> = neighbours.iter().map(|n| n.to_ascii()).collect(); + assert_eq!(ascii, vec![b"AAG".to_vec(), b"ACG".to_vec(), b"AGG".to_vec(), b"ATG".to_vec()]); + // The identity substitution (centre unchanged) must reproduce `ck`. + assert!(neighbours.contains(&ck)); + } + + #[test] + fn central_canonical_neighbors_identity_present_for_various_k() { + macro_rules! check { + ($n:expr) => {{ + let ck = KmerOf::>::from_ascii(&make_seq::<$n>()) + .unwrap() + .canonical(); + let neighbours = ck.central_canonical_neighbors(); + assert!( + neighbours.contains(&ck), + "identity substitution missing from central_canonical_neighbors for k={}", + $n + ); + // Every returned neighbour must itself already be canonical. + for n in &neighbours { + assert_eq!(n.into_kmer().canonical(), *n, "neighbour not canonical for k={}", $n); + } + }}; + } + check!(1); + check!(3); + check!(5); + check!(31); + } } diff --git a/src/obiskbuilder/src/lib.rs b/src/obiskbuilder/src/lib.rs index 2dd9c2f..65d15da 100644 --- a/src/obiskbuilder/src/lib.rs +++ b/src/obiskbuilder/src/lib.rs @@ -10,7 +10,8 @@ pub mod stream_iter; mod scratch; pub(crate) mod encoding; -pub(crate) mod rolling_stat; +#[allow(missing_docs)] +pub mod rolling_stat; pub use iter::SuperKmerIter; pub use scratch::SuperKmerScratch; -- 2.54.0 From ba990a48a0d8559b8ac93466eb7008d3b0de4da9 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 15:27:35 +0200 Subject: [PATCH 5/8] feat: add obipipeline for concurrent sibling annex stats Add the `obipipeline` crate and replace sequential scatter/gather logic with a concurrent pipeline using `Flat` and `Transform` stages. Introduce `SiblingAnnexStats` API to compute distributions, and add CLI flags to `distance.rs` for constructing the annex and exporting statistics as CSV. --- src/Cargo.lock | 1 + src/obikindex/Cargo.toml | 1 + src/obikindex/src/lib.rs | 1 + src/obikindex/src/siblings.rs | 291 +++++++++++++++++++++++++------- src/obikmer/src/cmd/distance.rs | 75 +++++++- 5 files changed, 310 insertions(+), 59 deletions(-) diff --git a/src/Cargo.lock b/src/Cargo.lock index 87ae4a4..aa92eee 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1701,6 +1701,7 @@ dependencies = [ "obikpartitionner", "obikseq", "obilayeredmap", + "obipipeline", "obiread", "obiskbuilder", "obiskio", diff --git a/src/obikindex/Cargo.toml b/src/obikindex/Cargo.toml index 05d3d77..20e1dd3 100644 --- a/src/obikindex/Cargo.toml +++ b/src/obikindex/Cargo.toml @@ -11,6 +11,7 @@ obisys = { path = "../obisys" } obicompactvec = { path = "../obicompactvec" } obilayeredmap = { path = "../obilayeredmap" } obiskbuilder = { path = "../obiskbuilder" } +obipipeline = { path = "../obipipeline" } ndarray = "0.16" rayon = "1" crossbeam-channel = "0.5" diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 1f769c6..4d22205 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -19,3 +19,4 @@ pub use merge::MergeMode; pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME}; pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED}; pub use stats::IndexBitsPerKmer; +pub use siblings::SiblingAnnexStats; diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index 73b09b9..4638bde 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -13,23 +13,22 @@ //! - whether this k-mer is the "minorant" of its family — the smallest //! canonical encoding among the members actually observed. //! -//! Implementation note (deviation from the fully staged `obipipeline` design -//! discussed at length in the doc): this first implementation processes each -//! layer with a straightforward sequential scatter (batch the layer's -//! outgoing variant queries by destination partition) then gather (one -//! `query_partition_with` call per destination partition) — not the -//! multi-stage elementary `obipipeline` pipeline the design settled on. The -//! external semantics (one annex file per layer, sequential outer loop over -//! layers, order-independent reconciliation) match the design exactly; only -//! the internal execution mechanism is simplified, as a scope trade-off. -//! Revisiting this to use `obipipeline` with elementary stages, as designed, -//! is a follow-up, not a behavioural change. +//! Per layer, the computation runs as an `obipipeline` pipeline with several +//! elementary stages (a `Flat` stage generating each k-mer's 3 central +//! variants, a `Transform` stage looking each variant up in its destination +//! partition), so the scheduler's shared worker pool interleaves this work +//! across many in-flight k-mers/variants rather than processing everything +//! on a single thread — see `docmd/theory/evolutionary_distances.md`, Step +//! 2b, "Mechanism", for why elementary stages were chosen deliberately over +//! a few coarse ones. -use std::collections::HashMap; use std::path::Path; +use std::sync::Arc; -use obicompactvec::{SiblingAnnexBuilder, SiblingInfo}; -use obikpartitionner::{KmerDesc, QueryHit}; +use obicompactvec::{ + PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder, SiblingInfo, +}; +use obikpartitionner::KmerPartition; use obikseq::{CanonicalKmer, Minimizer}; use obilayeredmap::{MphfLayer, OLMError}; use obilayeredmap::meta::PartitionMeta; @@ -90,6 +89,66 @@ impl Default for RunningState { } } +// ── obipipeline data types ───────────────────────────────────────────────── + +/// One distinct k-mer of the layer currently being processed, at its local +/// MPHF slot — the pipeline's source item. +#[derive(Clone, Copy)] +struct SourceItem { + slot: usize, + kmer: CanonicalKmer, +} + +/// One of a source k-mer's (up to 3) central-substitution variants, already +/// routed to its destination partition and carrying, decided here (both +/// encodings are already in hand — no need to wait for the lookup answer), +/// whether this specific variant would outrank the source as minorant. +#[derive(Clone, Copy)] +struct VariantQuery { + source_slot: usize, + dest_partition: usize, + variant: CanonicalKmer, + smaller: bool, +} + +/// Outcome of looking a [`VariantQuery`] up in its destination partition. +#[derive(Clone, Copy)] +struct AnswerMsg { + source_slot: usize, + hit: bool, + smaller: bool, +} + +#[derive(Clone, Copy)] +enum SibData { + Item(SourceItem), + Query(VariantQuery), + Answer(AnswerMsg), +} + +/// Existence-only lookup of `variant` in partition `dest_partition`: tries +/// each of the partition's layers in turn, stopping at the first hit — the +/// same "try every layer's MPHF" shape as `QueryLayer::find_slot` (private +/// to `obikpartitionner`), reimplemented here directly against the public +/// `MphfLayer::open`/`find` since only existence is needed, not a column +/// fetch. +fn lookup_exists(partition: &KmerPartition, dest_partition: usize, variant: CanonicalKmer) -> bool { + let index_dir = partition.part_dir(dest_partition).join(INDEX_SUBDIR); + if !index_dir.exists() { + return false; + } + let Ok(meta) = PartitionMeta::load(&index_dir) else { return false }; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + if let Ok(mphf) = MphfLayer::open(&layer_dir, &meta.mode) { + if mphf.find(variant).is_some() { + return true; + } + } + } + false +} + impl KmerIndex { /// Build the sibling-count/minorant annex for every layer of every /// partition of this (already built) index, writing one annex file per @@ -97,11 +156,32 @@ impl KmerIndex { /// (e.g. after a fresh `merge`) — each run simply overwrites the annex /// files of the index it is called on. /// + /// Construction only — no statistics gathered here on purpose: this is + /// meant to run routinely (it is the artefact the SNP-family distances + /// will consume), while the sibling-count distribution + /// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate, + /// occasional diagnostic pass over the result, not run every time. + /// /// Cross-partition/cross-layer lookups are required (a k-mer's siblings /// can live in any partition), but the layer loop itself — and thus the /// annex file this produces — stays local to one layer at a time. pub fn build_sibling_annex(&self) -> OKIResult<()> { let n_parts = self.n_partitions(); + let n_bits = n_parts.trailing_zeros() as usize; + + // A fresh, owned `KmerPartition` handle (read-only use only — no + // writers opened), wrapped in `Arc` so pipeline stage closures + // (which must be `'static`, running in spawned threads) can share + // it without borrowing `self`. + let partition = Arc::new( + KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?, + ); for part in 0..n_parts { let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); @@ -112,7 +192,7 @@ impl KmerIndex { for l in 0..meta.n_layers { let layer_dir = index_dir.join(format!("layer_{l}")); - self.build_layer_sibling_annex(&layer_dir, &meta, n_parts)?; + self.build_layer_sibling_annex(&layer_dir, n_parts, &partition)?; } } @@ -122,9 +202,11 @@ impl KmerIndex { fn build_layer_sibling_annex( &self, layer_dir: &Path, - meta: &PartitionMeta, n_parts: usize, + partition: &Arc, ) -> OKIResult<()> { + let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); + let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; let n_slots = mphf.n(); @@ -138,53 +220,55 @@ impl KmerIndex { } } - // ── Scatter: bucket outgoing variant queries by destination partition ── - // `KmerDesc.seq_idx` carries the origin (this layer's local slot); - // `KmerDesc.pos` is repurposed as a 0/1 flag: 1 iff this specific - // variant's own encoding is smaller than the source's — decided here, - // at scatter time, since both encodings are already in hand; the - // query response only needs to confirm existence (hit/miss). - let mut outgoing: Vec>> = - (0..n_parts).map(|_| HashMap::new()).collect(); + let sources: Vec = slot_kmer + .iter() + .enumerate() + .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| SourceItem { slot, kmer })) + .collect(); - for (slot, maybe_kmer) in slot_kmer.iter().enumerate() { - let Some(kmer) = maybe_kmer else { continue }; - for variant in kmer.central_canonical_neighbors() { - if variant == *kmer { - continue; // identity substitution — not a real variant + // ── obipipeline: Flat (generate variants) -> Transform (lookup) ───── + let n_workers = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4); + let capacity = 256; + + let partition_for_lookup = Arc::clone(partition); + let pipe = obipipeline::make_pipe! { + SibData : SourceItem => AnswerMsg, + || { + move |item: SourceItem| -> Vec { + let kmer = item.kmer; + kmer.central_canonical_neighbors() + .into_iter() + .filter(|variant| *variant != kmer) + .map(|variant| VariantQuery { + source_slot: item.slot, + dest_partition: partition_of(variant, n_parts), + variant, + smaller: variant.raw() < kmer.raw(), + }) + .collect::>() } - let smaller = variant.raw() < kmer.raw(); - let dest = partition_of(variant, n_parts); - outgoing[dest].entry(variant).or_default().push(KmerDesc { - seq_idx: slot as u32, - pos: smaller as u32, - }); - } - } + } : Item => Query, + | { + let partition_for_lookup = Arc::clone(&partition_for_lookup); + move |vq: VariantQuery| -> AnswerMsg { + let hit = lookup_exists(&partition_for_lookup, vq.dest_partition, vq.variant); + AnswerMsg { source_slot: vq.source_slot, hit, smaller: vq.smaller } + } + } : Query => Answer, + }; - // ── Gather + reconcile ─────────────────────────────────────────────── + // ── Reconciliation (the pipeline's sink): commutative fold of every + // answer into its origin k-mer's running state, in whatever order + // the pipeline delivers them. ───────────────────────────────────── let mut state = vec![RunningState::default(); n_slots]; - let n_genomes = self.meta.genomes.len(); - let with_counts = self.meta.config.with_counts; - - for (dest, kmers) in outgoing.iter().enumerate() { - if kmers.is_empty() { - continue; + for ans in pipe.apply(sources.into_iter(), n_workers, capacity) { + if ans.hit { + let st = &mut state[ans.source_slot]; + st.siblings = (st.siblings + 1).min(3); + if ans.smaller { + st.minorant = false; + } } - self.partition() - .query_partition_with(dest, kmers, n_genomes, with_counts, |hit| { - if let QueryHit::Found(descs) = hit { - for d in descs { - let slot = d.seq_idx as usize; - let st = &mut state[slot]; - st.siblings = (st.siblings + 1).min(3); - if d.pos == 1 { - st.minorant = false; - } - } - } - }) - .map_err(OKIError::Partition)?; } // ── Write the layer's annex file ───────────────────────────────────── @@ -202,6 +286,97 @@ impl KmerIndex { } } +/// Distribution of sibling counts (0-3), read back from an already-built +/// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's +/// presence/count data — a separate, occasional diagnostic pass, not fused +/// into construction. +#[derive(Debug, Clone, Default)] +pub struct SiblingAnnexStats { + /// `counts[s]` = number of k-mers (slots) with exactly `s` siblings, + /// counted once each regardless of how many genomes carry them. + pub counts: [u64; 4], + /// Of those, how many are minorant. + pub minorant_counts: [u64; 4], + /// `per_genome[g][s]` = number of k-mers with exactly `s` siblings that + /// genome `g` (index into `KmerIndex::meta().genomes`) carries. + pub per_genome: Vec<[u64; 4]>, +} + +impl KmerIndex { + /// Tally the sibling-count distribution of an already-built annex + /// (globally, and per genome). Errors if [`build_sibling_annex`] has not + /// been run on this index first. + /// + /// [`build_sibling_annex`]: Self::build_sibling_annex + pub fn sibling_annex_stats(&self) -> OKIResult { + let n_parts = self.n_partitions(); + let n_genomes = self.meta.genomes.len(); + let with_counts = self.meta.config.with_counts; + + let mut stats = SiblingAnnexStats { + per_genome: vec![[0u64; 4]; n_genomes], + ..Default::default() + }; + + for part in 0..n_parts { + let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + if !annex_path.exists() { + return Err(OKIError::InvalidInput(format!( + "no sibling annex at {} — run build_sibling_annex first", + annex_path.display() + ))); + } + let annex = SiblingAnnex::open(&annex_path)?; + let use_counts = with_counts && layer_dir.join("counts").exists(); + + // Opened once per layer, outside the slot loop. + enum Mat { + Count(PersistentCompactIntMatrix), + Presence(PersistentBitMatrix), + } + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(&layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(&layer_dir)?) + }; + let n_cols = match &mat { + Mat::Count(m) => m.n_cols(), + Mat::Presence(m) => m.n_cols(), + } + .min(n_genomes); + + for slot in 0..annex.len() { + let Some(info) = annex.get(slot) else { continue }; + let s = info.siblings as usize; + stats.counts[s] += 1; + if info.minorant { + stats.minorant_counts[s] += 1; + } + for g in 0..n_cols { + let carried = match &mat { + Mat::Count(m) => m.col_view(g).get(slot) != 0, + Mat::Presence(m) => m.get(g, slot) != 0, + }; + if carried { + stats.per_genome[g][s] += 1; + } + } + } + } + } + + Ok(stats) + } +} + #[cfg(test)] mod tests { use std::io::Write; diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index 999beba..3707387 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -3,7 +3,7 @@ use std::path::PathBuf; use clap::Args; use kodama::{Method, linkage}; -use obikindex::{DistanceMetric, KmerIndex}; +use obikindex::{DistanceMetric, KmerIndex, SiblingAnnexStats}; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; @@ -64,7 +64,21 @@ pub struct DistanceArgs { #[arg(long)] pub upgma: bool, + /// Build the sibling-count/minorant annex on this (multi-genome) index + /// — see `docmd/theory/evolutionary_distances.md`, Step 2b. Construction + /// only; does not by itself compute or write any statistics. + #[arg(long)] + pub sibling_annex: bool, + + /// Tally the sibling-count distribution (CSV) of an already-built annex + /// (run with `--sibling-annex` first, in this invocation or an earlier + /// one). A separate, occasional diagnostic pass — not run every time the + /// annex itself is (re)built. + #[arg(long)] + pub sibling_stats: bool, + /// Output prefix: _dist.csv, _shared.csv, + /// _siblings.csv, _siblings_per_genome.csv, /// _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] @@ -80,6 +94,26 @@ pub fn run(args: DistanceArgs) { let labels: Vec = idx.meta().genomes.iter().map(|g| g.label.clone()).collect(); let n = labels.len(); + + // ── Sibling-count/minorant annex (independent of the distance metric) ── + // Construction (`--sibling-annex`) and stats (`--sibling-stats`) are + // deliberately decoupled: the annex is meant to be (re)built routinely, + // the distribution only occasionally, on demand. + if args.sibling_annex { + info!("building sibling-count/minorant annex"); + idx.build_sibling_annex().unwrap_or_else(|e| { + eprintln!("error building sibling annex: {e}"); + std::process::exit(1); + }); + } + if args.sibling_stats { + let stats = idx.sibling_annex_stats().unwrap_or_else(|e| { + eprintln!("error computing sibling-annex stats: {e}"); + std::process::exit(1); + }); + write_sibling_stats_csv(&stats, &labels, &args.output); + } + info!( "computing {:?} distances for {} genome(s)", args.metric, n @@ -191,6 +225,45 @@ pub fn run(args: DistanceArgs) { } } +// ── Sibling-count distribution → CSV ──────────────────────────────────────── + +fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output: &Option) { + // Global histogram: one row per sibling count (0-3). + let global_path = output.as_ref() + .map(|p| format!("{}_siblings.csv", p.display())) + .unwrap_or_else(|| "siblings.csv".into()); + let mut f = BufWriter::new(std::fs::File::create(&global_path).unwrap_or_else(|e| { + eprintln!("error creating {global_path}: {e}"); + std::process::exit(1); + })); + writeln!(f, "siblings,slots,minorant_slots").unwrap(); + for s in 0..4 { + writeln!(f, "{s},{},{}", stats.counts[s], stats.minorant_counts[s]).unwrap(); + } + let total: u64 = stats.counts.iter().sum(); + info!("sibling-count distribution → {global_path} (total {total} slot(s))"); + + // Per-genome breakdown: one row per genome, 4 columns (0-3), + a total row. + let per_genome_path = output.as_ref() + .map(|p| format!("{}_siblings_per_genome.csv", p.display())) + .unwrap_or_else(|| "siblings_per_genome.csv".into()); + let mut f = BufWriter::new(std::fs::File::create(&per_genome_path).unwrap_or_else(|e| { + eprintln!("error creating {per_genome_path}: {e}"); + std::process::exit(1); + })); + writeln!(f, "genome,0,1,2,3").unwrap(); + let mut column_totals = [0u64; 4]; + for (label, counts) in labels.iter().zip(stats.per_genome.iter()) { + writeln!(f, "{label},{},{},{},{}", counts[0], counts[1], counts[2], counts[3]).unwrap(); + for s in 0..4 { column_totals[s] += counts[s]; } + } + writeln!( + f, "total,{},{},{},{}", + column_totals[0], column_totals[1], column_totals[2], column_totals[3], + ).unwrap(); + info!("per-genome sibling-count distribution → {per_genome_path}"); +} + // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── fn upgma_to_newick(dendro: &kodama::Dendrogram, names: &[String]) -> String { -- 2.54.0 From 1a470eab9e677bc4883ac19e18a4b6d00bf28cca Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 17:21:04 +0200 Subject: [PATCH 6/8] Refactor k-mer sibling tracking to compact bitmask and on-demand counts Replaces the explicit `SiblingInfo` struct and 3-bit minorant flags with a derived 4-bit presence mask (`FamilyMask`) that tracks observed bases per family. This eliminates redundant file I/O overhead by introducing a `PartitionCache` for batch lookups, simplifies serialization, and updates all downstream builders, stats computation, and tests to operate on the new bitmask representation. Adjusts CLI output to report deduplicated family sizes instead of histograms, ignores generated CSV files, and updates documentation to reflect the fixed canonical reference and new theory. --- .gitignore | 1 + docmd/theory/evolutionary_distances.md | 115 ++++- src/obicompactvec/src/lib.rs | 2 +- src/obicompactvec/src/siblingannex.rs | 165 ++++--- src/obikindex/src/siblings.rs | 604 ++++++++++++++++++------- src/obikmer/src/cmd/distance.rs | 61 +-- 6 files changed, 671 insertions(+), 277 deletions(-) diff --git a/.gitignore b/.gitignore index 793b0a9..d2d43e5 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,7 @@ data-stress ./**/*.json *.bin *.log +*.csv Betula_exilis--IGA-24-33 benchmark/genomes benchmark/simulated_data diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index 6ff266b..2c21a3c 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -65,6 +65,50 @@ 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. +### Definitions: family, and the canonical form of a family + +**Family.** The family of a k-mer `x` is the set of (up to) 4 k-mers sharing +`x`'s `2m` flanking bases, differing only at the central base `m`. Membership +is a property of the flank pattern, not of `x` itself: any of the 4 possible +central substitutions belongs to the same family. + +**`central_canonical_neighbors()`** (`obikseq`, `CanonicalKmerOf::central_canonical_neighbors`) +generates all 4 members from any one of them (observed or not), each +independently canonicalised (`.canonical()`, i.e. `min(kmer, revcomp(kmer))`). +This independent canonicalisation is necessary because a central substitution +can flip which orientation is lexicographically smaller — two members of the +same family can end up canonicalised in *different* orientations. Despite +that, the **set** of 4 resulting canonical k-mers is invariant: calling +`central_canonical_neighbors()` on any member of a family — present in the +index or not — yields the same 4 values. This is relied upon throughout the +rest of this document. + +**Canonical form of a family.** Because orientation can differ member to +member, "which of the 4 is the reference" cannot be defined relative to +*whichever member happened to be visited first*, nor relative to the +minorant (see below) — both are data-dependent (they depend on what is +actually observed), so using either as the reference would make the +reference itself vary depending on what happens to be present in a given +index. Instead: **the canonical form of a family is, by definition, the +member whose own central base — read in its own already-canonical +orientation — is `A`.** This is well-defined for every family, computed +purely from the flank pattern, whether or not that specific member (or any +member at all) is actually observed anywhere in the index. Concretely: call +`central_canonical_neighbors()` on any member (observed or not) to get the +family's 4 canonical forms; the one among them whose own centre nucleotide is +`A` is the family's canonical form. The other 3 (`C`, `G`, `T`) are labelled +relative to *that* fixed reference, not relative to the calling member's own +orientation. + +**Consequence for the minorant.** With this fixed A-referenced labelling, +`minorant` (the smallest raw encoding among the family's *observed* members, +introduced further below) becomes directly computable rather than needing to +be tracked as extra state: regenerate the family's 4 canonical forms from +any member's own k-mer (cheap, no lookup), compare the raw encodings of +whichever are marked present, and take the smallest. No separate stored bit +is required — see Step 2b below, where this replaces the earlier +minorant-bit design. + ## Locus eligibility: raw definition vs. paralogy filter For each k-mer `x` observed in genome A (source, one MPHF slot; the 3 @@ -482,31 +526,58 @@ here. ### Step 2b — sibling-count / minorant annex (consolidated plan) -Scope: only the precursor annex (sibling count 0-3 per slot, minorant -decided on demand) — not the SNP tally itself, whose Step 2 sweep remains -unresolved above. This piece is simpler than the sweep, because it writes to -an independent per-slot value, not a shared cross-k-mer accumulator, so it -needs no dedup/ownership logic at all at this stage. +Scope: only the precursor annex — not the SNP tally itself, whose Step 2 +sweep remains unresolved above. This piece is simpler than the sweep, +because it writes to an independent per-slot value, not a shared +cross-k-mer accumulator, so it needs no dedup/ownership logic at all at this +stage. + +**Revised annex encoding — 4-bit presence mask, not 3-bit (minorant + +count).** Superseded after settling the "canonical form of a family" +definition above. The 3-bit design (1 minorant bit + 2-bit sibling count, +§ below, kept for the historical record) had two problems: it discards +*which* variants are present (only how many), so any future consumer +(the SNP sweep, or a stats pass — see below) that needs to know which bases +exist still has to regenerate and blindly re-query all 3 candidates; and +the minorant bit's meaning was tied to whichever member was visited, not to +a fixed reference. Storing instead a **4-bit mask** — one bit per base +(A/C/G/T), set iff that member of the family (labelled relative to the +family's fixed canonical form, i.e. the member with `A` at the centre — see +above) is observed anywhere in the index — fixes both: +- **Sibling count is derived, not stored**: `siblings = popcount(mask) - 1`. +- **Minorant is derived, not stored**: regenerate the family's 4 canonical + forms from the slot's own k-mer (cheap, no lookup — see above), compare + the raw encodings of whichever bits are set in the mask, take the + smallest. +- **A future consumer knows exactly which variants to (re-)query** — + `popcount(mask) - 1` lookups instead of always 3, and it knows *which* + 3 (or fewer) to issue, not just how many hits to expect. +- The all-zero value (no base present at all) is still logically + unreachable as a real result — the slot's *own* base is always present in + its own family — so it remains available as a free "not yet computed" + sentinel, exactly as before. 1. **Primitive.** Reuse `central_canonical_neighbors()` from Step 0 - unchanged — the 3 canonicalised central-substitution variants of a k-mer. -2. **New annex type** (`obicompactvec`, alongside `bitmatrix.rs`): a 3-bit- - per-slot packed array, one per partition — same on-disk shape family as - `PersistentBitMatrix`'s `Packed` variant, but simpler (no per-genome - columns, a single derived read-only value per slot). 3 bits, not 2: - revised to also store minorant status alongside sibling count, since it - comes for free from the same lookups (point 3 below) — 5 real states + unchanged — the 3 canonicalised central-substitution variants of a k-mer + (plus the identity, i.e. all 4 members of the family — see "Definitions" + above). +2. **New annex type** (`obicompactvec`, alongside `bitmatrix.rs`): a 4-bit- + per-slot packed array (the presence mask above), one per partition — same + on-disk shape family as `PersistentBitMatrix`'s `Packed` variant, but + simpler (no per-genome columns, a single derived read-only value per + slot). +
Superseded 3-bit design (historical) + 3 bits, storing minorant status alongside sibling count directly, since + it came for free from the same lookups (point 3 below) — 5 real states (not-minorant; minorant with 0/1/2/3 siblings) fit in 3 bits (8 states, - 3 unused). This lets the future SNP sweep discard a non-minorant slot - **instantly**, with no lookup at all, instead of having to regenerate and - look up its siblings just to rediscover it isn't the designated writer — - moving that cost into this one-time, cached pass instead of repeating it - on every future sweep. The otherwise-unreachable combination - "not-minorant + 0 siblings" (impossible: 0 siblings always implies - minorant, see below) doubles as a free **"not yet computed" sentinel** — - annex files for all partitions/layers can be pre-initialised to this - value before the computation pass runs, distinguishing genuinely-computed - 0-sibling slots from not-yet-processed ones with no extra storage. + 3 unused). This let the future SNP sweep discard a non-minorant slot + instantly, with no lookup at all. The otherwise-unreachable combination + "not-minorant + 0 siblings" (0 siblings always implies minorant) doubled + as the "not yet computed" sentinel. Replaced by the 4-bit mask above, + which subsumes this benefit (minorant still derivable, now for free at + read time rather than stored) while also fixing the "which variant" + blindness. +
3. **Computation pass** (`obikindex`, new `siblings.rs`): **one `obipipeline` run per layer, iterated sequentially over the index's layers** — settled after two false starts, worth recording both. diff --git a/src/obicompactvec/src/lib.rs b/src/obicompactvec/src/lib.rs index 9a136bb..1f85a49 100644 --- a/src/obicompactvec/src/lib.rs +++ b/src/obicompactvec/src/lib.rs @@ -19,7 +19,7 @@ pub use builder::PersistentCompactIntVecBuilder; pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask}; pub use intmatrix::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, pack_compact_int_matrix}; pub use layer_meta::LayerMeta; -pub use siblingannex::{SiblingAnnex, SiblingAnnexBuilder, SiblingInfo}; +pub use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder}; pub use reader::{PersistentCompactIntVec, Iter as CompactIntVecIter}; pub use tempbitvec::{TempBitVec, TempBitVecBuilder}; pub use tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder}; diff --git a/src/obicompactvec/src/siblingannex.rs b/src/obicompactvec/src/siblingannex.rs index 94cb5cc..0de9733 100644 --- a/src/obicompactvec/src/siblingannex.rs +++ b/src/obicompactvec/src/siblingannex.rs @@ -1,27 +1,33 @@ -//! Sibling-count / minorant annex: a compact, read-only-after-build, per-slot +//! Family presence-mask annex: a compact, read-only-after-build, per-slot //! derived value used by the central-position SNP distance estimator (see -//! `docmd/theory/evolutionary_distances.md`, "Step 2b"). +//! `docmd/theory/evolutionary_distances.md`, "Step 2b" and "Definitions: +//! family, and the canonical form of a family"). //! -//! One byte is stored per MPHF slot of a partition/layer, encoding two -//! independent facts about the slot's k-mer's "family" (the up-to-4 k-mers -//! sharing the same flanks, differing only at the central base), both -//! properties of the whole current multi-genome index, not of any one -//! genome: +//! One byte is stored per MPHF slot of a partition/layer, its low 4 bits +//! encoding a **presence mask** for the slot's k-mer's "family" (the up to 4 +//! k-mers sharing the same flanks, differing only at the central base): +//! bit `b` (`b` = 0..3, in the fixed A/C/G/T = 0/1/2/3 encoding already used +//! for a single nucleotide) is set iff the family member whose *own* central +//! base — in its own canonical orientation — is `b`, is observed anywhere in +//! the current multi-genome index. This is a property of the whole index, +//! not of any one genome. //! -//! - bit 0: `minorant` — is this k-mer's canonical encoding the smallest -//! among the family members actually observed in the index? -//! - bits 1-2: `siblings` — how many *other* family members (0-3) are -//! observed anywhere in the index. +//! Both facts the earlier (superseded) 3-bit design stored explicitly are +//! derived from the mask instead, not stored: +//! - sibling count = `popcount(mask) - 1`; +//! - minorant = regenerate the family's 4 canonical forms from the slot's +//! own k-mer (`CanonicalKmerOf::central_canonical_neighbors`, cheap, no +//! lookup), compare the raw encodings of whichever are set in the mask, +//! take the smallest — see `obikindex::siblings`. //! -//! Byte value 0 (`minorant = false`, `siblings = 0`) is logically -//! unreachable as a real result (0 siblings always implies minorant — see -//! the design doc) and is reused as the "not yet computed" sentinel: annex -//! files are pre-initialised to all-zero, and a real value is only ever -//! written once, by the computation pass. +//! Mask value 0 is logically unreachable as a real result (a slot's own base +//! is always present in its own family) and is reused as the "not yet +//! computed" sentinel: annex files are pre-initialised to all-zero, and a +//! real value is only ever written once, by the computation pass. //! -//! Deliberately simpler than a true 3-bit pack (1 byte/slot instead of 3 -//! bits/slot): correctness and simplicity first: for a first implementation. -//! Packing to 3 bits/slot is a pure storage-density follow-up, not a +//! Deliberately simpler than a true 4-bit pack (1 byte/slot instead of 4 +//! bits/slot): correctness and simplicity first, for a first implementation. +//! Packing to 4 bits/slot is a pure storage-density follow-up, not a //! behavioural change, left for later. use std::fs::{File, OpenOptions}; @@ -35,30 +41,70 @@ const MAGIC: [u8; 4] = *b"PSIB"; // Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (1 byte/slot) follows. const HEADER_SIZE: usize = 16; -/// Decoded value of one slot's annex entry. +/// A family presence mask: bit `b` set iff the member whose own canonical +/// central base is `b` (0=A, 1=C, 2=G, 3=T) is observed in the index. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct SiblingInfo { - pub minorant: bool, - pub siblings: u8, // 0..=3 -} +pub struct FamilyMask(u8); + +impl FamilyMask { + /// The empty mask — never a valid *computed* result (a slot's own base + /// is always present in its own family) — used only to build up a mask + /// via repeated [`with`](Self::with) calls before storing it. + pub const EMPTY: FamilyMask = FamilyMask(0); + + /// Set bit `base` (0=A, 1=C, 2=G, 3=T). + #[inline] + pub fn with(self, base: u8) -> Self { + debug_assert!(base < 4, "base out of range: {base}"); + FamilyMask(self.0 | (1 << base)) + } + + /// Is the member with central base `base` (0..3) present? + #[inline] + pub fn has(self, base: u8) -> bool { + debug_assert!(base < 4, "base out of range: {base}"); + self.0 & (1 << base) != 0 + } + + /// Number of family members observed anywhere in the index (1..=4). + #[inline] + pub fn family_size(self) -> u32 { + self.0.count_ones() + } + + /// Number of *other* members observed (0..=3) — `family_size() - 1`. + #[inline] + pub fn siblings(self) -> u32 { + self.family_size() - 1 + } + + /// Raw bitmask (bit `b` = base `b` present) — for callers that build up + /// a mask via their own bit operations (e.g. concurrently, via an + /// `AtomicU8`) and only need the `FamilyMask` wrapper at the end. + #[inline] + pub fn bits(self) -> u8 { + self.0 + } + + /// Construct from a raw bitmask (only the low 4 bits are kept). + #[inline] + pub fn from_bits(bits: u8) -> Self { + FamilyMask(bits & 0b1111) + } -impl SiblingInfo { #[inline] fn encode(self) -> u8 { - (self.siblings << 1) | (self.minorant as u8) + self.0 } #[inline] fn decode(byte: u8) -> Option { if byte == 0 { - // The unreachable "not minorant + 0 siblings" combination — - // reserved as the "not yet computed" sentinel. + // Unreachable for a real result — reserved as the "not yet + // computed" sentinel. return None; } - Some(SiblingInfo { - minorant: byte & 1 != 0, - siblings: (byte >> 1) & 0b11, - }) + Some(FamilyMask(byte & 0b1111)) } } @@ -91,8 +137,8 @@ impl SiblingAnnex { pub fn is_empty(&self) -> bool { self.n == 0 } /// `None` means the slot has not (yet) been computed — see module docs. - pub fn get(&self, slot: usize) -> Option { - SiblingInfo::decode(self.mmap[HEADER_SIZE + slot]) + pub fn get(&self, slot: usize) -> Option { + FamilyMask::decode(self.mmap[HEADER_SIZE + slot]) } } @@ -124,17 +170,17 @@ impl SiblingAnnexBuilder { pub fn len(&self) -> usize { self.n } pub fn is_empty(&self) -> bool { self.n == 0 } - pub fn get(&self, slot: usize) -> Option { - SiblingInfo::decode(self.mmap[HEADER_SIZE + slot]) + pub fn get(&self, slot: usize) -> Option { + FamilyMask::decode(self.mmap[HEADER_SIZE + slot]) } - pub fn set(&mut self, slot: usize, info: SiblingInfo) { + pub fn set(&mut self, slot: usize, mask: FamilyMask) { // Redundant concurrent writes from independent recomputation paths // converge to the same encoded byte for a given slot, so a plain // store here is safe even without external synchronisation, as long // as the byte write itself is atomic (true for a single aligned // byte on every platform this project targets). - self.mmap[HEADER_SIZE + slot] = info.encode(); + self.mmap[HEADER_SIZE + slot] = mask.encode(); } pub fn close(self) -> io::Result<()> { self.mmap.flush() } @@ -163,36 +209,37 @@ mod tests { } #[test] - fn roundtrip_all_valid_states() { + fn roundtrip_all_valid_masks() { let dir = tempdir().unwrap(); let path = dir.path().join("test.psib"); - let mut builder = SiblingAnnexBuilder::new(5, &path).unwrap(); + let mut builder = SiblingAnnexBuilder::new(4, &path).unwrap(); - let cases = [ - SiblingInfo { minorant: true, siblings: 0 }, - SiblingInfo { minorant: true, siblings: 1 }, - SiblingInfo { minorant: true, siblings: 2 }, - SiblingInfo { minorant: true, siblings: 3 }, - SiblingInfo { minorant: false, siblings: 2 }, + let masks = [ + FamilyMask::EMPTY.with(0), // just A: family size 1 + FamilyMask::EMPTY.with(0).with(3), // A + T: size 2 + FamilyMask::EMPTY.with(1).with(2).with(3), // C+G+T: size 3 + FamilyMask::EMPTY.with(0).with(1).with(2).with(3), // all 4 ]; - for (slot, info) in cases.iter().enumerate() { - builder.set(slot, *info); + for (slot, mask) in masks.iter().enumerate() { + builder.set(slot, *mask); } let annex = builder.finish().unwrap(); - for (slot, info) in cases.iter().enumerate() { - assert_eq!(annex.get(slot), Some(*info)); + for (slot, mask) in masks.iter().enumerate() { + assert_eq!(annex.get(slot), Some(*mask)); } + assert_eq!(annex.get(0).unwrap().siblings(), 0); + assert_eq!(annex.get(1).unwrap().siblings(), 1); + assert_eq!(annex.get(2).unwrap().siblings(), 2); + assert_eq!(annex.get(3).unwrap().siblings(), 3); + assert_eq!(annex.get(3).unwrap().family_size(), 4); } #[test] - fn not_minorant_zero_siblings_is_unreachable_via_set_and_decodes_as_sentinel() { - // Documented invariant, not enforced by the type: callers must never - // construct this combination. If they do, it is indistinguishable - // from "not computed" — exercised here to pin the behaviour down. - let dir = tempdir().unwrap(); - let path = dir.path().join("test.psib"); - let mut builder = SiblingAnnexBuilder::new(1, &path).unwrap(); - builder.set(0, SiblingInfo { minorant: false, siblings: 0 }); - assert_eq!(builder.get(0), None); + fn has_reflects_individual_bits() { + let mask = FamilyMask::EMPTY.with(0).with(2); + assert!(mask.has(0)); + assert!(!mask.has(1)); + assert!(mask.has(2)); + assert!(!mask.has(3)); } } diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index 4638bde..e80bd67 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -1,39 +1,49 @@ -//! Sibling-count / minorant annex construction. +//! Family presence-mask annex construction. //! -//! See `docmd/theory/evolutionary_distances.md`, "Step 2b — sibling-count / -//! minorant annex", for the full design discussion this implements. +//! See `docmd/theory/evolutionary_distances.md`, "Definitions: family, and +//! the canonical form of a family" and "Step 2b", for the full design +//! discussion this implements. //! //! For each distinct k-mer of each layer of the (already built/merged) -//! index, computes two facts about its "family" (the up to 4 k-mers sharing -//! its flanks, differing only at the central base — well-defined for odd -//! k), both properties of the whole current multi-genome index rather than -//! of any one genome: -//! - how many *other* family members (0-3) are observed anywhere in the -//! index; -//! - whether this k-mer is the "minorant" of its family — the smallest -//! canonical encoding among the members actually observed. +//! index, computes a 4-bit presence mask for its "family" (the up to 4 +//! k-mers sharing its flanks, differing only at the central base — +//! well-defined for odd k): bit `b` set iff the family member whose own +//! canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed anywhere +//! in the current multi-genome index — a property of the whole index, not +//! of any one genome. Sibling count and minorant are *derived* from the +//! mask by callers, not stored (see `FamilyMask` and +//! [`sibling_annex_stats`](KmerIndex::sibling_annex_stats) below). //! -//! Per layer, the computation runs as an `obipipeline` pipeline with several -//! elementary stages (a `Flat` stage generating each k-mer's 3 central -//! variants, a `Transform` stage looking each variant up in its destination -//! partition), so the scheduler's shared worker pool interleaves this work -//! across many in-flight k-mers/variants rather than processing everything -//! on a single thread — see `docmd/theory/evolutionary_distances.md`, Step -//! 2b, "Mechanism", for why elementary stages were chosen deliberately over -//! a few coarse ones. +//! Per layer, an `obipipeline` `Flat` stage (throttled — see +//! `obipipeline::throttle`) generates each k-mer's 3 central variants, +//! interleaved across many in-flight k-mers by the scheduler's shared +//! worker pool rather than processed on a single thread. The actual +//! cross-partition lookup, though, reuses +//! `KmerPartition::query_partition_with` — the same partition-batching +//! mechanism `obikmer query` already uses (open a partition's files once, +//! answer a whole batch of queries against it) — rather than a per-item +//! pipeline stage: an earlier per-item design reopened/re-mmap'd every +//! target partition's files on every single lookup, which was fine at +//! toy scale but manifested as ~90% system time against a real index. +//! See `docmd/theory/evolutionary_distances.md`, Step 2b, "Mechanism". use std::path::Path; +use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; +use rayon::prelude::*; + use obicompactvec::{ - PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder, SiblingInfo, + FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder, }; use obikpartitionner::KmerPartition; +use obipipeline::ThrottleGuard; use obikseq::{CanonicalKmer, Minimizer}; use obilayeredmap::{MphfLayer, OLMError}; use obilayeredmap::meta::PartitionMeta; use obiskbuilder::rolling_stat::RollingStat; use obiskio::UnitigFileReader; +use obisys::progress_bar; use crate::error::{OKIError, OKIResult}; use crate::index::KmerIndex; @@ -48,6 +58,26 @@ fn olm_to_ok(e: OLMError) -> OKIError { } } +/// Central-position base of a canonical k-mer, in the fixed 0=A/1=C/2=G/3=T +/// encoding — the mask's bit index. `k` must be odd (project invariant). +#[inline] +fn central_base(kmer: CanonicalKmer, k: usize) -> u8 { + kmer.nucleotide((k - 1) / 2) +} + +/// Is `kmer` the minorant of its family, given the family's presence mask? +/// Regenerates the family's 4 canonical forms from `kmer` itself (cheap, no +/// lookup — see the design doc's "Definitions" section for why this is +/// always safe: the set of 4 forms is invariant regardless of which member +/// you start from), and compares the raw encodings of whichever are marked +/// present in `mask`. +fn is_minorant(kmer: CanonicalKmer, mask: FamilyMask, k: usize) -> bool { + kmer.central_canonical_neighbors().into_iter().all(|other| { + other == kmer || !mask.has(central_base(other, k)) || kmer.raw() <= other.raw() + }) +} + + /// Minimiser of a single, isolated canonical k-mer (not part of a streamed /// sequence). `RollingStat` computes minimisers incrementally along a /// sequence; this feeds one k-mer's bases through a fresh instance to get @@ -74,79 +104,144 @@ fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize { (lone_kmer_minimizer(kmer).seq_hash() & mask) as usize } -/// Running reconciliation state for one source k-mer, initialised to the -/// trivial "no siblings observed yet" state and folded incrementally (in any -/// order — commutative) as query answers come back. -#[derive(Clone, Copy)] -struct RunningState { - minorant: bool, - siblings: u8, -} - -impl Default for RunningState { - fn default() -> Self { - RunningState { minorant: true, siblings: 0 } - } -} - // ── obipipeline data types ───────────────────────────────────────────────── /// One distinct k-mer of the layer currently being processed, at its local -/// MPHF slot — the pipeline's source item. -#[derive(Clone, Copy)] +/// MPHF slot — the pipeline's source item. Carries a throttle slot (shared, +/// `Arc`-wrapped so it can be cloned into each of the (up to 3) items this +/// one fans out into via the `Flat` stage) that is only released once every +/// one of those descendants has been fully processed — see the module docs' +/// "Throttling" note for why this is required, not optional, once a `Flat` +/// stage is in the pipeline. struct SourceItem { slot: usize, kmer: CanonicalKmer, + _permit: Arc, } /// One of a source k-mer's (up to 3) central-substitution variants, already -/// routed to its destination partition and carrying, decided here (both -/// encodings are already in hand — no need to wait for the lookup answer), -/// whether this specific variant would outrank the source as minorant. -#[derive(Clone, Copy)] +/// routed to its destination partition and carrying its own central base +/// (0=A/1=C/2=G/3=T) — the mask bit it will set on a hit. Carries a clone of +/// the source item's throttle permit. struct VariantQuery { source_slot: usize, dest_partition: usize, variant: CanonicalKmer, - smaller: bool, + base: u8, + _permit: Arc, } -/// Outcome of looking a [`VariantQuery`] up in its destination partition. -#[derive(Clone, Copy)] -struct AnswerMsg { - source_slot: usize, - hit: bool, - smaller: bool, -} - -#[derive(Clone, Copy)] enum SibData { Item(SourceItem), Query(VariantQuery), - Answer(AnswerMsg), } -/// Existence-only lookup of `variant` in partition `dest_partition`: tries -/// each of the partition's layers in turn, stopping at the first hit — the -/// same "try every layer's MPHF" shape as `QueryLayer::find_slot` (private -/// to `obikpartitionner`), reimplemented here directly against the public -/// `MphfLayer::open`/`find` since only existence is needed, not a column -/// fetch. -fn lookup_exists(partition: &KmerPartition, dest_partition: usize, variant: CanonicalKmer) -> bool { - let index_dir = partition.part_dir(dest_partition).join(INDEX_SUBDIR); - if !index_dir.exists() { - return false; - } - let Ok(meta) = PartitionMeta::load(&index_dir) else { return false }; - for l in 0..meta.n_layers { - let layer_dir = index_dir.join(format!("layer_{l}")); - if let Ok(mphf) = MphfLayer::open(&layer_dir, &meta.mode) { - if mphf.find(variant).is_some() { - return true; - } +/// Every partition's already-open MPHF layers, built **once** for the whole +/// `build_sibling_annex` run and shared (read-only) across every lookup, in +/// every source layer, for the rest of the run — not reopened/re-mmap'd per +/// query, nor per source layer. +/// +/// Confirmed necessary by sampling a real run: routing lookups through +/// `KmerPartition::query_partition_with` (the same batching `obikmer query` +/// uses) still reopens+re-mmaps every target partition's files on every +/// call, and it is called once per destination partition **per source +/// layer** — for an index with many layers this repeats the same +/// `MphfLayer::open`/`Evidence::open`/`PersistentBitMatrix::open` work over +/// and over. Parallelising those calls (see the gather step below) spread +/// the redundant work across more cores but did not reduce it: sampling +/// showed Rayon workers spending their time inside repeated `open()` +/// syscalls, not computation. This cache amortises that cost to once per +/// partition for the entire run, regardless of how many source layers or +/// lookups follow. +/// A cached layer's opened presence/count matrix, alongside its `MphfLayer`. +enum Mat { + Count(PersistentCompactIntMatrix), + Presence(PersistentBitMatrix), +} + +impl Mat { + fn n_cols(&self) -> usize { + match self { + Mat::Count(m) => m.n_cols(), + Mat::Presence(m) => m.n_cols(), } } - false + fn carries(&self, g: usize, slot: usize) -> bool { + match self { + Mat::Count(m) => m.col_view(g).get(slot) != 0, + Mat::Presence(m) => m.get(g, slot) != 0, + } + } +} + +struct PartitionCache { + /// `layers[partition][layer]` = that partition's opened MPHF layers, + /// paired 1:1 with `mats[partition][layer]`; empty if the partition + /// directory doesn't exist. Used by both [`KmerIndex::build_sibling_annex`] + /// (`layers` only) and [`KmerIndex::sibling_annex_stats`] (both). + layers: Vec>, + mats: Vec>, +} + +impl PartitionCache { + fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult { + let pb = progress_bar("open_partitions", n_parts as u64, "partitions"); + let built: Vec<(Vec, Vec)> = (0..n_parts) + .into_par_iter() + .map(|part| -> OKIResult<(Vec, Vec)> { + let index_dir = partition.part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + pb.inc(1); + return Ok((Vec::new(), Vec::new())); + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + let mut layers = Vec::with_capacity(meta.n_layers); + let mut mats = Vec::with_capacity(meta.n_layers); + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let Ok(mphf) = MphfLayer::open(&layer_dir, &meta.mode) else { continue }; + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + PersistentCompactIntMatrix::open(&layer_dir).ok().map(Mat::Count) + } else { + PersistentBitMatrix::open(&layer_dir).ok().map(Mat::Presence) + }; + let Some(mat) = mat else { continue }; + layers.push(mphf); + mats.push(mat); + } + pb.inc(1); + Ok((layers, mats)) + }) + .collect::>>()?; + pb.finish_and_clear(); + let (layers, mats) = built.into_iter().unzip(); + Ok(Self { layers, mats }) + } + + /// Existence-only lookup of `variant` in partition `dest_partition`: + /// tries each of the partition's already-open layers in turn, stopping + /// at the first hit. + fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> bool { + self.layers + .get(dest_partition) + .is_some_and(|layers| layers.iter().any(|mphf| mphf.find(variant).is_some())) + } + + /// Per-genome presence vector for `variant` in partition `dest_partition` + /// (`true` iff that genome carries it), `None` on a miss. Same shape as + /// `find`, but also reads the cached matrix instead of just the MPHF. + fn find_presence(&self, dest_partition: usize, variant: CanonicalKmer, n_genomes: usize) -> Option> { + let layers = self.layers.get(dest_partition)?; + let mats = self.mats.get(dest_partition)?; + for (mphf, mat) in layers.iter().zip(mats.iter()) { + if let Some(slot) = mphf.find(variant) { + let n_cols = mat.n_cols().min(n_genomes); + return Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()); + } + } + None + } } impl KmerIndex { @@ -169,42 +264,50 @@ impl KmerIndex { let n_parts = self.n_partitions(); let n_bits = n_parts.trailing_zeros() as usize; - // A fresh, owned `KmerPartition` handle (read-only use only — no - // writers opened), wrapped in `Arc` so pipeline stage closures - // (which must be `'static`, running in spawned threads) can share - // it without borrowing `self`. - let partition = Arc::new( - KmerPartition::open_with_config( - &self.root_path, - self.kmer_size(), - self.minimizer_size(), - n_bits, - ) - .map_err(OKIError::Partition)?, - ); + let partition = KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep"); + let cache = Arc::new(PartitionCache::build(&partition, n_parts, self.meta.config.with_counts)?); + + let pb = progress_bar("sibling_annex", n_parts as u64, "partitions"); + let mut total_slots: u64 = 0; for part in 0..n_parts { let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); if !index_dir.exists() { + pb.inc(1); continue; } let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + let mut part_slots: u64 = 0; for l in 0..meta.n_layers { let layer_dir = index_dir.join(format!("layer_{l}")); - self.build_layer_sibling_annex(&layer_dir, n_parts, &partition)?; + part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?; } + total_slots += part_slots; + pb.inc(1); + pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)")); } + pb.finish_and_clear(); + tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions"); Ok(()) } + /// Returns the number of distinct k-mers (annex slots) processed, for + /// progress reporting. fn build_layer_sibling_annex( &self, layer_dir: &Path, n_parts: usize, - partition: &Arc, - ) -> OKIResult<()> { + cache: &Arc, + ) -> OKIResult { let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; @@ -220,22 +323,61 @@ impl KmerIndex { } } - let sources: Vec = slot_kmer - .iter() - .enumerate() - .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| SourceItem { slot, kmer })) - .collect(); + let k = self.kmer_size(); - // ── obipipeline: Flat (generate variants) -> Transform (lookup) ───── + // ── Reconciliation state, initialised with each slot's own base — + // that member is trivially present, no lookup needed. Built before + // the pipeline runs, from the same enumeration, since `sources` + // below is consumed as a throttled iterator, not collected. + // `AtomicU8`, not `FamilyMask`, because the gather phase below + // parallelises across destination partitions (independent + // `query_partition_with` calls, safe to run concurrently) and their + // `Found` hits can land on arbitrary, possibly-shared slots — a + // lock-free `fetch_or` avoids needing any synchronisation beyond + // that. ───────────────────────────────────────────────────────── + let mask: Vec = (0..n_slots).map(|_| AtomicU8::new(0)).collect(); + for (slot, kmer) in slot_kmer.iter().enumerate().filter_map(|(s, k)| k.map(|k| (s, k))) { + mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed); + } + + // ── obipipeline: Flat stage generates variants only — the actual + // cross-partition lookup reuses `KmerPartition::query_partition_with` + // (the same batching mechanism `obikmer query` already uses: open a + // partition's files once, answer a whole batch of queries against + // it) instead of one lookup per pipeline item. A per-item lookup + // (tried first) reopened/re-mmap'd every target partition's files on + // every single variant — fine at the scale of a handful of test + // k-mers, but with billions of lookups against a real index this + // manifested as ~90% system time, observed in practice. ─────────── let n_workers = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4); let capacity = 256; - let partition_for_lookup = Arc::clone(partition); + // Throttling is not optional once a `Flat` stage is in the pipeline + // (see `obipipeline::throttle`'s docs): without it, every worker can + // become a simultaneous `Flat` producer, saturate the shared output + // channel, and deadlock against the scheduler's own dispatch loop — + // also observed in practice. The permit acquired here for a source + // k-mer is held (via the `Arc`-shared guard carried through + // `SourceItem` -> `VariantQuery`) until every one of its (up to 3) + // descendants has been read out of the pipeline by the accumulation + // loop below, not just until the `Flat` stage itself returns. + let sources: Vec<(usize, CanonicalKmer)> = slot_kmer + .iter() + .enumerate() + .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer))) + .collect(); + let throttled = obipipeline::throttle(sources.into_iter(), n_workers).map(|t| SourceItem { + slot: t.item.0, + kmer: t.item.1, + _permit: Arc::new(t.guard), + }); + let pipe = obipipeline::make_pipe! { - SibData : SourceItem => AnswerMsg, + SibData : SourceItem => VariantQuery, || { move |item: SourceItem| -> Vec { let kmer = item.kmer; + let permit = item._permit; kmer.central_canonical_neighbors() .into_iter() .filter(|variant| *variant != kmer) @@ -243,88 +385,112 @@ impl KmerIndex { source_slot: item.slot, dest_partition: partition_of(variant, n_parts), variant, - smaller: variant.raw() < kmer.raw(), + base: central_base(variant, k), + _permit: Arc::clone(&permit), }) .collect::>() } } : Item => Query, - | { - let partition_for_lookup = Arc::clone(&partition_for_lookup); - move |vq: VariantQuery| -> AnswerMsg { - let hit = lookup_exists(&partition_for_lookup, vq.dest_partition, vq.variant); - AnswerMsg { source_slot: vq.source_slot, hit, smaller: vq.smaller } - } - } : Query => Answer, }; - // ── Reconciliation (the pipeline's sink): commutative fold of every - // answer into its origin k-mer's running state, in whatever order - // the pipeline delivers them. ───────────────────────────────────── - let mut state = vec![RunningState::default(); n_slots]; - for ans in pipe.apply(sources.into_iter(), n_workers, capacity) { - if ans.hit { - let st = &mut state[ans.source_slot]; - st.siblings = (st.siblings + 1).min(3); - if ans.smaller { - st.minorant = false; + // ── Group generated variants by destination partition. `cache` + // holds every partition already mmap'd (no more `open()` cost), but + // `mmap` pages are still loaded on demand and can be evicted — a + // lookup is not free just because the file isn't reopened. Grouping + // keeps one partition's pages hot while its whole batch is resolved, + // instead of faulting pages in and out as lookups jump between + // partitions in whatever order the `Flat` stage happens to produce + // them. The throttle permit drops here, once accumulated. ───────── + let mut outgoing: Vec> = (0..n_parts).map(|_| Vec::new()).collect(); + for vq in pipe.apply(throttled, n_workers, capacity) { + outgoing[vq.dest_partition].push((vq.variant, vq.source_slot, vq.base)); + } + + // ── Resolve each partition's batch against the cache in one + // contiguous pass; parallelised across partitions (independent, + // read-only) so this keeps using multiple cores without giving up + // the per-partition locality above. ───────────────────────────── + outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { + for &(variant, source_slot, base) in queries { + if cache.find(dest, variant) { + mask[source_slot].fetch_or(1 << base, Ordering::Relaxed); } } - } + }); // ── Write the layer's annex file ───────────────────────────────────── let annex_path = layer_dir.join(ANNEX_FILE_NAME); let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?; - for (slot, st) in state.iter().enumerate() { + for (slot, m) in mask.iter().enumerate() { if slot_kmer[slot].is_none() { continue; // unused MPHF slot, if any — leave at the sentinel } - builder.set(slot, SiblingInfo { minorant: st.minorant, siblings: st.siblings }); + builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed))); } builder.close()?; - Ok(()) + Ok(n_slots as u64) } } -/// Distribution of sibling counts (0-3), read back from an already-built +/// Distribution of family sizes (1-4), read back from an already-built /// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's /// presence/count data — a separate, occasional diagnostic pass, not fused /// into construction. +/// +/// Every count here is **per family, not per slot**: a family with `F` +/// members occupies `F` annex slots (one per observed member), all sharing +/// the same mask. Counting every slot would count each family up to 4 +/// times over; only the minorant's slot is tallied (minorant is derived on +/// the fly — see `is_minorant` — not stored, but cheap: no lookup, pure +/// bit arithmetic on already-in-hand data). #[derive(Debug, Clone, Default)] pub struct SiblingAnnexStats { - /// `counts[s]` = number of k-mers (slots) with exactly `s` siblings, - /// counted once each regardless of how many genomes carry them. + /// `counts[s]` = number of *families* of size `s + 1` (index 0 = size 1, + /// i.e. 0 siblings, ... index 3 = size 4, i.e. 3 siblings). pub counts: [u64; 4], - /// Of those, how many are minorant. - pub minorant_counts: [u64; 4], - /// `per_genome[g][s]` = number of k-mers with exactly `s` siblings that - /// genome `g` (index into `KmerIndex::meta().genomes`) carries. + /// `per_genome[g][s]` = number of families of size `s + 1` for which + /// genome `g` (index into `KmerIndex::meta().genomes`) carries at least + /// one member. pub per_genome: Vec<[u64; 4]>, } impl KmerIndex { - /// Tally the sibling-count distribution of an already-built annex - /// (globally, and per genome). Errors if [`build_sibling_annex`] has not - /// been run on this index first. + /// Tally the family-size distribution of an already-built annex + /// (globally, and per genome), counting each family once (at its + /// minorant slot). Errors if [`build_sibling_annex`] has not been run on + /// this index first. /// /// [`build_sibling_annex`]: Self::build_sibling_annex pub fn sibling_annex_stats(&self) -> OKIResult { let n_parts = self.n_partitions(); let n_genomes = self.meta.genomes.len(); let with_counts = self.meta.config.with_counts; + let k = self.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; - let mut stats = SiblingAnnexStats { - per_genome: vec![[0u64; 4]; n_genomes], - ..Default::default() - }; + // Same whole-run cache as `build_sibling_annex` — see its docs for + // why re-opening per lookup (or per call to a batching helper) is + // not good enough on a real index. + let partition = KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + // Gather the (partition, layer) pairs to process — cheap metadata + // reads only, checking every annex file exists up front so a + // missing one is reported before any real work starts. + let mut layer_dirs = Vec::new(); for part in 0..n_parts { let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); if !index_dir.exists() { continue; } let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; - for l in 0..meta.n_layers { let layer_dir = index_dir.join(format!("layer_{l}")); let annex_path = layer_dir.join(ANNEX_FILE_NAME); @@ -334,45 +500,111 @@ impl KmerIndex { annex_path.display() ))); } - let annex = SiblingAnnex::open(&annex_path)?; - let use_counts = with_counts && layer_dir.join("counts").exists(); + layer_dirs.push(layer_dir); + } + } - // Opened once per layer, outside the slot loop. - enum Mat { - Count(PersistentCompactIntMatrix), - Presence(PersistentBitMatrix), - } - let mat = if use_counts { - Mat::Count(PersistentCompactIntMatrix::open(&layer_dir)?) - } else { - Mat::Presence(PersistentBitMatrix::open(&layer_dir)?) + // One layer's worth of work, parallelised across layers with Rayon + // — independent, read-only, each producing its own partial tally + // merged at the end. + let pb = progress_bar("sibling_annex_stats", layer_dirs.len() as u64, "layers"); + let partials: Vec = layer_dirs + .par_iter() + .map(|layer_dir| -> OKIResult { + let mut stats = SiblingAnnexStats { + per_genome: vec![[0u64; 4]; n_genomes], + ..Default::default() }; - let n_cols = match &mat { - Mat::Count(m) => m.n_cols(), - Mat::Presence(m) => m.n_cols(), + + let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); + let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + + // Need each slot's own k-mer to derive minorant — same + // enumeration as construction. + let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + let mut slot_kmer: Vec> = vec![None; annex.len()]; + let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) + .map_err(OKIError::Partition)?; + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if let Some(slot) = mphf.find(kmer) { + slot_kmer[slot] = Some(kmer); + } } - .min(n_genomes); + + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(layer_dir)?) + }; + let n_cols = mat.n_cols().min(n_genomes); for slot in 0..annex.len() { - let Some(info) = annex.get(slot) else { continue }; - let s = info.siblings as usize; - stats.counts[s] += 1; - if info.minorant { - stats.minorant_counts[s] += 1; + let Some(mask) = annex.get(slot) else { continue }; + let Some(kmer) = slot_kmer[slot] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // this family is tallied at its minorant's slot only } + let s = mask.siblings() as usize; + stats.counts[s] += 1; + + // "Genome g represents this family" means g carries + // *any* of its members, not just the minorant's own — + // start from the minorant's own presence (already + // open, no lookup) and OR in every other present + // member's presence vector, resolved against the + // whole-run cache (no I/O) — exactly `mask.siblings()` + // of them, the mask tells us precisely which to fetch. + let mut carries = vec![false; n_cols]; for g in 0..n_cols { - let carried = match &mat { - Mat::Count(m) => m.col_view(g).get(slot) != 0, - Mat::Presence(m) => m.get(g, slot) != 0, - }; + carries[g] = mat.carries(g, slot); + } + for other in kmer.central_canonical_neighbors() { + if other == kmer { + continue; + } + let base = central_base(other, k); + if !mask.has(base) { + continue; + } + let dest = partition_of(other, n_parts); + if let Some(other_presence) = cache.find_presence(dest, other, n_genomes) { + for (g, &present) in other_presence.iter().enumerate() { + if present { + carries[g] = true; + } + } + } + } + + for (g, &carried) in carries.iter().enumerate() { if carried { stats.per_genome[g][s] += 1; } } } + + pb.inc(1); + Ok(stats) + }) + .collect::>>()?; + pb.finish_and_clear(); + + let mut stats = SiblingAnnexStats { + per_genome: vec![[0u64; 4]; n_genomes], + ..Default::default() + }; + for part in partials { + for s in 0..4 { + stats.counts[s] += part.counts[s]; + } + for g in 0..n_genomes { + for s in 0..4 { + stats.per_genome[g][s] += part.per_genome[g][s]; + } } } - Ok(stats) } } @@ -442,7 +674,7 @@ mod tests { /// Read back the annex entry for a given canonical k-mer from the merged /// index's (single) partition/layer, asserting it was found at all. - fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> SiblingInfo { + fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> FamilyMask { let index_dir = idx.partition().part_dir(0).join(INDEX_SUBDIR); let meta = PartitionMeta::load(&index_dir).unwrap(); for l in 0..meta.n_layers { @@ -474,24 +706,33 @@ mod tests { fn sibling_annex_one_sibling_each() { // k=11, centre = index 5 (0-based). Two genomes, each exactly one // k-mer, sharing every base except the centre: - // g1 = "AACCGCTTAAG" (centre 'C') - // g2 = "AACCGGTTAAG" (centre 'G') + // g1 = "AACCGCTTAAG" (centre 'C', base index 1) + // g2 = "AACCGGTTAAG" (centre 'G', base index 2) // Hand-verified: both stay forward-oriented under canonicalisation // (each is lexicographically smaller than its own reverse // complement, since both start with "AA"), and raw(g1) < raw(g2) // (only differing base: C=0b01 < G=0b10 at the centre) — so g1 is - // the minorant, g2 is not, and each is the other's one sibling. + // the minorant, g2 is not. The mask is a family-wide value: both + // slots must read back the *same* mask (bits 1 and 2 set). let dir = tempdir().unwrap(); let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG"); let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG"); let merged = merge_two(dir.path(), &g1, &g2); merged.build_sibling_annex().expect("build_sibling_annex"); - let a = annex_info_for(&merged, canonical(b"AACCGCTTAAG")); - assert_eq!(a, SiblingInfo { minorant: true, siblings: 1 }, "AACCGCTTAAG"); + let g1_kmer = canonical(b"AACCGCTTAAG"); + let g2_kmer = canonical(b"AACCGGTTAAG"); + let expected_mask = FamilyMask::EMPTY.with(1).with(2); - let b = annex_info_for(&merged, canonical(b"AACCGGTTAAG")); - assert_eq!(b, SiblingInfo { minorant: false, siblings: 1 }, "AACCGGTTAAG"); + let a = annex_info_for(&merged, g1_kmer); + assert_eq!(a, expected_mask, "AACCGCTTAAG"); + assert_eq!(a.siblings(), 1); + assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant"); + + let b = annex_info_for(&merged, g2_kmer); + assert_eq!(b, expected_mask, "AACCGGTTAAG"); + assert_eq!(b.siblings(), 1); + assert!(!is_minorant(g2_kmer, b, K), "g2 should not be the minorant"); } #[test] @@ -504,7 +745,36 @@ mod tests { let merged = merge_two(dir.path(), &g1, &g2); merged.build_sibling_annex().expect("build_sibling_annex"); - let info = annex_info_for(&merged, canonical(b"GATTACAGATC")); - assert_eq!(info, SiblingInfo { minorant: true, siblings: 0 }, "GATTACAGATC"); + let kmer = canonical(b"GATTACAGATC"); + let mask = annex_info_for(&merged, kmer); + assert_eq!(mask.siblings(), 0, "GATTACAGATC"); + assert_eq!(mask.family_size(), 1); + assert!(is_minorant(kmer, mask, K)); + } + + #[test] + fn sibling_annex_stats_counts_each_family_once_and_per_genome() { + // Reuses the one-sibling-each fixture: a single family of size 2 + // (g1's centre-C form + g2's centre-G form), each genome carrying + // exactly one of the two members. Stats must report exactly one + // family of size 2 (`counts[1] == 1`, since index 1 = size 2), not + // two (which naively summing both slots would give), and both + // genomes represented at size 2, neither at any other size. + let dir = tempdir().unwrap(); + let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG"); + let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG"); + let merged = merge_two(dir.path(), &g1, &g2); + merged.build_sibling_annex().expect("build_sibling_annex"); + + let stats = merged.sibling_annex_stats().expect("sibling_annex_stats"); + + assert_eq!(stats.counts, [0, 1, 0, 0], "one family of size 2, counted once"); + assert_eq!(stats.per_genome.len(), 2); + for g in 0..2 { + assert_eq!( + stats.per_genome[g], [0, 1, 0, 0], + "genome {g} should represent exactly one size-2 family" + ); + } } } diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index 3707387..7a5ad4e 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -78,8 +78,7 @@ pub struct DistanceArgs { pub sibling_stats: bool, /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _siblings_per_genome.csv, - /// _nj.nwk, _upgma.nwk. + /// _siblings.csv, _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -114,6 +113,17 @@ pub fn run(args: DistanceArgs) { write_sibling_stats_csv(&stats, &labels, &args.output); } + // `--sibling-annex`/`--sibling-stats` are their own operation, not a + // modifier on top of a distance-metric computation — a metric was + // never requested by asking for either of them, so there is nothing + // for the rest of this function to compute. Not a historical accident + // to keep: stop here rather than always also running a Jaccard (or + // whichever `--metric` defaults to) pass and printing an unrequested + // matrix. + if args.sibling_annex || args.sibling_stats { + return; + } + info!( "computing {:?} distances for {} genome(s)", args.metric, n @@ -225,43 +235,38 @@ pub fn run(args: DistanceArgs) { } } -// ── Sibling-count distribution → CSV ──────────────────────────────────────── +// ── Family-size distribution → CSV ────────────────────────────────────────── +// +// Each row is a family (the up-to-4 k-mers sharing flanks, differing only at +// the centre), counted once — at its minorant — regardless of how many of +// its members are observed. Family size 1..4 (not "sibling count" 0..3): +// see `docmd/theory/evolutionary_distances.md`, "Definitions". fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output: &Option) { - // Global histogram: one row per sibling count (0-3). - let global_path = output.as_ref() + // One row per genome (4 columns, family size 1-4: number of families of + // that size for which the genome carries at least one member), plus a + // `global` row — the actual deduplicated family-size histogram + // (`stats.counts`), NOT a sum of the per-genome columns (a family shared + // by several genomes would otherwise be counted once per genome it + // appears in, inflating the total beyond the real family count). + let path = output.as_ref() .map(|p| format!("{}_siblings.csv", p.display())) .unwrap_or_else(|| "siblings.csv".into()); - let mut f = BufWriter::new(std::fs::File::create(&global_path).unwrap_or_else(|e| { - eprintln!("error creating {global_path}: {e}"); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); std::process::exit(1); })); - writeln!(f, "siblings,slots,minorant_slots").unwrap(); - for s in 0..4 { - writeln!(f, "{s},{},{}", stats.counts[s], stats.minorant_counts[s]).unwrap(); - } - let total: u64 = stats.counts.iter().sum(); - info!("sibling-count distribution → {global_path} (total {total} slot(s))"); - - // Per-genome breakdown: one row per genome, 4 columns (0-3), + a total row. - let per_genome_path = output.as_ref() - .map(|p| format!("{}_siblings_per_genome.csv", p.display())) - .unwrap_or_else(|| "siblings_per_genome.csv".into()); - let mut f = BufWriter::new(std::fs::File::create(&per_genome_path).unwrap_or_else(|e| { - eprintln!("error creating {per_genome_path}: {e}"); - std::process::exit(1); - })); - writeln!(f, "genome,0,1,2,3").unwrap(); - let mut column_totals = [0u64; 4]; + writeln!(f, "genome,1,2,3,4").unwrap(); for (label, counts) in labels.iter().zip(stats.per_genome.iter()) { writeln!(f, "{label},{},{},{},{}", counts[0], counts[1], counts[2], counts[3]).unwrap(); - for s in 0..4 { column_totals[s] += counts[s]; } } writeln!( - f, "total,{},{},{},{}", - column_totals[0], column_totals[1], column_totals[2], column_totals[3], + f, "global,{},{},{},{}", + stats.counts[0], stats.counts[1], stats.counts[2], stats.counts[3], ).unwrap(); - info!("per-genome sibling-count distribution → {per_genome_path}"); + let total: u64 = stats.counts.iter().sum(); + info!("family-size distribution → {path} (total {total} famil{})", + if total == 1 { "y" } else { "ies" }); } // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── -- 2.54.0 From 49f329edd5db9bcac27e4626e1fb9ddd6df69dd8 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 17:31:48 +0200 Subject: [PATCH 7/8] feat: add raw SNP distance calculation and CLI flag Exposes RawSnpDistanceOutput and implements KmerIndex::raw_snp_distance() to compute pairwise single-copy locus counts under a paralogy-aware rule. The implementation leverages ndarray for parallel matrix aggregation, producing raw p-distance matrices for sanity-checking. A --raw-snp-distance CLI flag is added to export results as CSV, mapping zero-eligible pairs to NA. --- src/obikindex/src/lib.rs | 2 +- src/obikindex/src/siblings.rs | 172 ++++++++++++++++++++++++++++++++ src/obikmer/src/cmd/distance.rs | 71 +++++++++++-- 3 files changed, 235 insertions(+), 10 deletions(-) diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 4d22205..4b09bc3 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -19,4 +19,4 @@ pub use merge::MergeMode; pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME}; pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED}; pub use stats::IndexBitsPerKmer; -pub use siblings::SiblingAnnexStats; +pub use siblings::{RawSnpDistanceOutput, SiblingAnnexStats}; diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index e80bd67..6ab50a8 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -31,6 +31,7 @@ use std::path::Path; use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; +use ndarray::Array2; use rayon::prelude::*; use obicompactvec::{ @@ -609,6 +610,177 @@ impl KmerIndex { } } +/// Raw p-distance restricted to loci that are single-copy in **both** +/// genomes of a pair — the "stringent / paralogy-aware" locus eligibility +/// rule (`docmd/theory/evolutionary_distances.md`, "Locus eligibility"), +/// without the JC/K2P/LogDet finalisation math: `snp[i,j]` + `shared[i,j]` +/// is the number of eligible loci, `snp[i,j] / (snp[i,j] + shared[i,j])` is +/// `p_hat`. A quick, self-contained way to sanity-check the estimator +/// against a real index before the full `SnpTally` design is built. +/// +/// A locus (family, tallied once at its minorant) is eligible for pair +/// `(i, j)` iff genome `i` carries exactly one of the family's observed +/// forms **and** genome `j` carries exactly one (possibly a different one) +/// — presence-only: a genome carrying the same form twice (a same-allele +/// duplicate) is indistinguishable from carrying it once when only a +/// presence matrix is available, so such cases are not excluded here even +/// when a count index exists. See "Locus eligibility", stringent rule, for +/// why this matters and how a count index would close the gap — left as a +/// follow-up, not applied here. +pub struct RawSnpDistanceOutput { + /// n×n count of eligible loci where the two genomes' single forms differ. + pub snp: Array2, + /// n×n count of eligible loci where the two genomes' single forms agree. + pub shared: Array2, +} + +impl KmerIndex { + /// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex + /// (run [`build_sibling_annex`](Self::build_sibling_annex) first). + pub fn raw_snp_distance(&self) -> OKIResult { + let n_parts = self.n_partitions(); + let n_genomes = self.meta.genomes.len(); + let with_counts = self.meta.config.with_counts; + let k = self.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; + + let partition = KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + + let mut layer_dirs = Vec::new(); + for part in 0..n_parts { + let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + if !annex_path.exists() { + return Err(OKIError::InvalidInput(format!( + "no sibling annex at {} — run build_sibling_annex first", + annex_path.display() + ))); + } + layer_dirs.push(layer_dir); + } + } + + let pb = progress_bar("raw_snp_distance", layer_dirs.len() as u64, "layers"); + let partials: Vec<(Array2, Array2)> = layer_dirs + .par_iter() + .map(|layer_dir| -> OKIResult<(Array2, Array2)> { + let mut snp = Array2::::zeros((n_genomes, n_genomes)); + let mut shared = Array2::::zeros((n_genomes, n_genomes)); + + let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); + let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + + let mut slot_kmer: Vec> = vec![None; annex.len()]; + let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) + .map_err(OKIError::Partition)?; + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if let Some(slot) = mphf.find(kmer) { + slot_kmer[slot] = Some(kmer); + } + } + + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(layer_dir)?) + }; + let n_cols = mat.n_cols().min(n_genomes); + + // Per family: which single form (if exactly one) each genome + // carries — `None` once a second form is seen (ambiguous, + // not single-copy, ineligible for either side of a pair). + let mut single_form: Vec> = Vec::with_capacity(n_cols); + let mut ambiguous: Vec = Vec::with_capacity(n_cols); + + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + let Some(kmer) = slot_kmer[slot] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // family tallied once, at its minorant + } + + single_form.clear(); + single_form.resize(n_cols, None); + ambiguous.clear(); + ambiguous.resize(n_cols, false); + + for other in kmer.central_canonical_neighbors() { + let base = central_base(other, k); + if !mask.has(base) { + continue; + } + let presence: Option> = if other == kmer { + Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()) + } else { + let dest = partition_of(other, n_parts); + cache.find_presence(dest, other, n_genomes) + }; + let Some(presence) = presence else { continue }; + for (g, &present) in presence.iter().enumerate() { + if !present { + continue; + } + if single_form[g].is_some() { + ambiguous[g] = true; + } else { + single_form[g] = Some(base); + } + } + } + + for i in 0..n_cols { + if ambiguous[i] { + continue; + } + let Some(bi) = single_form[i] else { continue }; + for j in (i + 1)..n_cols { + if ambiguous[j] { + continue; + } + let Some(bj) = single_form[j] else { continue }; + if bi == bj { + shared[[i, j]] += 1; + shared[[j, i]] += 1; + } else { + snp[[i, j]] += 1; + snp[[j, i]] += 1; + } + } + } + } + + pb.inc(1); + Ok((snp, shared)) + }) + .collect::>>()?; + pb.finish_and_clear(); + + let mut snp = Array2::::zeros((n_genomes, n_genomes)); + let mut shared = Array2::::zeros((n_genomes, n_genomes)); + for (s, sh) in partials { + snp += &s; + shared += &sh; + } + Ok(RawSnpDistanceOutput { snp, shared }) + } +} + #[cfg(test)] mod tests { use std::io::Write; diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index 7a5ad4e..9e02552 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -3,7 +3,7 @@ use std::path::PathBuf; use clap::Args; use kodama::{Method, linkage}; -use obikindex::{DistanceMetric, KmerIndex, SiblingAnnexStats}; +use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats}; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; @@ -77,8 +77,17 @@ pub struct DistanceArgs { #[arg(long)] pub sibling_stats: bool, + /// Compute the raw p-distance restricted to loci that are single-copy + /// in both genomes of each pair (an already-built sibling annex is + /// required — run with `--sibling-annex` first, in this invocation or + /// an earlier one). A quick way to test the central-position SNP + /// estimator against a real index; not the full `SnpTally` design. + #[arg(long)] + pub raw_snp_distance: bool, + /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _nj.nwk, _upgma.nwk. + /// _siblings.csv, _rawsnp.csv, _nj.nwk, + /// _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -112,15 +121,22 @@ pub fn run(args: DistanceArgs) { }); write_sibling_stats_csv(&stats, &labels, &args.output); } + if args.raw_snp_distance { + let result = idx.raw_snp_distance().unwrap_or_else(|e| { + eprintln!("error computing raw SNP distance: {e}"); + std::process::exit(1); + }); + write_raw_snp_distance_csv(&result, &labels, &args.output); + } - // `--sibling-annex`/`--sibling-stats` are their own operation, not a - // modifier on top of a distance-metric computation — a metric was - // never requested by asking for either of them, so there is nothing - // for the rest of this function to compute. Not a historical accident - // to keep: stop here rather than always also running a Jaccard (or - // whichever `--metric` defaults to) pass and printing an unrequested + // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance` are their own + // operation, not a modifier on top of a distance-metric computation — a + // metric was never requested by asking for any of them, so there is + // nothing for the rest of this function to compute. Not a historical + // accident to keep: stop here rather than always also running a Jaccard + // (or whichever `--metric` defaults to) pass and printing an unrequested // matrix. - if args.sibling_annex || args.sibling_stats { + if args.sibling_annex || args.sibling_stats || args.raw_snp_distance { return; } @@ -269,6 +285,43 @@ fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output: if total == 1 { "y" } else { "ies" }); } +// ── Raw single-copy SNP distance → CSV ────────────────────────────────────── +// +// p_hat[i,j] = snp[i,j] / (snp[i,j] + shared[i,j]) over loci single-copy in +// both i and j — see `RawSnpDistanceOutput` / `KmerIndex::raw_snp_distance`. +// A single file: the distance matrix, with an eligible-loci count alongside +// each value so a 0/0 pair (no eligible locus at all) is distinguishable +// from a genuinely identical pair. + +fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_rawsnp.csv", p.display())) + .unwrap_or_else(|| "rawsnp.csv".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + let n = labels.len(); + write!(f, "genome").unwrap(); + for g in labels { write!(f, ",{g}").unwrap(); } + writeln!(f).unwrap(); + for (i, g) in labels.iter().enumerate() { + write!(f, "{g}").unwrap(); + for j in 0..n { + let snp = result.snp[[i, j]]; + let shared = result.shared[[i, j]]; + let eligible = snp + shared; + if eligible == 0 { + write!(f, ",NA").unwrap(); + } else { + write!(f, ",{:.6}", snp as f64 / eligible as f64).unwrap(); + } + } + writeln!(f).unwrap(); + } + info!("raw single-copy SNP distance matrix → {path}"); +} + // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── fn upgma_to_newick(dendro: &kodama::Dendrogram, names: &[String]) -> String { -- 2.54.0 From f5e508ed339b753f7d1ce253b87f13dc223ddd65 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 22:12:12 +0200 Subject: [PATCH 8/8] feat: add multi-genome SNP pseudo-alignment and CLI export Introduces a `SnpAlignment` struct and helper methods to construct per-genome SNP pseudo-alignments from sibling k-mer data, filtering monomorphic families and encoding bases as IUPAC ambiguity codes. Exposes the type at the crate root for simplified imports. Adds a `--snp` CLI flag to compute and export these alignments as an IUPAC-coded FASTA file. Updates theory documentation to propose a multi-genome framing approach for joint phylogenetic inference, resolving pairwise correspondence ambiguities through positional homology and partial coverage thresholds. Bumps crate version to 1.1.40. --- docmd/theory/evolutionary_distances.md | 57 +++++++++ src/Cargo.lock | 2 +- src/obikindex/src/lib.rs | 2 +- src/obikindex/src/siblings.rs | 169 +++++++++++++++++++++++++ src/obikmer/Cargo.toml | 2 +- src/obikmer/src/cmd/distance.rs | 64 ++++++++-- 6 files changed, 282 insertions(+), 14 deletions(-) diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index 2c21a3c..ba63278 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -146,6 +146,63 @@ A once multiplicity > 1 on either side. Any pairing rule invents a correspondence the data cannot support. Multiplicity > 1 is treated as non-identifiable, not as a puzzle to solve with a heuristic. +## Multi-genome framing: family as pseudo-alignment column + +**Idea.** Instead of resolving locus eligibility and correspondence one +genome pair at a time, treat a family as a column of a pseudo multiple +alignment across *all* genomes simultaneously: for each family, each genome +has either a net single-copy state (`A`/`C`/`G`/`T`, when the genome carries +exactly one of the 4 forms) or "missing" (`?`, multi-copy or absent). Flank +conservation (the `2m` bases fixed by construction) supplies positional +homology for free — the same role a real MSA would play, without alignment +software, gap penalties, or progressive-alignment approximations. Stacking +one such column per family, genomes as rows, produces a genuine SNP +pseudo-alignment matrix, not just a bag of pairwise distances. + +**Precedent.** This is the same principle behind reference-free +k-mer-based phylogenomics tools — SKA (Split K-mer Analysis, Harris 2018) and +kSNP: split the k-mer around a variable center, use flank identity to call +homologous columns across arbitrarily many genomes with no reference and no +MSA step, then feed the resulting pseudo-alignment to standard phylogenetic +tools. Landing on the same design independently is a good sign, not a +coincidence. + +**Resolves the pairwise-correspondence problem, properly.** The "Rejected: +parsimony-based multiset pairing" case above failed because, with only two +genomes' cardinalities to look at, there is no external constraint to justify +picking one correspondence between leftover alleles over another — `min(a,b)` +is a lower bound dressed up as a point estimate (see the follow-up discussion +on Felsenstein-style parsimony inconsistency: minimum-event explanations are +systematically biased low whenever homoplasy/multiplicity is real, not +noise-cancelling). With `N` genomes and many families jointly, the same +question can be answered the way real phylogenetics answers it: ancestral +state reconstruction / ML mapping over a tree estimated from the whole +column set. The tree supplies the missing constraint that two isolated +columns cannot — this is the principled way out, not a heuristic replacement +for one. + +**Relation to what's already implemented.** `KmerIndex::raw_snp_distance` +already computes, internally, per family, exactly this row — `single_form: +Vec>`, one entry per genome, `None` where ambiguous/absent — +before immediately collapsing it into pairwise `snp[i,j]`/`shared[i,j]` +tallies. The pivot this section proposes is small at the implementation +level: stop collapsing early, and surface the per-family row as a first-class +artifact (a `families x genomes` matrix). Pairwise raw p-distance becomes one +projection of that matrix (what's computed today), not the primary object; +downstream, the matrix itself could feed real phylogenetic tools (parsimony/ +ML, e.g. RAxML/IQ-TREE-style) instead of only NJ/UPGMA on a homemade +pairwise-distance matrix. + +**Caveat: column completeness shrinks with `N`.** The probability that a +family's flanks stay intact simultaneously across all `N` genomes decays with +`N` (same ascertainment-bias mechanism as Bias 1 above, compounded over more +genomes) — fully-resolved columns (no `?` anywhere) become rare as more +genomes are added. Same missing-data situation any real multi-species +alignment faces, and phylogenetic tools already handle it well; the practical +implication is that columns should be allowed partial coverage (>=2 resolved +genomes, not unanimous) rather than requiring every genome to be net +single-copy at that locus. + ## Heterozygosity, ploidy, and consensus-assembly inputs A within-genome multiplicity signal (more than one of the 4 central forms diff --git a/src/Cargo.lock b/src/Cargo.lock index aa92eee..25955a3 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1715,7 +1715,7 @@ dependencies = [ [[package]] name = "obikmer" -version = "1.1.39" +version = "1.1.40" dependencies = [ "clap", "csv", diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 4b09bc3..48daf54 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -19,4 +19,4 @@ pub use merge::MergeMode; pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME}; pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED}; pub use stats::IndexBitsPerKmer; -pub use siblings::{RawSnpDistanceOutput, SiblingAnnexStats}; +pub use siblings::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index 6ab50a8..da51270 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -781,6 +781,175 @@ impl KmerIndex { } } +/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set +/// iff the genome carries the member whose own central base is `b`): +/// single bit -> the plain base; 2 or 3 bits -> the matching IUPAC +/// ambiguity code (preserves partial information instead of collapsing to +/// `N`, the same convention used for diploid heterozygous VCF/FASTA sites); +/// all 4 bits -> `N`; no bits (genome carries none of the family's observed +/// members) -> `-` (no data at this locus for this genome). +fn iupac_code(mask: u8) -> u8 { + match mask & 0b1111 { + 0b0000 => b'-', + 0b0001 => b'A', + 0b0010 => b'C', + 0b0100 => b'G', + 0b1000 => b'T', + 0b0101 => b'R', // A/G + 0b1010 => b'Y', // C/T + 0b0110 => b'S', // C/G + 0b1001 => b'W', // A/T + 0b1100 => b'K', // G/T + 0b0011 => b'M', // A/C + 0b1110 => b'B', // C/G/T + 0b1101 => b'D', // A/G/T + 0b1011 => b'H', // A/C/T + 0b0111 => b'V', // A/C/G + 0b1111 => b'N', + _ => unreachable!("masked to 4 bits"), + } +} + +/// A SNP-only pseudo-alignment: one row (byte sequence, IUPAC-coded) per +/// genome, one column per variable family (`family_size() >= 2` — monomorphic +/// families carry no signal and are skipped, unlike `raw_snp_distance`'s +/// tally which does count them as `shared`). Column order is the same, +/// deterministic sweep order as the annex build (partition, then layer, then +/// slot) — arbitrary but stable and identical across genomes, which is all a +/// pseudo-alignment needs (there is no natural genomic coordinate to sort by +/// once flanks are dropped). See `docmd/theory/evolutionary_distances.md`, +/// "Multi-genome framing: family as pseudo-alignment column". +pub struct SnpAlignment { + /// `sequences[g]` = genome `g`'s IUPAC-coded row, same length for every + /// genome (`sequences.len()` columns). + pub sequences: Vec>, +} + +impl KmerIndex { + /// Build the SNP-only pseudo-alignment from an already-built sibling + /// annex (run [`build_sibling_annex`](Self::build_sibling_annex) first). + pub fn snp_pseudo_alignment(&self) -> OKIResult { + let n_parts = self.n_partitions(); + let n_genomes = self.meta.genomes.len(); + let with_counts = self.meta.config.with_counts; + let k = self.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; + + let partition = KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + + let mut layer_dirs = Vec::new(); + for part in 0..n_parts { + let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + if !annex_path.exists() { + return Err(OKIError::InvalidInput(format!( + "no sibling annex at {} — run build_sibling_annex first", + annex_path.display() + ))); + } + layer_dirs.push(layer_dir); + } + } + + let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers"); + // `Vec>` per layer, one entry (column) per variable family; + // `par_iter().map(...).collect()` on this indexed source preserves + // input order, so concatenating the results below in order gives a + // single deterministic column order across the whole index. + let partials: Vec>> = layer_dirs + .par_iter() + .map(|layer_dir| -> OKIResult>> { + let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); + let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + + let mut slot_kmer: Vec> = vec![None; annex.len()]; + let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) + .map_err(OKIError::Partition)?; + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if let Some(slot) = mphf.find(kmer) { + slot_kmer[slot] = Some(kmer); + } + } + + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(layer_dir)?) + }; + let n_cols = mat.n_cols().min(n_genomes); + + let mut columns: Vec> = Vec::new(); + let mut genome_mask: Vec = Vec::with_capacity(n_genomes); + + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + let Some(kmer) = slot_kmer[slot] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // family tallied once, at its minorant + } + if mask.family_size() < 2 { + continue; // monomorphic family — no signal, skip + } + + genome_mask.clear(); + genome_mask.resize(n_genomes, 0); + + for other in kmer.central_canonical_neighbors() { + let base = central_base(other, k); + if !mask.has(base) { + continue; + } + let presence: Option> = if other == kmer { + Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()) + } else { + let dest = partition_of(other, n_parts); + cache.find_presence(dest, other, n_genomes) + }; + let Some(presence) = presence else { continue }; + for (g, &present) in presence.iter().enumerate() { + if present { + genome_mask[g] |= 1 << base; + } + } + } + + columns.push(genome_mask.iter().map(|&m| iupac_code(m)).collect()); + } + + pb.inc(1); + Ok(columns) + }) + .collect::>>()?; + pb.finish_and_clear(); + + let mut sequences: Vec> = vec![Vec::new(); n_genomes]; + for layer_columns in partials { + for column in layer_columns { + for (g, &code) in column.iter().enumerate() { + sequences[g].push(code); + } + } + } + Ok(SnpAlignment { sequences }) + } +} + #[cfg(test)] mod tests { use std::io::Write; diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index 9d00ffd..8c0e42b 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "obikmer" -version = "1.1.39" +version = "1.1.40" edition = "2024" [[bin]] diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index 9e02552..8852478 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -3,7 +3,8 @@ use std::path::PathBuf; use clap::Args; use kodama::{Method, linkage}; -use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats}; +use obifastwrite::{JsonVal, write_record}; +use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; @@ -85,9 +86,17 @@ pub struct DistanceArgs { #[arg(long)] pub raw_snp_distance: bool, + /// Write a SNP-only pseudo-alignment (FASTA, IUPAC-coded) from an + /// already-built sibling annex — one row per genome, one column per + /// variable family (monomorphic families skipped), no flanking + /// sequence. See `docmd/theory/evolutionary_distances.md`, + /// "Multi-genome framing: family as pseudo-alignment column". + #[arg(long)] + pub snp: bool, + /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _rawsnp.csv, _nj.nwk, - /// _upgma.nwk. + /// _siblings.csv, _rawsnp.csv, _snp.fasta, + /// _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -128,15 +137,22 @@ pub fn run(args: DistanceArgs) { }); write_raw_snp_distance_csv(&result, &labels, &args.output); } + if args.snp { + let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { + eprintln!("error computing SNP pseudo-alignment: {e}"); + std::process::exit(1); + }); + write_snp_fasta(&alignment, &labels, &args.output); + } - // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance` are their own - // operation, not a modifier on top of a distance-metric computation — a - // metric was never requested by asking for any of them, so there is - // nothing for the rest of this function to compute. Not a historical - // accident to keep: stop here rather than always also running a Jaccard - // (or whichever `--metric` defaults to) pass and printing an unrequested - // matrix. - if args.sibling_annex || args.sibling_stats || args.raw_snp_distance { + // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp` are + // their own operation, not a modifier on top of a distance-metric + // computation — a metric was never requested by asking for any of them, + // so there is nothing for the rest of this function to compute. Not a + // historical accident to keep: stop here rather than always also + // running a Jaccard (or whichever `--metric` defaults to) pass and + // printing an unrequested matrix. + if args.sibling_annex || args.sibling_stats || args.raw_snp_distance || args.snp { return; } @@ -322,6 +338,32 @@ fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], info!("raw single-copy SNP distance matrix → {path}"); } +// ── SNP-only pseudo-alignment → FASTA ─────────────────────────────────────── +// +// One record per genome, IUPAC-coded, no flanking sequence — see +// `SnpAlignment` / `KmerIndex::snp_pseudo_alignment`. Uses the project's +// existing FASTA writer (`obifastwrite::write_record`) rather than +// hand-rolling one. + +fn write_snp_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_snp.fasta", p.display())) + .unwrap_or_else(|| "snp.fasta".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); + for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { + write_record(seq, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| { + eprintln!("error writing {path}: {e}"); + std::process::exit(1); + }); + } + info!("SNP pseudo-alignment → {path} ({n_sites} site{})", + if n_sites == 1 { "" } else { "s" }); +} + // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── fn upgma_to_newick(dendro: &kodama::Dendrogram, names: &[String]) -> String { -- 2.54.0