Refactor k-mer sibling tracking to compact bitmask and on-demand counts
Replaces the explicit `SiblingInfo` struct and 3-bit minorant flags with a derived 4-bit presence mask (`FamilyMask`) that tracks observed bases per family. This eliminates redundant file I/O overhead by introducing a `PartitionCache` for batch lookups, simplifies serialization, and updates all downstream builders, stats computation, and tests to operate on the new bitmask representation. Adjusts CLI output to report deduplicated family sizes instead of histograms, ignores generated CSV files, and updates documentation to reflect the fixed canonical reference and new theory.
This commit is contained in:
Eric Coissac
@@ -9,6 +9,7 @@ data-stress
|
||||
./**/*.json
|
||||
*.bin
|
||||
*.log
|
||||
*.csv
|
||||
Betula_exilis--IGA-24-33
|
||||
benchmark/genomes
|
||||
benchmark/simulated_data
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 ───────────────────────────────────────
|
||||
|
||||
Reference in New Issue
Block a user