From 23812d1af860e1dc8dfc9395089730007770b632 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Sat, 22 Aug 2026 13:56:50 +0200 Subject: [PATCH] feat: add col_weights API and refactor obikstats to use Algorithm trait Added a `col_weights` method to index layers for computing per-genome column sums or presence k-mer counts. Refactored `obikstats` to implement the `Algorithm` trait with a two-phase `new`/`run` model, replacing manual layer resolution with `IndexCache` for eager file I/O. Simplified per-genome counting logic and updated public exports and dependencies accordingly. --- src/Cargo.lock | 2 + src/obikindex/src/layer/content_layer.rs | 10 ++ src/obikindex/src/layer/typed_layer.rs | 15 +- src/obikstats/Cargo.toml | 2 + src/obikstats/src/lib.rs | 15 +- src/obikstats/src/stats.rs | 202 +++++++++++------------ 6 files changed, 137 insertions(+), 109 deletions(-) diff --git a/src/Cargo.lock b/src/Cargo.lock index 63fec439..7f35c589 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1767,6 +1767,8 @@ name = "obikstats" version = "0.1.0" dependencies = [ "obicompactvec", + "obikalgorithm", + "obikidxcache", "obikindex", "rayon", ] diff --git a/src/obikindex/src/layer/content_layer.rs b/src/obikindex/src/layer/content_layer.rs index 10e7dbc1..b61b1df2 100644 --- a/src/obikindex/src/layer/content_layer.rs +++ b/src/obikindex/src/layer/content_layer.rs @@ -217,6 +217,16 @@ impl KmerLayer { } } + /// Per-genome column weights — count sum or presence k-mer count, + /// depending on content (see `obicompactvec::ColumnWeights`). + pub fn col_weights(&self) -> ndarray::Array1 { + match self { + KmerLayer::Count { layer, .. } => layer.col_weights(), + KmerLayer::Presence { layer, .. } => layer.col_weights(), + KmerLayer::Empty { .. } => panic!("Layer::col_weights() called on an Empty layer"), + } + } + /// 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()`. diff --git a/src/obikindex/src/layer/typed_layer.rs b/src/obikindex/src/layer/typed_layer.rs index befb3d80..998294ca 100644 --- a/src/obikindex/src/layer/typed_layer.rs +++ b/src/obikindex/src/layer/typed_layer.rs @@ -364,6 +364,12 @@ impl TypedLayer { pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec]) { self.data.fill_sub_matrix(slots, out) } + + /// Per-genome column weights — for counts, the sum of values in each + /// column (see `obicompactvec::ColumnWeights::col_weights`). + pub fn col_weights(&self) -> ndarray::Array1 { + obicompactvec::ColumnWeights::col_weights(&self.data) + } } // ── Mode 3 — presence/absence matrix ───────────────────────────────────────── @@ -379,7 +385,7 @@ impl TypedLayer { // below: sparse matrices aren't built column-by-column, they're built // row-by-row from an already-built dense layer // (`PersistentSparseBitMatrixBuilder::build_from_dense`). -impl> + BinaryMatrix> TypedLayer { +impl> + BinaryMatrix + obicompactvec::ColumnWeights> TypedLayer { /// 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 @@ -408,6 +414,13 @@ impl> + BinaryMatrix> TypedLayer { pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec]) { self.data.fill_sub_matrix(slots, out) } + + /// Per-genome column weights — for presence/absence, the number of + /// k-mers each genome carries (see + /// `obicompactvec::ColumnWeights::col_weights`). + pub fn col_weights(&self) -> ndarray::Array1 { + obicompactvec::ColumnWeights::col_weights(&self.data) + } } impl TypedLayer { diff --git a/src/obikstats/Cargo.toml b/src/obikstats/Cargo.toml index d302450a..825b9a12 100644 --- a/src/obikstats/Cargo.toml +++ b/src/obikstats/Cargo.toml @@ -5,5 +5,7 @@ edition = "2024" [dependencies] obikindex = { path = "../obikindex" } +obikidxcache = { path = "../obikidxcache" } +obikalgorithm = { path = "../obikalgorithm" } obicompactvec = { path = "../obicompactvec" } rayon = "1" diff --git a/src/obikstats/src/lib.rs b/src/obikstats/src/lib.rs index 66baa672..5ebd0165 100644 --- a/src/obikstats/src/lib.rs +++ b/src/obikstats/src/lib.rs @@ -1,10 +1,15 @@ //! Read-only aggregation/statistics over an already-built //! `obikindex::KmerIndex`: bits-per-kmer breakdown by index component -//! (`bits_per_kmer`), per-genome k-mer counts (`genome_kmer_counts`). -//! Not part of the `Index { Partition { Layer } }` data model itself — -//! kept out of `obikindex` like every other read/write extension -//! (`obikquery`, `obikdump`, `obikselect`, `obikrebuild`, `obikmerge`). +//! (`BitsPerKmer`), per-genome k-mer counts (`GenomeKmerCounts`). Not part +//! of the `Index { Partition { Layer } }` data model itself — kept out of +//! `obikindex` like every other read/write extension (`obikquery`, +//! `obikdump`, `obikselect`, `obikrebuild`, `obikmerge`). +//! +//! Structured on the model of `obikindexer::algorithms`/`obikmerge`: each +//! computation is its own `obikalgorithm::Algorithm` (two-phase `new` then +//! `run`), built on `obikidxcache::IndexCache` rather than hand-opening +//! each layer's files. mod stats; -pub use stats::IndexBitsPerKmer; +pub use stats::{BitsPerKmer, GenomeKmerCounts, IndexBitsPerKmer}; diff --git a/src/obikstats/src/stats.rs b/src/obikstats/src/stats.rs index d32fb25e..a89b8878 100644 --- a/src/obikstats/src/stats.rs +++ b/src/obikstats/src/stats.rs @@ -1,32 +1,32 @@ use std::fs; use std::path::Path; -use obicompactvec::{LayerMeta, PersistentBitMatrix, PersistentCompactIntMatrix}; -use obicompactvec::traits::ColumnWeights; use rayon::prelude::*; -use obikindex::{OKIError, OKIResult}; +use obikalgorithm::Algorithm; +use obikidxcache::index_cache::IndexCache; use obikindex::KmerIndex; +use obikindex::layer::KmerLayer; /// Bits per kmer broken down by index component. pub struct IndexBitsPerKmer { /// Total distinct k-mers across all partitions and layers. - pub n_kmers: usize, + pub n_kmers: usize, /// Number of genomes in the index. - pub n_genomes: usize, + pub n_genomes: usize, /// Bits used by the minimal perfect hash function (`mphf.bin`). - pub mphf: f64, + pub mphf: f64, /// Bits used by the evidence files (`evidence.bin`, `unitigs.bin*`, /// `fingerprint.bin`). - pub evidence: f64, + pub evidence: f64, /// Bits used by the count/presence matrices (`counts/` and `presence/`), /// normalised by k-mers only. - pub matrix: f64, + pub matrix: f64, /// `matrix` divided by the number of genomes — intrinsic encoding /// efficiency, independent of index size. pub matrix_per_genome: f64, /// Sum of mphf + evidence + matrix. - pub total: f64, + pub total: f64, } // ── File-size helpers ───────────────────────────────────────────────────────── @@ -36,7 +36,9 @@ fn file_bytes(path: &Path) -> u64 { } fn dir_bytes(dir: &Path) -> u64 { - if !dir.exists() { return 0; } + if !dir.exists() { + return 0; + } fs::read_dir(dir) .map(|entries| { entries @@ -52,134 +54,128 @@ fn dir_bytes(dir: &Path) -> u64 { // ── Per-layer accounting ────────────────────────────────────────────────────── struct LayerBytes { - n_kmers: usize, - mphf: u64, + n_kmers: usize, + mphf: u64, evidence: u64, - matrix: u64, + matrix: u64, } -fn layer_bytes(layer_dir: &Path) -> LayerBytes { - let n_kmers = LayerMeta::load(layer_dir).map(|m| m.n).unwrap_or(0); +fn layer_bytes(layer: &KmerLayer) -> LayerBytes { + let mphf = file_bytes(&layer.mphf_path()); - let mphf = file_bytes(&layer_dir.join("mphf.bin")); + let evidence = file_bytes(&layer.unitigs_path()) + + file_bytes(&layer.dir().join("unitigs.bin.idx")) + + file_bytes(&layer.evidence_path()) + + file_bytes(&layer.fingerprint_path()); - let evidence = file_bytes(&layer_dir.join("unitigs.bin")) - + file_bytes(&layer_dir.join("unitigs.bin.idx")) - + file_bytes(&layer_dir.join("evidence.bin")) - + file_bytes(&layer_dir.join("fingerprint.bin")); + let matrix = dir_bytes(&layer.counts_dir()) + dir_bytes(&layer.presence_dir()); - let matrix = dir_bytes(&layer_dir.join("counts")) - + dir_bytes(&layer_dir.join("presence")); - - LayerBytes { n_kmers, mphf, evidence, matrix } + LayerBytes { + n_kmers: layer.n(), + mphf, + evidence, + matrix, + } } -// ── KmerIndex::bits_per_kmer ────────────────────────────────────────────────── +// ── BitsPerKmer ──────────────────────────────────────────────────────────────── -impl KmerIndex { - /// Compute bits-per-kmer statistics for the built index. - /// - /// File sizes are read directly from disk; kmer counts come from - /// `layer_meta.json` (no need to scan the MPHF or unitig files). - /// Computation is parallelised across partitions. - pub fn bits_per_kmer(&self) -> OKIResult { - let n = self.n_partitions(); - let n_genomes = self.meta().genomes().map_err(OKIError::Io)?.len().max(1); +/// Computes bits-per-kmer statistics for an already-built index. Two-phase +/// construction (`new`), same shape as the other pipeline algorithms in +/// `obikindexer::algorithms`/`obikmerge`. File sizes are read directly from +/// disk; kmer counts come from each layer's own `n()` (its MPHF slot +/// count). Built on `obikidxcache::IndexCache`, which opens every layer +/// once up front — this is diagnostic tooling, not a hot path, so trading +/// that eager open for not having to hand-reconstruct every layer's file +/// paths is the right tradeoff here. +pub struct BitsPerKmer<'a> { + index: &'a KmerIndex, +} - let (n_kmers, mphf_b, evidence_b, matrix_b) = (0..n) - .into_par_iter() - .map(|i| { - let index_dir = self.index_dir(i); - if !index_dir.exists() { return (0usize, 0u64, 0u64, 0u64); } +impl<'a> BitsPerKmer<'a> { + pub fn new(index: &'a KmerIndex) -> Self { + Self { index } + } +} - let n_layers = self.n_layers(i).unwrap_or(0); +impl Algorithm for BitsPerKmer<'_> { + type Output = IndexBitsPerKmer; - (0..n_layers).fold((0usize, 0u64, 0u64, 0u64), |acc, l| { - let Ok(layer_dir) = self.layer_dir(i, l) else { return acc }; - let lb = layer_bytes(&layer_dir); - (acc.0 + lb.n_kmers, acc.1 + lb.mphf, acc.2 + lb.evidence, acc.3 + lb.matrix) - }) - }) - .reduce(|| (0, 0, 0, 0), |a, b| (a.0 + b.0, a.1 + b.1, a.2 + b.2, a.3 + b.3)); + fn run(&mut self) -> obikalgorithm::Result { + let n_genomes = self.index.meta().genomes()?.len().max(1); + let cache = IndexCache::new(self.index, None); + let layers: Vec<&KmerLayer> = cache.iter().collect(); + + let (n_kmers, mphf_b, evidence_b, matrix_b) = layers + .par_iter() + .map(|layer| layer_bytes(layer)) + .map(|lb| (lb.n_kmers, lb.mphf, lb.evidence, lb.matrix)) + .reduce( + || (0usize, 0u64, 0u64, 0u64), + |a, b| (a.0 + b.0, a.1 + b.1, a.2 + b.2, a.3 + b.3), + ); if n_kmers == 0 { return Ok(IndexBitsPerKmer { - n_kmers: 0, n_genomes, - mphf: 0.0, evidence: 0.0, - matrix: 0.0, matrix_per_genome: 0.0, total: 0.0, + n_kmers: 0, + n_genomes, + mphf: 0.0, + evidence: 0.0, + matrix: 0.0, + matrix_per_genome: 0.0, + total: 0.0, }); } - let bpk = |bytes: u64| bytes as f64 * 8.0 / n_kmers as f64; + let bpk = |bytes: u64| bytes as f64 * 8.0 / n_kmers as f64; let matrix = bpk(matrix_b); Ok(IndexBitsPerKmer { n_kmers, n_genomes, - mphf: bpk(mphf_b), - evidence: bpk(evidence_b), + mphf: bpk(mphf_b), + evidence: bpk(evidence_b), matrix, matrix_per_genome: matrix / n_genomes as f64, - total: bpk(mphf_b + evidence_b + matrix_b), + total: bpk(mphf_b + evidence_b + matrix_b), }) } +} - /// Return `(total_distinct_kmers, per_genome_kmer_counts)`. - /// - /// For each genome, the count is the number of distinct k-mers for which - /// that genome has a non-zero value (presence = 1, count > 0). - /// Partitions are scanned in parallel; results are summed across partitions. - pub fn genome_kmer_counts(&self) -> OKIResult<(usize, Vec)> { - let n = self.n_partitions(); - let n_genomes = self.meta.genomes().map_err(OKIError::Io)?.len(); +// ── GenomeKmerCounts ───────────────────────────────────────────────────────── - let partials: Vec<(usize, Vec)> = (0..n) - .into_par_iter() - .map(|i| { - let mut counts = vec![0u64; n_genomes]; - let mut n_kmers = 0usize; +/// Computes, per genome, the number of distinct k-mers for which that +/// genome has a non-zero value (presence = 1, count > 0) — via +/// `KmerLayer::col_weights` (`obicompactvec::ColumnWeights`), one call per +/// layer instead of hand-opening each layer's matrix and summing rows. +pub struct GenomeKmerCounts<'a> { + index: &'a KmerIndex, +} - let index_dir = self.index_dir(i); - if !index_dir.exists() { return (0, counts); } +impl<'a> GenomeKmerCounts<'a> { + pub fn new(index: &'a KmerIndex) -> Self { + Self { index } + } +} - let n_layers = self.n_layers(i).unwrap_or(0); +impl Algorithm for GenomeKmerCounts<'_> { + /// `(total_distinct_kmers, per_genome_kmer_counts)`. + type Output = (usize, Vec); - for l in 0..n_layers { - let Ok(this_layer_dir) = self.layer_dir(i, l) else { continue }; - if !this_layer_dir.exists() { continue; } + fn run(&mut self) -> obikalgorithm::Result<(usize, Vec)> { + let n_genomes = self.index.meta().genomes()?.len(); + let cache = IndexCache::new(self.index, None); + let layers: Vec<&KmerLayer> = cache.iter().collect(); - n_kmers += LayerMeta::load(&this_layer_dir).map(|m| m.n).unwrap_or(0); + let total_kmers: usize = layers.iter().map(|l| l.n()).sum(); - let mat: Box = - if this_layer_dir.join("counts").exists() - && !this_layer_dir.join("presence").exists() - { - match obikindex::layer::open_data::(&index_dir, l) { - Ok(m) => Box::new(m), - Err(_) => continue, - } - } else { - match obikindex::layer::open_data::(&index_dir, l) { - Ok(m) => Box::new(m), - Err(_) => continue, - } - }; - let col_counts = mat.partial_kmer_counts(); - - for (c, &v) in col_counts.iter().enumerate() { - if c < n_genomes { counts[c] += v; } - } + let mut total_counts = vec![0u64; n_genomes]; + let per_layer_weights: Vec<_> = layers.par_iter().map(|l| l.col_weights()).collect(); + for weights in per_layer_weights { + for (g, &v) in weights.iter().enumerate() { + if g < n_genomes { + total_counts[g] += v; } - - (n_kmers, counts) - }) - .collect(); - - let total_kmers: usize = partials.iter().map(|(n, _)| n).sum(); - let mut total_counts = vec![0u64; n_genomes]; - for (_, counts) in partials { - for (i, v) in counts.into_iter().enumerate() { - total_counts[i] += v; } }