Add modular data structures, parallel pipelines, and system profiling

Establishes foundational infrastructure across multiple crates by introducing unified persistent bit matrix storage with columnar, packed, and implicit variants, alongside De Bruijn graph node encoding and unitig iteration logic. Adds a macro-driven parallel pipeline scheduler featuring NUMA-aware runners, bounded channels, and memory budgets to enforce concurrency limits. Implements streaming nucleotide parsers with pooled page buffers for FASTA, FASTQ, and Genbank formats, complemented by system resource monitoring, progress tracking, and stage profiling utilities. Collectively, these changes provide the core data models, execution frameworks, and I/O pipelines required for downstream k-mer indexing and analysis workloads.
This commit is contained in:
Eric Coissac
2026-08-14 14:20:08 +02:00
parent cc67023e2c
commit 79346c0c86
69 changed files with 7487 additions and 7060 deletions
+17
View File
@@ -0,0 +1,17 @@
//! NUMA-aware partition runner via hwlocality.
//!
//! Detects NUMA topology using hwloc (cross-platform: Linux, macOS, etc.) and
//! builds one Rayon ThreadPool per NUMA node with threads pinned to that node's
//! CPUs. Linux first-touch policy then places graph allocations in local DRAM
//! automatically — no explicit memory binding needed.
//!
//! UMA systems (single socket, Apple Silicon, etc.) are the degenerate case:
//! one synthetic node containing all cores, no pool, no pinning.
//!
//! Submodules: [`topology`] (NUMA detection, per-node pools, thread pinning),
//! [`runner`] ([`PartitionRunner`], the adaptive worker-activation scheduler).
mod runner;
mod topology;
pub use runner::PartitionRunner;
@@ -1,137 +1,11 @@
// NUMA-aware partition runner via hwlocality.
//
// Detects NUMA topology using hwloc (cross-platform: Linux, macOS, etc.) and
// builds one Rayon ThreadPool per NUMA node with threads pinned to that node's
// CPUs. Linux first-touch policy then places graph allocations in local DRAM
// automatically — no explicit memory binding needed.
//
// UMA systems (single socket, Apple Silicon, etc.) are the degenerate case:
// one synthetic node containing all cores, no pool, no pinning.
use std::sync::Arc;
use std::time::{Duration, Instant};
use crossbeam_channel::unbounded;
#[cfg(feature = "numa")]
use hwlocality::Topology;
#[cfg(feature = "numa")]
use hwlocality::cpu::binding::CpuBindingFlags;
#[cfg(feature = "numa")]
use hwlocality::cpu::cpuset::CpuSet;
#[cfg(feature = "numa")]
use hwlocality::object::types::ObjectType;
use obisys::{CpuSample, IoSample};
use tracing::debug;
// ── Public interface ──────────────────────────────────────────────────────────
pub struct NumaSetup {
/// One entry per NUMA node. `None` on UMA systems (no pool, no pinning).
pub pools: Vec<Option<Arc<rayon::ThreadPool>>>,
/// CPU indices for each NUMA node, in node order.
pub cpus_per_node: Vec<Vec<usize>>,
}
impl NumaSetup {
/// Maximum worker slots per node (one per physical core in the node).
pub fn workers_per_node(&self) -> usize {
self.cpus_per_node
.first()
.map(|c| c.len().max(1))
.unwrap_or(1)
}
}
/// Detect NUMA topology and build per-node Rayon pools.
/// Always succeeds: falls back to a single synthetic UMA node on failure.
#[cfg(feature = "numa")]
pub fn build() -> NumaSetup {
if let Ok(topology) = Topology::new() {
let nodes: Vec<Vec<usize>> = topology
.objects_with_type(ObjectType::NUMANode)
.filter_map(|obj| obj.cpuset())
.map(|cpuset| {
cpuset
.iter_set()
.map(|idx| usize::from(idx))
.collect::<Vec<_>>()
})
.filter(|v| !v.is_empty())
.collect();
if nodes.len() > 1 {
if let Some(pools) = nodes
.iter()
.map(|cpus| build_pool(cpus).map(|p| Some(Arc::new(p))))
.collect::<Option<Vec<_>>>()
{
debug!(
"NUMA topology: {} node(s), {} core(s)/node",
nodes.len(),
nodes.first().map_or(0, |v| v.len()),
);
return NumaSetup {
pools,
cpus_per_node: nodes,
};
}
}
}
// UMA fallback: single synthetic node, all cores, no pool, no pinning.
let n_cores = obisys::effective_parallelism();
debug!("UMA: single synthetic node, {} core(s)", n_cores);
NumaSetup {
pools: vec![None],
cpus_per_node: vec![(0..n_cores).collect()],
}
}
#[cfg(not(feature = "numa"))]
pub fn build() -> NumaSetup {
let n_cores = obisys::effective_parallelism();
debug!("UMA: single synthetic node, {} core(s)", n_cores);
NumaSetup {
pools: vec![None],
cpus_per_node: vec![(0..n_cores).collect()],
}
}
/// Bind the calling thread to `cpu_indices` using hwloc.
/// Silently returns on any error so the thread still runs, just unbound.
#[cfg(feature = "numa")]
pub fn pin_current_thread(cpu_indices: &[usize]) {
let Ok(topology) = Topology::new() else {
return;
};
let mut cpuset = CpuSet::new();
for &idx in cpu_indices {
cpuset.set(idx);
}
let _ = topology.bind_cpu(&cpuset, CpuBindingFlags::THREAD);
}
#[cfg(not(feature = "numa"))]
pub fn pin_current_thread(_cpu_indices: &[usize]) {}
// ── Internal helpers ──────────────────────────────────────────────────────────
#[cfg(feature = "numa")]
fn build_pool(cpus: &[usize]) -> Option<rayon::ThreadPool> {
let cpus = cpus.to_vec();
rayon::ThreadPoolBuilder::new()
.num_threads(cpus.len())
.spawn_handler(move |thread| {
let cpus = cpus.clone();
std::thread::Builder::new().spawn(move || {
pin_current_thread(&cpus);
thread.run();
})?;
Ok(())
})
.build()
.ok()
}
use super::topology::{build, pin_current_thread};
// ── PartitionRunner ─────────────────────────────────────────────────────────
+121
View File
@@ -0,0 +1,121 @@
use std::sync::Arc;
#[cfg(feature = "numa")]
use hwlocality::Topology;
#[cfg(feature = "numa")]
use hwlocality::cpu::binding::CpuBindingFlags;
#[cfg(feature = "numa")]
use hwlocality::cpu::cpuset::CpuSet;
#[cfg(feature = "numa")]
use hwlocality::object::types::ObjectType;
use tracing::debug;
// ── Public interface ──────────────────────────────────────────────────────────
pub struct NumaSetup {
/// One entry per NUMA node. `None` on UMA systems (no pool, no pinning).
pub pools: Vec<Option<Arc<rayon::ThreadPool>>>,
/// CPU indices for each NUMA node, in node order.
pub cpus_per_node: Vec<Vec<usize>>,
}
impl NumaSetup {
/// Maximum worker slots per node (one per physical core in the node).
pub fn workers_per_node(&self) -> usize {
self.cpus_per_node
.first()
.map(|c| c.len().max(1))
.unwrap_or(1)
}
}
/// Detect NUMA topology and build per-node Rayon pools.
/// Always succeeds: falls back to a single synthetic UMA node on failure.
#[cfg(feature = "numa")]
pub fn build() -> NumaSetup {
if let Ok(topology) = Topology::new() {
let nodes: Vec<Vec<usize>> = topology
.objects_with_type(ObjectType::NUMANode)
.filter_map(|obj| obj.cpuset())
.map(|cpuset| {
cpuset
.iter_set()
.map(|idx| usize::from(idx))
.collect::<Vec<_>>()
})
.filter(|v| !v.is_empty())
.collect();
if nodes.len() > 1 {
if let Some(pools) = nodes
.iter()
.map(|cpus| build_pool(cpus).map(|p| Some(Arc::new(p))))
.collect::<Option<Vec<_>>>()
{
debug!(
"NUMA topology: {} node(s), {} core(s)/node",
nodes.len(),
nodes.first().map_or(0, |v| v.len()),
);
return NumaSetup {
pools,
cpus_per_node: nodes,
};
}
}
}
// UMA fallback: single synthetic node, all cores, no pool, no pinning.
let n_cores = obisys::effective_parallelism();
debug!("UMA: single synthetic node, {} core(s)", n_cores);
NumaSetup {
pools: vec![None],
cpus_per_node: vec![(0..n_cores).collect()],
}
}
#[cfg(not(feature = "numa"))]
pub fn build() -> NumaSetup {
let n_cores = obisys::effective_parallelism();
debug!("UMA: single synthetic node, {} core(s)", n_cores);
NumaSetup {
pools: vec![None],
cpus_per_node: vec![(0..n_cores).collect()],
}
}
/// Bind the calling thread to `cpu_indices` using hwloc.
/// Silently returns on any error so the thread still runs, just unbound.
#[cfg(feature = "numa")]
pub fn pin_current_thread(cpu_indices: &[usize]) {
let Ok(topology) = Topology::new() else {
return;
};
let mut cpuset = CpuSet::new();
for &idx in cpu_indices {
cpuset.set(idx);
}
let _ = topology.bind_cpu(&cpuset, CpuBindingFlags::THREAD);
}
#[cfg(not(feature = "numa"))]
pub fn pin_current_thread(_cpu_indices: &[usize]) {}
// ── Internal helpers ──────────────────────────────────────────────────────────
#[cfg(feature = "numa")]
fn build_pool(cpus: &[usize]) -> Option<rayon::ThreadPool> {
let cpus = cpus.to_vec();
rayon::ThreadPoolBuilder::new()
.num_threads(cpus.len())
.spawn_handler(move |thread| {
let cpus = cpus.clone();
std::thread::Builder::new().spawn(move || {
pin_current_thread(&cpus);
thread.run();
})?;
Ok(())
})
.build()
.ok()
}
File diff suppressed because it is too large Load Diff
+185
View File
@@ -0,0 +1,185 @@
use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache};
use super::helpers::{central_base, is_minorant, partition_of};
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set
/// iff the genome carries the member whose own central base is `b`):
/// single bit -> the plain base; 2 or 3 bits -> the matching IUPAC
/// ambiguity code (preserves partial information instead of collapsing to
/// `N`, the same convention used for diploid heterozygous VCF/FASTA sites);
/// all 4 bits -> `N`; no bits (genome carries none of the family's observed
/// members) -> `-` (no data at this locus for this genome).
fn iupac_code(mask: u8) -> u8 {
match mask & 0b1111 {
0b0000 => b'-',
0b0001 => b'A',
0b0010 => b'C',
0b0100 => b'G',
0b1000 => b'T',
0b0101 => b'R', // A/G
0b1010 => b'Y', // C/T
0b0110 => b'S', // C/G
0b1001 => b'W', // A/T
0b1100 => b'K', // G/T
0b0011 => b'M', // A/C
0b1110 => b'B', // C/G/T
0b1101 => b'D', // A/G/T
0b1011 => b'H', // A/C/T
0b0111 => b'V', // A/C/G
0b1111 => b'N',
_ => unreachable!("masked to 4 bits"),
}
}
/// A SNP-only pseudo-alignment: one row (byte sequence, IUPAC-coded) per
/// genome, one column per variable family (`family_size() >= 2` — monomorphic
/// families carry no signal and are skipped, unlike `raw_snp_distance`'s
/// tally which does count them as `shared`). Column order is the same,
/// deterministic sweep order as the annex build (partition, then layer, then
/// slot) — arbitrary but stable and identical across genomes, which is all a
/// pseudo-alignment needs (there is no natural genomic coordinate to sort by
/// once flanks are dropped). See `docmd/theory/evolutionary_distances.md`,
/// "Multi-genome framing: family as pseudo-alignment column".
pub struct SnpAlignment {
/// `sequences[g]` = genome `g`'s IUPAC-coded row, same length for every
/// genome (`sequences.len()` columns).
pub sequences: Vec<Vec<u8>>,
}
impl KmerIndex {
/// Build the SNP-only pseudo-alignment from an already-built sibling
/// annex (run [`build_sibling_annex`](Self::build_sibling_annex) first).
pub fn snp_pseudo_alignment(&self) -> OKIResult<SnpAlignment> {
let n_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len();
let with_counts = self.meta.config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config(
&self.root_path,
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers");
// `Vec<Vec<u8>>` per layer, one entry (column) per variable family;
// `par_iter().map(...).collect()` on this indexed source preserves
// input order, so concatenating the results below in order gives a
// single deterministic column order across the whole index.
let partials: Vec<Vec<Vec<u8>>> = layer_dirs
.par_iter()
.map(|layer_dir| -> OKIResult<Vec<Vec<u8>>> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
let mut columns: Vec<Vec<u8>> = Vec::new();
let mut genome_mask: Vec<u8> = Vec::with_capacity(n_genomes);
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // family tallied once, at its minorant
}
if mask.family_size() < 2 {
continue; // monomorphic family — no signal, skip
}
genome_mask.clear();
genome_mask.resize(n_genomes, 0);
for other in kmer.central_canonical_neighbors() {
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let presence: Option<Vec<bool>> = if other == kmer {
Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
} else {
let dest = partition_of(other, n_parts);
cache.find_presence(dest, other, n_genomes)
};
let Some(presence) = presence else { continue };
for (g, &present) in presence.iter().enumerate() {
if present {
genome_mask[g] |= 1 << base;
}
}
}
columns.push(genome_mask.iter().map(|&m| iupac_code(m)).collect());
}
pb.inc(1);
Ok(columns)
})
.collect::<OKIResult<Vec<_>>>()?;
pb.finish_and_clear();
let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); n_genomes];
for layer_columns in partials {
for column in layer_columns {
for (g, &code) in column.iter().enumerate() {
sequences[g].push(code);
}
}
}
Ok(SnpAlignment { sequences })
}
}
+255
View File
@@ -0,0 +1,255 @@
use std::path::Path;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::Arc;
use rayon::prelude::*;
use obicompactvec::{FamilyMask, SiblingAnnexBuilder};
use obikpartitionner::KmerPartition;
use obipipeline::ThrottleGuard;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use super::cache::PartitionCache;
use super::helpers::{central_base, partition_of};
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
// ── obipipeline data types ─────────────────────────────────────────────────
/// A batch of this layer's distinct k-mers (local MPHF slot + k-mer), the
/// pipeline's source item — batched, not one k-mer per item, so that
/// pipeline messages and their synchronisation cost stay amortised over
/// thousands of lookups (see `build_layer_sibling_annex`'s comment on
/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved
/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform,
/// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
struct SourceBatch {
items: Vec<(usize, CanonicalKmer)>,
_permit: ThrottleGuard,
}
/// One batch's worth of central-substitution variants (up to 3 per source
/// k-mer), each already routed to its destination partition and carrying
/// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a
/// hit. `(dest_partition, variant, source_slot, base)` per entry.
struct VariantBatch {
items: Vec<(usize, CanonicalKmer, usize, u8)>,
_permit: ThrottleGuard,
}
enum SibData {
Batch(SourceBatch),
Variants(VariantBatch),
}
impl KmerIndex {
/// Build the sibling-count/minorant annex for every layer of every
/// partition of this (already built) index, writing one annex file per
/// layer alongside its existing index files. Safe to call again later
/// (e.g. after a fresh `merge`) — each run simply overwrites the annex
/// files of the index it is called on.
///
/// Construction only — no statistics gathered here on purpose: this is
/// meant to run routinely (it is the artefact the SNP-family distances
/// will consume), while the sibling-count distribution
/// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate,
/// occasional diagnostic pass over the result, not run every time.
///
/// Cross-partition/cross-layer lookups are required (a k-mer's siblings
/// can live in any partition), but the layer loop itself — and thus the
/// annex file this produces — stays local to one layer at a time.
pub fn build_sibling_annex(&self) -> OKIResult<()> {
let n_parts = self.n_partitions();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config(
&self.root_path,
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
let cache = Arc::new(PartitionCache::build(&partition, n_parts, self.meta.config.with_counts)?);
let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
let mut total_slots: u64 = 0;
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
pb.inc(1);
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
let mut part_slots: u64 = 0;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?;
}
total_slots += part_slots;
pb.inc(1);
pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)"));
}
pb.finish_and_clear();
tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions");
Ok(())
}
/// Returns the number of distinct k-mers (annex slots) processed, for
/// progress reporting.
fn build_layer_sibling_annex(
&self,
layer_dir: &Path,
n_parts: usize,
cache: &Arc<PartitionCache>,
) -> OKIResult<u64> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let n_slots = mphf.n();
// ── Enumerate this layer's distinct k-mers, one per slot ────────────
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; n_slots];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let k = self.kmer_size();
// ── Reconciliation state, initialised with each slot's own base —
// that member is trivially present, no lookup needed. Built before
// the pipeline runs, from the same enumeration, since `sources`
// below is consumed as a throttled iterator, not collected.
// `AtomicU8`, not `FamilyMask`, because the gather phase below
// parallelises across destination partitions (independent
// `query_partition_with` calls, safe to run concurrently) and their
// `Found` hits can land on arbitrary, possibly-shared slots — a
// lock-free `fetch_or` avoids needing any synchronisation beyond
// that. ─────────────────────────────────────────────────────────
let mask: Vec<AtomicU8> = (0..n_slots).map(|_| AtomicU8::new(0)).collect();
for (slot, kmer) in slot_kmer.iter().enumerate().filter_map(|(s, k)| k.map(|k| (s, k))) {
mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
}
// ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one —
// the actual cross-partition lookup reuses
// `KmerPartition::query_partition_with` (the same batching mechanism
// `obikmer query` already uses: open a partition's files once,
// answer a whole batch of queries against it) instead of one lookup
// per pipeline item. A per-item lookup (tried first) reopened/
// re-mmap'd every target partition's files on every single variant
// — fine at toy scale, but ~90% system time against a real index,
// observed in practice. A *later* attempt still pushed one pipeline
// message per generated variant (a `Flat` stage, `SourceItem` =>
// `VariantQuery`, one k-mer in => up to 3 variants out as separate
// messages) — cheaper than reopening files, but sampling a real run
// showed most wall-clock time going into per-message channel
// send/notify syscalls instead of the lookup itself: the pipeline's
// whole point is amortising synchronisation over a batch, and a
// single k-mer's ≤3 variants is far too fine a granularity for
// that. Batching `BATCH_SIZE` source k-mers into one pipeline item
// — a plain 1-to-1 (`|`, not `||`) transform, batch in, batch of
// variants out, one message either way — keeps the per-message
// synchronisation cost amortised over thousands of lookups instead
// of one to three. ──────────────────────────────────────────────
const BATCH_SIZE: usize = 4096;
let n_workers = obisys::effective_parallelism();
let capacity = 256;
// Throttling limits how many *batches* are in flight at once — the
// permit is acquired per batch (not per k-mer) in the source
// thread, and released once its `VariantBatch` has been read out of
// the pipeline by the accumulation loop below. See
// `obipipeline::throttle`'s docs for why this is required, not
// optional, once a `Flat`-style stage sits in the pipeline.
let sources: Vec<(usize, CanonicalKmer)> = slot_kmer
.iter()
.enumerate()
.filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer)))
.collect();
let batches: Vec<Vec<(usize, CanonicalKmer)>> = sources
.chunks(BATCH_SIZE)
.map(|chunk| chunk.to_vec())
.collect();
let throttled = obipipeline::throttle(batches.into_iter(), n_workers).map(|t| SourceBatch {
items: t.item,
_permit: t.guard,
});
let pipe = obipipeline::make_pipe! {
SibData : SourceBatch => VariantBatch,
| {
move |batch: SourceBatch| -> VariantBatch {
let mut items = Vec::with_capacity(batch.items.len() * 3);
for (slot, kmer) in batch.items {
for variant in kmer.central_canonical_neighbors() {
if variant == kmer {
continue;
}
items.push((
partition_of(variant, n_parts),
variant,
slot,
central_base(variant, k),
));
}
}
VariantBatch { items, _permit: batch._permit }
}
} : Batch => Variants,
};
// ── Group generated variants by destination partition. `cache`
// holds every partition already mmap'd (no more `open()` cost), but
// `mmap` pages are still loaded on demand and can be evicted — a
// lookup is not free just because the file isn't reopened. Grouping
// keeps one partition's pages hot while its whole batch is resolved,
// instead of faulting pages in and out as lookups jump between
// partitions in whatever order the pipeline happens to produce
// them. Each batch's throttle permit drops here, once accumulated.
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
for vb in pipe.apply(throttled, n_workers, capacity) {
for (dest_partition, variant, source_slot, base) in vb.items {
outgoing[dest_partition].push((variant, source_slot, base));
}
}
// ── Resolve each partition's batch against the cache in one
// contiguous pass; parallelised across partitions (independent,
// read-only) so this keeps using multiple cores without giving up
// the per-partition locality above. ─────────────────────────────
outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
for &(variant, source_slot, base) in queries {
if cache.find(dest, variant) {
mask[source_slot].fetch_or(1 << base, Ordering::Relaxed);
}
}
});
// ── Write the layer's annex file ─────────────────────────────────────
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?;
for (slot, m) in mask.iter().enumerate() {
if slot_kmer[slot].is_none() {
continue; // unused MPHF slot, if any — leave at the sentinel
}
builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed)));
}
builder.close()?;
Ok(n_slots as u64)
}
}
+121
View File
@@ -0,0 +1,121 @@
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()))
}
/// Per-genome presence vector for `variant` in partition `dest_partition`
/// (`true` iff that genome carries it), `None` on a miss. Same shape as
/// `find`, but also reads the cached matrix instead of just the MPHF.
pub(super) fn find_presence(&self, dest_partition: usize, variant: CanonicalKmer, n_genomes: usize) -> Option<Vec<bool>> {
let layers = self.layers.get(dest_partition)?;
let mats = self.mats.get(dest_partition)?;
for (mphf, mat) in layers.iter().zip(mats.iter()) {
if let Some(slot) = mphf.find(variant) {
let n_cols = mat.n_cols().min(n_genomes);
return Some((0..n_cols).map(|g| mat.carries(g, slot)).collect());
}
}
None
}
}
+204
View File
@@ -0,0 +1,204 @@
use ndarray::Array2;
use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache};
use super::distance::RawSnpDistanceOutput;
use super::helpers::{central_base, is_minorant, partition_of};
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// See [`KmerIndex::cardinality_tally`].
pub struct CardinalityTally {
/// `counts[a][b] == counts[b][a]` = number of family sites, pooled over
/// included genome pairs, where one genome's family cardinality
/// (popcount of its presence mask, `0..=4`) is `a` and the other's is
/// `b`. Diagonal is real data here (both genomes at the same
/// cardinality), unlike [`super::distance::BasePairTally::counts`].
pub counts: [[u64; 5]; 5],
}
impl KmerIndex {
/// 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
/// [`base_pair_tally`](Self::base_pair_tally). Unlike
/// `scan_family_pairs` (which resolves each genome to a single form and
/// silently drops any genome carrying more than one member of the
/// family), this needs the *full* per-genome presence mask — a family
/// member count of 2, 3 or 4 is exactly the signal being tallied, not
/// noise to discard — so it re-implements the traversal rather than
/// reusing that helper.
///
/// Restricted to variable families (`family_size() >= 2`), matching
/// `snp_pseudo_alignment`'s own scope — briefly removed, then
/// reinstated: without it, the diagonal is dominated by genome-wide
/// invariant background (family_size()<2 loci vastly outnumber the
/// ones that ever vary anywhere), which is inconsistent with the
/// `+ASC`-corrected alignment this matrix is ultimately used with —
/// `+ASC` exists specifically because the likelihood only ever sees
/// variable sites, so a rate model calibrated mostly from invariant
/// background sites doesn't describe the population it's applied to.
/// Verified empirically: removing the filter measurably worsened a
/// real IQ-TREE run (log-likelihood dropped, `NNI search needs
/// unusual large number of steps to converge` warnings appeared) — see
/// `docmd/theory/evolutionary_distances.md` for the full account.
/// [`base_pair_tally`](Self::base_pair_tally)'s own diagonal (`same`)
/// 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> {
let n_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len();
let with_counts = self.meta.config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
if i == j {
return false;
}
let snp = raw.snp[[i, j]];
let total = snp + raw.shared[[i, j]];
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
});
let partition = KmerPartition::open_with_config(
&self.root_path,
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers");
let partials: Vec<[[u64; 5]; 5]> = layer_dirs
.par_iter()
.map(|layer_dir| -> OKIResult<[[u64; 5]; 5]> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
let mut counts = [[0u64; 5]; 5];
let mut genome_mask: Vec<u8> = Vec::with_capacity(n_genomes);
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // family tallied once, at its minorant
}
if mask.family_size() < 2 {
// Fully invariant family (never varies anywhere in
// the index) — genome-wide background, not
// SNP-adjacent signal; would otherwise swamp the
// diagonal (`c=1/c=1` etc.), which needs to reflect
// the same variable-families-only population the
// `+ASC`-corrected alignment/likelihood actually
// models. See `base_pair_tally`'s `variable` gate
// on its own `same` diagonal for the matching fix.
continue;
}
genome_mask.clear();
genome_mask.resize(n_genomes, 0);
for other in kmer.central_canonical_neighbors() {
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let presence: Option<Vec<bool>> = if other == kmer {
Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
} else {
let dest = partition_of(other, n_parts);
cache.find_presence(dest, other, n_genomes)
};
let Some(presence) = presence else { continue };
for (g, &present) in presence.iter().enumerate() {
if present {
genome_mask[g] |= 1 << base;
}
}
}
for i in 0..n_genomes {
let card_i = genome_mask[i].count_ones() as usize;
for j in (i + 1)..n_genomes {
if !included[[i, j]] {
continue;
}
let card_j = genome_mask[j].count_ones() as usize;
counts[card_i][card_j] += 1;
if card_i != card_j {
counts[card_j][card_i] += 1;
}
}
}
}
pb.inc(1);
Ok(counts)
})
.collect::<OKIResult<Vec<_>>>()?;
pb.finish_and_clear();
let mut total = [[0u64; 5]; 5];
for partial in partials {
for a in 0..5 {
for b in 0..5 {
total[a][b] += partial[a][b];
}
}
}
Ok(CardinalityTally { counts: total })
}
}
+305
View File
@@ -0,0 +1,305 @@
use ndarray::Array2;
use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache};
use super::helpers::{central_base, is_minorant, partition_of};
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// Raw p-distance restricted to loci that are single-copy in **both**
/// genomes of a pair — the "stringent / paralogy-aware" locus eligibility
/// rule (`docmd/theory/evolutionary_distances.md`, "Locus eligibility"),
/// without the JC/K2P/LogDet finalisation math: `snp[i,j]` + `shared[i,j]`
/// is the number of eligible loci, `snp[i,j] / (snp[i,j] + shared[i,j])` is
/// `p_hat`. A quick, self-contained way to sanity-check the estimator
/// against a real index before the full `SnpTally` design is built.
///
/// A locus (family, tallied once at its minorant) is eligible for pair
/// `(i, j)` iff genome `i` carries exactly one of the family's observed
/// forms **and** genome `j` carries exactly one (possibly a different one)
/// — presence-only: a genome carrying the same form twice (a same-allele
/// duplicate) is indistinguishable from carrying it once when only a
/// presence matrix is available, so such cases are not excluded here even
/// when a count index exists. See "Locus eligibility", stringent rule, for
/// why this matters and how a count index would close the gap — left as a
/// follow-up, not applied here.
pub struct RawSnpDistanceOutput {
/// n×n count of eligible loci where the two genomes' single forms differ.
pub snp: Array2<u64>,
/// n×n count of eligible loci where the two genomes' single forms agree.
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 in parallel (rayon); each gets its own accumulator
/// from `zero()`, combined pairwise via `combine`.
fn scan_family_pairs<Acc, F, C>(
&self,
label: &str,
zero: impl Fn() -> Acc + Sync,
on_pair: F,
combine: C,
) -> 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_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config(
&self.root_path,
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
let partials: Vec<Acc> = layer_dirs
.par_iter()
.map(|layer_dir| -> OKIResult<Acc> {
let mut acc = zero();
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
// Per family: which single form (if exactly one) each genome
// carries — `None` once a second form is seen (ambiguous,
// not single-copy, ineligible for either side of a pair).
let mut single_form: Vec<Option<u8>> = Vec::with_capacity(n_cols);
let mut ambiguous: Vec<bool> = Vec::with_capacity(n_cols);
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // family tallied once, at its minorant
}
let variable = mask.family_size() >= 2;
single_form.clear();
single_form.resize(n_cols, None);
ambiguous.clear();
ambiguous.resize(n_cols, false);
for other in kmer.central_canonical_neighbors() {
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let presence: Option<Vec<bool>> = if other == kmer {
Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
} else {
let dest = partition_of(other, n_parts);
cache.find_presence(dest, other, n_genomes)
};
let Some(presence) = presence else { continue };
for (g, &present) in presence.iter().enumerate() {
if !present {
continue;
}
if single_form[g].is_some() {
ambiguous[g] = true;
} else {
single_form[g] = Some(base);
}
}
}
for i in 0..n_cols {
if ambiguous[i] {
continue;
}
let Some(bi) = single_form[i] else { continue };
for j in (i + 1)..n_cols {
if ambiguous[j] {
continue;
}
let Some(bj) = single_form[j] else { continue };
on_pair(&mut acc, i, j, bi, bj, variable);
}
}
}
pb.inc(1);
Ok(acc)
})
.collect::<OKIResult<Vec<_>>>()?;
pb.finish_and_clear();
let mut total = zero();
for partial in partials {
total = combine(total, partial);
}
Ok(total)
}
/// 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(
"raw_snp_distance",
|| (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::zeros((n_genomes, n_genomes))),
|(snp, shared), i, j, bi, bj, _variable| {
if bi == bj {
shared[[i, j]] += 1;
shared[[j, i]] += 1;
} else {
snp[[i, j]] += 1;
snp[[j, i]] += 1;
}
},
|(mut snp, mut shared), (s, sh)| {
snp += &s;
shared += &sh;
(snp, shared)
},
)?;
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();
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
if i == j {
return false;
}
let snp = raw.snp[[i, j]];
let total = snp + raw.shared[[i, j]];
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
});
let (counts, same) = self.scan_family_pairs(
"base_pair_tally",
|| ([[0u64; 4]; 4], [0u64; 4]),
|(counts, same), i, j, bi, bj, variable| {
if !included[[i, j]] {
return;
}
if bi != bj {
counts[bi as usize][bj as usize] += 1;
counts[bj as usize][bi as usize] += 1;
} else if variable {
// Only count "stayed the same" from families that vary
// *somewhere* in the index — a fully invariant family
// (never varies anywhere) isn't a SNP-adjacent
// agreement, it's genome-wide background, and would
// otherwise swamp the diagonal (see
// `docmd/theory/evolutionary_distances.md`, the
// ascertainment-bias regression this was reverting).
same[bi as usize] += 1;
}
},
|(mut counts, mut same), (partial_counts, partial_same)| {
for a in 0..4 {
same[a] += partial_same[a];
for b in 0..4 {
counts[a][b] += partial_counts[a][b];
}
}
(counts, same)
},
)?;
Ok(BasePairTally { counts, same })
}
}
/// See [`KmerIndex::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
/// `a` and `b` (0=A, 1=C, 2=G, 3=T). Diagonal always `0` — an `a == b`
/// locus is counted in `same`, not here.
pub counts: [[u64; 4]; 4],
/// `same[a]` = number of eligible loci, pooled over included genome
/// pairs, where both genomes' single forms are `a` — the diagonal
/// `counts` omits, needed to build a proper row-stochastic composition
/// probability matrix (the "stay the same base" entries), not just the
/// substitution-cost off-diagonal.
pub same: [u64; 4],
}
+49
View File
@@ -0,0 +1,49 @@
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
}
+73
View File
@@ -0,0 +1,73 @@
//! Family presence-mask annex construction.
//!
//! See `docmd/theory/evolutionary_distances.md`, "Definitions: family, and
//! the canonical form of a family" and "Step 2b", for the full design
//! discussion this implements.
//!
//! For each distinct k-mer of each layer of the (already built/merged)
//! index, computes a 4-bit presence mask for its "family" (the up to 4
//! k-mers sharing its flanks, differing only at the central base —
//! well-defined for odd k): bit `b` set iff the family member whose own
//! canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed anywhere
//! in the current multi-genome index — a property of the whole index, not
//! of any one genome. Sibling count and minorant are *derived* from the
//! mask by callers, not stored (see `FamilyMask` and
//! [`sibling_annex_stats`](crate::index::KmerIndex::sibling_annex_stats)
//! below).
//!
//! Per layer, an `obipipeline` batch transform (throttled — see
//! `obipipeline::throttle`) generates a whole batch's central variants at
//! once (`BATCH_SIZE` source k-mers in, that batch's variants out as one
//! pipeline message), interleaved across many in-flight batches by the
//! scheduler's shared worker pool rather than processed on a single
//! thread. The actual cross-partition lookup reuses a `PartitionCache` of
//! every partition's already-open MPHF layers, built once for the whole
//! `build_sibling_annex` run, rather than reopening files per lookup or
//! per source layer. Two earlier, coarser-grained designs were tried and
//! measured (not guessed) to be worse, in order: (1) reopening/re-mmap'ing
//! every target partition's files on every single lookup — fine at toy
//! scale, ~90% system time against a real index; (2) a `Flat` pipeline
//! stage pushing one message per generated *variant* (up to 3 per source
//! k-mer) — cheaper than reopening files, but sampling a real run showed
//! most wall-clock time going into per-message channel send/notify
//! syscalls rather than the lookup itself, because a single k-mer's ≤3
//! variants is far too fine a granularity to amortise a pipeline's
//! synchronisation cost over. See `docmd/theory/evolutionary_distances.md`,
//! Step 2b, "Mechanism".
//!
//! Submodules, in the order data flows through them: [`cache`] (shared
//! whole-run partition cache), [`helpers`] (small pure functions used
//! throughout), [`build`] (annex construction), [`stats`] (family-size
//! diagnostics), [`distance`] (raw SNP distance + base-pair tally),
//! [`cardinality`] (cardinality co-occurrence), [`alignment`] (SNP-only
//! pseudo-alignment).
mod alignment;
mod build;
mod cache;
mod cardinality;
mod distance;
mod helpers;
mod stats;
#[cfg(test)]
mod tests;
pub use alignment::SnpAlignment;
pub use cardinality::CardinalityTally;
pub use distance::{BasePairTally, RawSnpDistanceOutput};
pub use stats::SiblingAnnexStats;
use obilayeredmap::OLMError;
use crate::error::OKIError;
pub(super) const INDEX_SUBDIR: &str = "index";
pub(super) const ANNEX_FILE_NAME: &str = "siblings.psib";
pub(super) fn olm_to_ok(e: OLMError) -> OKIError {
match e {
OLMError::Io(e) => OKIError::Io(e),
other => OKIError::InvalidInput(format!("layered-map error: {other}")),
}
}
+192
View File
@@ -0,0 +1,192 @@
use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache};
use super::helpers::{central_base, is_minorant, partition_of};
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// Distribution of family sizes (1-4), read back from an already-built
/// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's
/// presence/count data — a separate, occasional diagnostic pass, not fused
/// into construction.
///
/// Every count here is **per family, not per slot**: a family with `F`
/// members occupies `F` annex slots (one per observed member), all sharing
/// the same mask. Counting every slot would count each family up to 4
/// times over; only the minorant's slot is tallied (minorant is derived on
/// the fly — see `is_minorant` — not stored, but cheap: no lookup, pure
/// bit arithmetic on already-in-hand data).
#[derive(Debug, Clone, Default)]
pub struct SiblingAnnexStats {
/// `counts[s]` = number of *families* of size `s + 1` (index 0 = size 1,
/// i.e. 0 siblings, ... index 3 = size 4, i.e. 3 siblings).
pub counts: [u64; 4],
/// `per_genome[g][s]` = number of families of size `s + 1` for which
/// genome `g` (index into `KmerIndex::meta().genomes`) carries at least
/// one member.
pub per_genome: Vec<[u64; 4]>,
}
impl KmerIndex {
/// Tally the family-size distribution of an already-built annex
/// (globally, and per genome), counting each family once (at its
/// minorant slot). Errors if [`build_sibling_annex`] has not been run on
/// this index first.
///
/// [`build_sibling_annex`]: Self::build_sibling_annex
pub fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats> {
let n_parts = self.n_partitions();
let n_genomes = self.meta.genomes.len();
let with_counts = self.meta.config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
// Same whole-run cache as `build_sibling_annex` — see its docs for
// why re-opening per lookup (or per call to a batching helper) is
// not good enough on a real index.
let partition = KmerPartition::open_with_config(
&self.root_path,
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
// Gather the (partition, layer) pairs to process — cheap metadata
// reads only, checking every annex file exists up front so a
// missing one is reported before any real work starts.
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
// One layer's worth of work, parallelised across layers with Rayon
// — independent, read-only, each producing its own partial tally
// merged at the end.
let pb = progress_bar("sibling_annex_stats", layer_dirs.len() as u64, "layers");
let partials: Vec<SiblingAnnexStats> = layer_dirs
.par_iter()
.map(|layer_dir| -> OKIResult<SiblingAnnexStats> {
let mut stats = SiblingAnnexStats {
per_genome: vec![[0u64; 4]; n_genomes],
..Default::default()
};
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
// Need each slot's own k-mer to derive minorant — same
// enumeration as construction.
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // this family is tallied at its minorant's slot only
}
let s = mask.siblings() as usize;
stats.counts[s] += 1;
// "Genome g represents this family" means g carries
// *any* of its members, not just the minorant's own —
// start from the minorant's own presence (already
// open, no lookup) and OR in every other present
// member's presence vector, resolved against the
// whole-run cache (no I/O) — exactly `mask.siblings()`
// of them, the mask tells us precisely which to fetch.
let mut carries = vec![false; n_cols];
for g in 0..n_cols {
carries[g] = mat.carries(g, slot);
}
for other in kmer.central_canonical_neighbors() {
if other == kmer {
continue;
}
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let dest = partition_of(other, n_parts);
if let Some(other_presence) = cache.find_presence(dest, other, n_genomes) {
for (g, &present) in other_presence.iter().enumerate() {
if present {
carries[g] = true;
}
}
}
}
for (g, &carried) in carries.iter().enumerate() {
if carried {
stats.per_genome[g][s] += 1;
}
}
}
pb.inc(1);
Ok(stats)
})
.collect::<OKIResult<Vec<_>>>()?;
pb.finish_and_clear();
let mut stats = SiblingAnnexStats {
per_genome: vec![[0u64; 4]; n_genomes],
..Default::default()
};
for part in partials {
for s in 0..4 {
stats.counts[s] += part.counts[s];
}
for g in 0..n_genomes {
for s in 0..4 {
stats.per_genome[g][s] += part.per_genome[g][s];
}
}
}
Ok(stats)
}
}
+186
View File
@@ -0,0 +1,186 @@
use std::io::Write;
use std::path::Path;
use obicompactvec::{FamilyMask, SiblingAnnex};
use obikseq::{CanonicalKmer, Kmer, Sequence};
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obisys::Reporter;
use tempfile::tempdir;
use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexConfig};
use crate::merge::MergeMode;
use super::helpers::is_minorant;
use super::{ANNEX_FILE_NAME, INDEX_SUBDIR};
// k must be >= 11 (project constraint, "k ∈ [11,31]"); k=11, level_max=1,
// theta=0.0 mirror `obiskbuilder`'s own tests (smaller k/level_max
// combinations trip an unrelated pre-existing bug in `obikentropy`'s
// sliding-window ring buffer — not this feature's concern).
const K: usize = 11;
const M: usize = 5;
/// Build a single-genome index from one in-memory FASTA sequence, driving
/// the same primitives `obikmer`'s `scatter` step uses (minus the
/// multi-file `obipipeline` wrapper — a single sequence needs none of
/// that): normalise -> build superkmers -> route -> write.
/// `cargo test` doesn't install a `tracing` subscriber the way `obikmer`'s
/// CLI does, so `debug!`/etc. are silent no-ops by default — including the
/// `PartitionRunner` instrumentation that would matter most for
/// re-diagnosing a hang here. `try_init` is idempotent across concurrently
/// running tests (later calls just find a subscriber already installed).
fn init_tracing() {
let _ = tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.with_writer(std::io::stderr)
.try_init();
}
fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
init_tracing();
let fasta_path = dir.join(format!("{label}.fasta"));
let mut f = std::fs::File::create(&fasta_path).unwrap();
writeln!(f, ">{label}").unwrap();
f.write_all(seq).unwrap();
writeln!(f).unwrap();
drop(f);
let index_path = dir.join(format!("{label}.idx"));
let config = IndexConfig {
kmer_size: K,
minimizer_size: M,
n_bits: 0, // 1 partition — keeps the test deterministic and simple
with_counts: false,
evidence: obilayeredmap::IndexMode::Exact,
block_bits: 0,
};
let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)), false)
.expect("create");
let mut rep = Reporter::new();
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
for page in stream {
let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
idx.partition_mut().write_batch(batch).expect("write_batch");
}
idx.partition_mut().close().expect("close partition writers");
idx.mark_scattered().expect("mark_scattered");
idx.dereplicate_and_count(false, &mut rep).expect("dereplicate_and_count");
idx.build_layers(1, None, false, &mut rep).expect("build_layers");
idx
}
fn canonical(ascii: &[u8]) -> CanonicalKmer {
Kmer::from_ascii(ascii).unwrap().canonical()
}
/// Read back the annex entry for a given canonical k-mer from the merged
/// index's (single) partition/layer, asserting it was found at all.
fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> FamilyMask {
let index_dir = idx.partition().part_dir(0).join(INDEX_SUBDIR);
let meta = PartitionMeta::load(&index_dir).unwrap();
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let mphf = MphfLayer::open(&layer_dir, &meta.mode).unwrap();
if let Some(slot) = mphf.find(kmer) {
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).unwrap();
return annex.get(slot).expect("slot must have a computed annex entry");
}
}
panic!("kmer not found in any layer of partition 0");
}
fn merge_two(dir: &Path, g1: &KmerIndex, g2: &KmerIndex) -> KmerIndex {
let mut rep = Reporter::new();
KmerIndex::merge(
&dir.join("merged.idx"),
&[g1, g2],
MergeMode::Presence,
false,
false,
1.0,
&mut rep,
)
.expect("merge")
}
#[test]
fn sibling_annex_one_sibling_each() {
// k=11, centre = index 5 (0-based). Two genomes, each exactly one
// k-mer, sharing every base except the centre:
// g1 = "AACCGCTTAAG" (centre 'C', base index 1)
// g2 = "AACCGGTTAAG" (centre 'G', base index 2)
// Hand-verified: both stay forward-oriented under canonicalisation
// (each is lexicographically smaller than its own reverse
// complement, since both start with "AA"), and raw(g1) < raw(g2)
// (only differing base: C=0b01 < G=0b10 at the centre) — so g1 is
// the minorant, g2 is not. The mask is a family-wide value: both
// slots must read back the *same* mask (bits 1 and 2 set).
let dir = tempdir().unwrap();
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
let g1_kmer = canonical(b"AACCGCTTAAG");
let g2_kmer = canonical(b"AACCGGTTAAG");
let expected_mask = FamilyMask::EMPTY.with(1).with(2);
let a = annex_info_for(&merged, g1_kmer);
assert_eq!(a, expected_mask, "AACCGCTTAAG");
assert_eq!(a.siblings(), 1);
assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant");
let b = annex_info_for(&merged, g2_kmer);
assert_eq!(b, expected_mask, "AACCGGTTAAG");
assert_eq!(b.siblings(), 1);
assert!(!is_minorant(g2_kmer, b, K), "g2 should not be the minorant");
}
#[test]
fn sibling_annex_zero_siblings_when_identical_and_no_variant_exists() {
// Same k-mer in both genomes, no other genome around to carry a
// variant -> 0 siblings, trivially its own minorant.
let dir = tempdir().unwrap();
let g1 = build_single_genome_index(dir.path(), "g1", b"GATTACAGATC");
let g2 = build_single_genome_index(dir.path(), "g2", b"GATTACAGATC");
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
let kmer = canonical(b"GATTACAGATC");
let mask = annex_info_for(&merged, kmer);
assert_eq!(mask.siblings(), 0, "GATTACAGATC");
assert_eq!(mask.family_size(), 1);
assert!(is_minorant(kmer, mask, K));
}
#[test]
fn sibling_annex_stats_counts_each_family_once_and_per_genome() {
// Reuses the one-sibling-each fixture: a single family of size 2
// (g1's centre-C form + g2's centre-G form), each genome carrying
// exactly one of the two members. Stats must report exactly one
// family of size 2 (`counts[1] == 1`, since index 1 = size 2), not
// two (which naively summing both slots would give), and both
// genomes represented at size 2, neither at any other size.
let dir = tempdir().unwrap();
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
let stats = merged.sibling_annex_stats().expect("sibling_annex_stats");
assert_eq!(stats.counts, [0, 1, 0, 0], "one family of size 2, counted once");
assert_eq!(stats.per_genome.len(), 2);
for g in 0..2 {
assert_eq!(
stats.per_genome[g], [0, 1, 0, 0],
"genome {g} should represent exactly one size-2 family"
);
}
}