OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 27 Wiki Activity

Extract phylogenetic sibling logic into new obikphylo crate

Browse Source 
Relocate the `siblings` and `cardcomp` modules from `obikindex` to a dedicated `obikphylo` workspace member. Convert inherent methods on `KmerIndex` into extension traits, update import paths across `obikmer`, and add supporting accessor methods to `obikseq` and `obilayeredmap`. This restructuring reduces the public API surface of `obikindex` while organizing phylogenetic iteration, caching, and distance calculation logic under a dedicated crate.
This commit is contained in:
Eric Coissac
2026-08-16 14:30:34 +02:00
parent 519195d4a1
commit 5997de6707
31 changed files with 831 additions and 547 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/Cargo.lock
Generated
+22 -2
View File
@@ -1768,9 +1768,7 @@ dependencies = [
"obikpartitionner",
"obikseq",
"obilayeredmap",
"obipipeline",
"obiread",
"obiskbuilder",
"obiskio",
"obisys",
"obitaxonomy",
@@ -1794,6 +1792,7 @@ dependencies = [
"obifastwrite",
"obikindex",
"obikpartitionner",
"obikphylo",
"obikrope",
"obikseq",
"obilayeredmap",
@@ -1843,6 +1842,27 @@ dependencies = [
"tracing",
]
[[package]]
name = "obikphylo"
version = "0.1.0"
dependencies = [
"ndarray",
"obicompactvec",
"obikindex",
"obikpartitionner",
"obikseq",
"obilayeredmap",
"obipipeline",
"obiread",
"obiskbuilder",
"obiskio",
"obisys",
"rayon",
"tempfile",
"tracing",
"tracing-subscriber",
]
[[package]]
name = "obikrope"
version = "0.1.0"
src/Cargo.toml
+1 -1
View File
@@ -1,5 +1,5 @@
[workspace]
resolver = "3"
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy"]
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy", "obikphylo"]
[profile.release]
debug = 1
src/obikindex/Cargo.toml
-2
View File
@@ -11,8 +11,6 @@ obiskio = { path = "../obiskio" }
obisys = { path = "../obisys" }
obicompactvec = { path = "../obicompactvec" }
obilayeredmap = { path = "../obilayeredmap" }
obiskbuilder = { path = "../obiskbuilder" }
obipipeline = { path = "../obipipeline" }
ndarray = "0.16"
rayon = "1"
crossbeam-channel = "0.5"
src/obikindex/src/index.rs
+4
View File
@@ -132,6 +132,10 @@ impl KmerIndex {
IndexState::detect(&self.root_path).unwrap_or(IndexState::Empty)
}
/// The index's root directory — needed by out-of-crate extension code
/// (e.g. `obikphylo`) that opens its own `KmerPartition` handle onto
/// the same on-disk index.
pub fn root_path(&self) -> &Path { &self.root_path }
pub fn meta(&self) -> &IndexMeta { &self.meta }
pub fn meta_mut(&mut self) -> &mut IndexMeta { &mut self.meta }
pub fn kmer_size(&self) -> usize { self.meta.config.kmer_size }
src/obikindex/src/lib.rs
-4
View File
@@ -2,7 +2,6 @@ pub mod error;
pub mod meta;
pub mod predicate;
pub mod state;
mod cardcomp;
mod distance;
mod dump;
mod index;
@@ -11,7 +10,6 @@ mod numa;
mod rebuild;
mod reindex;
mod select;
mod siblings;
mod stats;
pub use error::{OKIError, OKIResult};
@@ -22,5 +20,3 @@ pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME
pub use predicate::{GroupFilterParams, MetaPred};
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
pub use stats::IndexBitsPerKmer;
pub use siblings::{BasePairTally, CardinalityTally, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
src/obikindex/src/siblings/cache.rs
-166
View File
@@ -1,166 +0,0 @@
use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obisys::progress_bar;
use crate::error::OKIResult;
use super::{olm_to_ok, INDEX_SUBDIR};
/// 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`.
pub(super) enum Mat {
Count(PersistentCompactIntMatrix),
Presence(PersistentBitMatrix),
}
impl Mat {
pub(super) fn n_cols(&self) -> usize {
match self {
Mat::Count(m) => m.n_cols(),
Mat::Presence(m) => m.n_cols(),
}
}
pub(super) 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,
}
}
}
pub(super) 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
/// [`crate::index::KmerIndex::build_sibling_annex`] (`layers` only) and
/// [`crate::index::KmerIndex::sibling_annex_stats`] (both).
layers: Vec<Vec<MphfLayer>>,
mats: Vec<Vec<Mat>>,
}
impl PartitionCache {
pub(super) 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}"));
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.
pub(super) 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()))
}
/// Resolve many `(variant, family_idx, base)` queries against one
/// partition's matrices at once, calling `on_hit(family_idx, base, g)`
/// for every genome `g` that carries the resolved variant.
///
/// Genome-major, not query-major: `PersistentBitMatrix`/
/// `PersistentCompactIntMatrix` are stored one contiguous block per
/// genome (column), slot as the offset within it (see
/// `obicompactvec::bitmatrix::packed::PackedBitMatrix` — each column's
/// own `mmap` region). Resolving query-by-query (`for query { for genome
/// { mat.carries(genome, slot) } }`, this function's predecessor) visits
/// every genome's block once *per query* — for a batch of thousands of
/// queries against ~90 genomes, that is thousands of jumps into each of
/// ~90 widely separated multi-MB regions, in query order, not genome
/// order: the access pattern a column-major layout is least suited to.
/// Grouping first (by layer, since each layer's matrix is a separate
/// column set) and sorting each group by slot, then visiting genome by
/// genome, turns that into ~90 mostly-sequential sweeps through one
/// column's own bytes — the layout's fast axis — confirmed by sampling
/// a real run: `PersistentBitMatrix::get` dominated wall-clock time,
/// mostly blocked on page faults, even after every partition/batch
/// locality fix above this in the traversal.
pub(super) fn find_presence_batch(
&self,
dest_partition: usize,
queries: &[(CanonicalKmer, usize, u8)],
n_genomes: usize,
mut on_hit: impl FnMut(usize, u8, usize),
) {
let Some(layers) = self.layers.get(dest_partition) else { return };
let Some(mats) = self.mats.get(dest_partition) else { return };
// First hit wins, same semantics as the old per-query loop (a
// variant present in an earlier layer shadows later ones).
let mut by_layer: Vec<Vec<(usize, usize, u8)>> = vec![Vec::new(); layers.len()];
for &(variant, family_idx, base) in queries {
for (li, mphf) in layers.iter().enumerate() {
if let Some(slot) = mphf.find(variant) {
by_layer[li].push((slot, family_idx, base));
break;
}
}
}
for (li, mut hits) in by_layer.into_iter().enumerate() {
if hits.is_empty() {
continue;
}
hits.sort_unstable_by_key(|&(slot, _, _)| slot);
let mat = &mats[li];
let n_cols = mat.n_cols().min(n_genomes);
for g in 0..n_cols {
for &(slot, family_idx, base) in &hits {
if mat.carries(g, slot) {
on_hit(family_idx, base, g);
}
}
}
}
}
}
src/obikindex/src/siblings/helpers.rs
-49
View File
@@ -1,49 +0,0 @@
use obikseq::CanonicalKmer;
use obiskbuilder::rolling_stat::RollingStat;
use obicompactvec::FamilyMask;
/// 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]
pub(super) 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`.
pub(super) fn is_minorant(kmer: CanonicalKmer, mask: FamilyMask, k: usize) -> bool {
kmer.central_canonical_neighbors().into_iter().all(|other| {
other == kmer || !mask.has(central_base(other, k)) || kmer.raw() <= other.raw()
})
}
/// Minimiser of a single, isolated canonical k-mer (not part of a streamed
/// sequence). `RollingStat` computes minimisers incrementally along a
/// sequence; this feeds one k-mer's bases through a fresh instance to get
/// the same selection for a single, disconnected k-mer. Not the leanest
/// possible primitive (an O(1)-amortised dedicated scan, as originally
/// sketched in the design doc's Step 0, would avoid the ASCII round-trip and
/// `RollingStat` allocation) but correct and reuses already-tested logic;
/// left as a follow-up optimisation.
fn lone_kmer_minimizer(kmer: CanonicalKmer) -> obikseq::Minimizer {
let ascii = kmer.to_ascii();
let mut rs = RollingStat::new(0);
for b in ascii {
rs.push(b);
}
rs.canonical_minimizer()
.expect("RollingStat must be ready after k bases of a valid k-mer")
}
/// Destination partition for a (possibly synthetic) canonical k-mer, using
/// the same routing rule as the rest of the index (`minimiser.seq_hash() &
/// mask`, `n_partitions` is a power of two).
pub(super) fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize {
let mask = (n_partitions as u64) - 1;
(lone_kmer_minimizer(kmer).seq_hash() & mask) as usize
}
src/obikindex/src/siblings/iter.rs
-78
View File
@@ -1,78 +0,0 @@
//! Phylo/sibling-domain iteration over a layer — an extension trait, not a
//! new field on `MphfLayer`/`Layer<D>`: "family"/"minorant" are phylo
//! concepts, `obilayeredmap` stays kmer/slot-mapping only (see
//! `docmd/architecture/siblings.md`).
//!
//! The sibling annex is persisted in the same order as `iter_kmers()`
//! (`build_sibling_annex`, see `build.rs`), so pairing them is a plain zip —
//! no MPHF, no slot, no `kmer_at`. Both sides are already `Send + 'static`
//! (`KmerIter` owns an `Arc<UnitigFileReader>` clone; `SiblingAnnex` is
//! mmap-backed and handed in as an `Arc` by the caller), so `SiblingIter`
//! streams straight from disk and can be fed to `obipipeline` batch by
//! batch — never collected whole into memory (see the project's "no full
//! collect" rule).
use std::sync::Arc;
use obicompactvec::{FamilyMask, SiblingAnnex};
use obikseq::CanonicalKmer;
use obilayeredmap::{KmerIter, MphfLayer};
/// One layer entry: a k-mer's position in the layer's iteration order (the
/// same index the sibling annex is keyed on — not an MPHF slot), the k-mer
/// itself, and its family mask.
#[derive(Debug, Clone, Copy)]
pub struct SiblingEntry {
pub order: usize,
pub kmer: CanonicalKmer,
pub mask: FamilyMask,
}
/// Streams `(order, kmer, mask)` triples for one layer, in iteration order.
/// Produced by [`SiblingLayerExt::iter_siblings`].
pub struct SiblingIter {
kmers: KmerIter,
annex: Arc<SiblingAnnex>,
order: usize,
}
impl Iterator for SiblingIter {
type Item = SiblingEntry;
fn next(&mut self) -> Option<Self::Item> {
loop {
let kmer = self.kmers.next()?;
let order = self.order;
self.order += 1;
// `None` means "not yet computed" (see `SiblingAnnex` module
// docs) — shouldn't happen against a fully-built annex, but
// skip rather than misalign the two streams if it does.
if let Some(mask) = self.annex.get(order) {
return Some(SiblingEntry { order, kmer, mask });
}
}
}
}
/// Adds phylo/sibling iteration to `MphfLayer`.
pub trait SiblingLayerExt {
/// Zip this layer's k-mers with their sibling-annex entry, in iteration
/// order. `annex` must have been built from this same layer (its length
/// must match the layer's k-mer count).
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter;
/// Like [`iter_siblings`](Self::iter_siblings), filtered to the
/// minorant of each family — the common case, since a family is
/// tallied once, at its minorant.
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> impl Iterator<Item = SiblingEntry>;
}
impl SiblingLayerExt for MphfLayer {
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter {
SiblingIter { kmers: self.iter_kmers(), annex, order: 0 }
}
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> impl Iterator<Item = SiblingEntry> {
self.iter_siblings(annex).filter(|e| e.mask.is_minorant())
}
}
src/obikmer/Cargo.toml
+1
View File
@@ -19,6 +19,7 @@ obikpartitionner = { path = "../obikpartitionner" }
obisys = { path = "../obisys" }
obiskio = { path = "../obiskio" }
obikindex = { path = "../obikindex", default-features = false }
obikphylo = { path = "../obikphylo" }
obitaxonomy = { path = "../obitaxonomy" }
obilayeredmap = { path = "../obilayeredmap" }
clap = { version = "4", features = ["derive"] }
src/obikmer/src/cmd/phylo/family_overlap.rs
+2 -1
View File
@@ -1,7 +1,8 @@
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use obikindex::{KmerIndex, SnpAlignment};
use obikindex::KmerIndex;
use obikphylo::siblings::{SnpAlignment, SnpAlignmentExt};
use tracing::info;
// ── Family overlap: shared-family counts and the `--min-shared-family` /
src/obikmer/src/cmd/phylo/iqtree.rs
+1 -1
View File
@@ -2,7 +2,7 @@ use std::io::{BufWriter, Write};
use std::path::PathBuf;
use obifastwrite::{JsonVal, write_record};
use obikindex::SnpAlignment;
use obikphylo::siblings::SnpAlignment;
use tracing::info;
use super::sankoff::state_index_table;
src/obikmer/src/cmd/phylo/mod.rs
+6 -2
View File
@@ -9,9 +9,13 @@ mod tnt;
use std::io::{self, BufWriter, Write};
use kodama::{Method, linkage};
use obikindex::{
KmerIndex, RawSnpDistanceOutput, SnpAlignment,
use obikindex::KmerIndex;
use obikphylo::{
cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix,
siblings::{
CardinalityExt, DistanceExt, RawSnpDistanceOutput, SiblingAnnexBuildExt, SiblingStatsExt,
SnpAlignment, SnpAlignmentExt,
},
};
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
use tracing::info;
src/obikmer/src/cmd/phylo/outputs.rs
+1 -1
View File
@@ -2,7 +2,7 @@ use std::io::{BufWriter, Write};
use std::path::PathBuf;
use obifastwrite::{JsonVal, write_record};
use obikindex::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
use obikphylo::siblings::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
use tracing::info;
// ── Family-size distribution → CSV ──────────────────────────────────────────
src/obikmer/src/cmd/phylo/sankoff.rs
+1 -1
View File
@@ -2,7 +2,7 @@ use std::io::{BufWriter, Write};
use std::path::PathBuf;
use obifastwrite::{JsonVal, write_record};
use obikindex::{BasePairTally, CardinalityTally, SnpAlignment};
use obikphylo::siblings::{BasePairTally, CardinalityTally, SnpAlignment};
use tracing::info;
// ── Sankoff pseudo-alignment → FASTA ────────────────────────────────────────
src/obikmer/src/cmd/phylo/tnt.rs
+1 -1
View File
@@ -1,7 +1,7 @@
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use obikindex::SnpAlignment;
use obikphylo::siblings::SnpAlignment;
use tracing::info;
use super::sankoff::{scaled_metric_matrix, state_index_table};
src/obikphylo/Cargo.toml
+23
View File
@@ -0,0 +1,23 @@
[package]
name = "obikphylo"
version = "0.1.0"
edition = "2024"
[dependencies]
obikindex = { path = "../obikindex" }
obikseq = { path = "../obikseq" }
obikpartitionner = { path = "../obikpartitionner" }
obiskio = { path = "../obiskio" }
obisys = { path = "../obisys" }
obicompactvec = { path = "../obicompactvec" }
obilayeredmap = { path = "../obilayeredmap" }
obiskbuilder = { path = "../obiskbuilder" }
obipipeline = { path = "../obipipeline" }
ndarray = "0.16"
rayon = "1"
tracing = "0.1.44"
[dev-dependencies]
obiread = { path = "../obiread" }
tempfile = "3"
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
src/obikindex/src/cardcomp.rs → src/obikphylo/src/cardcomp.rs
+1 -1
View File
@@ -17,7 +17,7 @@
//! since TNT/PhyG search *unrooted* trees) if the cost matrix is
//! symmetric, the discrete-parsimony analogue of CTMC reversibility.
use crate::{BasePairTally, CardinalityTally};
use crate::siblings::{BasePairTally, CardinalityTally};
/// Row-stochastic 5×5 cardinality transition probabilities (`0..=4`),
/// diagonal included ("stay at the same cardinality"), from
src/obikphylo/src/lib.rs
+14
View File
@@ -0,0 +1,14 @@
//! Library-level phylogenetic functionality for `obikmer`, built as
//! extension traits over `obikindex::KmerIndex` and `obilayeredmap`'s
//! generic layer types — the `phylo` CLI command is a consumer of this
//! crate, not the owner of this logic (see `docmd/architecture/siblings.md`).
//!
//! Starts with [`siblings`] (family presence-mask annex, SNP distance,
//! cardinality, pseudo-alignment); further phylo-domain functionality
//! (currently `obikindex::distance`/`obikindex::cardcomp`) moves here
//! incrementally.
mod cardcomp;
pub mod siblings;
pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
src/obikindex/src/siblings/alignment.rs → src/obikphylo/src/siblings/alignment.rs
+16 -9
View File
@@ -3,8 +3,8 @@ use std::sync::Arc;
use obikpartitionner::KmerPartition;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;
use super::cache::PartitionCache;
use super::family_scan::scan_layer_families;
@@ -53,25 +53,32 @@ pub struct SnpAlignment {
pub sequences: Vec<Vec<u8>>,
}
impl KmerIndex {
/// Adds [`snp_pseudo_alignment`](Self::snp_pseudo_alignment) to `KmerIndex` —
/// phylo-domain functionality, kept out of `obikindex` itself (see
/// `docmd/architecture/siblings.md`).
pub trait SnpAlignmentExt {
/// Build the SNP-only pseudo-alignment from an already-built sibling
/// annex (run [`build_sibling_annex`](Self::build_sibling_annex) first).
pub fn snp_pseudo_alignment(&self) -> OKIResult<SnpAlignment> {
/// annex (run [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex) first).
fn snp_pseudo_alignment(&self) -> OKIResult<SnpAlignment>;
}
impl SnpAlignmentExt for KmerIndex {
fn snp_pseudo_alignment(&self) -> OKIResult<SnpAlignment> {
let n_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len();
let with_counts = self.meta.config.with_counts;
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config(
&self.root_path,
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let layer_dirs = self.sibling_layer_dirs()?;
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers");
// One layer at a time, not `par_iter()` over layers — same
src/obikindex/src/siblings/build.rs → src/obikphylo/src/siblings/build.rs
+25 -18
View File
@@ -12,11 +12,11 @@ use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;
use super::cache::PartitionCache;
use super::helpers::{central_base, is_minorant, partition_of};
use super::helpers::{central_base, is_minorant};
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
// ── obipipeline data types ─────────────────────────────────────────────────
@@ -53,7 +53,8 @@ enum SibData {
Variants(VariantBatch),
}
impl KmerIndex {
/// Adds [`build_sibling_annex`](Self::build_sibling_annex) to `KmerIndex`.
pub trait SiblingAnnexBuildExt {
/// Build the sibling-count/minorant annex for every layer of every
/// partition of this (already built) index, writing one annex file per
/// layer alongside its existing index files. Safe to call again later
@@ -63,18 +64,23 @@ impl KmerIndex {
/// Construction only — no statistics gathered here on purpose: this is
/// meant to run routinely (it is the artefact the SNP-family distances
/// will consume), while the sibling-count distribution
/// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate,
/// occasional diagnostic pass over the result, not run every time.
/// ([`sibling_annex_stats`](super::stats::SiblingStatsExt::sibling_annex_stats))
/// is a separate, occasional diagnostic pass over the result, not run
/// every time.
///
/// Cross-partition/cross-layer lookups are required (a k-mer's siblings
/// can live in any partition), but the layer loop itself — and thus the
/// annex file this produces — stays local to one layer at a time.
pub fn build_sibling_annex(&self) -> OKIResult<()> {
fn build_sibling_annex(&self) -> OKIResult<()>;
}
impl SiblingAnnexBuildExt for KmerIndex {
fn build_sibling_annex(&self) -> OKIResult<()> {
let n_parts = self.n_partitions();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config(
&self.root_path,
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
@@ -82,7 +88,7 @@ impl KmerIndex {
.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 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;
@@ -97,7 +103,7 @@ impl KmerIndex {
let mut part_slots: u64 = 0;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?;
part_slots += build_layer_sibling_annex(self, &layer_dir, n_parts, &cache)?;
}
total_slots += part_slots;
pb.inc(1);
@@ -108,19 +114,21 @@ impl KmerIndex {
Ok(())
}
}
/// Returns the number of distinct k-mers (annex entries) processed, for
/// progress reporting.
fn build_layer_sibling_annex(
&self,
/// Returns the number of distinct k-mers (annex entries) processed, for
/// progress reporting. A free function, not a `KmerIndex` method — called
/// only from `build_sibling_annex` above, in the same file.
fn build_layer_sibling_annex(
index: &KmerIndex,
layer_dir: &Path,
n_parts: usize,
cache: &Arc<PartitionCache>,
) -> OKIResult<u64> {
) -> 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)?;
let k = self.kmer_size();
let k = index.kmer_size();
let n = mphf.n();
// ── Reconciliation state, indexed by this layer's k-mer iteration
@@ -211,7 +219,7 @@ impl KmerIndex {
continue;
}
items.push((
partition_of(variant, n_parts),
variant.partition(n_parts),
variant,
order,
central_base(variant, k),
@@ -273,5 +281,4 @@ impl KmerIndex {
builder.close()?;
Ok(n as u64)
}
}
src/obikphylo/src/siblings/cache.rs
+225
View File
@@ -0,0 +1,225 @@
use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::Layer;
use obilayeredmap::meta::PartitionMeta;
use obisys::progress_bar;
use obikindex::OKIResult;
use super::iter::SiblingLayerExt;
use super::{olm_to_ok, INDEX_SUBDIR};
/// Every partition's already-open 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.
///
/// One `Layer<D>` per layer (MPHF + matrix bundled), not a separate
/// `MphfLayer` and a separate `PersistentCompactIntMatrix`/
/// `PersistentBitMatrix` in parallel arrays — `obilayeredmap::Layer` already
/// *is* that bundle, with `find_slot` (MPHF-only, no data read),
/// `n_cols`/`sub_matrix`/`fill_sub_matrix` (batched, sorted-internally
/// column access) on top of it. Reinventing that pairing here would just be
/// going back through the low-level pieces `Layer` already assembles.
pub(super) enum Mat {
Count(Layer<PersistentCompactIntMatrix>),
Presence(Layer<PersistentBitMatrix>),
}
impl Mat {
fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
match self {
Mat::Count(l) => l.find_slot(kmer),
Mat::Presence(l) => l.find_slot(kmer),
}
}
/// Raw MPHF batch lookup: kmer → slot, no membership check — for
/// callers that already know every kmer is a member of *this* layer
/// (e.g. it came from this layer's own `iter_minorants_batch`), so the
/// evidence check `find_slot`/`find` would perform is redundant work.
/// See `docmd/architecture/siblings.md`: iteration-pipeline kmers use
/// `index`, never `find`.
pub(super) fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
match self {
Mat::Count(l) => l.index_batch(kmers),
Mat::Presence(l) => l.index_batch(kmers),
}
}
/// This layer's own `SiblingLayerExt::iter_minorants_batch` — dispatch
/// only, both arms return the same concrete `MinorantBatchIter` (it
/// doesn't depend on `D`), so no boxing is needed.
pub(super) fn iter_minorants_batch(
&self,
annex: std::sync::Arc<obicompactvec::SiblingAnnex>,
batch_size: usize,
) -> super::iter::MinorantBatchIter {
match self {
Mat::Count(l) => l.iter_minorants_batch(annex, batch_size),
Mat::Presence(l) => l.iter_minorants_batch(annex, batch_size),
}
}
pub(super) fn n_cols(&self) -> usize {
match self {
Mat::Count(l) => l.n_cols(),
Mat::Presence(l) => l.n_cols(),
}
}
/// Batch, genome-major "carries" for a set of `slots` — `out[g][i]` =
/// whether genome `g` (0..`out.len()`) carries `slots[i]`. `out` must
/// have one entry per genome column, each resized to `slots.len()`.
///
/// Delegates entirely to `Layer<D>::fill_sub_matrix`, which sorts
/// `slots` once internally for a sequential mmap sweep per column, then
/// restores the original order — the same discipline this call site
/// (and `find_presence_batch`) used to hand-roll with its own sort +
/// genome-major loop. `Count` still needs one intermediate
/// `Vec<Vec<u32>>` fetch (the underlying store only has an int
/// sub-matrix, not a bool one), converted to presence (`!= 0`) in place
/// — the sort/sequential-access win is unaffected, just one extra
/// allocation pass over already-in-hand data.
pub(super) fn fill_sub_matrix_carries(&self, slots: &[usize], out: &mut [Vec<bool>]) {
match self {
Mat::Presence(l) => l.fill_sub_matrix(slots, out),
Mat::Count(l) => {
let mut counts: Vec<Vec<u32>> = out.iter().map(|_| Vec::new()).collect();
l.fill_sub_matrix(slots, &mut counts);
for (o, c) in out.iter_mut().zip(counts.iter()) {
o.clear();
o.extend(c.iter().map(|&v| v != 0));
}
}
}
}
}
pub(super) struct PartitionCache {
/// `mats[partition][layer]` = that partition's opened layers. Used by
/// both [`obikindex::KmerIndex::build_sibling_annex`] and
/// [`obikindex::KmerIndex::sibling_annex_stats`].
mats: Vec<Vec<Mat>>,
}
impl PartitionCache {
pub(super) 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<Mat>> = (0..n_parts)
.into_par_iter()
.map(|part| -> OKIResult<Vec<Mat>> {
let index_dir = partition.part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
pb.inc(1);
return Ok(Vec::new());
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
let mut mats = Vec::with_capacity(meta.n_layers);
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Layer::<PersistentCompactIntMatrix>::open(&layer_dir, &meta.mode).ok().map(Mat::Count)
} else {
Layer::<PersistentBitMatrix>::open(&layer_dir, &meta.mode).ok().map(Mat::Presence)
};
let Some(mat) = mat else { continue };
mats.push(mat);
}
pb.inc(1);
Ok(mats)
})
.collect::<OKIResult<Vec<_>>>()?;
pb.finish_and_clear();
Ok(Self { mats: built })
}
/// 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. `find_slot`, not `sub_matrix`/`carries` — no data
/// read needed for a plain existence check.
pub(super) fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> bool {
self.mats
.get(dest_partition)
.is_some_and(|mats| mats.iter().any(|mat| mat.find_slot(variant).is_some()))
}
/// Resolve many `(variant, family_idx, base)` queries against one
/// partition's matrices at once, calling `on_hit(family_idx, base, g)`
/// for every genome `g` that carries the resolved variant.
///
/// Genome-major, not query-major: `PersistentBitMatrix`/
/// `PersistentCompactIntMatrix` are stored one contiguous block per
/// genome (column), slot as the offset within it (see
/// `obicompactvec::bitmatrix::packed::PackedBitMatrix` — each column's
/// own `mmap` region). Resolving query-by-query (`for query { for genome
/// { mat.carries(genome, slot) } }`, this function's predecessor) visits
/// every genome's block once *per query* — for a batch of thousands of
/// queries against ~90 genomes, that is thousands of jumps into each of
/// ~90 widely separated multi-MB regions, in query order, not genome
/// order: the access pattern a column-major layout is least suited to.
/// Grouping first (by layer, since each layer's matrix is a separate
/// column set) and letting `fill_sub_matrix_carries` sort each group by
/// slot internally, then visiting genome by genome, turns that into ~90
/// mostly-sequential sweeps through one column's own bytes — the
/// layout's fast axis — confirmed by sampling a real run:
/// `PersistentBitMatrix::get` dominated wall-clock time, mostly blocked
/// on page faults, even after every partition/batch locality fix above
/// this in the traversal.
pub(super) fn find_presence_batch(
&self,
dest_partition: usize,
queries: &[(CanonicalKmer, usize, u8)],
n_genomes: usize,
mut on_hit: impl FnMut(usize, u8, usize),
) {
let Some(mats) = self.mats.get(dest_partition) else { return };
// First hit wins, same semantics as the old per-query loop (a
// variant present in an earlier layer shadows later ones).
let mut by_layer: Vec<Vec<(usize, usize, u8)>> = vec![Vec::new(); mats.len()];
for &(variant, family_idx, base) in queries {
for (li, mat) in mats.iter().enumerate() {
if let Some(slot) = mat.find_slot(variant) {
by_layer[li].push((slot, family_idx, base));
break;
}
}
}
for (li, hits) in by_layer.into_iter().enumerate() {
if hits.is_empty() {
continue;
}
let mat = &mats[li];
let n_cols = mat.n_cols().min(n_genomes);
let slots: Vec<usize> = hits.iter().map(|&(slot, _, _)| slot).collect();
let mut carries: Vec<Vec<bool>> = (0..n_cols).map(|_| Vec::new()).collect();
mat.fill_sub_matrix_carries(&slots, &mut carries);
for (g, col) in carries.iter().enumerate() {
for (&(_, family_idx, base), &carries_it) in hits.iter().zip(col.iter()) {
if carries_it {
on_hit(family_idx, base, g);
}
}
}
}
}
}
src/obikindex/src/siblings/cardinality.rs → src/obikphylo/src/siblings/cardinality.rs
+13 -8
View File
@@ -5,8 +5,8 @@ use ndarray::Array2;
use obikpartitionner::KmerPartition;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;
use super::cache::PartitionCache;
use super::distance::RawSnpDistanceOutput;
@@ -22,7 +22,8 @@ pub struct CardinalityTally {
pub counts: [[u64; 5]; 5],
}
impl KmerIndex {
/// Adds [`cardinality_tally`](Self::cardinality_tally) to `KmerIndex`.
pub trait CardinalityExt {
/// Cardinality co-occurrence, pooled only over genome pairs whose
/// overall SNP ratio in `raw` is at or below `ratio_ceiling` — same
/// saturation/no-data exclusion discipline as
@@ -51,10 +52,14 @@ impl KmerIndex {
/// gets the matching restriction via `scan_family_pairs`'s new
/// `variable` flag, rather than a `family_size()` check of its own (it
/// doesn't have direct access to the family's mask).
pub fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<CardinalityTally> {
fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<CardinalityTally>;
}
impl CardinalityExt for KmerIndex {
fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<CardinalityTally> {
let n_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len();
let with_counts = self.meta.config.with_counts;
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;
@@ -68,14 +73,14 @@ impl KmerIndex {
});
let partition = KmerPartition::open_with_config(
&self.root_path,
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let layer_dirs = self.sibling_layer_dirs()?;
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers");
// One layer at a time — see `snp_pseudo_alignment`'s comment for why
src/obikindex/src/siblings/distance.rs → src/obikphylo/src/siblings/distance.rs
+68 -58
View File
@@ -5,8 +5,8 @@ use ndarray::Array2;
use obikpartitionner::KmerPartition;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;
use super::cache::PartitionCache;
use super::family_scan::scan_layer_families;
@@ -35,54 +35,52 @@ pub struct RawSnpDistanceOutput {
pub shared: Array2<u64>,
}
impl KmerIndex {
/// Shared traversal behind [`raw_snp_distance`](Self::raw_snp_distance)
/// and [`base_pair_tally`](Self::base_pair_tally): for every family
/// (tallied once, at its minorant) of every layer of the already-built
/// sibling annex, resolves each genome's single observed form (`None`
/// if absent or ambiguous/multi-copy), then calls `on_pair(acc, i, j,
/// bi, bj, variable)` for every genome pair `(i, j)` where both are
/// unambiguous and single-copy (`bi == bj` means shared at that locus,
/// `bi != bj` means a SNP). `variable` is the family's own
/// `family_size() >= 2` (true if more than one member is observed
/// *anywhere* in the family, i.e. it isn't fully invariant across the
/// whole index) — `raw_snp_distance` ignores it (a fully-invariant
/// family is still legitimately "shared"), but callers whose diagonal
/// should only reflect genuine SNP-adjacent agreement, not the
/// genome-wide invariant background, need it (see
/// [`base_pair_tally`](Self::base_pair_tally)'s `same` field). Layers
/// are processed one at a time, not in parallel — see
/// `snp_pseudo_alignment`'s comment for why `par_iter()` over layers
/// would defeat `scan_layer_families`'s partition-grouped locality;
/// each layer gets its own accumulator from `zero()`, combined
/// pairwise via `combine`.
fn scan_family_pairs<Acc, F, C>(
&self,
/// Shared traversal behind [`DistanceExt::raw_snp_distance`] and
/// [`DistanceExt::base_pair_tally`]: for every family (tallied once, at its
/// minorant) of every layer of the already-built sibling annex, resolves
/// each genome's single observed form (`None` if absent or
/// ambiguous/multi-copy), then calls `on_pair(acc, i, j, bi, bj, variable)`
/// for every genome pair `(i, j)` where both are unambiguous and
/// single-copy (`bi == bj` means shared at that locus, `bi != bj` means a
/// SNP). `variable` is the family's own `family_size() >= 2` (true if more
/// than one member is observed *anywhere* in the family, i.e. it isn't
/// fully invariant across the whole index) — `raw_snp_distance` ignores it
/// (a fully-invariant family is still legitimately "shared"), but callers
/// whose diagonal should only reflect genuine SNP-adjacent agreement, not
/// the genome-wide invariant background, need it (see
/// [`DistanceExt::base_pair_tally`]'s `same` field). Layers are processed
/// one at a time, not in parallel — see `snp_pseudo_alignment`'s comment
/// for why `par_iter()` over layers would defeat `scan_layer_families`'s
/// partition-grouped locality; each layer gets its own accumulator from
/// `zero()`, combined pairwise via `combine`. A free function, not a
/// `KmerIndex` method — called only from this file's `DistanceExt` impl.
fn scan_family_pairs<Acc, F, C>(
index: &KmerIndex,
label: &str,
zero: impl Fn() -> Acc + Sync,
on_pair: F,
combine: C,
) -> OKIResult<Acc>
where
) -> OKIResult<Acc>
where
Acc: Send,
F: Fn(&mut Acc, usize, usize, u8, u8, bool) + Sync,
C: Fn(Acc, Acc) -> Acc,
{
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_parts = index.n_partitions();
let n_genomes = index.meta().genomes.len();
let with_counts = index.meta().config.with_counts;
let k = index.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config(
&self.root_path,
self.kmer_size(),
self.minimizer_size(),
index.root_path(),
index.kmer_size(),
index.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let layer_dirs = self.sibling_layer_dirs()?;
let layer_dirs = super::family_scan::sibling_layer_dirs(index)?;
let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
// One layer at a time — see `snp_pseudo_alignment`'s comment for why
@@ -121,11 +119,36 @@ impl KmerIndex {
Ok(total)
}
/// Adds [`raw_snp_distance`](Self::raw_snp_distance) and
/// [`base_pair_tally`](Self::base_pair_tally) to `KmerIndex`.
pub trait DistanceExt {
/// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex
/// (run [`build_sibling_annex`](Self::build_sibling_annex) first).
pub fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> {
let n_genomes = self.meta.genomes.len();
let (snp, shared) = self.scan_family_pairs(
/// (run [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex) first).
fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput>;
/// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG,
/// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs
/// whose overall SNP ratio in `raw` is at or below `ratio_ceiling` —
/// same saturation-exclusion discipline as
/// [`cardinality_tally`](super::cardinality::CardinalityExt::cardinality_tally),
/// for the same reason: a saturated pair's observed base-pair mix
/// trends toward neutral base composition, not the true point-mutation
/// spectrum.
///
/// A second full pass over the annex, sharing
/// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by
/// it, not a blind re-scan) — needed because `raw_snp_distance` only
/// keeps aggregate SNP/shared counts per genome pair, not which bases
/// were actually involved at each locus, and the ratio-ceiling filter
/// can only be evaluated once the aggregate counts are known.
fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally>;
}
impl DistanceExt for KmerIndex {
fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> {
let n_genomes = self.meta().genomes.len();
let (snp, shared) = scan_family_pairs(
self,
"raw_snp_distance",
|| (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::zeros((n_genomes, n_genomes))),
|(snp, shared), i, j, bi, bj, _variable| {
@@ -146,22 +169,8 @@ impl KmerIndex {
Ok(RawSnpDistanceOutput { snp, shared })
}
/// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG,
/// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs
/// whose overall SNP ratio in `raw` is at or below `ratio_ceiling` —
/// same saturation-exclusion discipline as
/// [`cardinality_tally`](Self::cardinality_tally), for the same reason:
/// a saturated pair's observed base-pair mix trends toward neutral base
/// composition, not the true point-mutation spectrum.
///
/// A second full pass over the annex, sharing
/// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by
/// it, not a blind re-scan) — needed because `raw_snp_distance` only
/// keeps aggregate SNP/shared counts per genome pair, not which bases
/// were actually involved at each locus, and the ratio-ceiling filter
/// can only be evaluated once the aggregate counts are known.
pub fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally> {
let n_genomes = self.meta.genomes.len();
fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally> {
let n_genomes = self.meta().genomes.len();
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
if i == j {
return false;
@@ -171,7 +180,8 @@ impl KmerIndex {
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
});
let (counts, same) = self.scan_family_pairs(
let (counts, same) = scan_family_pairs(
self,
"base_pair_tally",
|| ([[0u64; 4]; 4], [0u64; 4]),
|(counts, same), i, j, bi, bj, variable| {
@@ -206,7 +216,7 @@ impl KmerIndex {
}
}
/// See [`KmerIndex::base_pair_tally`].
/// See [`DistanceExt::base_pair_tally`].
pub struct BasePairTally {
/// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled
/// over included genome pairs, where the two genomes' single forms are
src/obikindex/src/siblings/family_scan.rs → src/obikphylo/src/siblings/family_scan.rs
+67 -77
View File
@@ -55,15 +55,16 @@ use rayon::prelude::*;
use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::Layer;
use obilayeredmap::meta::PartitionMeta;
use obipipeline::{ThrottleGuard, throttle};
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;
use super::cache::{Mat, PartitionCache};
use super::helpers::{central_base, partition_of};
use super::helpers::central_base;
use super::iter::SiblingEntry;
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// Families per batch — see the module docs for the memory-vs-per-partition-
@@ -73,14 +74,17 @@ use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// queries per batch to amortise against.
const FAMILY_BATCH: usize = 65536;
impl KmerIndex {
/// Every (partition, layer) directory carrying a sibling annex, checked
/// up front so a missing one is reported before any real work starts.
pub(super) fn sibling_layer_dirs(&self) -> OKIResult<Vec<PathBuf>> {
let n_parts = self.n_partitions();
/// Every (partition, layer) directory carrying a sibling annex, checked
/// up front so a missing one is reported before any real work starts.
/// Shared by every sibling-annex consumer's extension-trait impl (`alignment`,
/// `build`, `cardinality`, `distance`, `stats`) — a free function, not a
/// `KmerIndex` method, since it is crate-internal only and `KmerIndex` lives
/// in `obikindex`, a foreign crate from here (orphan rule).
pub(crate) fn sibling_layer_dirs(index: &KmerIndex) -> OKIResult<Vec<PathBuf>> {
let n_parts = index.n_partitions();
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
let index_dir = index.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
@@ -98,17 +102,15 @@ impl KmerIndex {
}
}
Ok(layer_dirs)
}
}
/// Read-only state shared (via `Arc`) across every pipeline worker
/// generating this layer's batches — opened once, not per batch. No `cache`
/// here: generation never touches the cross-partition cache, only this
/// layer's own already-open matrix. No `annex` either — the filter step
/// that builds `minorant_slots` already carries each slot's mask through,
/// so a worker never needs to re-read it.
/// layer's own already-open matrix. No separate MPHF/`slot_kmer` either —
/// `mat` (a `Layer<D>`) already bundles the MPHF, and each `SiblingEntry`
/// arrives with its kmer and mask already in hand from `iter_minorants_batch`.
struct LayerCtx {
slot_kmer: Vec<Option<CanonicalKmer>>,
mat: Mat,
n_parts: usize,
n_genomes: usize,
@@ -118,10 +120,11 @@ struct LayerCtx {
struct SourceBatch {
start_family_idx: usize,
/// `(slot, mask)` — the mask is carried through from the filter step
/// below rather than re-read from the annex per batch; it's the same
/// byte either way, just already in hand.
slots: Vec<(usize, FamilyMask)>,
/// One entry per minorant family in this batch, straight from
/// `iter_minorants_batch` — `order` (iteration-order index, not an MPHF
/// slot; see `docmd/architecture/siblings.md`), `kmer`, and `mask`
/// already carried through, no second annex read.
entries: Vec<SiblingEntry>,
_permit: ThrottleGuard,
}
@@ -143,10 +146,10 @@ enum FamData {
Generated(GeneratedBatch),
}
/// Visits every minorant family of one layer, in slot order, batched
/// Visits every minorant family of one layer, in iteration order, batched
/// [`FAMILY_BATCH`] at a time — see the module docs for why generation and
/// resolution use different concurrency. `on_family` is called once per
/// family, in slot order, with its own annex mask and its per-genome
/// family, in iteration order, with its own annex mask and its per-genome
/// base-presence (`genome_mask[g]`: bit `b` set iff genome `g` carries the
/// member whose own canonical central base is `b`), backed by a scratch
/// buffer reused across every call — callers that need to keep data past
@@ -162,59 +165,38 @@ pub(super) fn scan_layer_families(
) -> OKIResult<()> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
// `MphfLayer::kmer_at` — direct slot -> k-mer reconstruction (evidence +
// direct-access unitigs, no MPHF hashing, no file scan) instead of the
// previous scan-`unitigs.bin`-and-hash-every-k-mer-forward approach,
// which redundantly re-read the same data `mphf.find()` itself already
// reads through `evidence`/`unitigs` to verify each hit.
let slot_kmer: Vec<Option<CanonicalKmer>> = (0..annex.len()).map(|slot| mphf.kmer_at(slot)).collect();
let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?);
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
Mat::Count(Layer::<PersistentCompactIntMatrix>::open(layer_dir, &meta.mode).map_err(olm_to_ok)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
Mat::Presence(Layer::<PersistentBitMatrix>::open(layer_dir, &meta.mode).map_err(olm_to_ok)?)
};
let n_cols = mat.n_cols().min(n_genomes);
// Minorant slots, in order, mask carried along — the flag is already
// stored in the annex (set once, at construction — see
// `build_layer_sibling_annex`), no need to re-derive it from the kmer
// here, and no need for the pipeline workers below to re-read the same
// byte from the annex a second time. `slot_kmer[s]` is still required
// to exist (used later for `central_canonical_neighbors()`), but not
// for this check.
let minorant_slots: Vec<(usize, FamilyMask)> = (0..annex.len())
.filter_map(|s| {
let mask = annex.get(s)?;
(slot_kmer[s].is_some() && mask.is_minorant()).then_some((s, mask))
})
.collect();
let total_families = minorant_slots.len();
if total_families == 0 {
return Ok(());
}
let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k });
let ctx = Arc::new(LayerCtx { slot_kmer, mat, n_parts, n_genomes, n_cols, k });
let mut offset = 0usize;
let batches: Vec<(usize, Vec<(usize, FamilyMask)>)> = minorant_slots
.chunks(FAMILY_BATCH)
.map(|chunk| {
let start = offset;
offset += chunk.len();
(start, chunk.to_vec())
})
.collect();
// Streamed straight from `iter_minorants_batch` (zips this layer's own
// `iter_kmers()` with the annex, both in iteration order — never an
// MPHF slot; see `docmd/architecture/siblings.md`) — never collected
// into a `Vec` first: a layer can hold billions of k-mers, so
// materialising every minorant family up front is exactly the memory
// blowup an earlier version of this traversal was rewritten to avoid
// (see the module docs). `.scan()` computes each batch's starting
// family index lazily, mirroring what the eager `chunks()`+running
// `offset` used to do.
let batches = ctx.mat.iter_minorants_batch(annex, FAMILY_BATCH).scan(0usize, |offset, batch| {
let start = *offset;
*offset += batch.len();
Some((start, batch))
});
let n_workers = obisys::effective_parallelism();
let capacity = 4;
let throttled = throttle(batches.into_iter(), n_workers).map(|t| SourceBatch {
let throttled = throttle(batches, n_workers).map(|t| SourceBatch {
start_family_idx: t.item.0,
slots: t.item.1,
entries: t.item.1,
_permit: t.guard,
});
@@ -224,17 +206,18 @@ pub(super) fn scan_layer_families(
| {
move |batch: SourceBatch| -> GeneratedBatch {
let ctx = &worker_ctx;
let n = batch.slots.len();
let n = batch.entries.len();
let mut masks = Vec::with_capacity(n);
let mut bases = Vec::with_capacity(n);
let mut genome_mask = vec![0u8; n * ctx.n_genomes];
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..ctx.n_parts).map(|_| Vec::new()).collect();
// Pass 1: cheap, no matrix access — own base and this
// batch's cross-partition queries. Mask already in hand
// from the filter step, no second annex read.
for (i, &(fslot, mask)) in batch.slots.iter().enumerate() {
let kmer = ctx.slot_kmer[fslot].expect("minorant slot has a kmer");
// batch's cross-partition queries. Kmer and mask already in
// hand from `iter_minorants_batch`, no second annex read,
// no slot lookup needed for this pass.
for (i, entry) in batch.entries.iter().enumerate() {
let (kmer, mask) = (entry.kmer, entry.mask);
masks.push(mask);
let base = central_base(kmer, ctx.k);
bases.push(base);
@@ -246,7 +229,7 @@ pub(super) fn scan_layer_families(
if !mask.has(b) {
continue;
}
let dest = partition_of(other, ctx.n_parts);
let dest = other.partition(ctx.n_parts);
outgoing[dest].push((other, i, b));
}
}
@@ -254,14 +237,21 @@ pub(super) fn scan_layer_families(
// Pass 2: genome-major, not family-major — `mat` is stored
// one contiguous block per genome (column), slot as the
// offset within it (see `PartitionCache::find_presence_batch`'s
// docs for the full rationale). `batch.slots` is already
// sorted (a contiguous sub-range of the layer's sorted
// minorant-slot list), so this sweeps each column roughly
// in slot order instead of jumping between all `n_cols`
// columns once per family.
for g in 0..ctx.n_cols {
for (i, &(fslot, _)) in batch.slots.iter().enumerate() {
if ctx.mat.carries(g, fslot) {
// docs for the full rationale). `fill_sub_matrix_carries`
// sorts the slots internally for a sequential mmap sweep per
// column, then restores this batch's order — no hand-rolled
// sort/genome-major loop needed here. The presence/count
// matrix is still MPHF-slot-indexed (unlike the annex), so
// each entry's kmer is mapped to its slot via `index_batch`
// — a pure MPHF lookup, no evidence check, since these are
// this layer's own kmers, known members by construction.
let kmers: Vec<CanonicalKmer> = batch.entries.iter().map(|e| e.kmer).collect();
let slots = ctx.mat.index_batch(&kmers);
let mut carries: Vec<Vec<bool>> = (0..ctx.n_cols).map(|_| Vec::new()).collect();
ctx.mat.fill_sub_matrix_carries(&slots, &mut carries);
for (g, col) in carries.iter().enumerate() {
for (i, &carries_it) in col.iter().enumerate() {
if carries_it {
genome_mask[i * ctx.n_genomes + g] |= 1 << bases[i];
}
}
@@ -305,7 +295,7 @@ pub(super) fn scan_layer_families(
next_expected += n;
}
}
debug_assert_eq!(next_expected, total_families, "every batch must have been replayed");
debug_assert!(pending.is_empty(), "every generated batch must have been replayed");
Ok(())
}
src/obikphylo/src/siblings/helpers.rs
+23
View File
@@ -0,0 +1,23 @@
use obikseq::CanonicalKmer;
use obicompactvec::FamilyMask;
/// 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]
pub(super) 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`.
pub(super) 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()
})
}
src/obikphylo/src/siblings/iter.rs
+172
View File
@@ -0,0 +1,172 @@
//! Phylo/sibling-domain iteration over a layer — an extension trait, not a
//! new field on `MphfLayer`/`Layer<D>`: "family"/"minorant" are phylo
//! concepts, `obilayeredmap` stays kmer/slot-mapping only (see
//! `docmd/architecture/siblings.md`).
//!
//! The sibling annex is persisted in the same order as `iter_kmers()`
//! (`build_sibling_annex`, see `build.rs`), so pairing them is a plain zip —
//! no MPHF, no slot, no `kmer_at`. Both sides are already `Send + 'static`
//! (`KmerIter` owns an `Arc<UnitigFileReader>` clone; `SiblingAnnex` is
//! mmap-backed and handed in as an `Arc` by the caller), so `SiblingIter`
//! streams straight from disk and can be fed to `obipipeline` batch by
//! batch — never collected whole into memory (see the project's "no full
//! collect" rule).
//!
//! Four iterator types, deliberately mirroring `obilayeredmap`'s own
//! `KmerIter`/`KmerBatchIter` pair (single item vs. `Vec` batch) — plus the
//! minorant-filtered variant of each, since "all siblings" and "one row per
//! family" are both common cases:
//!
//! | | all entries | minorants only |
//! |------------|-------------------|---------------------|
//! | single | [`SiblingIter`] | [`MinorantIter`] |
//! | batch | [`SiblingBatchIter`] | [`MinorantBatchIter`] |
//!
//! No separate "enumerate" variant (unlike `KmerIter`/`enumerate_kmers`):
//! [`SiblingEntry`] already carries `order` for free, since pairing with
//! the annex requires it anyway.
use std::sync::Arc;
use obicompactvec::{FamilyMask, SiblingAnnex};
use obikseq::CanonicalKmer;
use obilayeredmap::{KmerIter, Layer, LayerData};
/// One layer entry: a k-mer's position in the layer's iteration order (the
/// same index the sibling annex is keyed on — not an MPHF slot), the k-mer
/// itself, and its family mask.
#[derive(Debug, Clone, Copy)]
pub struct SiblingEntry {
pub order: usize,
pub kmer: CanonicalKmer,
pub mask: FamilyMask,
}
/// Streams `(order, kmer, mask)` triples for one layer, in iteration order.
/// Produced by [`SiblingLayerExt::iter_siblings`].
pub struct SiblingIter {
kmers: KmerIter,
annex: Arc<SiblingAnnex>,
order: usize,
}
impl Iterator for SiblingIter {
type Item = SiblingEntry;
fn next(&mut self) -> Option<Self::Item> {
loop {
let kmer = self.kmers.next()?;
let order = self.order;
self.order += 1;
// `None` means "not yet computed" (see `SiblingAnnex` module
// docs) — shouldn't happen against a fully-built annex, but
// skip rather than misalign the two streams if it does.
if let Some(mask) = self.annex.get(order) {
return Some(SiblingEntry { order, kmer, mask });
}
}
}
}
/// Batches of [`SiblingIter`]'s entries, `batch_size` at a time — the last
/// batch may be shorter. Produced by [`SiblingLayerExt::iter_siblings_batch`].
pub struct SiblingBatchIter {
inner: SiblingIter,
batch_size: usize,
}
impl Iterator for SiblingBatchIter {
type Item = Vec<SiblingEntry>;
fn next(&mut self) -> Option<Self::Item> {
collect_batch(&mut self.inner, self.batch_size)
}
}
/// Like [`SiblingIter`], filtered to the minorant of each family — the
/// common case, since a family is tallied once, at its minorant. Produced
/// by [`SiblingLayerExt::iter_minorants`].
pub struct MinorantIter {
inner: SiblingIter,
}
impl Iterator for MinorantIter {
type Item = SiblingEntry;
fn next(&mut self) -> Option<Self::Item> {
self.inner.by_ref().find(|e| e.mask.is_minorant())
}
}
/// Batches of [`MinorantIter`]'s entries, `batch_size` at a time — the last
/// batch may be shorter. Produced by
/// [`SiblingLayerExt::iter_minorants_batch`].
pub struct MinorantBatchIter {
inner: MinorantIter,
batch_size: usize,
}
impl Iterator for MinorantBatchIter {
type Item = Vec<SiblingEntry>;
fn next(&mut self) -> Option<Self::Item> {
collect_batch(&mut self.inner, self.batch_size)
}
}
/// Shared by [`SiblingBatchIter`] and [`MinorantBatchIter`] — pull up to
/// `batch_size` items, `None` once the source is exhausted with nothing left.
fn collect_batch<I: Iterator>(inner: &mut I, batch_size: usize) -> Option<Vec<I::Item>> {
let mut batch = Vec::with_capacity(batch_size);
for _ in 0..batch_size {
match inner.next() {
Some(item) => batch.push(item),
None => break,
}
}
if batch.is_empty() { None } else { Some(batch) }
}
/// Adds phylo/sibling iteration to any `Layer<D>` — the extension that
/// turns a plain layer into a "sibling layer". Generic over `D`
/// (`LayerData`) rather than implemented once per matrix kind: kmer
/// iteration doesn't depend on the data payload, and `Layer<D>` already
/// delegates `iter_kmers`/`index`/`index_batch` to its inner MPHF for every
/// `D` — reusing that instead of going back through a separate `MphfLayer`.
pub trait SiblingLayerExt {
/// Zip this layer's k-mers with their sibling-annex entry, in iteration
/// order. `annex` must have been built from this same layer (its length
/// must match the layer's k-mer count).
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter;
/// Like [`iter_siblings`](Self::iter_siblings), yielding `batch_size`
/// entries at a time.
fn iter_siblings_batch(&self, annex: Arc<SiblingAnnex>, batch_size: usize) -> SiblingBatchIter;
/// Like [`iter_siblings`](Self::iter_siblings), filtered to the
/// minorant of each family — the common case, since a family is
/// tallied once, at its minorant.
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> MinorantIter;
/// Like [`iter_minorants`](Self::iter_minorants), yielding `batch_size`
/// minorants at a time.
fn iter_minorants_batch(&self, annex: Arc<SiblingAnnex>, batch_size: usize) -> MinorantBatchIter;
}
impl<D: LayerData> SiblingLayerExt for Layer<D> {
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter {
SiblingIter { kmers: self.iter_kmers(), annex, order: 0 }
}
fn iter_siblings_batch(&self, annex: Arc<SiblingAnnex>, batch_size: usize) -> SiblingBatchIter {
SiblingBatchIter { inner: self.iter_siblings(annex), batch_size }
}
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> MinorantIter {
MinorantIter { inner: self.iter_siblings(annex) }
}
fn iter_minorants_batch(&self, annex: Arc<SiblingAnnex>, batch_size: usize) -> MinorantBatchIter {
MinorantBatchIter { inner: self.iter_minorants(annex), batch_size }
}
}
src/obikindex/src/siblings/mod.rs → src/obikphylo/src/siblings/mod.rs
+8 -7
View File
@@ -12,7 +12,7 @@
//! 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`](crate::index::KmerIndex::sibling_annex_stats)
//! [`sibling_annex_stats`](stats::SiblingStatsExt::sibling_annex_stats)
//! below).
//!
//! Per layer, an `obipipeline` batch transform (throttled — see
@@ -55,15 +55,16 @@ mod stats;
#[cfg(test)]
mod tests;
pub use alignment::SnpAlignment;
pub use cardinality::CardinalityTally;
pub use distance::{BasePairTally, RawSnpDistanceOutput};
pub use iter::{SiblingEntry, SiblingIter, SiblingLayerExt};
pub use stats::SiblingAnnexStats;
pub use alignment::{SnpAlignment, SnpAlignmentExt};
pub use build::SiblingAnnexBuildExt;
pub use cardinality::{CardinalityExt, CardinalityTally};
pub use distance::{BasePairTally, DistanceExt, RawSnpDistanceOutput};
pub use iter::{MinorantBatchIter, MinorantIter, SiblingBatchIter, SiblingEntry, SiblingIter, SiblingLayerExt};
pub use stats::{SiblingAnnexStats, SiblingStatsExt};
use obilayeredmap::OLMError;
use crate::error::OKIError;
use obikindex::OKIError;
pub(super) const INDEX_SUBDIR: &str = "index";
pub(super) const ANNEX_FILE_NAME: &str = "siblings.psib";
src/obikindex/src/siblings/stats.rs → src/obikphylo/src/siblings/stats.rs
+26 -19
View File
@@ -4,17 +4,17 @@ use obicompactvec::SiblingAnnex;
use obikpartitionner::KmerPartition;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;
use super::ANNEX_FILE_NAME;
use super::cache::PartitionCache;
use super::family_scan::scan_layer_families;
/// 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.
/// annex (see [`super::build::SiblingAnnexBuildExt::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
@@ -33,7 +33,9 @@ pub struct SiblingAnnexStats {
pub per_genome: Vec<[u64; 4]>,
}
impl KmerIndex {
/// Adds [`sibling_family_size_histogram`](Self::sibling_family_size_histogram)
/// and [`sibling_annex_stats`](Self::sibling_annex_stats) to `KmerIndex`.
pub trait SiblingStatsExt {
/// The global family-size histogram alone (`SiblingAnnexStats::counts`,
/// no `per_genome`) — reads only the already-built annex (`mask.siblings()`
/// + `mask.is_minorant()`, 1 byte/slot, mmap'd), nothing else: no
@@ -52,8 +54,19 @@ impl KmerIndex {
/// cost to get there too; use this when only the global histogram is
/// needed. Requires an annex built after the minorant flag was added —
/// re-run `build_sibling_annex` if this reads as all-zero on an older one.
pub fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]> {
let layer_dirs = self.sibling_layer_dirs()?;
fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]>;
/// 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`](super::build::SiblingAnnexBuildExt::build_sibling_annex)
/// has not been run on this index first.
fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats>;
}
impl SiblingStatsExt for KmerIndex {
fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]> {
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let mut counts = [0u64; 4];
for layer_dir in &layer_dirs {
@@ -70,16 +83,10 @@ impl KmerIndex {
Ok(counts)
}
/// 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> {
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 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;
@@ -87,14 +94,14 @@ impl KmerIndex {
// 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.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let layer_dirs = self.sibling_layer_dirs()?;
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
// One layer at a time, not parallelised across layers — see
// `snp_pseudo_alignment`'s comment for why `par_iter()` over layers
src/obikindex/src/siblings/tests.rs → src/obikphylo/src/siblings/tests.rs
+7 -4
View File
@@ -8,11 +8,14 @@ use obilayeredmap::meta::PartitionMeta;
use obisys::Reporter;
use tempfile::tempdir;
use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexConfig};
use crate::merge::MergeMode;
use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode};
use super::alignment::SnpAlignmentExt;
use super::build::SiblingAnnexBuildExt;
use super::cardinality::CardinalityExt;
use super::distance::DistanceExt;
use super::helpers::is_minorant;
use super::stats::SiblingStatsExt;
use super::{ANNEX_FILE_NAME, INDEX_SUBDIR};
// k must be >= 11 (project constraint, "k ∈ [11,31]"); k=11, level_max=1,
@@ -222,7 +225,7 @@ fn family_scan_consumers_agree_on_one_sibling_each() {
// Merge doesn't promise to preserve source order, so resolve each
// genome's index by label rather than assuming g1 -> 0, g2 -> 1.
let idx_of = |label: &str| merged.meta.genomes.iter().position(|g| g.label == label).unwrap();
let idx_of = |label: &str| merged.meta().genomes.iter().position(|g| g.label == label).unwrap();
let (i1, i2) = (idx_of("g1"), idx_of("g2"));
// snp_pseudo_alignment: one variable family, one column — g1's row
src/obikseq/src/kmer.rs
+44
View File
@@ -367,6 +367,50 @@ impl<L: KmerLength> CanonicalKmerOf<L> {
pub fn into_kmer(self) -> KmerOf<L> {
KmerOf(self.0, PhantomData)
}
/// This k-mer's own minimiser — as a standalone unit, not part of a
/// streamed sequence. Enumerates all `L::len() - MLen::len() + 1`
/// windows directly on the packed 2-bit representation (no ASCII
/// round-trip, no rolling state, no entropy tracking) and keeps the one
/// with the lowest [`hash_kmer`] — the same selection rule
/// `obiskbuilder::rolling_stat::RollingStat` computes incrementally
/// along a sequence, replicated here in `O(k)` pure bit arithmetic for
/// a single isolated k-mer (the case a streaming rolling scan is
/// unnecessary machinery for). Ties (equal hash) keep the last
/// (highest-position) window, matching `RollingStat`'s monotonic-deque
/// eviction rule (`>=` pops the previous record) — a formal detail, not
/// a practical concern: a `mix64` collision between two distinct
/// m-mers is vanishingly unlikely.
pub fn minimizer(&self) -> Minimizer {
let k = L::len();
let ml = MLen::len();
let rc = self.revcomp().raw();
let mut best_canon: RawKmer = 0;
let mut best_hash = u64::MAX;
for p in 0..=(k - ml) {
let fwd = (self.0 << (2 * p)) >> (KMER_BITS - 2 * ml);
let rev = (rc << (2 * (k - ml - p))) >> (KMER_BITS - 2 * ml);
let canon = fwd.min(rev);
let hash = hash_kmer(canon << (KMER_BITS - 2 * ml));
if hash < best_hash {
best_hash = hash;
best_canon = canon;
}
}
Minimizer::from_raw_unchecked(best_canon << (KMER_BITS - 2 * ml))
}
/// Destination partition for this k-mer, using the project-wide routing
/// rule (`minimizer().seq_hash() & mask`) — the same rule
/// `KmerPartition`/`RoutableSuperKmer` apply to k-mers read from a
/// streamed sequence, here for a standalone k-mer (e.g. a synthetic
/// variant generated outside any sequence). `n_partitions` must be a
/// power of two.
#[inline]
pub fn partition(&self, n_partitions: usize) -> usize {
let mask = (n_partitions as u64) - 1;
(self.minimizer().seq_hash() & mask) as usize
}
}
impl<L: KmerLength> Sequence for CanonicalKmerOf<L> {
src/obilayeredmap/src/layer.rs
+22
View File
@@ -72,6 +72,15 @@ impl<D: LayerData> Layer<D> {
self.mphf.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot) })
}
/// MPHF + evidence membership check only — no data read. For callers
/// that batch many lookups before touching the matrix at all (e.g. to
/// group hits by column for a later `sub_matrix`/`fill_sub_matrix`
/// sweep), so a plain `query` reading — and discarding — a full row
/// per lookup would be wasted work.
pub fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
self.mphf.find(kmer)
}
pub fn n(&self) -> usize { self.mphf.n() }
/// Raw MPHF lookup: kmer → slot, no membership check.
@@ -187,6 +196,11 @@ impl Layer<PersistentCompactIntMatrix> {
.map_err(OLMError::Io)
}
/// Number of genome columns in this layer's count matrix.
pub fn n_cols(&self) -> usize {
self.data.n_cols()
}
/// Extract a sub-matrix of counts for the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<u32>>`: one inner `Vec` per genome
@@ -218,6 +232,14 @@ impl Layer<PersistentBitMatrix> {
.map_err(OLMError::Io)
}
/// Number of genome columns in this layer's presence matrix — see
/// `PersistentBitMatrix::n_cols`'s docs for the `Implicit` mono-genome
/// special case (always reports `1`, regardless of the index's real
/// genome count).
pub fn n_cols(&self) -> usize {
self.data.n_cols()
}
/// Extract a sub-matrix of presence/absence for the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<bool>>`: one inner `Vec` per genome