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

feat: add col_weights API and refactor obikstats to use Algorithm trait

Browse Source 
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.
This commit is contained in:
Eric Coissac
2026-08-22 17:28:51 +02:00
parent fba9c65b1a
commit 23812d1af8
6 changed files with 137 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/Cargo.lock
Generated
+2
View File
@@ -1767,6 +1767,8 @@ name = "obikstats"
version = "0.1.0"
dependencies = [
"obicompactvec",
"obikalgorithm",
"obikidxcache",
"obikindex",
"rayon",
]
src/obikindex/src/layer/content_layer.rs
+10
View File
@@ -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<u64> {
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()`.
src/obikindex/src/layer/typed_layer.rs
+14 -1
View File
@@ -364,6 +364,12 @@ impl TypedLayer<PersistentCompactIntMatrix> {
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<u32>]) {
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<u64> {
obicompactvec::ColumnWeights::col_weights(&self.data)
}
}
// ── Mode 3 — presence/absence matrix ─────────────────────────────────────────
@@ -379,7 +385,7 @@ impl TypedLayer<PersistentCompactIntMatrix> {
// 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<D: LayerData<Item = Box<[bool]>> + BinaryMatrix> TypedLayer<D> {
impl<D: LayerData<Item = Box<[bool]>> + BinaryMatrix + obicompactvec::ColumnWeights> TypedLayer<D> {
/// 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<D: LayerData<Item = Box<[bool]>> + BinaryMatrix> TypedLayer<D> {
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<bool>]) {
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<u64> {
obicompactvec::ColumnWeights::col_weights(&self.data)
}
}
impl TypedLayer<PersistentBitMatrix> {
src/obikstats/Cargo.toml
+2
View File
@@ -5,5 +5,7 @@ edition = "2024"
[dependencies]
obikindex = { path = "../obikindex" }
obikidxcache = { path = "../obikidxcache" }
obikalgorithm = { path = "../obikalgorithm" }
obicompactvec = { path = "../obicompactvec" }
rayon = "1"
src/obikstats/src/lib.rs
+10 -5
View File
@@ -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};
src/obikstats/src/stats.rs
+99 -103
View File
@@ -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<IndexBitsPerKmer> {
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<IndexBitsPerKmer> {
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<u64>)> {
let n = self.n_partitions();
let n_genomes = self.meta.genomes().map_err(OKIError::Io)?.len();
// ── GenomeKmerCounts ─────────────────────────────────────────────────────────
let partials: Vec<(usize, Vec<u64>)> = (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<u64>);
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<u64>)> {
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<dyn ColumnWeights> =
if this_layer_dir.join("counts").exists()
&& !this_layer_dir.join("presence").exists()
{
match obikindex::layer::open_data::<PersistentCompactIntMatrix>(&index_dir, l) {
Ok(m) => Box::new(m),
Err(_) => continue,
}
} else {
match obikindex::layer::open_data::<PersistentBitMatrix>(&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;
}
}