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/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/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/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md new file mode 100644 index 0000000..ba63278 --- /dev/null +++ b/docmd/theory/evolutionary_distances.md @@ -0,0 +1,820 @@ +# 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 + +**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` + +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 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)`. `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. + +### 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 +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. + +## 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 +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 +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 (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 +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**. + 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; + 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. 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 + +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), +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`) + +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 + 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 + 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 + 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 +``` + +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). + +### 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 — 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 + (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 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. + - *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` 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 + 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 diff --git a/src/Cargo.lock b/src/Cargo.lock index 0dc087d..25955a3 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1701,17 +1701,21 @@ dependencies = [ "obikpartitionner", "obikseq", "obilayeredmap", + "obipipeline", + "obiread", + "obiskbuilder", "obiskio", "obisys", "rayon", "serde", "serde_json", + "tempfile", "tracing", ] [[package]] name = "obikmer" -version = "1.1.39" +version = "1.1.40" dependencies = [ "clap", "csv", diff --git a/src/obicompactvec/src/lib.rs b/src/obicompactvec/src/lib.rs index 9041ab7..1f85a49 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::{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 new file mode 100644 index 0000000..0de9733 --- /dev/null +++ b/src/obicompactvec/src/siblingannex.rs @@ -0,0 +1,245 @@ +//! 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" and "Definitions: +//! family, and the canonical form of a family"). +//! +//! 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. +//! +//! 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`. +//! +//! 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 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}; +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; + +/// 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 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) + } + + #[inline] + fn encode(self) -> u8 { + self.0 + } + + #[inline] + fn decode(byte: u8) -> Option { + if byte == 0 { + // Unreachable for a real result — reserved as the "not yet + // computed" sentinel. + return None; + } + Some(FamilyMask(byte & 0b1111)) + } +} + +// ── 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 { + FamilyMask::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 { + FamilyMask::decode(self.mmap[HEADER_SIZE + slot]) + } + + 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] = mask.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_masks() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test.psib"); + let mut builder = SiblingAnnexBuilder::new(4, &path).unwrap(); + + 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, mask) in masks.iter().enumerate() { + builder.set(slot, *mask); + } + let annex = builder.finish().unwrap(); + 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 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/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/Cargo.toml b/src/obikindex/Cargo.toml index c8fef8e..20e1dd3 100644 --- a/src/obikindex/Cargo.toml +++ b/src/obikindex/Cargo.toml @@ -10,6 +10,8 @@ obiskio = { path = "../obiskio" } obisys = { path = "../obisys" } obicompactvec = { path = "../obicompactvec" } obilayeredmap = { path = "../obilayeredmap" } +obiskbuilder = { path = "../obiskbuilder" } +obipipeline = { path = "../obipipeline" } ndarray = "0.16" rayon = "1" crossbeam-channel = "0.5" @@ -19,6 +21,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/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/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 371a05c..48daf54 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}; @@ -18,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::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs new file mode 100644 index 0000000..da51270 --- /dev/null +++ b/src/obikindex/src/siblings.rs @@ -0,0 +1,1121 @@ +//! Family presence-mask annex construction. +//! +//! 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 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, 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 ndarray::Array2; +use rayon::prelude::*; + +use obicompactvec::{ + 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; + +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}")), + } +} + +/// 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 +/// 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 +} + +// ── obipipeline data types ───────────────────────────────────────────────── + +/// One distinct k-mer of the layer currently being processed, at its local +/// MPHF slot — the pipeline's source item. Carries a throttle slot (shared, +/// `Arc`-wrapped so it can be cloned into each of the (up to 3) items this +/// one fans out into via the `Flat` stage) that is only released once every +/// one of those descendants has been fully processed — see the module docs' +/// "Throttling" note for why this is required, not optional, once a `Flat` +/// stage is in the pipeline. +struct SourceItem { + slot: usize, + kmer: CanonicalKmer, + _permit: Arc, +} + +/// One of a source k-mer's (up to 3) central-substitution variants, already +/// routed to its destination partition and carrying its own central base +/// (0=A/1=C/2=G/3=T) — the mask bit it will set on a hit. Carries a clone of +/// the source item's throttle permit. +struct VariantQuery { + source_slot: usize, + dest_partition: usize, + variant: CanonicalKmer, + base: u8, + _permit: Arc, +} + +enum SibData { + Item(SourceItem), + Query(VariantQuery), +} + +/// 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(), + } + } + 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 { + /// 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. + /// + /// 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; + + 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}")); + 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, + 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)?; + 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); + } + } + + let k = self.kmer_size(); + + // ── 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; + + // 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 => VariantQuery, + || { + move |item: SourceItem| -> Vec { + let kmer = item.kmer; + let permit = item._permit; + kmer.central_canonical_neighbors() + .into_iter() + .filter(|variant| *variant != kmer) + .map(|variant| VariantQuery { + source_slot: item.slot, + dest_partition: partition_of(variant, n_parts), + variant, + base: central_base(variant, k), + _permit: Arc::clone(&permit), + }) + .collect::>() + } + } : Item => Query, + }; + + // ── 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, m) in mask.iter().enumerate() { + if slot_kmer[slot].is_none() { + continue; // unused MPHF slot, if any — leave at the sentinel + } + builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed))); + } + builder.close()?; + + Ok(n_slots as u64) + } +} + +/// 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 *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], + /// `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 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; + + // 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); + 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); + } + } + + // 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 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); + } + } + + 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(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 { + 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) + } +} + +/// 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 }) + } +} + +/// 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; + 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) -> 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 { + 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', 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. 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 g1_kmer = canonical(b"AACCGCTTAAG"); + let g2_kmer = canonical(b"AACCGGTTAAG"); + let expected_mask = FamilyMask::EMPTY.with(1).with(2); + + 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] + 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 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/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 f66d78f..8852478 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -3,13 +3,15 @@ use std::path::PathBuf; use clap::Args; use kodama::{Method, linkage}; -use obikindex::{DistanceMetric, KmerIndex}; +use obifastwrite::{JsonVal, write_record}; +use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; #[derive(clap::ValueEnum, Clone, Copy, Debug)] pub enum MetricArg { Jaccard, + Mash, Hamming, BrayCurtis, #[value(name = "relfreq-bray-curtis")] @@ -26,6 +28,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, @@ -62,7 +65,37 @@ 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, + + /// 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, + + /// 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, _snp.fasta, /// _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] @@ -78,6 +111,51 @@ 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); + } + 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); + } + 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`/`--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; + } + info!( "computing {:?} distances for {} genome(s)", args.metric, n @@ -189,6 +267,103 @@ pub fn run(args: DistanceArgs) { } } +// ── 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) { + // 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(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + 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(); + } + writeln!( + f, "global,{},{},{},{}", + stats.counts[0], stats.counts[1], stats.counts[2], stats.counts[3], + ).unwrap(); + let total: u64 = stats.counts.iter().sum(); + info!("family-size distribution → {path} (total {total} famil{})", + 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}"); +} + +// ── 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 { 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); 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; 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;