Push lsqnpxrxuvpp #62

Merged
coissac merged 8 commits from push-lsqnpxrxuvpp into main 2026-08-11 09:09:24 +00:00
6 changed files with 671 additions and 277 deletions
Showing only changes of commit 1a470eab9e - Show all commits
+1
View File
@@ -9,6 +9,7 @@ data-stress
./**/*.json
*.bin
*.log
*.csv
Betula_exilis--IGA-24-33
benchmark/genomes
benchmark/simulated_data
+93 -22
View File
@@ -65,6 +65,50 @@ reverse-complement (`m -> k-1-m = m`, base complemented). A transition maps to
a transition, a transversion to a transversion — the transition/transversion
split is well-defined in canonical space.
### Definitions: family, and the canonical form of a family
**Family.** The family of a k-mer `x` is the set of (up to) 4 k-mers sharing
`x`'s `2m` flanking bases, differing only at the central base `m`. Membership
is a property of the flank pattern, not of `x` itself: any of the 4 possible
central substitutions belongs to the same family.
**`central_canonical_neighbors()`** (`obikseq`, `CanonicalKmerOf::central_canonical_neighbors`)
generates all 4 members from any one of them (observed or not), each
independently canonicalised (`.canonical()`, i.e. `min(kmer, revcomp(kmer))`).
This independent canonicalisation is necessary because a central substitution
can flip which orientation is lexicographically smaller — two members of the
same family can end up canonicalised in *different* orientations. Despite
that, the **set** of 4 resulting canonical k-mers is invariant: calling
`central_canonical_neighbors()` on any member of a family — present in the
index or not — yields the same 4 values. This is relied upon throughout the
rest of this document.
**Canonical form of a family.** Because orientation can differ member to
member, "which of the 4 is the reference" cannot be defined relative to
*whichever member happened to be visited first*, nor relative to the
minorant (see below) — both are data-dependent (they depend on what is
actually observed), so using either as the reference would make the
reference itself vary depending on what happens to be present in a given
index. Instead: **the canonical form of a family is, by definition, the
member whose own central base — read in its own already-canonical
orientation — is `A`.** This is well-defined for every family, computed
purely from the flank pattern, whether or not that specific member (or any
member at all) is actually observed anywhere in the index. Concretely: call
`central_canonical_neighbors()` on any member (observed or not) to get the
family's 4 canonical forms; the one among them whose own centre nucleotide is
`A` is the family's canonical form. The other 3 (`C`, `G`, `T`) are labelled
relative to *that* fixed reference, not relative to the calling member's own
orientation.
**Consequence for the minorant.** With this fixed A-referenced labelling,
`minorant` (the smallest raw encoding among the family's *observed* members,
introduced further below) becomes directly computable rather than needing to
be tracked as extra state: regenerate the family's 4 canonical forms from
any member's own k-mer (cheap, no lookup), compare the raw encodings of
whichever are marked present, and take the smallest. No separate stored bit
is required — see Step 2b below, where this replaces the earlier
minorant-bit design.
## Locus eligibility: raw definition vs. paralogy filter
For each k-mer `x` observed in genome A (source, one MPHF slot; the 3
@@ -482,31 +526,58 @@ here.
### Step 2b — sibling-count / minorant annex (consolidated plan)
Scope: only the precursor annex (sibling count 0-3 per slot, minorant
decided on demand) — not the SNP tally itself, whose Step 2 sweep remains
unresolved above. This piece is simpler than the sweep, because it writes to
an independent per-slot value, not a shared cross-k-mer accumulator, so it
needs no dedup/ownership logic at all at this stage.
Scope: only the precursor annex — not the SNP tally itself, whose Step 2
sweep remains unresolved above. This piece is simpler than the sweep,
because it writes to an independent per-slot value, not a shared
cross-k-mer accumulator, so it needs no dedup/ownership logic at all at this
stage.
**Revised annex encoding — 4-bit presence mask, not 3-bit (minorant +
count).** Superseded after settling the "canonical form of a family"
definition above. The 3-bit design (1 minorant bit + 2-bit sibling count,
§ below, kept for the historical record) had two problems: it discards
*which* variants are present (only how many), so any future consumer
(the SNP sweep, or a stats pass — see below) that needs to know which bases
exist still has to regenerate and blindly re-query all 3 candidates; and
the minorant bit's meaning was tied to whichever member was visited, not to
a fixed reference. Storing instead a **4-bit mask** — one bit per base
(A/C/G/T), set iff that member of the family (labelled relative to the
family's fixed canonical form, i.e. the member with `A` at the centre — see
above) is observed anywhere in the index — fixes both:
- **Sibling count is derived, not stored**: `siblings = popcount(mask) - 1`.
- **Minorant is derived, not stored**: regenerate the family's 4 canonical
forms from the slot's own k-mer (cheap, no lookup — see above), compare
the raw encodings of whichever bits are set in the mask, take the
smallest.
- **A future consumer knows exactly which variants to (re-)query** —
`popcount(mask) - 1` lookups instead of always 3, and it knows *which*
3 (or fewer) to issue, not just how many hits to expect.
- The all-zero value (no base present at all) is still logically
unreachable as a real result — the slot's *own* base is always present in
its own family — so it remains available as a free "not yet computed"
sentinel, exactly as before.
1. **Primitive.** Reuse `central_canonical_neighbors()` from Step 0
unchanged — the 3 canonicalised central-substitution variants of a k-mer.
2. **New annex type** (`obicompactvec`, alongside `bitmatrix.rs`): a 3-bit-
per-slot packed array, one per partition — same on-disk shape family as
`PersistentBitMatrix`'s `Packed` variant, but simpler (no per-genome
columns, a single derived read-only value per slot). 3 bits, not 2:
revised to also store minorant status alongside sibling count, since it
comes for free from the same lookups (point 3 below) — 5 real states
unchanged — the 3 canonicalised central-substitution variants of a k-mer
(plus the identity, i.e. all 4 members of the family — see "Definitions"
above).
2. **New annex type** (`obicompactvec`, alongside `bitmatrix.rs`): a 4-bit-
per-slot packed array (the presence mask above), one per partition — same
on-disk shape family as `PersistentBitMatrix`'s `Packed` variant, but
simpler (no per-genome columns, a single derived read-only value per
slot).
<details><summary>Superseded 3-bit design (historical)</summary>
3 bits, storing minorant status alongside sibling count directly, since
it came for free from the same lookups (point 3 below) — 5 real states
(not-minorant; minorant with 0/1/2/3 siblings) fit in 3 bits (8 states,
3 unused). This lets the future SNP sweep discard a non-minorant slot
**instantly**, with no lookup at all, instead of having to regenerate and
look up its siblings just to rediscover it isn't the designated writer —
moving that cost into this one-time, cached pass instead of repeating it
on every future sweep. The otherwise-unreachable combination
"not-minorant + 0 siblings" (impossible: 0 siblings always implies
minorant, see below) doubles as a free **"not yet computed" sentinel** —
annex files for all partitions/layers can be pre-initialised to this
value before the computation pass runs, distinguishing genuinely-computed
0-sibling slots from not-yet-processed ones with no extra storage.
3 unused). This let the future SNP sweep discard a non-minorant slot
instantly, with no lookup at all. The otherwise-unreachable combination
"not-minorant + 0 siblings" (0 siblings always implies minorant) doubled
as the "not yet computed" sentinel. Replaced by the 4-bit mask above,
which subsumes this benefit (minorant still derivable, now for free at
read time rather than stored) while also fixing the "which variant"
blindness.
</details>
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.
+1 -1
View File
@@ -19,7 +19,7 @@ pub use builder::PersistentCompactIntVecBuilder;
pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask};
pub use intmatrix::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, pack_compact_int_matrix};
pub use layer_meta::LayerMeta;
pub use siblingannex::{SiblingAnnex, SiblingAnnexBuilder, SiblingInfo};
pub use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
pub use reader::{PersistentCompactIntVec, Iter as CompactIntVecIter};
pub use tempbitvec::{TempBitVec, TempBitVecBuilder};
pub use tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
+105 -58
View File
@@ -1,27 +1,33 @@
//! Sibling-count / minorant annex: a compact, read-only-after-build, per-slot
//! Family presence-mask annex: a compact, read-only-after-build, per-slot
//! derived value used by the central-position SNP distance estimator (see
//! `docmd/theory/evolutionary_distances.md`, "Step 2b").
//! `docmd/theory/evolutionary_distances.md`, "Step 2b" and "Definitions:
//! family, and the canonical form of a family").
//!
//! One byte is stored per MPHF slot of a partition/layer, encoding two
//! independent facts about the slot's k-mer's "family" (the up-to-4 k-mers
//! sharing the same flanks, differing only at the central base), both
//! properties of the whole current multi-genome index, not of any one
//! genome:
//! One byte is stored per MPHF slot of a partition/layer, its low 4 bits
//! encoding a **presence mask** for the slot's k-mer's "family" (the up to 4
//! k-mers sharing the same flanks, differing only at the central base):
//! bit `b` (`b` = 0..3, in the fixed A/C/G/T = 0/1/2/3 encoding already used
//! for a single nucleotide) is set iff the family member whose *own* central
//! base — in its own canonical orientation — is `b`, is observed anywhere in
//! the current multi-genome index. This is a property of the whole index,
//! not of any one genome.
//!
//! - bit 0: `minorant` — is this k-mer's canonical encoding the smallest
//! among the family members actually observed in the index?
//! - bits 1-2: `siblings` — how many *other* family members (0-3) are
//! observed anywhere in the index.
//! Both facts the earlier (superseded) 3-bit design stored explicitly are
//! derived from the mask instead, not stored:
//! - sibling count = `popcount(mask) - 1`;
//! - minorant = regenerate the family's 4 canonical forms from the slot's
//! own k-mer (`CanonicalKmerOf::central_canonical_neighbors`, cheap, no
//! lookup), compare the raw encodings of whichever are set in the mask,
//! take the smallest — see `obikindex::siblings`.
//!
//! Byte value 0 (`minorant = false`, `siblings = 0`) is logically
//! unreachable as a real result (0 siblings always implies minorant — see
//! the design doc) and is reused as the "not yet computed" sentinel: annex
//! files are pre-initialised to all-zero, and a real value is only ever
//! written once, by the computation pass.
//! Mask value 0 is logically unreachable as a real result (a slot's own base
//! is always present in its own family) and is reused as the "not yet
//! computed" sentinel: annex files are pre-initialised to all-zero, and a
//! real value is only ever written once, by the computation pass.
//!
//! Deliberately simpler than a true 3-bit pack (1 byte/slot instead of 3
//! bits/slot): correctness and simplicity first: for a first implementation.
//! Packing to 3 bits/slot is a pure storage-density follow-up, not a
//! Deliberately simpler than a true 4-bit pack (1 byte/slot instead of 4
//! bits/slot): correctness and simplicity first, for a first implementation.
//! Packing to 4 bits/slot is a pure storage-density follow-up, not a
//! behavioural change, left for later.
use std::fs::{File, OpenOptions};
@@ -35,30 +41,70 @@ const MAGIC: [u8; 4] = *b"PSIB";
// Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (1 byte/slot) follows.
const HEADER_SIZE: usize = 16;
/// Decoded value of one slot's annex entry.
/// A family presence mask: bit `b` set iff the member whose own canonical
/// central base is `b` (0=A, 1=C, 2=G, 3=T) is observed in the index.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SiblingInfo {
pub minorant: bool,
pub siblings: u8, // 0..=3
pub struct FamilyMask(u8);
impl FamilyMask {
/// The empty mask — never a valid *computed* result (a slot's own base
/// is always present in its own family) — used only to build up a mask
/// via repeated [`with`](Self::with) calls before storing it.
pub const EMPTY: FamilyMask = FamilyMask(0);
/// Set bit `base` (0=A, 1=C, 2=G, 3=T).
#[inline]
pub fn with(self, base: u8) -> Self {
debug_assert!(base < 4, "base out of range: {base}");
FamilyMask(self.0 | (1 << base))
}
/// Is the member with central base `base` (0..3) present?
#[inline]
pub fn has(self, base: u8) -> bool {
debug_assert!(base < 4, "base out of range: {base}");
self.0 & (1 << base) != 0
}
/// Number of family members observed anywhere in the index (1..=4).
#[inline]
pub fn family_size(self) -> u32 {
self.0.count_ones()
}
/// Number of *other* members observed (0..=3) — `family_size() - 1`.
#[inline]
pub fn siblings(self) -> u32 {
self.family_size() - 1
}
/// Raw bitmask (bit `b` = base `b` present) — for callers that build up
/// a mask via their own bit operations (e.g. concurrently, via an
/// `AtomicU8`) and only need the `FamilyMask` wrapper at the end.
#[inline]
pub fn bits(self) -> u8 {
self.0
}
/// Construct from a raw bitmask (only the low 4 bits are kept).
#[inline]
pub fn from_bits(bits: u8) -> Self {
FamilyMask(bits & 0b1111)
}
impl SiblingInfo {
#[inline]
fn encode(self) -> u8 {
(self.siblings << 1) | (self.minorant as u8)
self.0
}
#[inline]
fn decode(byte: u8) -> Option<Self> {
if byte == 0 {
// The unreachable "not minorant + 0 siblings" combination —
// reserved as the "not yet computed" sentinel.
// Unreachable for a real result — reserved as the "not yet
// computed" sentinel.
return None;
}
Some(SiblingInfo {
minorant: byte & 1 != 0,
siblings: (byte >> 1) & 0b11,
})
Some(FamilyMask(byte & 0b1111))
}
}
@@ -91,8 +137,8 @@ impl SiblingAnnex {
pub fn is_empty(&self) -> bool { self.n == 0 }
/// `None` means the slot has not (yet) been computed — see module docs.
pub fn get(&self, slot: usize) -> Option<SiblingInfo> {
SiblingInfo::decode(self.mmap[HEADER_SIZE + slot])
pub fn get(&self, slot: usize) -> Option<FamilyMask> {
FamilyMask::decode(self.mmap[HEADER_SIZE + slot])
}
}
@@ -124,17 +170,17 @@ impl SiblingAnnexBuilder {
pub fn len(&self) -> usize { self.n }
pub fn is_empty(&self) -> bool { self.n == 0 }
pub fn get(&self, slot: usize) -> Option<SiblingInfo> {
SiblingInfo::decode(self.mmap[HEADER_SIZE + slot])
pub fn get(&self, slot: usize) -> Option<FamilyMask> {
FamilyMask::decode(self.mmap[HEADER_SIZE + slot])
}
pub fn set(&mut self, slot: usize, info: SiblingInfo) {
pub fn set(&mut self, slot: usize, mask: FamilyMask) {
// Redundant concurrent writes from independent recomputation paths
// converge to the same encoded byte for a given slot, so a plain
// store here is safe even without external synchronisation, as long
// as the byte write itself is atomic (true for a single aligned
// byte on every platform this project targets).
self.mmap[HEADER_SIZE + slot] = info.encode();
self.mmap[HEADER_SIZE + slot] = mask.encode();
}
pub fn close(self) -> io::Result<()> { self.mmap.flush() }
@@ -163,36 +209,37 @@ mod tests {
}
#[test]
fn roundtrip_all_valid_states() {
fn roundtrip_all_valid_masks() {
let dir = tempdir().unwrap();
let path = dir.path().join("test.psib");
let mut builder = SiblingAnnexBuilder::new(5, &path).unwrap();
let mut builder = SiblingAnnexBuilder::new(4, &path).unwrap();
let cases = [
SiblingInfo { minorant: true, siblings: 0 },
SiblingInfo { minorant: true, siblings: 1 },
SiblingInfo { minorant: true, siblings: 2 },
SiblingInfo { minorant: true, siblings: 3 },
SiblingInfo { minorant: false, siblings: 2 },
let masks = [
FamilyMask::EMPTY.with(0), // just A: family size 1
FamilyMask::EMPTY.with(0).with(3), // A + T: size 2
FamilyMask::EMPTY.with(1).with(2).with(3), // C+G+T: size 3
FamilyMask::EMPTY.with(0).with(1).with(2).with(3), // all 4
];
for (slot, info) in cases.iter().enumerate() {
builder.set(slot, *info);
for (slot, mask) in masks.iter().enumerate() {
builder.set(slot, *mask);
}
let annex = builder.finish().unwrap();
for (slot, info) in cases.iter().enumerate() {
assert_eq!(annex.get(slot), Some(*info));
for (slot, mask) in masks.iter().enumerate() {
assert_eq!(annex.get(slot), Some(*mask));
}
assert_eq!(annex.get(0).unwrap().siblings(), 0);
assert_eq!(annex.get(1).unwrap().siblings(), 1);
assert_eq!(annex.get(2).unwrap().siblings(), 2);
assert_eq!(annex.get(3).unwrap().siblings(), 3);
assert_eq!(annex.get(3).unwrap().family_size(), 4);
}
#[test]
fn not_minorant_zero_siblings_is_unreachable_via_set_and_decodes_as_sentinel() {
// Documented invariant, not enforced by the type: callers must never
// construct this combination. If they do, it is indistinguishable
// from "not computed" — exercised here to pin the behaviour down.
let dir = tempdir().unwrap();
let path = dir.path().join("test.psib");
let mut builder = SiblingAnnexBuilder::new(1, &path).unwrap();
builder.set(0, SiblingInfo { minorant: false, siblings: 0 });
assert_eq!(builder.get(0), None);
fn has_reflects_individual_bits() {
let mask = FamilyMask::EMPTY.with(0).with(2);
assert!(mask.has(0));
assert!(!mask.has(1));
assert!(mask.has(2));
assert!(!mask.has(3));
}
}
+428 -158
View File
@@ -1,39 +1,49 @@
//! Sibling-count / minorant annex construction.
//! Family presence-mask annex construction.
//!
//! See `docmd/theory/evolutionary_distances.md`, "Step 2b — sibling-count /
//! minorant annex", for the full design discussion this implements.
//! See `docmd/theory/evolutionary_distances.md`, "Definitions: family, and
//! the canonical form of a family" and "Step 2b", for the full design
//! discussion this implements.
//!
//! For each distinct k-mer of each layer of the (already built/merged)
//! index, computes two facts about its "family" (the up to 4 k-mers sharing
//! its flanks, differing only at the central base — well-defined for odd
//! k), both properties of the whole current multi-genome index rather than
//! of any one genome:
//! - how many *other* family members (0-3) are observed anywhere in the
//! index;
//! - whether this k-mer is the "minorant" of its family — the smallest
//! canonical encoding among the members actually observed.
//! index, computes a 4-bit presence mask for its "family" (the up to 4
//! k-mers sharing its flanks, differing only at the central base —
//! well-defined for odd k): bit `b` set iff the family member whose own
//! canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed anywhere
//! in the current multi-genome index — a property of the whole index, not
//! of any one genome. Sibling count and minorant are *derived* from the
//! mask by callers, not stored (see `FamilyMask` and
//! [`sibling_annex_stats`](KmerIndex::sibling_annex_stats) below).
//!
//! Per layer, the computation runs as an `obipipeline` pipeline with several
//! elementary stages (a `Flat` stage generating each k-mer's 3 central
//! variants, a `Transform` stage looking each variant up in its destination
//! partition), so the scheduler's shared worker pool interleaves this work
//! across many in-flight k-mers/variants rather than processing everything
//! on a single thread — see `docmd/theory/evolutionary_distances.md`, Step
//! 2b, "Mechanism", for why elementary stages were chosen deliberately over
//! a few coarse ones.
//! Per layer, an `obipipeline` `Flat` stage (throttled — see
//! `obipipeline::throttle`) generates each k-mer's 3 central variants,
//! interleaved across many in-flight k-mers by the scheduler's shared
//! worker pool rather than processed on a single thread. The actual
//! cross-partition lookup, though, reuses
//! `KmerPartition::query_partition_with` — the same partition-batching
//! mechanism `obikmer query` already uses (open a partition's files once,
//! answer a whole batch of queries against it) — rather than a per-item
//! pipeline stage: an earlier per-item design reopened/re-mmap'd every
//! target partition's files on every single lookup, which was fine at
//! toy scale but manifested as ~90% system time against a real index.
//! See `docmd/theory/evolutionary_distances.md`, Step 2b, "Mechanism".
use std::path::Path;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::Arc;
use rayon::prelude::*;
use obicompactvec::{
PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder, SiblingInfo,
FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder,
};
use obikpartitionner::KmerPartition;
use obipipeline::ThrottleGuard;
use obikseq::{CanonicalKmer, Minimizer};
use obilayeredmap::{MphfLayer, OLMError};
use obilayeredmap::meta::PartitionMeta;
use obiskbuilder::rolling_stat::RollingStat;
use obiskio::UnitigFileReader;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
@@ -48,6 +58,26 @@ fn olm_to_ok(e: OLMError) -> OKIError {
}
}
/// Central-position base of a canonical k-mer, in the fixed 0=A/1=C/2=G/3=T
/// encoding — the mask's bit index. `k` must be odd (project invariant).
#[inline]
fn central_base(kmer: CanonicalKmer, k: usize) -> u8 {
kmer.nucleotide((k - 1) / 2)
}
/// Is `kmer` the minorant of its family, given the family's presence mask?
/// Regenerates the family's 4 canonical forms from `kmer` itself (cheap, no
/// lookup — see the design doc's "Definitions" section for why this is
/// always safe: the set of 4 forms is invariant regardless of which member
/// you start from), and compares the raw encodings of whichever are marked
/// present in `mask`.
fn is_minorant(kmer: CanonicalKmer, mask: FamilyMask, k: usize) -> bool {
kmer.central_canonical_neighbors().into_iter().all(|other| {
other == kmer || !mask.has(central_base(other, k)) || kmer.raw() <= other.raw()
})
}
/// Minimiser of a single, isolated canonical k-mer (not part of a streamed
/// sequence). `RollingStat` computes minimisers incrementally along a
/// sequence; this feeds one k-mer's bases through a fresh instance to get
@@ -74,79 +104,144 @@ fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize {
(lone_kmer_minimizer(kmer).seq_hash() & mask) as usize
}
/// Running reconciliation state for one source k-mer, initialised to the
/// trivial "no siblings observed yet" state and folded incrementally (in any
/// order — commutative) as query answers come back.
#[derive(Clone, Copy)]
struct RunningState {
minorant: bool,
siblings: u8,
}
impl Default for RunningState {
fn default() -> Self {
RunningState { minorant: true, siblings: 0 }
}
}
// ── obipipeline data types ─────────────────────────────────────────────────
/// One distinct k-mer of the layer currently being processed, at its local
/// MPHF slot — the pipeline's source item.
#[derive(Clone, Copy)]
/// MPHF slot — the pipeline's source item. Carries a throttle slot (shared,
/// `Arc`-wrapped so it can be cloned into each of the (up to 3) items this
/// one fans out into via the `Flat` stage) that is only released once every
/// one of those descendants has been fully processed — see the module docs'
/// "Throttling" note for why this is required, not optional, once a `Flat`
/// stage is in the pipeline.
struct SourceItem {
slot: usize,
kmer: CanonicalKmer,
_permit: Arc<ThrottleGuard>,
}
/// One of a source k-mer's (up to 3) central-substitution variants, already
/// routed to its destination partition and carrying, decided here (both
/// encodings are already in hand — no need to wait for the lookup answer),
/// whether this specific variant would outrank the source as minorant.
#[derive(Clone, Copy)]
/// routed to its destination partition and carrying its own central base
/// (0=A/1=C/2=G/3=T) — the mask bit it will set on a hit. Carries a clone of
/// the source item's throttle permit.
struct VariantQuery {
source_slot: usize,
dest_partition: usize,
variant: CanonicalKmer,
smaller: bool,
base: u8,
_permit: Arc<ThrottleGuard>,
}
/// Outcome of looking a [`VariantQuery`] up in its destination partition.
#[derive(Clone, Copy)]
struct AnswerMsg {
source_slot: usize,
hit: bool,
smaller: bool,
}
#[derive(Clone, Copy)]
enum SibData {
Item(SourceItem),
Query(VariantQuery),
Answer(AnswerMsg),
}
/// Existence-only lookup of `variant` in partition `dest_partition`: tries
/// each of the partition's layers in turn, stopping at the first hit — the
/// same "try every layer's MPHF" shape as `QueryLayer::find_slot` (private
/// to `obikpartitionner`), reimplemented here directly against the public
/// `MphfLayer::open`/`find` since only existence is needed, not a column
/// fetch.
fn lookup_exists(partition: &KmerPartition, dest_partition: usize, variant: CanonicalKmer) -> bool {
let index_dir = partition.part_dir(dest_partition).join(INDEX_SUBDIR);
if !index_dir.exists() {
return false;
/// 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),
}
let Ok(meta) = PartitionMeta::load(&index_dir) else { return false };
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<Vec<MphfLayer>>,
mats: Vec<Vec<Mat>>,
}
impl PartitionCache {
fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult<Self> {
let pb = progress_bar("open_partitions", n_parts as u64, "partitions");
let built: Vec<(Vec<MphfLayer>, Vec<Mat>)> = (0..n_parts)
.into_par_iter()
.map(|part| -> OKIResult<(Vec<MphfLayer>, Vec<Mat>)> {
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}"));
if let Ok(mphf) = MphfLayer::open(&layer_dir, &meta.mode) {
if mphf.find(variant).is_some() {
return true;
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::<OKIResult<Vec<_>>>()?;
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<Vec<bool>> {
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
}
false
}
impl KmerIndex {
@@ -169,42 +264,50 @@ impl KmerIndex {
let n_parts = self.n_partitions();
let n_bits = n_parts.trailing_zeros() as usize;
// A fresh, owned `KmerPartition` handle (read-only use only — no
// writers opened), wrapped in `Arc` so pipeline stage closures
// (which must be `'static`, running in spawned threads) can share
// it without borrowing `self`.
let partition = Arc::new(
KmerPartition::open_with_config(
let partition = KmerPartition::open_with_config(
&self.root_path,
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?,
);
.map_err(OKIError::Partition)?;
tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
let cache = Arc::new(PartitionCache::build(&partition, n_parts, self.meta.config.with_counts)?);
let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
let mut total_slots: u64 = 0;
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
pb.inc(1);
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
let mut part_slots: u64 = 0;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
self.build_layer_sibling_annex(&layer_dir, n_parts, &partition)?;
part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?;
}
total_slots += part_slots;
pb.inc(1);
pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)"));
}
pb.finish_and_clear();
tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions");
Ok(())
}
/// Returns the number of distinct k-mers (annex slots) processed, for
/// progress reporting.
fn build_layer_sibling_annex(
&self,
layer_dir: &Path,
n_parts: usize,
partition: &Arc<KmerPartition>,
) -> OKIResult<()> {
cache: &Arc<PartitionCache>,
) -> OKIResult<u64> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
@@ -220,22 +323,61 @@ impl KmerIndex {
}
}
let sources: Vec<SourceItem> = slot_kmer
.iter()
.enumerate()
.filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| SourceItem { slot, kmer }))
.collect();
let k = self.kmer_size();
// ── obipipeline: Flat (generate variants) -> Transform (lookup) ─────
// ── Reconciliation state, initialised with each slot's own base —
// that member is trivially present, no lookup needed. Built before
// the pipeline runs, from the same enumeration, since `sources`
// below is consumed as a throttled iterator, not collected.
// `AtomicU8`, not `FamilyMask`, because the gather phase below
// parallelises across destination partitions (independent
// `query_partition_with` calls, safe to run concurrently) and their
// `Found` hits can land on arbitrary, possibly-shared slots — a
// lock-free `fetch_or` avoids needing any synchronisation beyond
// that. ─────────────────────────────────────────────────────────
let mask: Vec<AtomicU8> = (0..n_slots).map(|_| AtomicU8::new(0)).collect();
for (slot, kmer) in slot_kmer.iter().enumerate().filter_map(|(s, k)| k.map(|k| (s, k))) {
mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
}
// ── obipipeline: Flat stage generates variants only — the actual
// cross-partition lookup reuses `KmerPartition::query_partition_with`
// (the same batching mechanism `obikmer query` already uses: open a
// partition's files once, answer a whole batch of queries against
// it) instead of one lookup per pipeline item. A per-item lookup
// (tried first) reopened/re-mmap'd every target partition's files on
// every single variant — fine at the scale of a handful of test
// k-mers, but with billions of lookups against a real index this
// manifested as ~90% system time, observed in practice. ───────────
let n_workers = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4);
let capacity = 256;
let partition_for_lookup = Arc::clone(partition);
// Throttling is not optional once a `Flat` stage is in the pipeline
// (see `obipipeline::throttle`'s docs): without it, every worker can
// become a simultaneous `Flat` producer, saturate the shared output
// channel, and deadlock against the scheduler's own dispatch loop —
// also observed in practice. The permit acquired here for a source
// k-mer is held (via the `Arc`-shared guard carried through
// `SourceItem` -> `VariantQuery`) until every one of its (up to 3)
// descendants has been read out of the pipeline by the accumulation
// loop below, not just until the `Flat` stage itself returns.
let sources: Vec<(usize, CanonicalKmer)> = slot_kmer
.iter()
.enumerate()
.filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer)))
.collect();
let throttled = obipipeline::throttle(sources.into_iter(), n_workers).map(|t| SourceItem {
slot: t.item.0,
kmer: t.item.1,
_permit: Arc::new(t.guard),
});
let pipe = obipipeline::make_pipe! {
SibData : SourceItem => AnswerMsg,
SibData : SourceItem => VariantQuery,
|| {
move |item: SourceItem| -> Vec<VariantQuery> {
let kmer = item.kmer;
let permit = item._permit;
kmer.central_canonical_neighbors()
.into_iter()
.filter(|variant| *variant != kmer)
@@ -243,88 +385,112 @@ impl KmerIndex {
source_slot: item.slot,
dest_partition: partition_of(variant, n_parts),
variant,
smaller: variant.raw() < kmer.raw(),
base: central_base(variant, k),
_permit: Arc::clone(&permit),
})
.collect::<Vec<_>>()
}
} : Item => Query,
| {
let partition_for_lookup = Arc::clone(&partition_for_lookup);
move |vq: VariantQuery| -> AnswerMsg {
let hit = lookup_exists(&partition_for_lookup, vq.dest_partition, vq.variant);
AnswerMsg { source_slot: vq.source_slot, hit, smaller: vq.smaller }
}
} : Query => Answer,
};
// ── Reconciliation (the pipeline's sink): commutative fold of every
// answer into its origin k-mer's running state, in whatever order
// the pipeline delivers them. ─────────────────────────────────────
let mut state = vec![RunningState::default(); n_slots];
for ans in pipe.apply(sources.into_iter(), n_workers, capacity) {
if ans.hit {
let st = &mut state[ans.source_slot];
st.siblings = (st.siblings + 1).min(3);
if ans.smaller {
st.minorant = false;
// ── Group generated variants by destination partition. `cache`
// holds every partition already mmap'd (no more `open()` cost), but
// `mmap` pages are still loaded on demand and can be evicted — a
// lookup is not free just because the file isn't reopened. Grouping
// keeps one partition's pages hot while its whole batch is resolved,
// instead of faulting pages in and out as lookups jump between
// partitions in whatever order the `Flat` stage happens to produce
// them. The throttle permit drops here, once accumulated. ─────────
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
for vq in pipe.apply(throttled, n_workers, capacity) {
outgoing[vq.dest_partition].push((vq.variant, vq.source_slot, vq.base));
}
// ── Resolve each partition's batch against the cache in one
// contiguous pass; parallelised across partitions (independent,
// read-only) so this keeps using multiple cores without giving up
// the per-partition locality above. ─────────────────────────────
outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
for &(variant, source_slot, base) in queries {
if cache.find(dest, variant) {
mask[source_slot].fetch_or(1 << base, Ordering::Relaxed);
}
}
});
// ── Write the layer's annex file ─────────────────────────────────────
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?;
for (slot, st) in state.iter().enumerate() {
for (slot, m) in mask.iter().enumerate() {
if slot_kmer[slot].is_none() {
continue; // unused MPHF slot, if any — leave at the sentinel
}
builder.set(slot, SiblingInfo { minorant: st.minorant, siblings: st.siblings });
builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed)));
}
builder.close()?;
Ok(())
Ok(n_slots as u64)
}
}
/// Distribution of sibling counts (0-3), read back from an already-built
/// Distribution of family sizes (1-4), read back from an already-built
/// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's
/// presence/count data — a separate, occasional diagnostic pass, not fused
/// into construction.
///
/// Every count here is **per family, not per slot**: a family with `F`
/// members occupies `F` annex slots (one per observed member), all sharing
/// the same mask. Counting every slot would count each family up to 4
/// times over; only the minorant's slot is tallied (minorant is derived on
/// the fly — see `is_minorant` — not stored, but cheap: no lookup, pure
/// bit arithmetic on already-in-hand data).
#[derive(Debug, Clone, Default)]
pub struct SiblingAnnexStats {
/// `counts[s]` = number of k-mers (slots) with exactly `s` siblings,
/// counted once each regardless of how many genomes carry them.
/// `counts[s]` = number of *families* of size `s + 1` (index 0 = size 1,
/// i.e. 0 siblings, ... index 3 = size 4, i.e. 3 siblings).
pub counts: [u64; 4],
/// Of those, how many are minorant.
pub minorant_counts: [u64; 4],
/// `per_genome[g][s]` = number of k-mers with exactly `s` siblings that
/// genome `g` (index into `KmerIndex::meta().genomes`) carries.
/// `per_genome[g][s]` = number of families of size `s + 1` for which
/// genome `g` (index into `KmerIndex::meta().genomes`) carries at least
/// one member.
pub per_genome: Vec<[u64; 4]>,
}
impl KmerIndex {
/// Tally the sibling-count distribution of an already-built annex
/// (globally, and per genome). Errors if [`build_sibling_annex`] has not
/// been run on this index first.
/// Tally the family-size distribution of an already-built annex
/// (globally, and per genome), counting each family once (at its
/// minorant slot). Errors if [`build_sibling_annex`] has not been run on
/// this index first.
///
/// [`build_sibling_annex`]: Self::build_sibling_annex
pub fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats> {
let n_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len();
let with_counts = self.meta.config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let mut stats = SiblingAnnexStats {
per_genome: vec![[0u64; 4]; n_genomes],
..Default::default()
};
// Same whole-run cache as `build_sibling_annex` — see its docs for
// why re-opening per lookup (or per call to a batching helper) is
// not good enough on a real index.
let partition = KmerPartition::open_with_config(
&self.root_path,
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
// Gather the (partition, layer) pairs to process — cheap metadata
// reads only, checking every annex file exists up front so a
// missing one is reported before any real work starts.
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
@@ -334,45 +500,111 @@ impl KmerIndex {
annex_path.display()
)));
}
let annex = SiblingAnnex::open(&annex_path)?;
let use_counts = with_counts && layer_dir.join("counts").exists();
layer_dirs.push(layer_dir);
}
}
// Opened once per layer, outside the slot loop.
enum Mat {
Count(PersistentCompactIntMatrix),
Presence(PersistentBitMatrix),
}
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(&layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(&layer_dir)?)
// One layer's worth of work, parallelised across layers with Rayon
// — independent, read-only, each producing its own partial tally
// merged at the end.
let pb = progress_bar("sibling_annex_stats", layer_dirs.len() as u64, "layers");
let partials: Vec<SiblingAnnexStats> = layer_dirs
.par_iter()
.map(|layer_dir| -> OKIResult<SiblingAnnexStats> {
let mut stats = SiblingAnnexStats {
per_genome: vec![[0u64; 4]; n_genomes],
..Default::default()
};
let n_cols = match &mat {
Mat::Count(m) => m.n_cols(),
Mat::Presence(m) => m.n_cols(),
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
// Need each slot's own k-mer to derive minorant — same
// enumeration as construction.
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
.min(n_genomes);
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
for slot in 0..annex.len() {
let Some(info) = annex.get(slot) else { continue };
let s = info.siblings as usize;
stats.counts[s] += 1;
if info.minorant {
stats.minorant_counts[s] += 1;
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // this family is tallied at its minorant's slot only
}
let s = mask.siblings() as usize;
stats.counts[s] += 1;
// "Genome g represents this family" means g carries
// *any* of its members, not just the minorant's own —
// start from the minorant's own presence (already
// open, no lookup) and OR in every other present
// member's presence vector, resolved against the
// whole-run cache (no I/O) — exactly `mask.siblings()`
// of them, the mask tells us precisely which to fetch.
let mut carries = vec![false; n_cols];
for g in 0..n_cols {
let carried = match &mat {
Mat::Count(m) => m.col_view(g).get(slot) != 0,
Mat::Presence(m) => m.get(g, slot) != 0,
};
carries[g] = mat.carries(g, slot);
}
for other in kmer.central_canonical_neighbors() {
if other == kmer {
continue;
}
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let dest = partition_of(other, n_parts);
if let Some(other_presence) = cache.find_presence(dest, other, n_genomes) {
for (g, &present) in other_presence.iter().enumerate() {
if present {
carries[g] = true;
}
}
}
}
for (g, &carried) in carries.iter().enumerate() {
if carried {
stats.per_genome[g][s] += 1;
}
}
}
}
}
pb.inc(1);
Ok(stats)
})
.collect::<OKIResult<Vec<_>>>()?;
pb.finish_and_clear();
let mut stats = SiblingAnnexStats {
per_genome: vec![[0u64; 4]; n_genomes],
..Default::default()
};
for part in partials {
for s in 0..4 {
stats.counts[s] += part.counts[s];
}
for g in 0..n_genomes {
for s in 0..4 {
stats.per_genome[g][s] += part.per_genome[g][s];
}
}
}
Ok(stats)
}
}
@@ -442,7 +674,7 @@ mod tests {
/// Read back the annex entry for a given canonical k-mer from the merged
/// index's (single) partition/layer, asserting it was found at all.
fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> SiblingInfo {
fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> FamilyMask {
let index_dir = idx.partition().part_dir(0).join(INDEX_SUBDIR);
let meta = PartitionMeta::load(&index_dir).unwrap();
for l in 0..meta.n_layers {
@@ -474,24 +706,33 @@ mod tests {
fn sibling_annex_one_sibling_each() {
// k=11, centre = index 5 (0-based). Two genomes, each exactly one
// k-mer, sharing every base except the centre:
// g1 = "AACCGCTTAAG" (centre 'C')
// g2 = "AACCGGTTAAG" (centre 'G')
// g1 = "AACCGCTTAAG" (centre 'C', base index 1)
// g2 = "AACCGGTTAAG" (centre 'G', base index 2)
// Hand-verified: both stay forward-oriented under canonicalisation
// (each is lexicographically smaller than its own reverse
// complement, since both start with "AA"), and raw(g1) < raw(g2)
// (only differing base: C=0b01 < G=0b10 at the centre) — so g1 is
// the minorant, g2 is not, and each is the other's one sibling.
// the minorant, g2 is not. The mask is a family-wide value: both
// slots must read back the *same* mask (bits 1 and 2 set).
let dir = tempdir().unwrap();
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
let a = annex_info_for(&merged, canonical(b"AACCGCTTAAG"));
assert_eq!(a, SiblingInfo { minorant: true, siblings: 1 }, "AACCGCTTAAG");
let g1_kmer = canonical(b"AACCGCTTAAG");
let g2_kmer = canonical(b"AACCGGTTAAG");
let expected_mask = FamilyMask::EMPTY.with(1).with(2);
let b = annex_info_for(&merged, canonical(b"AACCGGTTAAG"));
assert_eq!(b, SiblingInfo { minorant: false, siblings: 1 }, "AACCGGTTAAG");
let a = annex_info_for(&merged, g1_kmer);
assert_eq!(a, expected_mask, "AACCGCTTAAG");
assert_eq!(a.siblings(), 1);
assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant");
let b = annex_info_for(&merged, g2_kmer);
assert_eq!(b, expected_mask, "AACCGGTTAAG");
assert_eq!(b.siblings(), 1);
assert!(!is_minorant(g2_kmer, b, K), "g2 should not be the minorant");
}
#[test]
@@ -504,7 +745,36 @@ mod tests {
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
let info = annex_info_for(&merged, canonical(b"GATTACAGATC"));
assert_eq!(info, SiblingInfo { minorant: true, siblings: 0 }, "GATTACAGATC");
let kmer = canonical(b"GATTACAGATC");
let mask = annex_info_for(&merged, kmer);
assert_eq!(mask.siblings(), 0, "GATTACAGATC");
assert_eq!(mask.family_size(), 1);
assert!(is_minorant(kmer, mask, K));
}
#[test]
fn sibling_annex_stats_counts_each_family_once_and_per_genome() {
// Reuses the one-sibling-each fixture: a single family of size 2
// (g1's centre-C form + g2's centre-G form), each genome carrying
// exactly one of the two members. Stats must report exactly one
// family of size 2 (`counts[1] == 1`, since index 1 = size 2), not
// two (which naively summing both slots would give), and both
// genomes represented at size 2, neither at any other size.
let dir = tempdir().unwrap();
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
let stats = merged.sibling_annex_stats().expect("sibling_annex_stats");
assert_eq!(stats.counts, [0, 1, 0, 0], "one family of size 2, counted once");
assert_eq!(stats.per_genome.len(), 2);
for g in 0..2 {
assert_eq!(
stats.per_genome[g], [0, 1, 0, 0],
"genome {g} should represent exactly one size-2 family"
);
}
}
}
+33 -28
View File
@@ -78,8 +78,7 @@ pub struct DistanceArgs {
pub sibling_stats: bool,
/// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
/// <prefix>_siblings.csv, <prefix>_siblings_per_genome.csv,
/// <prefix>_nj.nwk, <prefix>_upgma.nwk.
/// <prefix>_siblings.csv, <prefix>_nj.nwk, <prefix>_upgma.nwk.
/// If omitted, the distance matrix is written to stdout.
#[arg(short, long)]
pub output: Option<PathBuf>,
@@ -114,6 +113,17 @@ pub fn run(args: DistanceArgs) {
write_sibling_stats_csv(&stats, &labels, &args.output);
}
// `--sibling-annex`/`--sibling-stats` are their own operation, not a
// modifier on top of a distance-metric computation — a metric was
// never requested by asking for either of them, so there is nothing
// for the rest of this function to compute. Not a historical accident
// to keep: stop here rather than always also running a Jaccard (or
// whichever `--metric` defaults to) pass and printing an unrequested
// matrix.
if args.sibling_annex || args.sibling_stats {
return;
}
info!(
"computing {:?} distances for {} genome(s)",
args.metric, n
@@ -225,43 +235,38 @@ pub fn run(args: DistanceArgs) {
}
}
// ── Sibling-count distribution → CSV ────────────────────────────────────────
// ── Family-size distribution → CSV ──────────────────────────────────────────
//
// Each row is a family (the up-to-4 k-mers sharing flanks, differing only at
// the centre), counted once — at its minorant — regardless of how many of
// its members are observed. Family size 1..4 (not "sibling count" 0..3):
// see `docmd/theory/evolutionary_distances.md`, "Definitions".
fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output: &Option<PathBuf>) {
// Global histogram: one row per sibling count (0-3).
let global_path = output.as_ref()
// One row per genome (4 columns, family size 1-4: number of families of
// that size for which the genome carries at least one member), plus a
// `global` row — the actual deduplicated family-size histogram
// (`stats.counts`), NOT a sum of the per-genome columns (a family shared
// by several genomes would otherwise be counted once per genome it
// appears in, inflating the total beyond the real family count).
let path = output.as_ref()
.map(|p| format!("{}_siblings.csv", p.display()))
.unwrap_or_else(|| "siblings.csv".into());
let mut f = BufWriter::new(std::fs::File::create(&global_path).unwrap_or_else(|e| {
eprintln!("error creating {global_path}: {e}");
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
eprintln!("error creating {path}: {e}");
std::process::exit(1);
}));
writeln!(f, "siblings,slots,minorant_slots").unwrap();
for s in 0..4 {
writeln!(f, "{s},{},{}", stats.counts[s], stats.minorant_counts[s]).unwrap();
}
let total: u64 = stats.counts.iter().sum();
info!("sibling-count distribution → {global_path} (total {total} slot(s))");
// Per-genome breakdown: one row per genome, 4 columns (0-3), + a total row.
let per_genome_path = output.as_ref()
.map(|p| format!("{}_siblings_per_genome.csv", p.display()))
.unwrap_or_else(|| "siblings_per_genome.csv".into());
let mut f = BufWriter::new(std::fs::File::create(&per_genome_path).unwrap_or_else(|e| {
eprintln!("error creating {per_genome_path}: {e}");
std::process::exit(1);
}));
writeln!(f, "genome,0,1,2,3").unwrap();
let mut column_totals = [0u64; 4];
writeln!(f, "genome,1,2,3,4").unwrap();
for (label, counts) in labels.iter().zip(stats.per_genome.iter()) {
writeln!(f, "{label},{},{},{},{}", counts[0], counts[1], counts[2], counts[3]).unwrap();
for s in 0..4 { column_totals[s] += counts[s]; }
}
writeln!(
f, "total,{},{},{},{}",
column_totals[0], column_totals[1], column_totals[2], column_totals[3],
f, "global,{},{},{},{}",
stats.counts[0], stats.counts[1], stats.counts[2], stats.counts[3],
).unwrap();
info!("per-genome sibling-count distribution → {per_genome_path}");
let total: u64 = stats.counts.iter().sum();
info!("family-size distribution → {path} (total {total} famil{})",
if total == 1 { "y" } else { "ies" });
}
// ── UPGMA Newick from kodama dendrogram ───────────────────────────────────────