Extract phylogenetic sibling logic into new obikphylo crate

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
+22 -2
View File
@@ -1768,9 +1768,7 @@ dependencies = [
"obikpartitionner", "obikpartitionner",
"obikseq", "obikseq",
"obilayeredmap", "obilayeredmap",
"obipipeline",
"obiread", "obiread",
"obiskbuilder",
"obiskio", "obiskio",
"obisys", "obisys",
"obitaxonomy", "obitaxonomy",
@@ -1794,6 +1792,7 @@ dependencies = [
"obifastwrite", "obifastwrite",
"obikindex", "obikindex",
"obikpartitionner", "obikpartitionner",
"obikphylo",
"obikrope", "obikrope",
"obikseq", "obikseq",
"obilayeredmap", "obilayeredmap",
@@ -1843,6 +1842,27 @@ dependencies = [
"tracing", "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]] [[package]]
name = "obikrope" name = "obikrope"
version = "0.1.0" version = "0.1.0"
+1 -1
View File
@@ -1,5 +1,5 @@
[workspace] [workspace]
resolver = "3" 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] [profile.release]
debug = 1 debug = 1
-2
View File
@@ -11,8 +11,6 @@ obiskio = { path = "../obiskio" }
obisys = { path = "../obisys" } obisys = { path = "../obisys" }
obicompactvec = { path = "../obicompactvec" } obicompactvec = { path = "../obicompactvec" }
obilayeredmap = { path = "../obilayeredmap" } obilayeredmap = { path = "../obilayeredmap" }
obiskbuilder = { path = "../obiskbuilder" }
obipipeline = { path = "../obipipeline" }
ndarray = "0.16" ndarray = "0.16"
rayon = "1" rayon = "1"
crossbeam-channel = "0.5" crossbeam-channel = "0.5"
+4
View File
@@ -132,6 +132,10 @@ impl KmerIndex {
IndexState::detect(&self.root_path).unwrap_or(IndexState::Empty) 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(&self) -> &IndexMeta { &self.meta }
pub fn meta_mut(&mut self) -> &mut IndexMeta { &mut self.meta } pub fn meta_mut(&mut self) -> &mut IndexMeta { &mut self.meta }
pub fn kmer_size(&self) -> usize { self.meta.config.kmer_size } pub fn kmer_size(&self) -> usize { self.meta.config.kmer_size }
-4
View File
@@ -2,7 +2,6 @@ pub mod error;
pub mod meta; pub mod meta;
pub mod predicate; pub mod predicate;
pub mod state; pub mod state;
mod cardcomp;
mod distance; mod distance;
mod dump; mod dump;
mod index; mod index;
@@ -11,7 +10,6 @@ mod numa;
mod rebuild; mod rebuild;
mod reindex; mod reindex;
mod select; mod select;
mod siblings;
mod stats; mod stats;
pub use error::{OKIError, OKIResult}; 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 predicate::{GroupFilterParams, MetaPred};
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED}; pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
pub use stats::IndexBitsPerKmer; pub use stats::IndexBitsPerKmer;
pub use siblings::{BasePairTally, CardinalityTally, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
-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);
}
}
}
}
}
}
-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
}
-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())
}
}
+1
View File
@@ -19,6 +19,7 @@ obikpartitionner = { path = "../obikpartitionner" }
obisys = { path = "../obisys" } obisys = { path = "../obisys" }
obiskio = { path = "../obiskio" } obiskio = { path = "../obiskio" }
obikindex = { path = "../obikindex", default-features = false } obikindex = { path = "../obikindex", default-features = false }
obikphylo = { path = "../obikphylo" }
obitaxonomy = { path = "../obitaxonomy" } obitaxonomy = { path = "../obitaxonomy" }
obilayeredmap = { path = "../obilayeredmap" } obilayeredmap = { path = "../obilayeredmap" }
clap = { version = "4", features = ["derive"] } clap = { version = "4", features = ["derive"] }
+2 -1
View File
@@ -1,7 +1,8 @@
use std::io::{BufWriter, Write}; use std::io::{BufWriter, Write};
use std::path::PathBuf; use std::path::PathBuf;
use obikindex::{KmerIndex, SnpAlignment}; use obikindex::KmerIndex;
use obikphylo::siblings::{SnpAlignment, SnpAlignmentExt};
use tracing::info; use tracing::info;
// ── Family overlap: shared-family counts and the `--min-shared-family` / // ── Family overlap: shared-family counts and the `--min-shared-family` /
+1 -1
View File
@@ -2,7 +2,7 @@ use std::io::{BufWriter, Write};
use std::path::PathBuf; use std::path::PathBuf;
use obifastwrite::{JsonVal, write_record}; use obifastwrite::{JsonVal, write_record};
use obikindex::SnpAlignment; use obikphylo::siblings::SnpAlignment;
use tracing::info; use tracing::info;
use super::sankoff::state_index_table; use super::sankoff::state_index_table;
+6 -2
View File
@@ -9,9 +9,13 @@ mod tnt;
use std::io::{self, BufWriter, Write}; use std::io::{self, BufWriter, Write};
use kodama::{Method, linkage}; use kodama::{Method, linkage};
use obikindex::{ use obikindex::KmerIndex;
KmerIndex, RawSnpDistanceOutput, SnpAlignment, use obikphylo::{
cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix, 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 speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
use tracing::info; use tracing::info;
+1 -1
View File
@@ -2,7 +2,7 @@ use std::io::{BufWriter, Write};
use std::path::PathBuf; use std::path::PathBuf;
use obifastwrite::{JsonVal, write_record}; use obifastwrite::{JsonVal, write_record};
use obikindex::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; use obikphylo::siblings::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
use tracing::info; use tracing::info;
// ── Family-size distribution → CSV ────────────────────────────────────────── // ── Family-size distribution → CSV ──────────────────────────────────────────
+1 -1
View File
@@ -2,7 +2,7 @@ use std::io::{BufWriter, Write};
use std::path::PathBuf; use std::path::PathBuf;
use obifastwrite::{JsonVal, write_record}; use obifastwrite::{JsonVal, write_record};
use obikindex::{BasePairTally, CardinalityTally, SnpAlignment}; use obikphylo::siblings::{BasePairTally, CardinalityTally, SnpAlignment};
use tracing::info; use tracing::info;
// ── Sankoff pseudo-alignment → FASTA ──────────────────────────────────────── // ── Sankoff pseudo-alignment → FASTA ────────────────────────────────────────
+1 -1
View File
@@ -1,7 +1,7 @@
use std::io::{BufWriter, Write}; use std::io::{BufWriter, Write};
use std::path::PathBuf; use std::path::PathBuf;
use obikindex::SnpAlignment; use obikphylo::siblings::SnpAlignment;
use tracing::info; use tracing::info;
use super::sankoff::{scaled_metric_matrix, state_index_table}; use super::sankoff::{scaled_metric_matrix, state_index_table};
+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"] }
@@ -17,7 +17,7 @@
//! since TNT/PhyG search *unrooted* trees) if the cost matrix is //! since TNT/PhyG search *unrooted* trees) if the cost matrix is
//! symmetric, the discrete-parsimony analogue of CTMC reversibility. //! 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`), /// Row-stochastic 5×5 cardinality transition probabilities (`0..=4`),
/// diagonal included ("stay at the same cardinality"), from /// diagonal included ("stay at the same cardinality"), from
+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};
@@ -3,8 +3,8 @@ use std::sync::Arc;
use obikpartitionner::KmerPartition; use obikpartitionner::KmerPartition;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use obikindex::{OKIError, OKIResult};
use crate::index::KmerIndex; use obikindex::KmerIndex;
use super::cache::PartitionCache; use super::cache::PartitionCache;
use super::family_scan::scan_layer_families; use super::family_scan::scan_layer_families;
@@ -53,25 +53,32 @@ pub struct SnpAlignment {
pub sequences: Vec<Vec<u8>>, 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 /// Build the SNP-only pseudo-alignment from an already-built sibling
/// annex (run [`build_sibling_annex`](Self::build_sibling_annex) first). /// annex (run [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex) first).
pub fn snp_pseudo_alignment(&self) -> OKIResult<SnpAlignment> { 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_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len(); let n_genomes = self.meta().genomes.len();
let with_counts = self.meta.config.with_counts; let with_counts = self.meta().config.with_counts;
let k = self.kmer_size(); let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize; let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config( let partition = KmerPartition::open_with_config(
&self.root_path, self.root_path(),
self.kmer_size(), self.kmer_size(),
self.minimizer_size(), self.minimizer_size(),
n_bits, n_bits,
) )
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); 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"); let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers");
// One layer at a time, not `par_iter()` over layers — same // One layer at a time, not `par_iter()` over layers — same
@@ -12,11 +12,11 @@ use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta; use obilayeredmap::meta::PartitionMeta;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use obikindex::{OKIError, OKIResult};
use crate::index::KmerIndex; use obikindex::KmerIndex;
use super::cache::PartitionCache; 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}; use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
// ── obipipeline data types ───────────────────────────────────────────────── // ── obipipeline data types ─────────────────────────────────────────────────
@@ -53,7 +53,8 @@ enum SibData {
Variants(VariantBatch), 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 /// Build the sibling-count/minorant annex for every layer of every
/// partition of this (already built) index, writing one annex file per /// partition of this (already built) index, writing one annex file per
/// layer alongside its existing index files. Safe to call again later /// 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 /// Construction only — no statistics gathered here on purpose: this is
/// meant to run routinely (it is the artefact the SNP-family distances /// meant to run routinely (it is the artefact the SNP-family distances
/// will consume), while the sibling-count distribution /// will consume), while the sibling-count distribution
/// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate, /// ([`sibling_annex_stats`](super::stats::SiblingStatsExt::sibling_annex_stats))
/// occasional diagnostic pass over the result, not run every time. /// 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 /// 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 /// 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. /// 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_parts = self.n_partitions();
let n_bits = n_parts.trailing_zeros() as usize; let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config( let partition = KmerPartition::open_with_config(
&self.root_path, self.root_path(),
self.kmer_size(), self.kmer_size(),
self.minimizer_size(), self.minimizer_size(),
n_bits, n_bits,
@@ -82,7 +88,7 @@ impl KmerIndex {
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep"); 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 pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
let mut total_slots: u64 = 0; let mut total_slots: u64 = 0;
@@ -97,7 +103,7 @@ impl KmerIndex {
let mut part_slots: u64 = 0; let mut part_slots: u64 = 0;
for l in 0..meta.n_layers { for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}")); 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; total_slots += part_slots;
pb.inc(1); pb.inc(1);
@@ -108,20 +114,22 @@ impl KmerIndex {
Ok(()) Ok(())
} }
}
/// Returns the number of distinct k-mers (annex entries) processed, for /// Returns the number of distinct k-mers (annex entries) processed, for
/// progress reporting. /// progress reporting. A free function, not a `KmerIndex` method — called
fn build_layer_sibling_annex( /// only from `build_sibling_annex` above, in the same file.
&self, fn build_layer_sibling_annex(
layer_dir: &Path, index: &KmerIndex,
n_parts: usize, layer_dir: &Path,
cache: &Arc<PartitionCache>, n_parts: usize,
) -> OKIResult<u64> { cache: &Arc<PartitionCache>,
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); ) -> OKIResult<u64> {
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let k = self.kmer_size(); let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let n = mphf.n(); let k = index.kmer_size();
let n = mphf.n();
// ── Reconciliation state, indexed by this layer's k-mer iteration // ── Reconciliation state, indexed by this layer's k-mer iteration
// order (the physical layout of `unitigs.bin`), never by MPHF slot // order (the physical layout of `unitigs.bin`), never by MPHF slot
@@ -211,7 +219,7 @@ impl KmerIndex {
continue; continue;
} }
items.push(( items.push((
partition_of(variant, n_parts), variant.partition(n_parts),
variant, variant,
order, order,
central_base(variant, k), central_base(variant, k),
@@ -272,6 +280,5 @@ impl KmerIndex {
} }
builder.close()?; builder.close()?;
Ok(n as u64) Ok(n as u64)
}
} }
+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);
}
}
}
}
}
}
@@ -5,8 +5,8 @@ use ndarray::Array2;
use obikpartitionner::KmerPartition; use obikpartitionner::KmerPartition;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use obikindex::{OKIError, OKIResult};
use crate::index::KmerIndex; use obikindex::KmerIndex;
use super::cache::PartitionCache; use super::cache::PartitionCache;
use super::distance::RawSnpDistanceOutput; use super::distance::RawSnpDistanceOutput;
@@ -22,7 +22,8 @@ pub struct CardinalityTally {
pub counts: [[u64; 5]; 5], 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 /// Cardinality co-occurrence, pooled only over genome pairs whose
/// overall SNP ratio in `raw` is at or below `ratio_ceiling` — same /// overall SNP ratio in `raw` is at or below `ratio_ceiling` — same
/// saturation/no-data exclusion discipline as /// saturation/no-data exclusion discipline as
@@ -51,10 +52,14 @@ impl KmerIndex {
/// gets the matching restriction via `scan_family_pairs`'s new /// gets the matching restriction via `scan_family_pairs`'s new
/// `variable` flag, rather than a `family_size()` check of its own (it /// `variable` flag, rather than a `family_size()` check of its own (it
/// doesn't have direct access to the family's mask). /// 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_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len(); let n_genomes = self.meta().genomes.len();
let with_counts = self.meta.config.with_counts; let with_counts = self.meta().config.with_counts;
let k = self.kmer_size(); let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize; let n_bits = n_parts.trailing_zeros() as usize;
@@ -68,14 +73,14 @@ impl KmerIndex {
}); });
let partition = KmerPartition::open_with_config( let partition = KmerPartition::open_with_config(
&self.root_path, self.root_path(),
self.kmer_size(), self.kmer_size(),
self.minimizer_size(), self.minimizer_size(),
n_bits, n_bits,
) )
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); 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"); 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 // One layer at a time — see `snp_pseudo_alignment`'s comment for why
@@ -5,8 +5,8 @@ use ndarray::Array2;
use obikpartitionner::KmerPartition; use obikpartitionner::KmerPartition;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use obikindex::{OKIError, OKIResult};
use crate::index::KmerIndex; use obikindex::KmerIndex;
use super::cache::PartitionCache; use super::cache::PartitionCache;
use super::family_scan::scan_layer_families; use super::family_scan::scan_layer_families;
@@ -35,54 +35,52 @@ pub struct RawSnpDistanceOutput {
pub shared: Array2<u64>, pub shared: Array2<u64>,
} }
impl KmerIndex { /// Shared traversal behind [`DistanceExt::raw_snp_distance`] and
/// Shared traversal behind [`raw_snp_distance`](Self::raw_snp_distance) /// [`DistanceExt::base_pair_tally`]: for every family (tallied once, at its
/// and [`base_pair_tally`](Self::base_pair_tally): for every family /// minorant) of every layer of the already-built sibling annex, resolves
/// (tallied once, at its minorant) of every layer of the already-built /// each genome's single observed form (`None` if absent or
/// sibling annex, resolves each genome's single observed form (`None` /// ambiguous/multi-copy), then calls `on_pair(acc, i, j, bi, bj, variable)`
/// if absent or ambiguous/multi-copy), then calls `on_pair(acc, i, j, /// for every genome pair `(i, j)` where both are unambiguous and
/// bi, bj, variable)` for every genome pair `(i, j)` where both are /// single-copy (`bi == bj` means shared at that locus, `bi != bj` means a
/// unambiguous and single-copy (`bi == bj` means shared at that locus, /// SNP). `variable` is the family's own `family_size() >= 2` (true if more
/// `bi != bj` means a SNP). `variable` is the family's own /// than one member is observed *anywhere* in the family, i.e. it isn't
/// `family_size() >= 2` (true if more than one member is observed /// fully invariant across the whole index) — `raw_snp_distance` ignores it
/// *anywhere* in the family, i.e. it isn't fully invariant across the /// (a fully-invariant family is still legitimately "shared"), but callers
/// whole index) — `raw_snp_distance` ignores it (a fully-invariant /// whose diagonal should only reflect genuine SNP-adjacent agreement, not
/// family is still legitimately "shared"), but callers whose diagonal /// the genome-wide invariant background, need it (see
/// should only reflect genuine SNP-adjacent agreement, not the /// [`DistanceExt::base_pair_tally`]'s `same` field). Layers are processed
/// genome-wide invariant background, need it (see /// one at a time, not in parallel — see `snp_pseudo_alignment`'s comment
/// [`base_pair_tally`](Self::base_pair_tally)'s `same` field). Layers /// for why `par_iter()` over layers would defeat `scan_layer_families`'s
/// are processed one at a time, not in parallel — see /// partition-grouped locality; each layer gets its own accumulator from
/// `snp_pseudo_alignment`'s comment for why `par_iter()` over layers /// `zero()`, combined pairwise via `combine`. A free function, not a
/// would defeat `scan_layer_families`'s partition-grouped locality; /// `KmerIndex` method — called only from this file's `DistanceExt` impl.
/// each layer gets its own accumulator from `zero()`, combined fn scan_family_pairs<Acc, F, C>(
/// pairwise via `combine`. index: &KmerIndex,
fn scan_family_pairs<Acc, F, C>( label: &str,
&self, zero: impl Fn() -> Acc + Sync,
label: &str, on_pair: F,
zero: impl Fn() -> Acc + Sync, combine: C,
on_pair: F, ) -> OKIResult<Acc>
combine: C, where
) -> OKIResult<Acc> Acc: Send,
where F: Fn(&mut Acc, usize, usize, u8, u8, bool) + Sync,
Acc: Send, C: Fn(Acc, Acc) -> Acc,
F: Fn(&mut Acc, usize, usize, u8, u8, bool) + Sync, {
C: Fn(Acc, Acc) -> Acc, let n_parts = index.n_partitions();
{ let n_genomes = index.meta().genomes.len();
let n_parts = self.n_partitions(); let with_counts = index.meta().config.with_counts;
let n_genomes = self.meta.genomes.len(); let k = index.kmer_size();
let with_counts = self.meta.config.with_counts; let n_bits = n_parts.trailing_zeros() as usize;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config( let partition = KmerPartition::open_with_config(
&self.root_path, index.root_path(),
self.kmer_size(), index.kmer_size(),
self.minimizer_size(), index.minimizer_size(),
n_bits, n_bits,
) )
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); 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"); let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
// One layer at a time — see `snp_pseudo_alignment`'s comment for why // One layer at a time — see `snp_pseudo_alignment`'s comment for why
@@ -121,11 +119,36 @@ impl KmerIndex {
Ok(total) 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 /// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex
/// (run [`build_sibling_annex`](Self::build_sibling_annex) first). /// (run [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex) first).
pub fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> { fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput>;
let n_genomes = self.meta.genomes.len();
let (snp, shared) = self.scan_family_pairs( /// 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", "raw_snp_distance",
|| (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::zeros((n_genomes, n_genomes))), || (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::zeros((n_genomes, n_genomes))),
|(snp, shared), i, j, bi, bj, _variable| { |(snp, shared), i, j, bi, bj, _variable| {
@@ -146,22 +169,8 @@ impl KmerIndex {
Ok(RawSnpDistanceOutput { snp, shared }) Ok(RawSnpDistanceOutput { snp, shared })
} }
/// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG, fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally> {
/// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs let n_genomes = self.meta().genomes.len();
/// 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();
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| { let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
if i == j { if i == j {
return false; return false;
@@ -171,7 +180,8 @@ impl KmerIndex {
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling 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", "base_pair_tally",
|| ([[0u64; 4]; 4], [0u64; 4]), || ([[0u64; 4]; 4], [0u64; 4]),
|(counts, same), i, j, bi, bj, variable| { |(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 { pub struct BasePairTally {
/// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled /// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled
/// over included genome pairs, where the two genomes' single forms are /// over included genome pairs, where the two genomes' single forms are
@@ -55,15 +55,16 @@ use rayon::prelude::*;
use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex}; use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikseq::CanonicalKmer; use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer; use obilayeredmap::Layer;
use obilayeredmap::meta::PartitionMeta; use obilayeredmap::meta::PartitionMeta;
use obipipeline::{ThrottleGuard, throttle}; use obipipeline::{ThrottleGuard, throttle};
use crate::error::{OKIError, OKIResult}; use obikindex::{OKIError, OKIResult};
use crate::index::KmerIndex; use obikindex::KmerIndex;
use super::cache::{Mat, PartitionCache}; 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}; use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// Families per batch — see the module docs for the memory-vs-per-partition- /// Families per batch — see the module docs for the memory-vs-per-partition-
@@ -73,42 +74,43 @@ use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// queries per batch to amortise against. /// queries per batch to amortise against.
const FAMILY_BATCH: usize = 65536; const FAMILY_BATCH: usize = 65536;
impl KmerIndex { /// Every (partition, layer) directory carrying a sibling annex, checked
/// Every (partition, layer) directory carrying a sibling annex, checked /// up front so a missing one is reported before any real work starts.
/// up front so a missing one is reported before any real work starts. /// Shared by every sibling-annex consumer's extension-trait impl (`alignment`,
pub(super) fn sibling_layer_dirs(&self) -> OKIResult<Vec<PathBuf>> { /// `build`, `cardinality`, `distance`, `stats`) — a free function, not a
let n_parts = self.n_partitions(); /// `KmerIndex` method, since it is crate-internal only and `KmerIndex` lives
let mut layer_dirs = Vec::new(); /// in `obikindex`, a foreign crate from here (orphan rule).
for part in 0..n_parts { pub(crate) fn sibling_layer_dirs(index: &KmerIndex) -> OKIResult<Vec<PathBuf>> {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); let n_parts = index.n_partitions();
if !index_dir.exists() { let mut layer_dirs = Vec::new();
continue; for part in 0..n_parts {
let index_dir = index.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
} }
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; layer_dirs.push(layer_dir);
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
} }
Ok(layer_dirs)
} }
Ok(layer_dirs)
} }
/// Read-only state shared (via `Arc`) across every pipeline worker /// Read-only state shared (via `Arc`) across every pipeline worker
/// generating this layer's batches — opened once, not per batch. No `cache` /// generating this layer's batches — opened once, not per batch. No `cache`
/// here: generation never touches the cross-partition cache, only this /// here: generation never touches the cross-partition cache, only this
/// layer's own already-open matrix. No `annex` either — the filter step /// layer's own already-open matrix. No separate MPHF/`slot_kmer` either —
/// that builds `minorant_slots` already carries each slot's mask through, /// `mat` (a `Layer<D>`) already bundles the MPHF, and each `SiblingEntry`
/// so a worker never needs to re-read it. /// arrives with its kmer and mask already in hand from `iter_minorants_batch`.
struct LayerCtx { struct LayerCtx {
slot_kmer: Vec<Option<CanonicalKmer>>,
mat: Mat, mat: Mat,
n_parts: usize, n_parts: usize,
n_genomes: usize, n_genomes: usize,
@@ -118,10 +120,11 @@ struct LayerCtx {
struct SourceBatch { struct SourceBatch {
start_family_idx: usize, start_family_idx: usize,
/// `(slot, mask)` — the mask is carried through from the filter step /// One entry per minorant family in this batch, straight from
/// below rather than re-read from the annex per batch; it's the same /// `iter_minorants_batch` — `order` (iteration-order index, not an MPHF
/// byte either way, just already in hand. /// slot; see `docmd/architecture/siblings.md`), `kmer`, and `mask`
slots: Vec<(usize, FamilyMask)>, /// already carried through, no second annex read.
entries: Vec<SiblingEntry>,
_permit: ThrottleGuard, _permit: ThrottleGuard,
} }
@@ -143,10 +146,10 @@ enum FamData {
Generated(GeneratedBatch), 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 /// [`FAMILY_BATCH`] at a time — see the module docs for why generation and
/// resolution use different concurrency. `on_family` is called once per /// 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 /// 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 /// member whose own canonical central base is `b`), backed by a scratch
/// buffer reused across every call — callers that need to keep data past /// buffer reused across every call — callers that need to keep data past
@@ -162,59 +165,38 @@ pub(super) fn scan_layer_families(
) -> OKIResult<()> { ) -> OKIResult<()> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); 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 meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; let annex = Arc::new(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 use_counts = with_counts && layer_dir.join("counts").exists(); let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts { 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 { } 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); let n_cols = mat.n_cols().min(n_genomes);
// Minorant slots, in order, mask carried along — the flag is already let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k });
// 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 { slot_kmer, mat, n_parts, n_genomes, n_cols, k }); // Streamed straight from `iter_minorants_batch` (zips this layer's own
// `iter_kmers()` with the annex, both in iteration order — never an
let mut offset = 0usize; // MPHF slot; see `docmd/architecture/siblings.md`) — never collected
let batches: Vec<(usize, Vec<(usize, FamilyMask)>)> = minorant_slots // into a `Vec` first: a layer can hold billions of k-mers, so
.chunks(FAMILY_BATCH) // materialising every minorant family up front is exactly the memory
.map(|chunk| { // blowup an earlier version of this traversal was rewritten to avoid
let start = offset; // (see the module docs). `.scan()` computes each batch's starting
offset += chunk.len(); // family index lazily, mirroring what the eager `chunks()`+running
(start, chunk.to_vec()) // `offset` used to do.
}) let batches = ctx.mat.iter_minorants_batch(annex, FAMILY_BATCH).scan(0usize, |offset, batch| {
.collect(); let start = *offset;
*offset += batch.len();
Some((start, batch))
});
let n_workers = obisys::effective_parallelism(); let n_workers = obisys::effective_parallelism();
let capacity = 4; 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, start_family_idx: t.item.0,
slots: t.item.1, entries: t.item.1,
_permit: t.guard, _permit: t.guard,
}); });
@@ -224,17 +206,18 @@ pub(super) fn scan_layer_families(
| { | {
move |batch: SourceBatch| -> GeneratedBatch { move |batch: SourceBatch| -> GeneratedBatch {
let ctx = &worker_ctx; let ctx = &worker_ctx;
let n = batch.slots.len(); let n = batch.entries.len();
let mut masks = Vec::with_capacity(n); let mut masks = Vec::with_capacity(n);
let mut bases = Vec::with_capacity(n); let mut bases = Vec::with_capacity(n);
let mut genome_mask = vec![0u8; n * ctx.n_genomes]; 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(); 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 // Pass 1: cheap, no matrix access — own base and this
// batch's cross-partition queries. Mask already in hand // batch's cross-partition queries. Kmer and mask already in
// from the filter step, no second annex read. // hand from `iter_minorants_batch`, no second annex read,
for (i, &(fslot, mask)) in batch.slots.iter().enumerate() { // no slot lookup needed for this pass.
let kmer = ctx.slot_kmer[fslot].expect("minorant slot has a kmer"); for (i, entry) in batch.entries.iter().enumerate() {
let (kmer, mask) = (entry.kmer, entry.mask);
masks.push(mask); masks.push(mask);
let base = central_base(kmer, ctx.k); let base = central_base(kmer, ctx.k);
bases.push(base); bases.push(base);
@@ -246,7 +229,7 @@ pub(super) fn scan_layer_families(
if !mask.has(b) { if !mask.has(b) {
continue; continue;
} }
let dest = partition_of(other, ctx.n_parts); let dest = other.partition(ctx.n_parts);
outgoing[dest].push((other, i, b)); 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 // Pass 2: genome-major, not family-major — `mat` is stored
// one contiguous block per genome (column), slot as the // one contiguous block per genome (column), slot as the
// offset within it (see `PartitionCache::find_presence_batch`'s // offset within it (see `PartitionCache::find_presence_batch`'s
// docs for the full rationale). `batch.slots` is already // docs for the full rationale). `fill_sub_matrix_carries`
// sorted (a contiguous sub-range of the layer's sorted // sorts the slots internally for a sequential mmap sweep per
// minorant-slot list), so this sweeps each column roughly // column, then restores this batch's order — no hand-rolled
// in slot order instead of jumping between all `n_cols` // sort/genome-major loop needed here. The presence/count
// columns once per family. // matrix is still MPHF-slot-indexed (unlike the annex), so
for g in 0..ctx.n_cols { // each entry's kmer is mapped to its slot via `index_batch`
for (i, &(fslot, _)) in batch.slots.iter().enumerate() { // — a pure MPHF lookup, no evidence check, since these are
if ctx.mat.carries(g, fslot) { // 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]; genome_mask[i * ctx.n_genomes + g] |= 1 << bases[i];
} }
} }
@@ -305,7 +295,7 @@ pub(super) fn scan_layer_families(
next_expected += n; 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(()) Ok(())
} }
+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()
})
}
+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 }
}
}
@@ -12,7 +12,7 @@
//! in the current multi-genome index — a property of the whole index, not //! 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 //! of any one genome. Sibling count and minorant are *derived* from the
//! mask by callers, not stored (see `FamilyMask` and //! 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). //! below).
//! //!
//! Per layer, an `obipipeline` batch transform (throttled — see //! Per layer, an `obipipeline` batch transform (throttled — see
@@ -55,15 +55,16 @@ mod stats;
#[cfg(test)] #[cfg(test)]
mod tests; mod tests;
pub use alignment::SnpAlignment; pub use alignment::{SnpAlignment, SnpAlignmentExt};
pub use cardinality::CardinalityTally; pub use build::SiblingAnnexBuildExt;
pub use distance::{BasePairTally, RawSnpDistanceOutput}; pub use cardinality::{CardinalityExt, CardinalityTally};
pub use iter::{SiblingEntry, SiblingIter, SiblingLayerExt}; pub use distance::{BasePairTally, DistanceExt, RawSnpDistanceOutput};
pub use stats::SiblingAnnexStats; pub use iter::{MinorantBatchIter, MinorantIter, SiblingBatchIter, SiblingEntry, SiblingIter, SiblingLayerExt};
pub use stats::{SiblingAnnexStats, SiblingStatsExt};
use obilayeredmap::OLMError; use obilayeredmap::OLMError;
use crate::error::OKIError; use obikindex::OKIError;
pub(super) const INDEX_SUBDIR: &str = "index"; pub(super) const INDEX_SUBDIR: &str = "index";
pub(super) const ANNEX_FILE_NAME: &str = "siblings.psib"; pub(super) const ANNEX_FILE_NAME: &str = "siblings.psib";
@@ -4,17 +4,17 @@ use obicompactvec::SiblingAnnex;
use obikpartitionner::KmerPartition; use obikpartitionner::KmerPartition;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use obikindex::{OKIError, OKIResult};
use crate::index::KmerIndex; use obikindex::KmerIndex;
use super::ANNEX_FILE_NAME; use super::ANNEX_FILE_NAME;
use super::cache::PartitionCache; use super::cache::PartitionCache;
use super::family_scan::scan_layer_families; use super::family_scan::scan_layer_families;
/// Distribution of family sizes (1-4), 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 /// annex (see [`super::build::SiblingAnnexBuildExt::build_sibling_annex`])
/// presence/count data — a separate, occasional diagnostic pass, not fused /// plus the index's presence/count data — a separate, occasional diagnostic
/// into construction. /// pass, not fused into construction.
/// ///
/// Every count here is **per family, not per slot**: a family with `F` /// Every count here is **per family, not per slot**: a family with `F`
/// members occupies `F` annex slots (one per observed member), all sharing /// members occupies `F` annex slots (one per observed member), all sharing
@@ -33,7 +33,9 @@ pub struct SiblingAnnexStats {
pub per_genome: Vec<[u64; 4]>, 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`, /// The global family-size histogram alone (`SiblingAnnexStats::counts`,
/// no `per_genome`) — reads only the already-built annex (`mask.siblings()` /// no `per_genome`) — reads only the already-built annex (`mask.siblings()`
/// + `mask.is_minorant()`, 1 byte/slot, mmap'd), nothing else: no /// + `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 /// cost to get there too; use this when only the global histogram is
/// needed. Requires an annex built after the minorant flag was added — /// 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. /// 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]> { fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]>;
let layer_dirs = self.sibling_layer_dirs()?;
/// 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]; let mut counts = [0u64; 4];
for layer_dir in &layer_dirs { for layer_dir in &layer_dirs {
@@ -70,16 +83,10 @@ impl KmerIndex {
Ok(counts) Ok(counts)
} }
/// Tally the family-size distribution of an already-built annex fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats> {
/// (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_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len(); let n_genomes = self.meta().genomes.len();
let with_counts = self.meta.config.with_counts; let with_counts = self.meta().config.with_counts;
let k = self.kmer_size(); let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize; 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 // why re-opening per lookup (or per call to a batching helper) is
// not good enough on a real index. // not good enough on a real index.
let partition = KmerPartition::open_with_config( let partition = KmerPartition::open_with_config(
&self.root_path, self.root_path(),
self.kmer_size(), self.kmer_size(),
self.minimizer_size(), self.minimizer_size(),
n_bits, n_bits,
) )
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); 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 // One layer at a time, not parallelised across layers — see
// `snp_pseudo_alignment`'s comment for why `par_iter()` over layers // `snp_pseudo_alignment`'s comment for why `par_iter()` over layers
@@ -8,11 +8,14 @@ use obilayeredmap::meta::PartitionMeta;
use obisys::Reporter; use obisys::Reporter;
use tempfile::tempdir; use tempfile::tempdir;
use crate::index::KmerIndex; use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode};
use crate::meta::{GenomeInfo, IndexConfig};
use crate::merge::MergeMode;
use super::alignment::SnpAlignmentExt;
use super::build::SiblingAnnexBuildExt;
use super::cardinality::CardinalityExt;
use super::distance::DistanceExt;
use super::helpers::is_minorant; use super::helpers::is_minorant;
use super::stats::SiblingStatsExt;
use super::{ANNEX_FILE_NAME, INDEX_SUBDIR}; use super::{ANNEX_FILE_NAME, INDEX_SUBDIR};
// k must be >= 11 (project constraint, "k ∈ [11,31]"); k=11, level_max=1, // 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 // Merge doesn't promise to preserve source order, so resolve each
// genome's index by label rather than assuming g1 -> 0, g2 -> 1. // 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")); let (i1, i2) = (idx_of("g1"), idx_of("g2"));
// snp_pseudo_alignment: one variable family, one column — g1's row // snp_pseudo_alignment: one variable family, one column — g1's row
+44
View File
@@ -367,6 +367,50 @@ impl<L: KmerLength> CanonicalKmerOf<L> {
pub fn into_kmer(self) -> KmerOf<L> { pub fn into_kmer(self) -> KmerOf<L> {
KmerOf(self.0, PhantomData) 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> { impl<L: KmerLength> Sequence for CanonicalKmerOf<L> {
+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) }) 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() } pub fn n(&self) -> usize { self.mphf.n() }
/// Raw MPHF lookup: kmer → slot, no membership check. /// Raw MPHF lookup: kmer → slot, no membership check.
@@ -187,6 +196,11 @@ impl Layer<PersistentCompactIntMatrix> {
.map_err(OLMError::Io) .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`. /// Extract a sub-matrix of counts for the rows at `slots`.
/// ///
/// Returns a column-first `Vec<Vec<u32>>`: one inner `Vec` per genome /// Returns a column-first `Vec<Vec<u32>>`: one inner `Vec` per genome
@@ -218,6 +232,14 @@ impl Layer<PersistentBitMatrix> {
.map_err(OLMError::Io) .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`. /// Extract a sub-matrix of presence/absence for the rows at `slots`.
/// ///
/// Returns a column-first `Vec<Vec<bool>>`: one inner `Vec` per genome /// Returns a column-first `Vec<Vec<bool>>`: one inner `Vec` per genome