From f9ef6b8391b128efc30e6bd1f33864a21a17b507 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 20 Aug 2026 10:11:46 +0200 Subject: [PATCH] rename KmerPartition to KmerPartitions and update Mat enum Rename the KmerPartition type to KmerPartitions across obikindex, obikpartitionner, and obikphylo/siblings to reflect an updated data model. Update the Mat enum in siblings/cache.rs to add a SparsePresence variant and simplify opening logic by delegating sparse versus dense detection to PersistentBitMatrix. Apply consistent code formatting, import reordering, and multi-line refactoring throughout the affected modules. --- .../implementation/partition_layer_cache.md | 83 +++- src/obikindex/src/index.rs | 142 ++++-- src/obikindex/src/merge.rs | 43 +- src/obikindex/src/reindex.rs | 31 +- src/obikindex/src/select.rs | 78 +++- src/obikmer/src/steps/scatter.rs | 35 +- src/obikpartitionner/src/distance.rs | 14 +- src/obikpartitionner/src/dump_layer.rs | 67 ++- src/obikpartitionner/src/index_layer.rs | 34 +- src/obikpartitionner/src/lib.rs | 4 +- src/obikpartitionner/src/merge_layer/mod.rs | 197 +++++--- .../src/partition/kmer_partition.rs | 16 +- src/obikpartitionner/src/partition/mod.rs | 2 +- src/obikpartitionner/src/partition/tests.rs | 21 +- src/obikpartitionner/src/query_layer.rs | 20 +- src/obikpartitionner/src/rebuild_layer.rs | 57 ++- src/obikpartitionner/src/select_layer.rs | 162 +++++-- src/obikpartitionner/src/tests/query_layer.rs | 8 +- src/obikphylo/src/siblings/alignment.rs | 48 +- src/obikphylo/src/siblings/build.rs | 420 ++++++++++-------- src/obikphylo/src/siblings/cache.rs | 81 ++-- src/obikphylo/src/siblings/cardinality.rs | 77 ++-- src/obikphylo/src/siblings/distance.rs | 66 ++- src/obikphylo/src/siblings/entropy.rs | 113 +++-- src/obikphylo/src/siblings/sankoff_bundle.rs | 161 ++++--- src/obikphylo/src/siblings/stats.rs | 43 +- 26 files changed, 1290 insertions(+), 733 deletions(-) diff --git a/DevDocMD/implementation/partition_layer_cache.md b/DevDocMD/implementation/partition_layer_cache.md index e154be9b..32df7e8c 100644 --- a/DevDocMD/implementation/partition_layer_cache.md +++ b/DevDocMD/implementation/partition_layer_cache.md @@ -5,6 +5,70 @@ themselves not started; ownership split (below) still open. Preparatory encapsulation work (partition/layer path and metadata accessors on `KmerPartition`) landed the same day — see "Preparatory work done" below. +## Type-to-concept mapping: Index / Partition / Layer + +The conceptual nesting `Index { Partition { Layer { MPHF, Evidence, Matrix +} } } }` does **not** have one Rust type per level — worth stating +explicitly, since two of the names below are misleading. + +- **Index** = `obikindex::KmerIndex` — `{ root_path, meta: IndexMeta, + partition: KmerPartition }`. The field is named `partition` (singular), + but it holds the *whole* multi-partition structure below — the name + suggests "one partition", the value is all of them. +- **Partition, the collection (not one partition)** = + `obikpartitionner::KmerPartition`. Despite the singular name, this owns + *every* partition of the index: `root_path`, `n_partitions`, and + per-partition accessors that all take an explicit index `i` + (`partition_dir(i)`, `index_dir(i)`, `layer_dir(i, l)`, + `partition_meta(i)`, `n_layers(i)`, `index_mode(i)` — see the + `part_dir`/`layer_dir` and `PartitionMeta`-encapsulation work above). + It never holds one partition's layers open in memory — no `Vec>` + here, only path arithmetic and metadata reads. A more honest name would + be `KmerPartitionSet` or `KmerPartitioner`; renaming is out of scope for + now, just worth knowing it's not "a partition." +- **Partition, one of them** = **no dedicated type exists today**. The + structural equivalent of "one partition, its layers held open" is + `obilayeredmap::LayeredMap` — `{ root, meta: PartitionMeta, layers: + Vec> }` — but `LayeredMap` itself doesn't know it's playing that + role: it operates purely on whatever `root` path it's given and has no + notion of `KmerPartition`, `n_partitions`, or a partition index `i` (see + "Layering: who owns what" below — established independently while + fixing `open_data`/`layer_dir`). Every caller that wants "one partition, + opened" builds the path itself + (`kmer_partition.index_dir(i)`/`.layer_dir(i, l)`) and either passes it + to `LayeredMap::open` or — as `obikphylo::siblings::cache` does — + bypasses `LayeredMap` entirely and opens each `Layer` directly. +- **Layer** = `obilayeredmap::Layer` — `{ mphf: MphfLayer, data: D }`. + - **MPHF** = `MphfLayer.mphf: MemCase` — kmer → slot. + - **Evidence** = `MphfLayer.ev: LayerEvidence` (`Exact`/`Approx`/ + `Hybrid` — `evidence.bin`/`fingerprint.bin`; see `EvidenceKind`). + - **Matrix** = `Layer.data: D` — `PersistentBitMatrix` / + `PersistentCompactIntMatrix` / `PersistentSparseBitMatrix` / `()`. + +Today's actual nesting, in Rust terms: + +``` +KmerIndex + └─ partition: KmerPartition (all N partitions — paths + metadata only) + └─ (opened on demand, per i) LayeredMap ← "one partition", unnamed as such + └─ layers: Vec> + └─ Layer { mphf: MphfLayer, data: D } + ├─ mphf.mphf → MPHF + ├─ mphf.ev → Evidence + └─ data → Matrix +``` + +`obikphylo::siblings::cache::PartitionCache` — the thing (2) is meant to +generalise — skips the middle of this nesting entirely: it doesn't build a +`Vec>`, it builds `mats: Vec>` directly — outer +index = partition `i` (via `KmerPartition`), inner index = layer `l` (via +`Layer`/`Mat`) — reconstructing "one partition, its layers open" as a +bare nested `Vec` because no owning type for that concept exists to reuse. +That gap is exactly what (2) needs to fill, and (1) is exactly what would +let a per-partition type hold layers of mixed `D` (today `LayeredMap` +can't, being monomorphic — see "The gap in `obilayeredmap`'s existing +cache" below). + ## The problem Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps @@ -217,13 +281,18 @@ storage decisions, now a *third* copy of the same logic to keep in sync) about the code. Lesson: for a `pub enum` whose variant list matters, read the definition unfiltered, don't grep for the variant names you expect to find. -- One real consequence of that same correction: `obikphylo::siblings:: - cache::Mat::SparsePresence(Layer)` looks - redundant now — `Mat::Presence(Layer)` alone - would already handle sparse layers transparently, since - `PersistentBitMatrix` absorbs `Sparse` internally. Likely `Mat` predates - `PersistentBitMatrix` growing native sparse support. Signalled, not - removed — no mandate to touch `obikphylo` for this. +- One real consequence of that same correction, **fixed (2026-08-20)**: + `obikphylo::siblings::cache::Mat::SparsePresence(Layer< + PersistentSparseBitMatrix>)` was redundant — `Mat::Presence(Layer< + PersistentBitMatrix>)` alone already handles sparse layers + transparently, since `PersistentBitMatrix` absorbs `Sparse` internally. + Removed the variant, the `presence/is_multi.prsb` probe in `Mat::open` + (now just opens `Layer::` unconditionally for the + non-count case — sparse-vs-dense is `PersistentBitMatrix::open`'s own + concern), and every now-single-armed match in `find_slot`/`index_batch`/ + `iter_minorants_batch`/`n_cols`/`fill_sub_matrix_carries`. Full + workspace test suite green after, including the 27 `obikphylo::siblings` + tests that exercise `pack_matrices(true)`/sparse through `Mat`. ## Remaining instance of the PartitionMeta-encapsulation problem diff --git a/src/obikindex/src/index.rs b/src/obikindex/src/index.rs index 355483d0..db471c21 100644 --- a/src/obikindex/src/index.rs +++ b/src/obikindex/src/index.rs @@ -2,7 +2,7 @@ use std::collections::BTreeMap; use std::fs; use std::path::{Path, PathBuf}; -use obikpartitionner::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR}; +use obikpartitionner::{KmerPartitions, KmerSpectrum, PARTITIONS_SUBDIR}; use obisys::{Reporter, Stage, progress_bar}; use rayon::prelude::*; use tracing::info; @@ -16,7 +16,7 @@ use crate::state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCAT pub struct KmerIndex { pub(crate) root_path: PathBuf, pub(crate) meta: IndexMeta, - pub(crate) partition: KmerPartition, + pub(crate) partition: KmerPartitions, } impl KmerIndex { @@ -31,7 +31,7 @@ impl KmerIndex { force: bool, ) -> OKIResult { let root_path = path.as_ref().to_owned(); - let partition = KmerPartition::create( + let partition = KmerPartitions::create( &root_path, config.n_bits, config.kmer_size, @@ -45,7 +45,11 @@ impl KmerIndex { meta.genomes.push(info); } meta.write(&root_path)?; - Ok(Self { root_path, meta, partition }) + Ok(Self { + root_path, + meta, + partition, + }) } pub fn open>(path: P) -> OKIResult { @@ -53,13 +57,17 @@ impl KmerIndex { let meta = IndexMeta::read(&root_path).map_err(OKIError::Io)?; set_k(meta.config.kmer_size); set_m(meta.config.minimizer_size); - let partition = KmerPartition::open_with_config( + let partition = KmerPartitions::open_with_config( &root_path, meta.config.kmer_size, meta.config.minimizer_size, meta.config.n_bits, )?; - Ok(Self { root_path, meta, partition }) + Ok(Self { + root_path, + meta, + partition, + }) } /// Return `true` if `path` contains an `index.meta` file. @@ -96,12 +104,12 @@ impl KmerIndex { pub(crate) fn create_skeleton>( output: P, meta: &IndexMeta, - ) -> OKIResult { + ) -> OKIResult { let output = output.as_ref(); fs::create_dir_all(output).map_err(OKIError::Io)?; meta.write(output).map_err(OKIError::Io)?; fs::create_dir_all(output.join(PARTITIONS_SUBDIR)).map_err(OKIError::Io)?; - Ok(KmerPartition::open_with_config( + Ok(KmerPartitions::open_with_config( output, meta.config.kmer_size, meta.config.minimizer_size, @@ -134,12 +142,24 @@ impl KmerIndex { /// The index's root directory — needed by out-of-crate extension code /// (e.g. `obikphylo`) that opens its own `KmerPartition` handle onto /// the same on-disk index. - pub fn root_path(&self) -> &Path { &self.root_path } - pub fn meta(&self) -> &IndexMeta { &self.meta } - pub fn meta_mut(&mut self) -> &mut IndexMeta { &mut self.meta } - pub fn kmer_size(&self) -> usize { self.meta.config.kmer_size } - pub fn minimizer_size(&self) -> usize { self.meta.config.minimizer_size } - pub fn n_partitions(&self) -> usize { self.partition.n_partitions() } + pub fn root_path(&self) -> &Path { + &self.root_path + } + pub fn meta(&self) -> &IndexMeta { + &self.meta + } + pub fn meta_mut(&mut self) -> &mut IndexMeta { + &mut self.meta + } + pub fn kmer_size(&self) -> usize { + self.meta.config.kmer_size + } + pub fn minimizer_size(&self) -> usize { + self.meta.config.minimizer_size + } + pub fn n_partitions(&self) -> usize { + self.partition.n_partitions() + } /// Number of layers per partition. /// @@ -152,7 +172,7 @@ impl KmerIndex { /// Expose the inner partition so the caller can run scatter into it. /// Call `mark_scattered` once scatter is complete. - pub fn partition_mut(&mut self) -> &mut KmerPartition { + pub fn partition_mut(&mut self) -> &mut KmerPartitions { &mut self.partition } @@ -174,7 +194,11 @@ impl KmerIndex { /// /// Writes `spectrums/{label}.json` and touches `count.done` upon completion. /// Per-partition spectrum files are removed unless `keep_intermediate` is true. - pub fn dereplicate_and_count(&self, keep_intermediate: bool, rep: &mut Reporter) -> OKIResult<()> { + pub fn dereplicate_and_count( + &self, + keep_intermediate: bool, + rep: &mut Reporter, + ) -> OKIResult<()> { let t = Stage::start("dereplicate"); self.partition.dereplicate()?; rep.push(t.stop()); @@ -189,11 +213,17 @@ impl KmerIndex { } fn write_spectrum(&self, sp: &KmerSpectrum) -> OKIResult<()> { - let label = self.meta.genomes.first().map(|g| g.label.as_str()).unwrap_or("unknown"); + let label = self + .meta + .genomes + .first() + .map(|g| g.label.as_str()) + .unwrap_or("unknown"); let spectrums_dir = self.root_path.join("spectrums"); fs::create_dir_all(&spectrums_dir)?; let path = spectrums_dir.join(format!("{label}.json")); - let spectrum_map: BTreeMap = sp.counts + let spectrum_map: BTreeMap = sp + .counts .iter() .map(|(&c, &f)| (format!("{c:010}"), f)) .collect(); @@ -226,17 +256,28 @@ impl KmerIndex { let order: Vec = (0..n).collect(); let runner = crate::numa::PartitionRunner::new(); - runner.run( - &order, - |i| self.partition.build_index_layer(i, min_ab, max_ab, with_counts, &evidence, block_bits), - |i, n_kmers, _| { - if n_kmers > 0 { - total_kmers += n_kmers; - pb.inc(1); - pb.set_message(format!("{i}: {n_kmers} kmers")); - } - }, - ).map_err(OKIError::Partition)?; + runner + .run( + &order, + |i| { + self.partition.build_index_layer( + i, + min_ab, + max_ab, + with_counts, + &evidence, + block_bits, + ) + }, + |i, n_kmers, _| { + if n_kmers > 0 { + total_kmers += n_kmers; + pb.inc(1); + pb.set_message(format!("{i}: {n_kmers} kmers")); + } + }, + ) + .map_err(OKIError::Partition)?; pb.finish_and_clear(); info!("done — {} total kmers indexed", total_kmers); @@ -253,7 +294,7 @@ impl KmerIndex { } /// Borrow the inner partition for direct superkmer-level queries. - pub fn partition(&self) -> &KmerPartition { + pub fn partition(&self) -> &KmerPartitions { &self.partition } @@ -282,12 +323,14 @@ impl KmerIndex { &order, |i| -> OKIResult<()> { let index_dir = self.partition.index_dir(i); - if !index_dir.exists() { return Ok(()); } + if !index_dir.exists() { + return Ok(()); + } let n_layers = self.partition.n_layers(i)?; for l in 0..n_layers { let layer_dir = self.partition.layer_dir(i, l); let presence_dir = layer_dir.join("presence"); - let counts_dir = layer_dir.join("counts"); + let counts_dir = layer_dir.join("counts"); if presence_dir.exists() { if sparse { pack_sparse_bit_matrix(&presence_dir).map_err(OKIError::Io)?; @@ -295,11 +338,15 @@ impl KmerIndex { pack_bit_matrix(&presence_dir).map_err(OKIError::Io)?; } } - if counts_dir.exists() { pack_compact_int_matrix(&counts_dir).map_err(OKIError::Io)?; } + if counts_dir.exists() { + pack_compact_int_matrix(&counts_dir).map_err(OKIError::Io)?; + } } Ok(()) }, - |_, _, _| { pb.inc(1); }, + |_, _, _| { + pb.inc(1); + }, )?; pb.finish_and_clear(); Ok(()) @@ -318,18 +365,27 @@ impl KmerIndex { .into_par_iter() .filter_map(|i| { let index_dir = self.partition.index_dir(i); - if !index_dir.exists() { return None; } + if !index_dir.exists() { + return None; + } let n_layers = match self.partition.n_layers(i) { Ok(n) => n, - Err(e) => return Some(OKIError::Io(std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))), + Err(e) => { + return Some(OKIError::Io(std::io::Error::new( + std::io::ErrorKind::Other, + e.to_string(), + ))); + } }; for l in 0..n_layers { let layer_dir = self.partition.layer_dir(i, l); let meta_path = layer_dir.join(LayerMeta::FILENAME); - if meta_path.exists() { continue; } + if meta_path.exists() { + continue; + } let unitigs_path = layer_dir.join("unitigs.bin"); let n_kmers = match UnitigFileReader::open_sequential(&unitigs_path) { - Ok(r) => r.n_kmers(), + Ok(r) => r.n_kmers(), Err(e) => return Some(OKIError::Partition(e)), }; if let Err(e) = LayerMeta::save(&layer_dir, n_kmers) { @@ -340,7 +396,9 @@ impl KmerIndex { }) .collect(); - if let Some(e) = errors.into_iter().next() { return Err(e); } + if let Some(e) = errors.into_iter().next() { + return Err(e); + } Ok(()) } } @@ -355,7 +413,11 @@ fn label_from_path(path: &Path) -> String { while let Some(pos) = s.rfind('.') { s.truncate(pos); } - if s.is_empty() { "unknown".to_string() } else { s } + if s.is_empty() { + "unknown".to_string() + } else { + s + } } fn touch(path: &Path) -> Result<(), std::io::Error> { diff --git a/src/obikindex/src/merge.rs b/src/obikindex/src/merge.rs index 5139d72e..444dc487 100644 --- a/src/obikindex/src/merge.rs +++ b/src/obikindex/src/merge.rs @@ -193,7 +193,7 @@ impl KmerIndex { let block_bits = dst.meta.config.block_bits; // Pre-build source list once (avoid rebuilding per partition) - let srcs: Vec<(&obikpartitionner::KmerPartition, usize)> = remaining_sources + let srcs: Vec<(&obikpartitionner::KmerPartitions, usize)> = remaining_sources .iter() .map(|s| (&s.partition, s.meta.genomes.len())) .collect(); @@ -221,19 +221,34 @@ impl KmerIndex { let runner = crate::numa::PartitionRunner::new(); let mut part_stats: Vec = Vec::with_capacity(n_partitions); - runner.run( - &order, - |i| dst_partition.merge_partition(i, srcs, mode, n_dst_genomes, block_bits, evidence), - |i, g_len, dur| { - pb.inc(1); - debug!( - "partition {i}: done in {:.1}s — {} new kmers", - dur.as_secs_f64(), - g_len, - ); - part_stats.push(PartStat { id: i, unitig_bytes: partition_sizes[i], g_len }); - }, - ).map_err(OKIError::Partition)?; + runner + .run( + &order, + |i| { + dst_partition.merge_partition( + i, + srcs, + mode, + n_dst_genomes, + block_bits, + evidence, + ) + }, + |i, g_len, dur| { + pb.inc(1); + debug!( + "partition {i}: done in {:.1}s — {} new kmers", + dur.as_secs_f64(), + g_len, + ); + part_stats.push(PartStat { + id: i, + unitig_bytes: partition_sizes[i], + g_len, + }); + }, + ) + .map_err(OKIError::Partition)?; pb.finish_and_clear(); diff --git a/src/obikindex/src/reindex.rs b/src/obikindex/src/reindex.rs index 932fee23..4f43be01 100644 --- a/src/obikindex/src/reindex.rs +++ b/src/obikindex/src/reindex.rs @@ -1,17 +1,17 @@ -use std::fs; -use std::path::Path; -use obikpartitionner::KmerPartition; +use obikpartitionner::KmerPartitions; use obilayeredmap::{IndexMode, layer::Layer}; use obisys::{Reporter, Stage, progress_bar}; +use std::fs; +use std::path::Path; use tracing::info; use crate::error::{OKIError, OKIResult}; use crate::index::KmerIndex; use crate::state::IndexState; -const EVIDENCE_FILE: &str = "evidence.bin"; +const EVIDENCE_FILE: &str = "evidence.bin"; const FINGERPRINT_FILE: &str = "fingerprint.bin"; -const UNITIG_IDX_FILE: &str = "unitigs.bin.idx"; +const UNITIG_IDX_FILE: &str = "unitigs.bin.idx"; fn olm_to_oki(e: obilayeredmap::OLMError) -> OKIError { OKIError::InvalidInput(e.to_string()) @@ -48,9 +48,13 @@ impl KmerIndex { let runner = crate::numa::PartitionRunner::new(); runner.run( &order, - |i| reindex_partition(&self.partition, i, &target, block_bits) - .map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}"))), - |_, _, _| { pb.inc(1); }, + |i| { + reindex_partition(&self.partition, i, &target, block_bits) + .map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}"))) + }, + |_, _, _| { + pb.inc(1); + }, )?; pb.finish_and_clear(); @@ -66,11 +70,18 @@ impl KmerIndex { } /// Process all layers of one partition's index directory. -fn reindex_partition(partition: &KmerPartition, i: usize, target: &IndexMode, block_bits: u8) -> OKIResult<()> { +fn reindex_partition( + partition: &KmerPartitions + , i: usize + , target: &IndexMode + , block_bits: u, +8) -> OKIResult<()> { if !partition.index_dir(i).exists() { return Ok(()); } - let n_layers = partition.n_layers(i).map_err(|e| OKIError::InvalidInput(e.to_string()))?; + let n_layers = partition + .n_layers(i) + .map_err(|e| OKIError::InvalidInput(e.to_string()))?; for layer_idx in 0..n_layers { reindex_layer(&partition.layer_dir(i, layer_idx), target, block_bits)?; } diff --git a/src/obikindex/src/select.rs b/src/obikindex/src/select.rs index bc42a3a1..1f6591f5 100644 --- a/src/obikindex/src/select.rs +++ b/src/obikindex/src/select.rs @@ -1,6 +1,6 @@ use std::path::Path; -use obikpartitionner::{KmerPartition, OutputCol}; +use obikpartitionner::{KmerPartitions, OutputCol}; use obisys::{Reporter, Stage, progress_bar}; use tracing::info; @@ -35,31 +35,48 @@ impl KmerIndex { let mut meta = IndexMeta::new(src.meta.config.clone()); meta.config.with_counts = !output_presence; - meta.genomes = specs.iter() + meta.genomes = specs + .iter() .map(|s| GenomeInfo::new(s.label.clone())) .collect(); - let n_src_genomes = src.meta.genomes.len(); - let n_partitions = src.partition.n_partitions(); + let n_src_genomes = src.meta.genomes.len(); + let n_partitions = src.partition.n_partitions(); let dst_partition = KmerIndex::create_skeleton(output, &meta)?; info!( "select: {} partition(s), {} source genome(s) → {} output column(s)", - n_partitions, n_src_genomes, specs.len(), + n_partitions, + n_src_genomes, + specs.len(), ); - let t = Stage::start("select"); - let pb = progress_bar("select", n_partitions as u64, "partitions"); + let t = Stage::start("select"); + let pb = progress_bar("select", n_partitions as u64, "partitions"); let src_partition = &src.partition; let order: Vec = (0..n_partitions).collect(); let runner = crate::numa::PartitionRunner::new(); - runner.run( - &order, - |i| dst_partition.select_partition(src_partition, i, specs, n_src_genomes, threshold, output_presence, false), - |_, _, _| { pb.inc(1); }, - ).map_err(OKIError::Partition)?; + runner + .run( + &order, + |i| { + dst_partition.select_partition( + src_partition, + i, + specs, + n_src_genomes, + threshold, + output_presence, + false, + ) + }, + |_, _, _| { + pb.inc(1); + }, + ) + .map_err(OKIError::Partition)?; pb.finish_and_clear(); rep.push(t.stop()); @@ -82,9 +99,9 @@ impl KmerIndex { } let n_src_genomes = self.meta.genomes.len(); - let n_partitions = self.partition.n_partitions(); + let n_partitions = self.partition.n_partitions(); - let src_partition = KmerPartition::open_with_config( + let src_partition = KmerPartitions::open_with_config( &self.root_path, self.meta.config.kmer_size, self.meta.config.minimizer_size, @@ -93,26 +110,43 @@ impl KmerIndex { info!( "select (in-place): {} partition(s), {} source genome(s) → {} output column(s)", - n_partitions, n_src_genomes, specs.len(), + n_partitions, + n_src_genomes, + specs.len(), ); - let t = Stage::start("select"); + let t = Stage::start("select"); let pb = progress_bar("select", n_partitions as u64, "partitions"); let partition = &self.partition; let order: Vec = (0..n_partitions).collect(); let runner = crate::numa::PartitionRunner::new(); - runner.run( - &order, - |i| partition.select_partition(&src_partition, i, specs, n_src_genomes, threshold, output_presence, true), - |_, _, _| { pb.inc(1); }, - ).map_err(OKIError::Partition)?; + runner + .run( + &order, + |i| { + partition.select_partition( + &src_partition, + i, + specs, + n_src_genomes, + threshold, + output_presence, + true, + ) + }, + |_, _, _| { + pb.inc(1); + }, + ) + .map_err(OKIError::Partition)?; pb.finish_and_clear(); rep.push(t.stop()); self.meta.config.with_counts = !output_presence; - self.meta.genomes = specs.iter() + self.meta.genomes = specs + .iter() .map(|s| GenomeInfo::new(s.label.clone())) .collect(); self.meta.write(&self.root_path)?; diff --git a/src/obikmer/src/steps/scatter.rs b/src/obikmer/src/steps/scatter.rs index 28bd90a9..a6b3b2f2 100644 --- a/src/obikmer/src/steps/scatter.rs +++ b/src/obikmer/src/steps/scatter.rs @@ -1,12 +1,12 @@ use std::path::PathBuf; -use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::time::Instant; -use obisys::spinner; -use obiread::NucPage; -use obikpartitionner::KmerPartition; +use obikpartitionner::KmerPartitions; use obipipeline::{ThrottleGuard, Throttled, throttle}; +use obiread::NucPage; +use obisys::spinner; use obisys::{Reporter, Stage}; use tracing::info; @@ -15,8 +15,8 @@ use crate::cli::PipelineData; // ── Iterator that keeps the slot guard alive until the file is exhausted ────── struct GuardedIter { - inner: Box + Send>, - _guard: ThrottleGuard, + inner: Box + Send>, + _guard: ThrottleGuard, flat_active: Arc, } @@ -38,7 +38,7 @@ impl Drop for GuardedIter { /// Run scatter: normalise → build superkmers → route to partition → close. /// Reports the "scatter" stage to `rep`. pub fn scatter( - kp: &mut KmerPartition, + kp: &mut KmerPartitions, path_source: impl Iterator + Send + 'static, k: usize, level_max: usize, @@ -52,8 +52,8 @@ pub fn scatter( // Throttle in the source thread — never in a worker — to prevent deadlock. let throttled = throttle(path_source, max_open); - let file_count = Arc::new(AtomicU64::new(0)); - let flat_active = Arc::new(AtomicU32::new(0)); + let file_count = Arc::new(AtomicU64::new(0)); + let flat_active = Arc::new(AtomicU32::new(0)); let transform_active = Arc::new(AtomicU32::new(0)); let t = Stage::start("scatter"); @@ -95,7 +95,8 @@ pub fn scatter( const ALPHA: f64 = 0.15; for batch in pipe.apply(throttled, n_workers, 1) { - total_bases += batch.iter() + total_bases += batch + .iter() .map(|sk| (sk.seql() as u64).saturating_sub(kmer_overlap)) .sum::(); let now = Instant::now(); @@ -107,14 +108,22 @@ pub fn scatter( last_bases = total_bases; let bp = total_bases as f64; let (count_str, rate_str) = if bp >= 1e9 { - (format!("{:.2} Gbp", bp / 1e9), format!("{:.0} Mbp/s", ema_rate / 1e6)) + ( + format!("{:.2} Gbp", bp / 1e9), + format!("{:.0} Mbp/s", ema_rate / 1e6), + ) } else { - (format!("{:.0} Mbp", bp / 1e6), format!("{:.0} Mbp/s", ema_rate / 1e6)) + ( + format!("{:.0} Mbp", bp / 1e6), + format!("{:.0} Mbp/s", ema_rate / 1e6), + ) }; let n_files = file_count.load(Ordering::Relaxed); let r = flat_active.load(Ordering::Relaxed); let c = transform_active.load(Ordering::Relaxed); - pb.set_message(format!("{count_str} {rate_str} {n_files} files [R:{r} C:{c}]")); + pb.set_message(format!( + "{count_str} {rate_str} {n_files} files [R:{r} C:{c}]" + )); } kp.write_batch(batch).unwrap_or_else(|e| { eprintln!("error: {e}"); diff --git a/src/obikpartitionner/src/distance.rs b/src/obikpartitionner/src/distance.rs index 837d2cdc..291f3b05 100644 --- a/src/obikpartitionner/src/distance.rs +++ b/src/obikpartitionner/src/distance.rs @@ -1,11 +1,11 @@ use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; -use obilayeredmap::{open_data, LayeredStore}; +use obilayeredmap::{LayeredStore, open_data}; use obiskio::SKResult; use crate::common::{load_meta, olm_to_sk}; -use crate::partition::KmerPartition; +use crate::partition::KmerPartitions; -impl KmerPartition { +impl KmerPartitions { /// Open all count matrices for partition `part`, one per layer. /// Layers without a `counts/` directory are skipped. pub fn count_store(&self, part: usize) -> SKResult> { @@ -16,7 +16,9 @@ impl KmerPartition { let n_layers = load_meta(&index_dir, "distance")?.n_layers; let matrices = (0..n_layers) .filter_map(|l| { - self.layer_dir(part, l).join("counts").exists() + self.layer_dir(part, l) + .join("counts") + .exists() .then(|| open_data(&index_dir, l).map_err(|e| olm_to_sk(e, "distance"))) }) .collect::>>()?; @@ -33,7 +35,9 @@ impl KmerPartition { let n_layers = load_meta(&index_dir, "distance")?.n_layers; let matrices = (0..n_layers) .filter_map(|l| { - self.layer_dir(part, l).join("presence").exists() + self.layer_dir(part, l) + .join("presence") + .exists() .then(|| open_data(&index_dir, l).map_err(|e| olm_to_sk(e, "distance"))) }) .collect::>>()?; diff --git a/src/obikpartitionner/src/dump_layer.rs b/src/obikpartitionner/src/dump_layer.rs index eb18af0e..119c71be 100644 --- a/src/obikpartitionner/src/dump_layer.rs +++ b/src/obikpartitionner/src/dump_layer.rs @@ -1,19 +1,22 @@ use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; use obikseq::CanonicalKmer; -use obiskio::{SKError, SKResult, UnitigFileReader}; use obilayeredmap::{IndexMode, MphfLayer, OLMError}; +use obiskio::{SKError, SKResult, UnitigFileReader}; use crate::filter::{KmerFilter, passes_all}; -use crate::partition::KmerPartition; +use crate::partition::KmerPartitions; fn olm_to_sk(e: OLMError) -> SKError { match e { OLMError::Io(e) => SKError::Io(e), - other => SKError::InvalidData { context: "dump", detail: other.to_string() }, + other => SKError::InvalidData { + context: "dump", + detail: other.to_string(), + }, } } -impl KmerPartition { +impl KmerPartitions { /// Iterate all indexed kmers in partition `part`, calling `cb(kmer, row)` for each /// kmer that passes every filter in `filters`. /// @@ -43,12 +46,14 @@ impl KmerPartition { let mut l = 0; loop { let layer_dir = self.layer_dir(part, l); - if !layer_dir.exists() { break; } + if !layer_dir.exists() { + break; + } l += 1; - let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?; + let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?; let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?; - let counts_dir = layer_dir.join("counts"); + let counts_dir = layer_dir.join("counts"); let presence_dir = layer_dir.join("presence"); let cont = if use_counts && counts_dir.exists() { @@ -59,7 +64,9 @@ impl KmerPartition { let row = mat.row(slot); if passes_all(filters, kmer, &row, n_genomes) { cont = cb(kmer, row); - if !cont { break; } + if !cont { + break; + } } } } @@ -72,7 +79,9 @@ impl KmerPartition { let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect(); if passes_all(filters, kmer, &row, n_genomes) { cont = cb(kmer, row); - if !cont { break; } + if !cont { + break; + } } } } @@ -86,15 +95,21 @@ impl KmerPartition { let all_present: Box<[u32]> = vec![1u32; n_genomes].into(); let mut cont = true; for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { - if mphf.find(kmer).is_some() && passes_all(filters, kmer, &all_present, n_genomes) { + if mphf.find(kmer).is_some() + && passes_all(filters, kmer, &all_present, n_genomes) + { cont = cb(kmer, all_present.clone()); - if !cont { break; } + if !cont { + break; + } } } cont }; - if !cont { return Ok(false); } + if !cont { + return Ok(false); + } } Ok(true) @@ -121,11 +136,13 @@ impl KmerPartition { let mut layer = 0; loop { let layer_dir = self.layer_dir(part, layer); - if !layer_dir.exists() { break; } - let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?; + if !layer_dir.exists() { + break; + } + let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?; let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?; - let counts_dir = layer_dir.join("counts"); + let counts_dir = layer_dir.join("counts"); let presence_dir = layer_dir.join("presence"); let cont = if use_counts && counts_dir.exists() { @@ -136,7 +153,9 @@ impl KmerPartition { let row = mat.row(slot); if passes_all(filters, kmer, &row, n_genomes) { cont = cb(part, layer, kmer, row); - if !cont { break; } + if !cont { + break; + } } } } @@ -149,7 +168,9 @@ impl KmerPartition { let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect(); if passes_all(filters, kmer, &row, n_genomes) { cont = cb(part, layer, kmer, row); - if !cont { break; } + if !cont { + break; + } } } } @@ -161,15 +182,21 @@ impl KmerPartition { let all_present: Box<[u32]> = vec![1u32; n_genomes].into(); let mut cont = true; for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { - if mphf.find(kmer).is_some() && passes_all(filters, kmer, &all_present, n_genomes) { + if mphf.find(kmer).is_some() + && passes_all(filters, kmer, &all_present, n_genomes) + { cont = cb(part, layer, kmer, all_present.clone()); - if !cont { break; } + if !cont { + break; + } } } cont }; - if !cont { return Ok(false); } + if !cont { + return Ok(false); + } layer += 1; } diff --git a/src/obikpartitionner/src/index_layer.rs b/src/obikpartitionner/src/index_layer.rs index 9217df69..8c6a77cc 100644 --- a/src/obikpartitionner/src/index_layer.rs +++ b/src/obikpartitionner/src/index_layer.rs @@ -12,7 +12,7 @@ use ptr_hash::{PtrHash, bucket_fn::CubicEps, hash::Xx64}; use crate::common::olm_to_sk; use crate::graph_pipeline::{materialize_layer, write_graph_as_unitigs}; -use crate::partition::KmerPartition; +use crate::partition::KmerPartitions; type Mphf = PtrHash>, Xx64, Vec>; @@ -24,7 +24,7 @@ fn remove_if_exists(path: &std::path::Path) { } } -impl KmerPartition { +impl KmerPartitions { /// Build the layered MPHF index for partition `i`. /// /// Returns the number of canonical k-mers indexed, or 0 if the partition @@ -93,22 +93,20 @@ impl KmerPartition { } } - let n_kmers = if with_counts { - let n = write_graph_as_unitigs(g, &layer_dir)?; - Layer::::build( - &layer_dir, - block_bits, - mode, - |kmer| match (&mphf1_opt, &counts1_opt) { - (Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())), - _ => 1, - }, - ) - .map_err(|e| olm_to_sk(e, "layer build"))?; - n - } else { - materialize_layer(g, &layer_dir, block_bits, mode)? - }; + let n_kmers = + if with_counts { + let n = write_graph_as_unitigs(g, &layer_dir)?; + Layer::::build(&layer_dir, block_bits, mode, |kmer| { + match (&mphf1_opt, &counts1_opt) { + (Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())), + _ => 1, + } + }) + .map_err(|e| olm_to_sk(e, "layer build"))?; + n + } else { + materialize_layer(g, &layer_dir, block_bits, mode)? + }; let index_dir = layer_dir.parent().expect("layer_dir has a parent"); PartitionMeta { diff --git a/src/obikpartitionner/src/lib.rs b/src/obikpartitionner/src/lib.rs index 5e7801c3..7c59f8ce 100644 --- a/src/obikpartitionner/src/lib.rs +++ b/src/obikpartitionner/src/lib.rs @@ -1,7 +1,7 @@ -pub mod filter; mod common; mod distance; mod dump_layer; +pub mod filter; mod graph_pipeline; mod index_layer; mod kmer_sort; @@ -13,6 +13,6 @@ mod select_layer; pub use filter::{GroupQuorumFilter, KmerFilter, passes_all}; pub use merge_layer::MergeMode; -pub use partition::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR}; +pub use partition::{KmerPartitions, KmerSpectrum, PARTITIONS_SUBDIR}; pub use query_layer::{KmerDesc, QueryHit, QueryStats}; pub use select_layer::{AggOp, OutputCol}; diff --git a/src/obikpartitionner/src/merge_layer/mod.rs b/src/obikpartitionner/src/merge_layer/mod.rs index 1fc68811..77e68e29 100644 --- a/src/obikpartitionner/src/merge_layer/mod.rs +++ b/src/obikpartitionner/src/merge_layer/mod.rs @@ -12,21 +12,20 @@ use std::io; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; -use tracing::debug; use obipipeline::{ - Pipeline, PipelineError, PipelineSender, SharedFlatFn, Stage, WorkerPool, - ThrottleGuard, throttle, - make_sink, make_source, make_transform, + Pipeline, PipelineError, PipelineSender, SharedFlatFn, Stage, ThrottleGuard, WorkerPool, + make_sink, make_source, make_transform, throttle, }; +use tracing::debug; use obicompactvec::{PersistentBitMatrixBuilder, PersistentCompactIntMatrixBuilder}; use obikseq::CanonicalKmer; -use obilayeredmap::{layer_dir, IndexMode, Layer, LayeredMap, MphfOnly}; +use obilayeredmap::{IndexMode, Layer, LayeredMap, MphfOnly, layer_dir}; use obiskio::{SKError, SKResult, UnitigFileReader}; use crate::common::{ColBuilder, load_meta, olm_to_sk}; use crate::graph_pipeline::{build_graph, materialize_layer}; -use crate::partition::KmerPartition; +use crate::partition::KmerPartitions; mod src_layer; @@ -58,14 +57,14 @@ impl MatrixBuilder { fn new(mode: MergeMode, n: usize, dir: &Path) -> io::Result { Ok(match mode { MergeMode::Presence => MatrixBuilder::Bit(PersistentBitMatrixBuilder::new(n, dir)?), - MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::new(n, dir)?), + MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::new(n, dir)?), }) } fn resume(mode: MergeMode, dir: &Path) -> io::Result { Ok(match mode { MergeMode::Presence => MatrixBuilder::Bit(PersistentBitMatrixBuilder::resume(dir)?), - MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::resume(dir)?), + MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::resume(dir)?), }) } @@ -117,7 +116,11 @@ mod matrix_builder_tests { // One source column, filled like pass 2 would. let mut col = mb.add_col().unwrap(); match &mut col { - ColBuilder::Bit(b) => { b.set(0, true); b.set(1, false); b.set(2, true); } + ColBuilder::Bit(b) => { + b.set(0, true); + b.set(1, false); + b.set(2, true); + } ColBuilder::Int(_) => unreachable!(), } col.close().unwrap(); @@ -140,7 +143,10 @@ mod matrix_builder_tests { let mut col = mb.add_col().unwrap(); match &mut col { - ColBuilder::Int(b) => { b.set(0, 7); b.set(1, 42); } + ColBuilder::Int(b) => { + b.set(0, 7); + b.set(1, 42); + } ColBuilder::Bit(_) => unreachable!(), } col.close().unwrap(); @@ -171,7 +177,11 @@ mod matrix_builder_tests { for vals in [[true, false, true], [false, true, false]] { let mut col = mb.add_col().unwrap(); match &mut col { - ColBuilder::Bit(b) => for (slot, v) in vals.into_iter().enumerate() { b.set(slot, v); }, + ColBuilder::Bit(b) => { + for (slot, v) in vals.into_iter().enumerate() { + b.set(slot, v); + } + } ColBuilder::Int(_) => unreachable!(), } col.close().unwrap(); @@ -188,7 +198,7 @@ mod matrix_builder_tests { // ── KmerPartition::merge_partition ──────────────────────────────────────────── -impl KmerPartition { +impl KmerPartitions { /// Merge `sources` into destination partition `i`. /// /// Each entry in `sources` is `(partition, n_genomes)` where `n_genomes` is @@ -201,7 +211,7 @@ impl KmerPartition { pub fn merge_partition( &self, i: usize, - sources: &[(&KmerPartition, usize)], + sources: &[(&KmerPartitions, usize)], mode: MergeMode, n_dst_genomes: usize, block_bits: u8, @@ -213,7 +223,8 @@ impl KmerPartition { } load_meta(&dst_index_dir, "merge")?; // ensure meta.json exists before LayeredMap::open - let dst_map = Arc::new(LayeredMap::<()>::open(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?); + let dst_map = + Arc::new(LayeredMap::<()>::open(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?); let n_dst_layers = dst_map.n_layers(); let n_src_total: usize = sources.iter().map(|(_, n)| *n).sum(); @@ -221,8 +232,11 @@ impl KmerPartition { // (all slots true — every kmer in those layers belongs to genome_0). if n_dst_genomes == 1 && mode == MergeMode::Presence { for l in 0..n_dst_layers { - Layer::<()>::init_presence_matrix(&layer_dir(&dst_index_dir, l), dst_map.layer(l).n()) - .map_err(|e| olm_to_sk(e, "merge"))?; + Layer::<()>::init_presence_matrix( + &layer_dir(&dst_index_dir, l), + dst_map.layer(l).n(), + ) + .map_err(|e| olm_to_sk(e, "merge"))?; } } @@ -283,7 +297,10 @@ impl KmerPartition { )?; let any_new = g.len() > 0; - debug!("partition {i}: de Bruijn graph done — {} new kmers", g.len()); + debug!( + "partition {i}: de Bruijn graph done — {} new kmers", + g.len() + ); // Build new layer from de Bruijn graph if there are new kmers. let new_layer_idx = n_dst_layers; @@ -301,11 +318,16 @@ impl KmerPartition { let t_open = std::time::Instant::now(); let new_mphf: Option> = if any_new { - Some(Arc::new(MphfOnly::open(&new_layer_dir).map_err(|e| olm_to_sk(e, "merge"))?)) + Some(Arc::new( + MphfOnly::open(&new_layer_dir).map_err(|e| olm_to_sk(e, "merge"))?, + )) } else { None }; - debug!("partition {i}: MPHF open in {:.3}s", t_open.elapsed().as_secs_f64()); + debug!( + "partition {i}: MPHF open in {:.3}s", + t_open.elapsed().as_secs_f64() + ); // ── Prepare matrix directories for the new layer ────────────────────── // Absent columns (dst genomes) get an all-zero/false column. Source-genome @@ -351,7 +373,10 @@ impl KmerPartition { exist_builders.push(cols); } - debug!("partition {i}: builders ready in {:.3}s", t_builders.elapsed().as_secs_f64()); + debug!( + "partition {i}: builders ready in {:.3}s", + t_builders.elapsed().as_secs_f64() + ); // ── Pass 2: fill builders (pipeline) ───────────────────────────────── let t_pass2 = std::time::Instant::now(); @@ -391,35 +416,41 @@ impl KmerPartition { .into_iter() .map(|b| Arc::new(Mutex::new(b))) .collect(); - let exist_sink: Vec>>> = exist_locked.iter() + let exist_sink: Vec>>> = exist_locked + .iter() .map(|layer| layer.iter().map(Arc::clone).collect()) .collect(); let new_sink: Vec>> = new_locked.iter().map(Arc::clone).collect(); - let dst_map_t2 = Arc::clone(&dst_map); - let new_mphf_t2 = new_mphf.clone(); + let dst_map_t2 = Arc::clone(&dst_map); + let new_mphf_t2 = new_mphf.clone(); let pass2_err: Arc>> = Arc::new(Mutex::new(None)); - let err_cap2 = Arc::clone(&pass2_err); + let err_cap2 = Arc::clone(&pass2_err); let capacity = 2; let throttled_pass2 = throttle(pass2_items.into_iter(), max_open); let pipeline2 = Pipeline::new( - make_source!(Pass2Data, throttled_pass2.map(|t| { - let (col_offset, src_n, src_layer_dir) = t.item; - (col_offset, src_n, src_layer_dir, t.guard) - }), SrcLayer), + make_source!( + Pass2Data, + throttled_pass2.map(|t| { + let (col_offset, src_n, src_layer_dir) = t.item; + (col_offset, src_n, src_layer_dir, t.guard) + }), + SrcLayer + ), vec![ Stage::Flat(Arc::new( move |data: Pass2Data, - push: &PipelineSender>, - delta: &PipelineSender| - { - if let Pass2Data::SrcLayer((col_offset, src_n, src_layer_dir, _guard)) = data { + push: &PipelineSender>, + delta: &PipelineSender| { + if let Pass2Data::SrcLayer((col_offset, src_n, src_layer_dir, _guard)) = + data + { // _guard dropped at end of block, releasing the slot. let reader = match UnitigFileReader::open_sequential( &src_layer_dir.join("unitigs.bin"), ) { - Ok(r) => r, + Ok(r) => r, Err(e) => { *err_cap2.lock().unwrap() = Some(e.to_string()); delta.send(-1).ok(); @@ -427,7 +458,7 @@ impl KmerPartition { } }; let src_data = match SrcLayerData::open(&src_layer_dir, mode) { - Ok(d) => Arc::new(d), + Ok(d) => Arc::new(d), Err(e) => { *err_cap2.lock().unwrap() = Some(e.to_string()); delta.send(-1).ok(); @@ -440,66 +471,91 @@ impl KmerPartition { for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { batch.push(kmer); if batch.len() == BATCH { - let b = std::mem::replace(&mut batch, Vec::with_capacity(BATCH)); + let b = + std::mem::replace(&mut batch, Vec::with_capacity(BATCH)); push.send(Ok(Pass2Data::RawBatch(( - col_offset, src_n, Arc::clone(&src_data), b, - )))).ok(); + col_offset, + src_n, + Arc::clone(&src_data), + b, + )))) + .ok(); count += 1; } } if !batch.is_empty() { push.send(Ok(Pass2Data::RawBatch(( col_offset, src_n, src_data, batch, - )))).ok(); + )))) + .ok(); count += 1; } delta.send(count - 1).ok(); } - } + }, ) as SharedFlatFn), - make_transform!(Pass2Data, { - move |(col_offset, src_n, src_data, kmers): (usize, usize, Arc, Vec)| - -> Vec<(Option, usize, usize, u32)> + make_transform!( + Pass2Data, { - let mut ops: Vec<(Option, usize, usize, u32)> = Vec::new(); - for kmer in kmers { - let values = src_data.lookup(kmer, src_n); - if let Some((dst_layer, hit)) = dst_map_t2.query(kmer) { - for (g, val) in values.into_iter().enumerate() { - ops.push((Some(dst_layer), col_offset + g, hit.slot, val)); - } - } else if let Some(ref mphf) = new_mphf_t2 { - let slot = mphf.index(kmer); - for (g, val) in values.into_iter().enumerate() { - ops.push((None, col_offset + g, slot, val)); + move |(col_offset, src_n, src_data, kmers): ( + usize, + usize, + Arc, + Vec, + )| + -> Vec<(Option, usize, usize, u32)> { + let mut ops: Vec<(Option, usize, usize, u32)> = Vec::new(); + for kmer in kmers { + let values = src_data.lookup(kmer, src_n); + if let Some((dst_layer, hit)) = dst_map_t2.query(kmer) { + for (g, val) in values.into_iter().enumerate() { + ops.push((Some(dst_layer), col_offset + g, hit.slot, val)); + } + } else if let Some(ref mphf) = new_mphf_t2 { + let slot = mphf.index(kmer); + for (g, val) in values.into_iter().enumerate() { + ops.push((None, col_offset + g, slot, val)); + } } } + ops } - ops - } - }, RawBatch, WriteBatch), + }, + RawBatch, + WriteBatch + ), ], - make_sink!(Pass2Data, { - move |ops: Vec<(Option, usize, usize, u32)>| { - for (layer_opt, col, slot, val) in ops { - match layer_opt { - Some(l) => exist_sink[l][col].lock().unwrap().set_val(slot, val), - None => new_sink[col].lock().unwrap().set_val(slot, val), + make_sink!( + Pass2Data, + { + move |ops: Vec<(Option, usize, usize, u32)>| { + for (layer_opt, col, slot, val) in ops { + match layer_opt { + Some(l) => exist_sink[l][col].lock().unwrap().set_val(slot, val), + None => new_sink[col].lock().unwrap().set_val(slot, val), + } } } - } - }, WriteBatch), + }, + WriteBatch + ), ); WorkerPool::new(pipeline2, n_workers, capacity).run(); - debug!("partition {i}: pass2 pipeline done in {:.3}s", t_pass2.elapsed().as_secs_f64()); + debug!( + "partition {i}: pass2 pipeline done in {:.3}s", + t_pass2.elapsed().as_secs_f64() + ); if let Some(msg) = Arc::try_unwrap(pass2_err) .unwrap_or_else(|_| panic!("pass2: pass2_err not uniquely owned")) .into_inner() .unwrap_or_else(|e| e.into_inner()) { - return Err(SKError::InvalidData { context: "merge pass2", detail: msg }); + return Err(SKError::InvalidData { + context: "merge pass2", + detail: msg, + }); } let t_close = std::time::Instant::now(); @@ -527,10 +583,15 @@ impl KmerPartition { let mut part_meta = self.partition_meta(i)?; part_meta.n_layers = new_layer_idx + 1; - part_meta.save(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?; + part_meta + .save(&dst_index_dir) + .map_err(|e| olm_to_sk(e, "merge"))?; } - debug!("partition {i}: builders closed in {:.3}s", t_close.elapsed().as_secs_f64()); + debug!( + "partition {i}: builders closed in {:.3}s", + t_close.elapsed().as_secs_f64() + ); Ok(n_new) } diff --git a/src/obikpartitionner/src/partition/kmer_partition.rs b/src/obikpartitionner/src/partition/kmer_partition.rs index 57a202c2..aad39b4f 100644 --- a/src/obikpartitionner/src/partition/kmer_partition.rs +++ b/src/obikpartitionner/src/partition/kmer_partition.rs @@ -7,9 +7,9 @@ use std::time::Instant; use obisys::progress_bar; use obikseq::RoutableSuperKmer; -use obiskio::SKResult; use obilayeredmap::IndexMode; use obilayeredmap::meta::PartitionMeta; +use obiskio::SKResult; use rayon::prelude::*; use remove_dir_all::remove_dir_all; use sysinfo::System; @@ -33,12 +33,12 @@ use super::{PARTITIONS_SUBDIR, SK_EXT}; const INDEX_SUBDIR: &str = "index"; pub struct KmerSpectrum { - pub f0: u64, - pub f1: u64, + pub f0: u64, + pub f1: u64, pub counts: BTreeMap, } -pub struct KmerPartition { +pub struct KmerPartitions { root_path: PathBuf, n_partitions: usize, partitions_mask: u64, @@ -49,7 +49,7 @@ pub struct KmerPartition { closed: bool, } -impl KmerPartition { +impl KmerPartitions { pub fn create>( path: P, n_bits: usize, @@ -179,7 +179,9 @@ impl KmerPartition { /// Path of partition `i` directory. pub fn partition_dir(&self, i: usize) -> PathBuf { - self.root_path.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}")) + self.root_path + .join(PARTITIONS_SUBDIR) + .join(format!("part_{i:05}")) } /// Path of partition `i`'s layered-index directory (`/index`). @@ -380,7 +382,7 @@ impl KmerPartition { } } -impl Drop for KmerPartition { +impl Drop for KmerPartitions { fn drop(&mut self) { let _ = self.close(); } diff --git a/src/obikpartitionner/src/partition/mod.rs b/src/obikpartitionner/src/partition/mod.rs index 22ee849b..883c83a3 100644 --- a/src/obikpartitionner/src/partition/mod.rs +++ b/src/obikpartitionner/src/partition/mod.rs @@ -13,7 +13,7 @@ mod kmer_partition; #[cfg(test)] mod tests; -pub use kmer_partition::{KmerPartition, KmerSpectrum}; +pub use kmer_partition::{KmerPartitions, KmerSpectrum}; const SK_EXT: &str = "skmer.zst"; pub const PARTITIONS_SUBDIR: &str = "partitions"; diff --git a/src/obikpartitionner/src/partition/tests.rs b/src/obikpartitionner/src/partition/tests.rs index a98878bb..38d44728 100644 --- a/src/obikpartitionner/src/partition/tests.rs +++ b/src/obikpartitionner/src/partition/tests.rs @@ -6,7 +6,7 @@ use obikseq::SuperKmer; use obiskbuilder::build_superkmers; use super::count::count_partition; -use super::{KmerPartition, PARTITIONS_SUBDIR}; +use super::{KmerPartitions, PARTITIONS_SUBDIR}; const K: usize = 11; const M: usize = 5; @@ -46,7 +46,7 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) { let superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0); let dir = tempfile::tempdir().unwrap(); - let mut kp = KmerPartition::create(dir.path(), 0, K, M, true).unwrap(); + let mut kp = KmerPartitions::create(dir.path(), 0, K, M, true).unwrap(); kp.write_batch(superkmers).unwrap(); kp.close().unwrap(); kp.dereplicate().unwrap(); @@ -58,9 +58,9 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) { } count_partition(&part_dir, &dedup_path, 1 << 20).unwrap(); - let spec: serde_json::Value = serde_json::from_reader( - fs::File::open(part_dir.join("kmer_spectrum_raw.json")).unwrap(), - ).unwrap(); + let spec: serde_json::Value = + serde_json::from_reader(fs::File::open(part_dir.join("kmer_spectrum_raw.json")).unwrap()) + .unwrap(); let f0 = spec["f0"].as_u64().unwrap_or(0); let f1 = spec["f1"].as_u64().unwrap_or(0); (f0, f1) @@ -77,10 +77,7 @@ fn single_sequence_f0_f1_match() { #[test] fn two_sequences_f0_f1_match() { - let seqs: &[&[u8]] = &[ - b"ACGTACGTACGTACGTACGT", - b"TGCATGCATGCATGCATGCA", - ]; + let seqs: &[&[u8]] = &[b"ACGTACGTACGTACGTACGT", b"TGCATGCATGCATGCATGCA"]; let (ef0, ef1) = direct_counts(seqs); let (gf0, gf1) = pipeline_counts(seqs); assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}"); @@ -103,7 +100,11 @@ fn many_sequences_f0_f1_match() { // multiple minimizer boundaries per sequence. let bases = b"ACGT"; let seqs: Vec> = (0..20u32) - .map(|i| (0..40).map(|j| bases[((i * 7 + j * 3) % 4) as usize]).collect()) + .map(|i| { + (0..40) + .map(|j| bases[((i * 7 + j * 3) % 4) as usize]) + .collect() + }) .collect(); let seq_refs: Vec<&[u8]> = seqs.iter().map(|v| v.as_slice()).collect(); let (ef0, ef1) = direct_counts(&seq_refs); diff --git a/src/obikpartitionner/src/query_layer.rs b/src/obikpartitionner/src/query_layer.rs index fdd0ac84..8c4b1588 100644 --- a/src/obikpartitionner/src/query_layer.rs +++ b/src/obikpartitionner/src/query_layer.rs @@ -3,15 +3,18 @@ use std::path::Path; use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; use obikseq::CanonicalKmer; -use obiskio::{SKError, SKResult}; use obilayeredmap::{IndexMode, MphfLayer, OLMError}; +use obiskio::{SKError, SKResult}; -use crate::partition::KmerPartition; +use crate::partition::KmerPartitions; fn olm_to_sk(e: OLMError) -> SKError { match e { OLMError::Io(io_err) => SKError::Io(io_err), - other => SKError::InvalidData { context: "query", detail: other.to_string() }, + other => SKError::InvalidData { + context: "query", + detail: other.to_string(), + }, } } @@ -24,8 +27,8 @@ enum QueryLayer { impl QueryLayer { fn open(layer_dir: &Path, with_counts: bool, mode: &IndexMode) -> SKResult { - let mphf = MphfLayer::open(layer_dir, mode).map_err(olm_to_sk)?; - let counts_dir = layer_dir.join("counts"); + let mphf = MphfLayer::open(layer_dir, mode).map_err(olm_to_sk)?; + let counts_dir = layer_dir.join("counts"); let presence_dir = layer_dir.join("presence"); if with_counts && counts_dir.exists() { @@ -70,7 +73,10 @@ impl QueryLayer { /// slot)` `col_value` point lookup, which this layer's `Sparse` /// presence matrices paid for badly: each such lookup rebuilt the /// entire row just to return one cell. - fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> Box + 'a> { + fn nonzero_iter<'a>( + &'a self, + slots: &'a [usize], + ) -> Box + 'a> { match self { QueryLayer::Presence(_, mat) => mat.nonzero_iter(slots), QueryLayer::Count(_, mat) => Box::new(mat.nonzero_iter(slots)), @@ -137,7 +143,7 @@ pub enum QueryHit<'a> { // ── KmerPartition::query_partition_with ────────────────────────────────────── -impl KmerPartition { +impl KmerPartitions { /// Query a single partition for a pre-deduplicated map of canonical /// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch. /// diff --git a/src/obikpartitionner/src/rebuild_layer.rs b/src/obikpartitionner/src/rebuild_layer.rs index 3d23915e..9873ce84 100644 --- a/src/obikpartitionner/src/rebuild_layer.rs +++ b/src/obikpartitionner/src/rebuild_layer.rs @@ -1,27 +1,29 @@ use std::path::Path; use obicompactvec::{ - FilterMask, eval_filter_mask, - PersistentBitMatrixBuilder, PersistentBitVecBuilder, - PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder, + FilterMask, PersistentBitMatrixBuilder, PersistentBitVecBuilder, + PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder, eval_filter_mask, }; use obidebruinj::GraphDeBruijn; use obikseq::CanonicalKmer; use obilayeredmap::meta::PartitionMeta; -use obilayeredmap::{layer_dir, IndexMode, MphfLayer}; +use obilayeredmap::{IndexMode, MphfLayer, layer_dir}; use obiskio::{SKError, SKResult, UnitigFileReader}; use crate::common::{load_meta, olm_to_sk}; use crate::filter::KmerFilter; use crate::graph_pipeline::materialize_layer; use crate::merge_layer::{MergeMode, SrcLayerData}; -use crate::partition::KmerPartition; +use crate::partition::KmerPartitions; // ── Builders — pair matrix builder + column builders for one mode ───────────── enum Builders { Presence(PersistentBitMatrixBuilder, Vec), - Count(PersistentCompactIntMatrixBuilder, Vec), + Count( + PersistentCompactIntMatrixBuilder, + Vec, + ), } impl Builders { @@ -30,13 +32,18 @@ impl Builders { MergeMode::Presence => { let mut mat = PersistentBitMatrixBuilder::new(n, dir).map_err(SKError::Io)?; let mut cols = Vec::with_capacity(n_genomes); - for _ in 0..n_genomes { cols.push(mat.add_col().map_err(SKError::Io)?); } + for _ in 0..n_genomes { + cols.push(mat.add_col().map_err(SKError::Io)?); + } Ok(Builders::Presence(mat, cols)) } MergeMode::Count => { - let mut mat = PersistentCompactIntMatrixBuilder::new(n, dir).map_err(SKError::Io)?; + let mut mat = + PersistentCompactIntMatrixBuilder::new(n, dir).map_err(SKError::Io)?; let mut cols = Vec::with_capacity(n_genomes); - for _ in 0..n_genomes { cols.push(mat.add_col().map_err(SKError::Io)?); } + for _ in 0..n_genomes { + cols.push(mat.add_col().map_err(SKError::Io)?); + } Ok(Builders::Count(mat, cols)) } } @@ -45,18 +52,22 @@ impl Builders { fn set_val(&mut self, col: usize, slot: usize, value: u32) { match self { Builders::Presence(_, cols) => cols[col].set(slot, value > 0), - Builders::Count(_, cols) => cols[col].set(slot, value), + Builders::Count(_, cols) => cols[col].set(slot, value), } } fn close(self) -> SKResult<()> { match self { Builders::Presence(mat, cols) => { - for b in cols { b.close().map_err(SKError::Io)?; } + for b in cols { + b.close().map_err(SKError::Io)?; + } mat.close().map_err(SKError::Io) } Builders::Count(mat, cols) => { - for b in cols { b.close().map_err(SKError::Io)?; } + for b in cols { + b.close().map_err(SKError::Io)?; + } mat.close().map_err(SKError::Io) } } @@ -112,7 +123,9 @@ fn iter_src_kmers_masked( for l in 0..src_meta.n_layers { let src_layer_dir = layer_dir(src_index_dir, l); let unitigs_path = src_layer_dir.join("unitigs.bin"); - if !unitigs_path.exists() { continue; } + if !unitigs_path.exists() { + continue; + } let src_data = SrcLayerData::open(&src_layer_dir, mode)?; let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?; @@ -127,7 +140,9 @@ fn iter_src_kmers_masked( filters.iter().all(|f| f.passes(kmer, &row, n_genomes)) } }; - if passes { cb(kmer); } + if passes { + cb(kmer); + } } } Ok(()) @@ -149,7 +164,9 @@ fn iter_src_layers( for l in 0..src_meta.n_layers { let src_layer_dir = layer_dir(src_index_dir, l); let unitigs_path = src_layer_dir.join("unitigs.bin"); - if !unitigs_path.exists() { continue; } + if !unitigs_path.exists() { + continue; + } let src_data = SrcLayerData::open(&src_layer_dir, mode)?; let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?; @@ -158,7 +175,9 @@ fn iter_src_layers( for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { let slot = src_data.slot(kmer); if let Some(ref m) = mask { - if !m.get(slot) { continue; } + if !m.get(slot) { + continue; + } let row = src_data.fill_row_by_slot(slot, n_genomes); cb(kmer, row.into_boxed_slice()); } else { @@ -174,7 +193,7 @@ fn iter_src_layers( // ── KmerPartition::rebuild_partition ───────────────────────────────────────── -impl KmerPartition { +impl KmerPartitions { /// Rebuild partition `i` from `src` into `self` (an empty destination partition). /// /// Only k-mers whose per-genome row passes all `filters` are written. @@ -184,7 +203,7 @@ impl KmerPartition { /// `n_genomes` is the number of genome columns in the source (and destination). pub fn rebuild_partition( &self, - src: &KmerPartition, + src: &KmerPartitions, i: usize, filters: &[Box], mode: MergeMode, @@ -222,7 +241,7 @@ impl KmerPartition { // ── Prepare matrix builders (one column per genome) ─────────────────── let data_dir = match mode { MergeMode::Presence => dst_layer_dir.join("presence"), - MergeMode::Count => dst_layer_dir.join("counts"), + MergeMode::Count => dst_layer_dir.join("counts"), }; std::fs::create_dir_all(&data_dir)?; let mut builders = Builders::new(mode, n_new, &data_dir, n_genomes)?; diff --git a/src/obikpartitionner/src/select_layer.rs b/src/obikpartitionner/src/select_layer.rs index 158ab93c..bd66883e 100644 --- a/src/obikpartitionner/src/select_layer.rs +++ b/src/obikpartitionner/src/select_layer.rs @@ -3,14 +3,13 @@ use std::io; use std::path::{Path, PathBuf}; use obicompactvec::{ - ColGroup, MatrixGroupOps, - PersistentBitMatrix, PersistentBitMatrixBuilder, + ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, }; use obilayeredmap::OLMError; use obiskio::{SKError, SKResult}; -use crate::partition::KmerPartition; +use crate::partition::KmerPartitions; // ── AggOp ───────────────────────────────────────────────────────────────────── @@ -33,9 +32,9 @@ impl AggOp { // ── OutputCol ───────────────────────────────────────────────────────────────── pub struct OutputCol { - pub label: String, + pub label: String, pub indices: Vec, - pub op: AggOp, + pub op: AggOp, } // ── Helpers ─────────────────────────────────────────────────────────────────── @@ -43,7 +42,10 @@ pub struct OutputCol { fn olm_to_sk(e: OLMError) -> SKError { match e { OLMError::Io(e) => SKError::Io(e), - other => SKError::InvalidData { context: "select", detail: other.to_string() }, + other => SKError::InvalidData { + context: "select", + detail: other.to_string(), + }, } } @@ -77,23 +79,43 @@ fn fill_builders( if output_presence { let b = dst_bit.as_deref_mut().unwrap(); match spec.op { - AggOp::Any => b.add_col_from (&mat.partial_group_any (&g, threshold).map_err(SKError::Io)?), - AggOp::All => b.add_col_from (&mat.partial_group_all (&g, threshold).map_err(SKError::Io)?), - AggOp::None => b.add_col_from (&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?), - AggOp::Sum => b.add_col_from_int(&mat.partial_group_sum (&g).map_err(SKError::Io)?), - AggOp::Min => b.add_col_from_int(&mat.partial_group_min (&g).map_err(SKError::Io)?), - AggOp::Max => b.add_col_from_int(&mat.partial_group_max (&g).map_err(SKError::Io)?), - }.map_err(SKError::Io)?; + AggOp::Any => { + b.add_col_from(&mat.partial_group_any(&g, threshold).map_err(SKError::Io)?) + } + AggOp::All => { + b.add_col_from(&mat.partial_group_all(&g, threshold).map_err(SKError::Io)?) + } + AggOp::None => { + b.add_col_from(&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?) + } + AggOp::Sum => { + b.add_col_from_int(&mat.partial_group_sum(&g).map_err(SKError::Io)?) + } + AggOp::Min => { + b.add_col_from_int(&mat.partial_group_min(&g).map_err(SKError::Io)?) + } + AggOp::Max => { + b.add_col_from_int(&mat.partial_group_max(&g).map_err(SKError::Io)?) + } + } + .map_err(SKError::Io)?; } else { let b = dst_int.as_deref_mut().unwrap(); match spec.op { - AggOp::Sum => b.add_col_from (&mat.partial_group_sum (&g).map_err(SKError::Io)?), - AggOp::Min => b.add_col_from (&mat.partial_group_min (&g).map_err(SKError::Io)?), - AggOp::Max => b.add_col_from (&mat.partial_group_max (&g).map_err(SKError::Io)?), - AggOp::Any => b.add_col_from_bit(&mat.partial_group_any (&g, threshold).map_err(SKError::Io)?), - AggOp::All => b.add_col_from_bit(&mat.partial_group_all (&g, threshold).map_err(SKError::Io)?), - AggOp::None => b.add_col_from_bit(&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?), - }.map_err(SKError::Io)?; + AggOp::Sum => b.add_col_from(&mat.partial_group_sum(&g).map_err(SKError::Io)?), + AggOp::Min => b.add_col_from(&mat.partial_group_min(&g).map_err(SKError::Io)?), + AggOp::Max => b.add_col_from(&mat.partial_group_max(&g).map_err(SKError::Io)?), + AggOp::Any => b.add_col_from_bit( + &mat.partial_group_any(&g, threshold).map_err(SKError::Io)?, + ), + AggOp::All => b.add_col_from_bit( + &mat.partial_group_all(&g, threshold).map_err(SKError::Io)?, + ), + AggOp::None => b.add_col_from_bit( + &mat.partial_group_none(&g, threshold).map_err(SKError::Io)?, + ), + } + .map_err(SKError::Io)?; } } } else { @@ -103,23 +125,43 @@ fn fill_builders( if output_presence { let b = dst_bit.as_deref_mut().unwrap(); match spec.op { - AggOp::Any => b.add_col_from (&mat.partial_group_any (&g, 1).map_err(SKError::Io)?), - AggOp::All => b.add_col_from (&mat.partial_group_all (&g, 1).map_err(SKError::Io)?), - AggOp::None => b.add_col_from (&mat.partial_group_none(&g, 1).map_err(SKError::Io)?), - AggOp::Sum => b.add_col_from_int(&mat.partial_group_sum (&g).map_err(SKError::Io)?), - AggOp::Min => b.add_col_from_int(&mat.partial_group_min (&g).map_err(SKError::Io)?), - AggOp::Max => b.add_col_from_int(&mat.partial_group_max (&g).map_err(SKError::Io)?), - }.map_err(SKError::Io)?; + AggOp::Any => { + b.add_col_from(&mat.partial_group_any(&g, 1).map_err(SKError::Io)?) + } + AggOp::All => { + b.add_col_from(&mat.partial_group_all(&g, 1).map_err(SKError::Io)?) + } + AggOp::None => { + b.add_col_from(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?) + } + AggOp::Sum => { + b.add_col_from_int(&mat.partial_group_sum(&g).map_err(SKError::Io)?) + } + AggOp::Min => { + b.add_col_from_int(&mat.partial_group_min(&g).map_err(SKError::Io)?) + } + AggOp::Max => { + b.add_col_from_int(&mat.partial_group_max(&g).map_err(SKError::Io)?) + } + } + .map_err(SKError::Io)?; } else { let b = dst_int.as_deref_mut().unwrap(); match spec.op { - AggOp::Sum => b.add_col_from (&mat.partial_group_sum (&g).map_err(SKError::Io)?), - AggOp::Min => b.add_col_from (&mat.partial_group_min (&g).map_err(SKError::Io)?), - AggOp::Max => b.add_col_from (&mat.partial_group_max (&g).map_err(SKError::Io)?), - AggOp::Any => b.add_col_from_bit(&mat.partial_group_any (&g, 1).map_err(SKError::Io)?), - AggOp::All => b.add_col_from_bit(&mat.partial_group_all (&g, 1).map_err(SKError::Io)?), - AggOp::None => b.add_col_from_bit(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?), - }.map_err(SKError::Io)?; + AggOp::Sum => b.add_col_from(&mat.partial_group_sum(&g).map_err(SKError::Io)?), + AggOp::Min => b.add_col_from(&mat.partial_group_min(&g).map_err(SKError::Io)?), + AggOp::Max => b.add_col_from(&mat.partial_group_max(&g).map_err(SKError::Io)?), + AggOp::Any => { + b.add_col_from_bit(&mat.partial_group_any(&g, 1).map_err(SKError::Io)?) + } + AggOp::All => { + b.add_col_from_bit(&mat.partial_group_all(&g, 1).map_err(SKError::Io)?) + } + AggOp::None => { + b.add_col_from_bit(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?) + } + } + .map_err(SKError::Io)?; } } } @@ -128,7 +170,7 @@ fn fill_builders( // ── KmerPartition::select_partition ────────────────────────────────────────── -impl KmerPartition { +impl KmerPartitions { /// Rewrite the data matrices of partition `i` in `src` into `self`. /// /// `specs` defines the output columns (projection/aggregation). @@ -136,7 +178,7 @@ impl KmerPartition { /// `in_place` — `self` and `src` share the same root; write to temp dirs then swap. pub fn select_partition( &self, - src: &KmerPartition, + src: &KmerPartitions, i: usize, specs: &[OutputCol], _n_src_genomes: usize, @@ -159,23 +201,33 @@ impl KmerPartition { fs::create_dir_all(&dst_index_dir)?; } - let data_subdir = if output_presence { "presence" } else { "counts" }; + let data_subdir = if output_presence { + "presence" + } else { + "counts" + }; for l in 0..n_src_layers { let src_layer_dir = src.layer_dir(i, l); - if !src_layer_dir.exists() { continue; } + if !src_layer_dir.exists() { + continue; + } let dst_layer_dir = self.layer_dir(i, l); - let counts_dir = src_layer_dir.join("counts"); + let counts_dir = src_layer_dir.join("counts"); let presence_dir = src_layer_dir.join("presence"); let src_is_count = counts_dir.exists() && !presence_dir.exists(); // Determine number of slots and detect implicit layers. let n = if counts_dir.exists() { - PersistentCompactIntMatrix::open(&src_layer_dir).map_err(SKError::Io)?.n() + PersistentCompactIntMatrix::open(&src_layer_dir) + .map_err(SKError::Io)? + .n() } else if presence_dir.exists() { - PersistentBitMatrix::open(&src_layer_dir).map_err(SKError::Io)?.n() + PersistentBitMatrix::open(&src_layer_dir) + .map_err(SKError::Io)? + .n() } else { // Implicit single-genome layer: no data matrix needed in output either. if !in_place { @@ -187,7 +239,7 @@ impl KmerPartition { // Choose the output data directory (temp name for in-place). let (dst_data_dir, final_data_dir): (PathBuf, PathBuf) = if in_place { - let tmp = dst_layer_dir.join(format!("{data_subdir}_new")); + let tmp = dst_layer_dir.join(format!("{data_subdir}_new")); let perm = dst_layer_dir.join(data_subdir); (tmp, perm) } else { @@ -202,14 +254,28 @@ impl KmerPartition { fs::create_dir_all(&dst_data_dir)?; let (mut dst_bit, mut dst_int) = if output_presence { - (Some(PersistentBitMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?), None) + ( + Some(PersistentBitMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?), + None, + ) } else { - (None, Some(PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?)) + ( + None, + Some( + PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir) + .map_err(SKError::Io)?, + ), + ) }; fill_builders( - specs, &src_layer_dir, src_is_count, threshold, output_presence, - dst_bit.as_mut(), dst_int.as_mut(), + specs, + &src_layer_dir, + src_is_count, + threshold, + output_presence, + dst_bit.as_mut(), + dst_int.as_mut(), )?; if output_presence { @@ -233,7 +299,9 @@ impl KmerPartition { } if !in_place { - src.partition_meta(i)?.save(&dst_index_dir).map_err(olm_to_sk)?; + src.partition_meta(i)? + .save(&dst_index_dir) + .map_err(olm_to_sk)?; } Ok(()) diff --git a/src/obikpartitionner/src/tests/query_layer.rs b/src/obikpartitionner/src/tests/query_layer.rs index ee6f170a..8a07e95f 100644 --- a/src/obikpartitionner/src/tests/query_layer.rs +++ b/src/obikpartitionner/src/tests/query_layer.rs @@ -44,8 +44,8 @@ fn query_stats_default_is_zero() { #[test] fn query_partition_with_missing_index_dir_returns_default_stats() { let tmp = tempfile::tempdir().expect("tempdir"); - let partition = KmerPartition::create(tmp.path().join("idx"), 2, 21, 9, false) - .expect("create partition"); + let partition = + KmerPartitions::create(tmp.path().join("idx"), 2, 21, 9, false.expect("create partition"); let mut kmers: HashMap> = HashMap::new(); // Any well-formed canonical k-mer works here — the call must return @@ -65,8 +65,8 @@ fn query_partition_with_missing_index_dir_returns_default_stats() { #[test] fn query_partition_with_empty_kmers_is_a_noop() { let tmp = tempfile::tempdir().expect("tempdir"); - let partition = KmerPartition::create(tmp.path().join("idx"), 2, 21, 9, false) - .expect("create partition"); + let partition = + KmerPartitions::create(tmp.path().join("idx"), 2, 21, 9, false.expect("create partition"); let kmers: HashMap> = HashMap::new(); let stats = partition diff --git a/src/obikphylo/src/siblings/alignment.rs b/src/obikphylo/src/siblings/alignment.rs index 8aae451c..754fd577 100644 --- a/src/obikphylo/src/siblings/alignment.rs +++ b/src/obikphylo/src/siblings/alignment.rs @@ -1,13 +1,13 @@ use std::sync::Arc; -use obikpartitionner::KmerPartition; +use obikpartitionner::KmerPartitions; use obisys::progress_bar; -use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; +use obikindex::{OKIError, OKIResult}; use super::cache::PartitionCache; -use super::family_scan::{scan_layer_families, Selection}; +use super::family_scan::{Selection, scan_layer_families}; use super::subsample::EntropyBias; /// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set @@ -70,18 +70,26 @@ pub trait SnpAlignmentExt { /// that sample (or, with `subsample = None`, filter the whole index) /// by a Gaussian kernel on each family's entropy — see /// `DevDocMD/architecture/siblings.md`, "Entropy-biased selection". - fn snp_pseudo_alignment(&self, subsample: Option, entropy_bias: Option) -> OKIResult; + fn snp_pseudo_alignment( + &self, + subsample: Option, + entropy_bias: Option, + ) -> OKIResult; } impl SnpAlignmentExt for KmerIndex { - fn snp_pseudo_alignment(&self, subsample: Option, entropy_bias: Option) -> OKIResult { + fn snp_pseudo_alignment( + &self, + subsample: Option, + entropy_bias: Option, + ) -> OKIResult { 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( + let partition = KmerPartitions::open_with_config( self.root_path(), self.kmer_size(), self.minimizer_size(), @@ -90,7 +98,8 @@ impl SnpAlignmentExt for KmerIndex { .map_err(OKIError::Partition)?; let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); let layer_dirs = super::family_scan::sibling_layer_dirs(self)?; - let selections = super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?; + let selections = + super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?; let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers"); // One layer at a time, not `par_iter()` over layers — same @@ -109,14 +118,23 @@ impl SnpAlignmentExt for KmerIndex { None => Selection::All, Some(set) => Selection::Some(set), }; - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |_family_idx, mask, genome_mask| { - if mask.family_size() < 2 { - return; // monomorphic family — no signal, skip - } - for (g, &m) in genome_mask.iter().enumerate() { - sequences[g].push(iupac_code(m)); - } - })?; + scan_layer_families( + layer_dir, + n_parts, + n_genomes, + with_counts, + k, + &cache, + &selection, + |_family_idx, mask, genome_mask| { + if mask.family_size() < 2 { + return; // monomorphic family — no signal, skip + } + for (g, &m) in genome_mask.iter().enumerate() { + sequences[g].push(iupac_code(m)); + } + }, + )?; pb.inc(1); } pb.finish_and_clear(); diff --git a/src/obikphylo/src/siblings/build.rs b/src/obikphylo/src/siblings/build.rs index 358b0778..6d125e19 100644 --- a/src/obikphylo/src/siblings/build.rs +++ b/src/obikphylo/src/siblings/build.rs @@ -1,22 +1,22 @@ use std::path::Path; -use std::sync::atomic::Ordering; use std::sync::Arc; +use std::sync::atomic::Ordering; use rayon::prelude::*; -use obikpartitionner::KmerPartition; -use obipipeline::ThrottleGuard; +use obikpartitionner::KmerPartitions; use obikseq::CanonicalKmer; use obilayeredmap::MphfLayer; use obilayeredmap::meta::IndexMode; +use obipipeline::ThrottleGuard; use obisys::progress_bar; -use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; +use obikindex::{OKIError, OKIResult}; use super::cache::PartitionCache; use super::helpers::{central_base, is_minorant}; -use super::{olm_to_ok, FamilyMask, SiblingAnnexBuilder, ANNEX_FILE_NAME}; +use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnexBuilder, olm_to_ok}; // ── obipipeline data types ───────────────────────────────────────────────── @@ -78,7 +78,7 @@ impl SiblingAnnexBuildExt for KmerIndex { let n_parts = self.n_partitions(); let n_bits = n_parts.trailing_zeros() as usize; - let partition = KmerPartition::open_with_config( + let partition = KmerPartitions::open_with_config( self.root_path(), self.kmer_size(), self.minimizer_size(), @@ -87,7 +87,11 @@ impl SiblingAnnexBuildExt for KmerIndex { .map_err(OKIError::Partition)?; tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep"); - let cache = Arc::new(PartitionCache::build(&partition, n_parts, self.meta().config.with_counts)?); + let cache = Arc::new(PartitionCache::build( + &partition, + n_parts, + self.meta().config.with_counts, + )?); let pb = progress_bar("sibling_annex", n_parts as u64, "partitions"); let mut total_slots: u64 = 0; @@ -102,12 +106,19 @@ impl SiblingAnnexBuildExt for KmerIndex { let mut part_slots: u64 = 0; for l in 0..meta.n_layers { part_slots += build_layer_sibling_annex( - self, &self.partition().layer_dir(part, l), &meta.mode, n_parts, l, &cache, + self, + &self.partition().layer_dir(part, l), + &meta.mode, + n_parts, + l, + &cache, )?; } total_slots += part_slots; pb.inc(1); - pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)")); + 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"); @@ -139,213 +150,232 @@ fn build_layer_sibling_annex( // agree; see `PartitionCache::fast_mode`'s docs. let fast_mode = cache.fast_mode(); - // ── The annex file itself is the reconciliation state, indexed by - // this layer's k-mer iteration order (the physical layout of - // `unitigs.bin`), never by MPHF slot — this layer's own k-mers are - // known members by construction, so no evidence check, no MPHF - // slot, and no slot -> k-mer reconstruction is needed or legitimate - // here (see `DevDocMD/architecture/siblings.md`: the MPHF is not - // invertible, and evidence answers membership, not identity). No - // separate in-memory accumulator: every concurrent write goes - // straight through `SiblingAnnexBuilder::atomic_slot` into the - // mmap — a layer can hold billions of k-mers, so a `Vec` shadowing - // the whole file in RAM just to copy it out again at the end (the - // previous design) doubles memory for no benefit once the builder - // itself can be written concurrently. - // - // `Arc`, not a bare `SiblingAnnexBuilder` — - // `Pipe::apply` requires its source iterator to be `Send + 'static` - // (its items are dispatched to worker threads that outlive this - // call), so the batch-generating closure below needs an owned - // handle it can move in, not a borrow of a local. `atomic_slot`, - // not `set`, for every concurrent write below: the gather phase - // parallelises across destination partitions (independent - // `cache.find` calls, safe to run concurrently) and their hits can - // land on arbitrary, possibly-shared source entries — a lock-free - // `fetch_or` avoids needing any synchronisation beyond that. - // Reclaimed as a plain owned value (`Arc::try_unwrap`) once every - // concurrent phase below is done, for the sequential finalisation - // pass. ───────────────────────────────────────────────────────── - let annex_path = layer_dir.join(ANNEX_FILE_NAME); - let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?); + // ── The annex file itself is the reconciliation state, indexed by + // this layer's k-mer iteration order (the physical layout of + // `unitigs.bin`), never by MPHF slot — this layer's own k-mers are + // known members by construction, so no evidence check, no MPHF + // slot, and no slot -> k-mer reconstruction is needed or legitimate + // here (see `DevDocMD/architecture/siblings.md`: the MPHF is not + // invertible, and evidence answers membership, not identity). No + // separate in-memory accumulator: every concurrent write goes + // straight through `SiblingAnnexBuilder::atomic_slot` into the + // mmap — a layer can hold billions of k-mers, so a `Vec` shadowing + // the whole file in RAM just to copy it out again at the end (the + // previous design) doubles memory for no benefit once the builder + // itself can be written concurrently. + // + // `Arc`, not a bare `SiblingAnnexBuilder` — + // `Pipe::apply` requires its source iterator to be `Send + 'static` + // (its items are dispatched to worker threads that outlive this + // call), so the batch-generating closure below needs an owned + // handle it can move in, not a borrow of a local. `atomic_slot`, + // not `set`, for every concurrent write below: the gather phase + // parallelises across destination partitions (independent + // `cache.find` calls, safe to run concurrently) and their hits can + // land on arbitrary, possibly-shared source entries — a lock-free + // `fetch_or` avoids needing any synchronisation beyond that. + // Reclaimed as a plain owned value (`Arc::try_unwrap`) once every + // concurrent phase below is done, for the sequential finalisation + // pass. ───────────────────────────────────────────────────────── + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?); - // ── 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. `BATCH_SIZE` was originally tuned back when the - // per-k-mer cost inside this stage (minimiser via `RollingStat`) - // was ~3x higher than it is now (see `CanonicalKmerOf::minimizer`, - // a direct O(k) bit-arithmetic replacement) — at the old per-k-mer - // cost, dispatch overhead (the scheduler's single-threaded `Select` - // loop: one channel round-trip per batch) was negligible next to - // the work each batch represented; now that the work itself is - // cheaper, that fixed per-batch overhead is proportionally larger, - // so a bigger batch amortises it over more k-mers again. ───────── - const BATCH_SIZE: usize = 32768; - let n_workers = obisys::effective_parallelism(); - let capacity = 256; + // ── 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. `BATCH_SIZE` was originally tuned back when the + // per-k-mer cost inside this stage (minimiser via `RollingStat`) + // was ~3x higher than it is now (see `CanonicalKmerOf::minimizer`, + // a direct O(k) bit-arithmetic replacement) — at the old per-k-mer + // cost, dispatch overhead (the scheduler's single-threaded `Select` + // loop: one channel round-trip per batch) was negligible next to + // the work each batch represented; now that the work itself is + // cheaper, that fixed per-batch overhead is proportionally larger, + // so a bigger batch amortises it over more k-mers again. ───────── + const BATCH_SIZE: usize = 32768; + 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. - // - // Batches stream straight from `unitigs.bin` via - // `enumerate_kmers_batch` — never a full-layer `Vec` collect (a - // layer can hold billions of k-mers; see the "no full collect" - // rule). Each k-mer's own base is seeded straight into the annex in - // this same pass, since this is exactly the iteration order the - // annex is keyed on — no separate seeding pass needed. - let seed_builder = Arc::clone(&builder); - let batches = mphf.enumerate_kmers_batch(BATCH_SIZE).map(move |(start, kmers)| { + // 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. + // + // Batches stream straight from `unitigs.bin` via + // `enumerate_kmers_batch` — never a full-layer `Vec` collect (a + // layer can hold billions of k-mers; see the "no full collect" + // rule). Each k-mer's own base is seeded straight into the annex in + // this same pass, since this is exactly the iteration order the + // annex is keyed on — no separate seeding pass needed. + let seed_builder = Arc::clone(&builder); + let batches = mphf + .enumerate_kmers_batch(BATCH_SIZE) + .map(move |(start, kmers)| { kmers .into_iter() .enumerate() .map(|(i, kmer)| { let base = central_base(kmer, k); - let bits = if fast_mode { FamilyMask::layer_bits(base, l) } else { FamilyMask::presence_bits(base) }; - seed_builder.atomic_slot(start + i).fetch_or(bits, Ordering::Relaxed); + let bits = if fast_mode { + FamilyMask::layer_bits(base, l) + } else { + FamilyMask::presence_bits(base) + }; + seed_builder + .atomic_slot(start + i) + .fetch_or(bits, Ordering::Relaxed); (start + i, kmer) }) .collect::>() }); - // Diagnostic: count still-empty annex slots after seeding + cross-partition - // resolution. A non-zero count here means some k-mer's own base was never - // written (should be impossible after the batch_offset fix above). - let check_empty = |label: &str| { - let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count(); - if empty > 0 { - tracing::warn!("{label}: {empty} mask slots still empty after resolution (layer {layer_dir:?})"); - } - }; - let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch { - items: t.item, - _permit: t.guard, - }); + // Diagnostic: count still-empty annex slots after seeding + cross-partition + // resolution. A non-zero count here means some k-mer's own base was never + // written (should be impossible after the batch_offset fix above). + let check_empty = |label: &str| { + let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count(); + if empty > 0 { + tracing::warn!( + "{label}: {empty} mask slots still empty after resolution (layer {layer_dir:?})" + ); + } + }; + let throttled = obipipeline::throttle(batches, 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 (order, kmer) in batch.items { - for variant in kmer.central_canonical_neighbors() { - if variant == kmer { - continue; - } - items.push(( - variant.partition(n_parts), - variant, - order, - central_base(variant, k), - )); + let pipe = obipipeline::make_pipe! { + SibData : SourceBatch => VariantBatch, + | { + move |batch: SourceBatch| -> VariantBatch { + let mut items = Vec::with_capacity(batch.items.len() * 3); + for (order, kmer) in batch.items { + for variant in kmer.central_canonical_neighbors() { + if variant == kmer { + continue; } + items.push(( + variant.partition(n_parts), + variant, + order, + central_base(variant, k), + )); } - VariantBatch { items, _permit: batch._permit } } - } : Batch => Variants, - }; + 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> = (0..n_parts).map(|_| Vec::new()).collect(); - for vb in pipe.apply(throttled, n_workers, capacity) { - for (dest_partition, variant, source_order, base) in vb.items { - outgoing[dest_partition].push((variant, source_order, base)); + // ── 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> = + (0..n_parts).map(|_| Vec::new()).collect(); + for vb in pipe.apply(throttled, n_workers, capacity) { + for (dest_partition, variant, source_order, base) in vb.items { + outgoing[dest_partition].push((variant, source_order, base)); + } + } + + check_empty("after_seeding"); + + // ── Resolve each partition's batch against the cache, parallelised + // over roughly equal-sized *chunks*, not over partitions — a plain + // `outgoing.par_iter()` over `n_parts` buckets gives one thread the + // whole of one partition's bucket, however large, and a lot of + // buckets are far from equal: a central-base substitution changes + // the minimiser (and thus routes to a different partition) only + // when the winning minimiser window overlaps the central position. + // For k=31/m=11 that is ~11 of the 21 possible windows — so *most* + // of the remaining ~10/21 send the variant right back to the + // partition already being built. That self-partition bucket ends up + // far larger than any other, so the naive per-partition split + // leaves one thread grinding through it alone long after every + // other partition's (tiny) bucket is done — confirmed by sampling a + // real run: one thread solely in `MphfLayer::find` while the rest + // of the pool sits idle. Splitting each non-empty bucket into + // `chunk_size`-sized pieces first keeps this pass's per-partition + // locality (each chunk is still contiguous within one partition, + // still resolved via `cache.find` grouped by that partition) while + // letting Rayon spread a single oversized bucket across several + // threads instead of pinning it to one. ────────────────────────── + let chunk_size = ((outgoing.iter().map(Vec::len).sum::() / n_workers).max(1)).min(4096); + let work: Vec<(usize, &[(CanonicalKmer, usize, u8)])> = outgoing + .iter() + .enumerate() + .filter(|(_, q)| !q.is_empty()) + .flat_map(|(dest, q)| q.chunks(chunk_size).map(move |c| (dest, c))) + .collect(); + work.par_iter().for_each(|&(dest, chunk)| { + for &(variant, source_order, base) in chunk { + if let Some(layer) = cache.find(dest, variant) { + let bits = if fast_mode { + FamilyMask::layer_bits(base, layer) + } else { + FamilyMask::presence_bits(base) + }; + builder + .atomic_slot(source_order) + .fetch_or(bits, Ordering::Relaxed); } } + }); - check_empty("after_seeding"); + check_empty("after_resolution"); - // ── Resolve each partition's batch against the cache, parallelised - // over roughly equal-sized *chunks*, not over partitions — a plain - // `outgoing.par_iter()` over `n_parts` buckets gives one thread the - // whole of one partition's bucket, however large, and a lot of - // buckets are far from equal: a central-base substitution changes - // the minimiser (and thus routes to a different partition) only - // when the winning minimiser window overlaps the central position. - // For k=31/m=11 that is ~11 of the 21 possible windows — so *most* - // of the remaining ~10/21 send the variant right back to the - // partition already being built. That self-partition bucket ends up - // far larger than any other, so the naive per-partition split - // leaves one thread grinding through it alone long after every - // other partition's (tiny) bucket is done — confirmed by sampling a - // real run: one thread solely in `MphfLayer::find` while the rest - // of the pool sits idle. Splitting each non-empty bucket into - // `chunk_size`-sized pieces first keeps this pass's per-partition - // locality (each chunk is still contiguous within one partition, - // still resolved via `cache.find` grouped by that partition) while - // letting Rayon spread a single oversized bucket across several - // threads instead of pinning it to one. ────────────────────────── - let chunk_size = ((outgoing.iter().map(Vec::len).sum::() / n_workers).max(1)).min(4096); - let work: Vec<(usize, &[(CanonicalKmer, usize, u8)])> = outgoing - .iter() - .enumerate() - .filter(|(_, q)| !q.is_empty()) - .flat_map(|(dest, q)| q.chunks(chunk_size).map(move |c| (dest, c))) - .collect(); - work.par_iter().for_each(|&(dest, chunk)| { - for &(variant, source_order, base) in chunk { - if let Some(layer) = cache.find(dest, variant) { - let bits = if fast_mode { FamilyMask::layer_bits(base, layer) } else { FamilyMask::presence_bits(base) }; - builder.atomic_slot(source_order).fetch_or(bits, Ordering::Relaxed); - } - } - }); - - check_empty("after_resolution"); - - // ── Write the layer's annex file, indexed by iteration order — a - // second streamed pass over `unitigs.bin` (via `enumerate_kmers`), - // now that every entry's mask is final; never a `Vec` hold of the - // whole layer. The minorant flag is computed here, not by every - // later reader: this is the one place the whole family's *final* - // mask and this entry's own k-mer (already in hand, no extra - // lookup) are both available together. Every consumer that only - // needs to know "is this k-mer the family's minorant" — the common - // case, since a family is tallied once, at its minorant — reads - // the bit straight back instead of re-deriving it (re-scanning - // `unitigs.bin` and re-hashing through the MPHF, the cost - // `is_minorant` was cheap to *compute* but expensive to *get the - // inputs for* every time). - // Every concurrent phase above is done and its `Arc` clone (the - // seeding pass's `seed_builder`) has already been dropped along - // with the now-fully-drained `batches` iterator — this is the only - // remaining strong reference, so reclaiming plain ownership for the - // sequential `set` calls below is guaranteed to succeed. - let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| panic!("sibling-annex builder still shared after every concurrent phase completed")); - for (order, kmer) in mphf.enumerate_kmers() { - let final_mask = builder.get(order).expect("seeded by construction"); - let minorant = is_minorant(kmer, final_mask, k); - builder.set(order, final_mask.with_minorant(minorant)); - } - builder.close()?; + // ── Write the layer's annex file, indexed by iteration order — a + // second streamed pass over `unitigs.bin` (via `enumerate_kmers`), + // now that every entry's mask is final; never a `Vec` hold of the + // whole layer. The minorant flag is computed here, not by every + // later reader: this is the one place the whole family's *final* + // mask and this entry's own k-mer (already in hand, no extra + // lookup) are both available together. Every consumer that only + // needs to know "is this k-mer the family's minorant" — the common + // case, since a family is tallied once, at its minorant — reads + // the bit straight back instead of re-deriving it (re-scanning + // `unitigs.bin` and re-hashing through the MPHF, the cost + // `is_minorant` was cheap to *compute* but expensive to *get the + // inputs for* every time). + // Every concurrent phase above is done and its `Arc` clone (the + // seeding pass's `seed_builder`) has already been dropped along + // with the now-fully-drained `batches` iterator — this is the only + // remaining strong reference, so reclaiming plain ownership for the + // sequential `set` calls below is guaranteed to succeed. + let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| { + panic!("sibling-annex builder still shared after every concurrent phase completed") + }); + for (order, kmer) in mphf.enumerate_kmers() { + let final_mask = builder.get(order).expect("seeded by construction"); + let minorant = is_minorant(kmer, final_mask, k); + builder.set(order, final_mask.with_minorant(minorant)); + } + builder.close()?; Ok(n as u64) } diff --git a/src/obikphylo/src/siblings/cache.rs b/src/obikphylo/src/siblings/cache.rs index 315c0660..8bc8f4f2 100644 --- a/src/obikphylo/src/siblings/cache.rs +++ b/src/obikphylo/src/siblings/cache.rs @@ -2,17 +2,17 @@ use rayon::prelude::*; use std::path::Path; -use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, PersistentSparseBitMatrix}; -use obikpartitionner::KmerPartition; +use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; +use obikpartitionner::KmerPartitions; use obikseq::CanonicalKmer; -use obilayeredmap::{Layer, OLMResult}; use obilayeredmap::meta::IndexMode; +use obilayeredmap::{Layer, OLMResult}; use obisys::progress_bar; use obikindex::OKIResult; -use super::iter::SiblingLayerExt; use super::SiblingAnnex; +use super::iter::SiblingLayerExt; /// Every partition's already-open layers, built **once** for the whole /// `build_sibling_annex` run and shared (read-only) across every lookup, in @@ -41,47 +41,40 @@ use super::SiblingAnnex; /// going back through the low-level pieces `Layer` already assembles. pub(super) enum Mat { Count(Layer), + /// Dense *or* sparse (`pack --sparse`d) presence — `PersistentBitMatrix` + /// itself is a 4-way enum (`Columnar`/`Packed`/`Sparse`/`Implicit`) that + /// already auto-detects sparse storage in its own `open()` (checks + /// `presence/sparse_meta.json`) and dispatches every method + /// (`row`/`fill_row`/`nonzero_iter`/...) across all four internally — + /// see `DevDocMD/implementation/partition_layer_cache.md`. A separate + /// `SparsePresence(Layer)` arm used to exist + /// here, opened via its own `presence/is_multi.prsb` check; it + /// predated `PersistentBitMatrix` growing native sparse support and + /// was pure duplication by the time it was removed — every method + /// below dispatched it identically to this arm. Presence(Layer), - /// Same role as `Presence`, over a layer packed by `pack --sparse` - /// (`PersistentSparseBitMatrix`) instead of the dense form — see - /// `DevDocMD/architecture/siblings.md`'s sparse-matrix section. Every - /// `Mat` method below dispatches to the same `Layer` generic code - /// (`D: LayerData`) as `Presence`, so this arm is purely a storage - /// choice, not a behavioural difference. - SparsePresence(Layer), } impl Mat { - /// Open one layer's matrix, auto-detecting count vs. dense-presence vs. - /// sparse-presence from what is actually on disk — the single source of - /// truth for *every* caller that opens a layer's own matrix (this - /// module's cross-partition [`PartitionCache::build`] and - /// `family_scan::scan_layer_families`'s own-layer lookup alike), so the - /// two can never disagree about which format a layer's `presence/` was - /// packed to. Before this existed, `scan_layer_families` open-coded its - /// own (non-sparse-aware) copy of this decision: on a `pack --sparse`d - /// layer, `presence/matrix.pbmx` no longer exists (removed by - /// `pack_sparse_bit_matrix`), so a caller that only checks for it and - /// otherwise unconditionally opens `PersistentBitMatrix` doesn't error — - /// `PersistentBitMatrix::open`'s own auto-detection falls all the way - /// through to the `Implicit` (mono-genome, all-present) case, silently - /// corrupting every read. + /// Open one layer's matrix, auto-detecting count vs. presence from what + /// is actually on disk — the single source of truth for *every* caller + /// that opens a layer's own matrix (this module's cross-partition + /// [`PartitionCache::build`] and `family_scan::scan_layer_families`'s + /// own-layer lookup alike), so the two can never disagree. Sparse vs. + /// dense presence storage is `PersistentBitMatrix::open`'s own concern + /// (see the [`Presence`](Mat::Presence) variant's docs), not decided + /// here. pub(super) fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult { if with_counts && layer_dir.join("counts").exists() { return Layer::::open(layer_dir, mode).map(Mat::Count); } - if layer_dir.join("presence").join("is_multi.prsb").exists() { - Layer::::open(layer_dir, mode).map(Mat::SparsePresence) - } else { - Layer::::open(layer_dir, mode).map(Mat::Presence) - } + Layer::::open(layer_dir, mode).map(Mat::Presence) } fn find_slot(&self, kmer: CanonicalKmer) -> Option { match self { Mat::Count(l) => l.find_slot(kmer), Mat::Presence(l) => l.find_slot(kmer), - Mat::SparsePresence(l) => l.find_slot(kmer), } } @@ -95,7 +88,6 @@ impl Mat { match self { Mat::Count(l) => l.index_batch(kmers), Mat::Presence(l) => l.index_batch(kmers), - Mat::SparsePresence(l) => l.index_batch(kmers), } } @@ -110,7 +102,6 @@ impl Mat { match self { Mat::Count(l) => l.iter_minorants_batch(annex, batch_size), Mat::Presence(l) => l.iter_minorants_batch(annex, batch_size), - Mat::SparsePresence(l) => l.iter_minorants_batch(annex, batch_size), } } @@ -118,7 +109,6 @@ impl Mat { match self { Mat::Count(l) => l.n_cols(), Mat::Presence(l) => l.n_cols(), - Mat::SparsePresence(l) => l.n_cols(), } } @@ -138,7 +128,6 @@ impl Mat { pub(super) fn fill_sub_matrix_carries(&self, slots: &[usize], out: &mut [Vec]) { match self { Mat::Presence(l) => l.fill_sub_matrix(slots, out), - Mat::SparsePresence(l) => l.fill_sub_matrix(slots, out), Mat::Count(l) => { let mut counts: Vec> = out.iter().map(|_| Vec::new()).collect(); l.fill_sub_matrix(slots, &mut counts); @@ -169,7 +158,11 @@ pub(super) struct PartitionCache { } impl PartitionCache { - pub(super) fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult { + pub(super) fn build( + partition: &KmerPartitions, + n_parts: usize, + with_counts: bool, + ) -> OKIResult { let pb = progress_bar("open_partitions", n_parts as u64, "partitions"); let built: Vec<(Vec, usize)> = (0..n_parts) .into_par_iter() @@ -182,7 +175,10 @@ impl PartitionCache { let meta = partition.partition_meta(part)?; let mut mats = Vec::with_capacity(meta.n_layers); for l in 0..meta.n_layers { - let Ok(mat) = Mat::open(&partition.layer_dir(part, l), &meta.mode, with_counts) else { continue }; + let Ok(mat) = Mat::open(&partition.layer_dir(part, l), &meta.mode, with_counts) + else { + continue; + }; mats.push(mat); } pb.inc(1); @@ -257,7 +253,9 @@ impl PartitionCache { n_genomes: usize, mut on_hit: impl FnMut(usize, u8, usize), ) { - let Some(mats) = self.mats.get(dest_partition) else { return }; + let Some(mats) = self.mats.get(dest_partition) else { + return; + }; // First hit wins, same semantics as the old per-query loop (a // variant present in an earlier layer shadows later ones). @@ -297,7 +295,9 @@ impl PartitionCache { n_genomes: usize, mut on_hit: impl FnMut(usize, u8, usize), ) { - let Some(mats) = self.mats.get(dest_partition) else { return }; + let Some(mats) = self.mats.get(dest_partition) else { + return; + }; let mut by_layer: Vec> = vec![Vec::new(); mats.len()]; for &(variant, family_idx, base, layer) in queries { @@ -311,7 +311,8 @@ impl PartitionCache { continue; } let mat = &mats[li]; - let variants: Vec = entries.iter().map(|&(variant, _, _)| variant).collect(); + let variants: Vec = + entries.iter().map(|&(variant, _, _)| variant).collect(); let slots = mat.index_batch(&variants); let hits: Vec<(usize, usize, u8)> = slots .into_iter() diff --git a/src/obikphylo/src/siblings/cardinality.rs b/src/obikphylo/src/siblings/cardinality.rs index 1392bc3e..211a1f63 100644 --- a/src/obikphylo/src/siblings/cardinality.rs +++ b/src/obikphylo/src/siblings/cardinality.rs @@ -2,15 +2,15 @@ use std::sync::Arc; use ndarray::Array2; -use obikpartitionner::KmerPartition; +use obikpartitionner::KmerPartitions; use obisys::progress_bar; -use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; +use obikindex::{OKIError, OKIResult}; use super::cache::PartitionCache; use super::distance::RawSnpDistanceOutput; -use super::family_scan::{scan_layer_families, Selection}; +use super::family_scan::{Selection, scan_layer_families}; /// See [`KmerIndex::cardinality_tally`]. pub struct CardinalityTally { @@ -52,11 +52,19 @@ pub trait CardinalityExt { /// 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). - fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult; + fn cardinality_tally( + &self, + raw: &RawSnpDistanceOutput, + ratio_ceiling: f64, + ) -> OKIResult; } impl CardinalityExt for KmerIndex { - fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult { + fn cardinality_tally( + &self, + raw: &RawSnpDistanceOutput, + ratio_ceiling: f64, + ) -> OKIResult { let n_parts = self.n_partitions(); let n_genomes = self.meta().genomes.len(); let with_counts = self.meta().config.with_counts; @@ -72,7 +80,7 @@ impl CardinalityExt for KmerIndex { total > 0 && (snp as f64 / total as f64) <= ratio_ceiling }); - let partition = KmerPartition::open_with_config( + let partition = KmerPartitions::open_with_config( self.root_path(), self.kmer_size(), self.minimizer_size(), @@ -89,33 +97,42 @@ impl CardinalityExt for KmerIndex { // family comes back — no per-layer buffer. let mut total = [[0u64; 5]; 5]; for layer_dir in &layer_dirs { - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| { - 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. - return; - } + scan_layer_families( + layer_dir, + n_parts, + n_genomes, + with_counts, + k, + &cache, + &Selection::All, + |_family_idx, mask, genome_mask| { + 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. + return; + } - 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; - total[card_i][card_j] += 1; - if card_i != card_j { - total[card_j][card_i] += 1; + 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; + total[card_i][card_j] += 1; + if card_i != card_j { + total[card_j][card_i] += 1; + } } } - } - })?; + }, + )?; pb.inc(1); } pb.finish_and_clear(); diff --git a/src/obikphylo/src/siblings/distance.rs b/src/obikphylo/src/siblings/distance.rs index 34150300..8b8b15e6 100644 --- a/src/obikphylo/src/siblings/distance.rs +++ b/src/obikphylo/src/siblings/distance.rs @@ -2,14 +2,14 @@ use std::sync::Arc; use ndarray::Array2; -use obikpartitionner::KmerPartition; +use obikpartitionner::KmerPartitions; use obisys::progress_bar; -use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; +use obikindex::{OKIError, OKIResult}; use super::cache::PartitionCache; -use super::family_scan::{scan_layer_families, Selection}; +use super::family_scan::{Selection, scan_layer_families}; /// Raw p-distance restricted to loci that are single-copy in **both** /// genomes of a pair — the "stringent / paralogy-aware" locus eligibility @@ -72,7 +72,7 @@ where let k = index.kmer_size(); let n_bits = n_parts.trailing_zeros() as usize; - let partition = KmerPartition::open_with_config( + let partition = KmerPartitions::open_with_config( index.root_path(), index.kmer_size(), index.minimizer_size(), @@ -82,15 +82,23 @@ where let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); let layer_dirs = super::family_scan::sibling_layer_dirs(index)?; - let pb = progress_bar(label, layer_dirs.len() as u64, "layers"); - // One layer at a time — see `snp_pseudo_alignment`'s comment for why - // `par_iter()` over layers would defeat `scan_layer_families`'s - // partition-grouped locality. - let mut total = zero(); - for layer_dir in &layer_dirs { - let mut acc = zero(); + let pb = progress_bar(label, layer_dirs.len() as u64, "layers"); + // One layer at a time — see `snp_pseudo_alignment`'s comment for why + // `par_iter()` over layers would defeat `scan_layer_families`'s + // partition-grouped locality. + let mut total = zero(); + for layer_dir in &layer_dirs { + let mut acc = zero(); - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| { + scan_layer_families( + layer_dir, + n_parts, + n_genomes, + with_counts, + k, + &cache, + &Selection::All, + |_family_idx, mask, genome_mask| { let variable = mask.family_size() >= 2; // Per genome: which single form (if exactly one) it @@ -109,15 +117,16 @@ where on_pair(&mut acc, i, j, bi, bj, variable); } } - })?; + }, + )?; - total = combine(total, acc); - pb.inc(1); - } - pb.finish_and_clear(); - - Ok(total) + total = combine(total, acc); + pb.inc(1); } + pb.finish_and_clear(); + + Ok(total) +} /// Adds [`raw_snp_distance`](Self::raw_snp_distance) and /// [`base_pair_tally`](Self::base_pair_tally) to `KmerIndex`. @@ -141,7 +150,11 @@ pub trait DistanceExt { /// keeps aggregate SNP/shared counts per genome pair, not which bases /// were actually involved at each locus, and the ratio-ceiling filter /// can only be evaluated once the aggregate counts are known. - fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult; + fn base_pair_tally( + &self, + raw: &RawSnpDistanceOutput, + ratio_ceiling: f64, + ) -> OKIResult; } impl DistanceExt for KmerIndex { @@ -150,7 +163,12 @@ impl DistanceExt for KmerIndex { let (snp, shared) = scan_family_pairs( self, "raw_snp_distance", - || (Array2::::zeros((n_genomes, n_genomes)), Array2::::zeros((n_genomes, n_genomes))), + || { + ( + Array2::::zeros((n_genomes, n_genomes)), + Array2::::zeros((n_genomes, n_genomes)), + ) + }, |(snp, shared), i, j, bi, bj, _variable| { if bi == bj { shared[[i, j]] += 1; @@ -169,7 +187,11 @@ impl DistanceExt for KmerIndex { Ok(RawSnpDistanceOutput { snp, shared }) } - fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult { + fn base_pair_tally( + &self, + raw: &RawSnpDistanceOutput, + ratio_ceiling: f64, + ) -> OKIResult { let n_genomes = self.meta().genomes.len(); let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| { if i == j { diff --git a/src/obikphylo/src/siblings/entropy.rs b/src/obikphylo/src/siblings/entropy.rs index 18f7e15b..16aef5e3 100644 --- a/src/obikphylo/src/siblings/entropy.rs +++ b/src/obikphylo/src/siblings/entropy.rs @@ -13,14 +13,14 @@ use std::io::{BufWriter, Write}; use std::path::{Path, PathBuf}; use std::sync::Arc; -use obikpartitionner::KmerPartition; +use obikpartitionner::KmerPartitions; use obisys::progress_bar; -use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; +use obikindex::{OKIError, OKIResult}; use super::cache::PartitionCache; -use super::entropy_annex::{EntropyAnnexBuilder, ENTROPY_ANNEX_FILE_NAME}; +use super::entropy_annex::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnexBuilder}; use super::family_scan::{Selection, scan_layer_families}; use super::subsample::EntropyBias; @@ -120,18 +120,28 @@ pub trait ShannonEntropyExt { /// `DevDocMD/architecture/siblings.md`, "Entropy-biased selection") — /// useful here specifically to see the biased sample's own entropy /// distribution, not just to speed up a distance computation. - fn shannon_entropy_csv(&self, path: &Path, subsample: Option, entropy_bias: Option) -> OKIResult<()>; + fn shannon_entropy_csv( + &self, + path: &Path, + subsample: Option, + entropy_bias: Option, + ) -> OKIResult<()>; } impl ShannonEntropyExt for KmerIndex { - fn shannon_entropy_csv(&self, path: &Path, subsample: Option, entropy_bias: Option) -> OKIResult<()> { + fn shannon_entropy_csv( + &self, + path: &Path, + subsample: Option, + entropy_bias: Option, + ) -> OKIResult<()> { 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( + let partition = KmerPartitions::open_with_config( self.root_path(), self.kmer_size(), self.minimizer_size(), @@ -140,31 +150,53 @@ impl ShannonEntropyExt for KmerIndex { .map_err(OKIError::Partition)?; let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); let layer_dirs = super::family_scan::sibling_layer_dirs(self)?; - let selections = super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?; + let selections = + super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?; let mut f = BufWriter::new(std::fs::File::create(path).map_err(OKIError::Io)?); - writeln!(f, "layer,family_idx,entropy15,entropy4,family_size,n_genomes_present").map_err(OKIError::Io)?; + writeln!( + f, + "layer,family_idx,entropy15,entropy4,family_size,n_genomes_present" + ) + .map_err(OKIError::Io)?; let pb = progress_bar("shannon_entropy", layer_dirs.len() as u64, "layers"); - for (layer_ord, (layer_dir, layer_selection)) in layer_dirs.iter().zip(selections.iter()).enumerate() { + for (layer_ord, (layer_dir, layer_selection)) in + layer_dirs.iter().zip(selections.iter()).enumerate() + { let selection = match layer_selection { None => Selection::All, Some(set) => Selection::Some(set), }; let mut write_err = None; - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |family_idx, mask, genome_mask| { - if write_err.is_some() { - return; - } - if mask.family_size() < 2 { - return; // monomorphic family — entropy trivially 0, no signal, skip - } - let Some((h15, n_present)) = family_entropy(genome_mask) else { return }; - let h4 = family_entropy_4(genome_mask).map_or(0.0, |(h, _)| h); - if let Err(e) = writeln!(f, "{layer_ord},{family_idx},{h15:.6},{h4:.6},{},{n_present}", mask.family_size()) { - write_err = Some(e); - } - })?; + scan_layer_families( + layer_dir, + n_parts, + n_genomes, + with_counts, + k, + &cache, + &selection, + |family_idx, mask, genome_mask| { + if write_err.is_some() { + return; + } + if mask.family_size() < 2 { + return; // monomorphic family — entropy trivially 0, no signal, skip + } + let Some((h15, n_present)) = family_entropy(genome_mask) else { + return; + }; + let h4 = family_entropy_4(genome_mask).map_or(0.0, |(h, _)| h); + if let Err(e) = writeln!( + f, + "{layer_ord},{family_idx},{h15:.6},{h4:.6},{},{n_present}", + mask.family_size() + ) { + write_err = Some(e); + } + }, + )?; if let Some(e) = write_err { return Err(OKIError::Io(e)); } @@ -194,7 +226,9 @@ impl ShannonEntropyExt for KmerIndex { /// `scan_layer_families`'s expensive per-genome resolution for the ~98% /// of minorants that are monomorphic only to discard the result. pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf]) -> OKIResult<()> { - let missing: Vec = layer_dirs.iter().enumerate() + let missing: Vec = layer_dirs + .iter() + .enumerate() .filter(|(_, dir)| !dir.join(ENTROPY_ANNEX_FILE_NAME).exists()) .map(|(i, _)| i) .collect(); @@ -208,7 +242,7 @@ pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf]) let k = index.kmer_size(); let n_bits = n_parts.trailing_zeros() as usize; - let partition = KmerPartition::open_with_config( + let partition = KmerPartitions::open_with_config( index.root_path(), index.kmer_size(), index.minimizer_size(), @@ -221,15 +255,30 @@ pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf]) let pb = progress_bar("entropy_annex_build", missing.len() as u64, "layers"); for &li in &missing { let layer_dir = &layer_dirs[li]; - let mut builder = EntropyAnnexBuilder::new(minorant_counts[li] as usize, &layer_dir.join(ENTROPY_ANNEX_FILE_NAME)) - .map_err(OKIError::Io)?; + let mut builder = EntropyAnnexBuilder::new( + minorant_counts[li] as usize, + &layer_dir.join(ENTROPY_ANNEX_FILE_NAME), + ) + .map_err(OKIError::Io)?; let eligible = super::subsample::non_monomorphic_selection_layer(layer_dir)?; - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::Some(&eligible), |family_idx, mask, genome_mask| { - debug_assert!(mask.family_size() >= 2, "Selection::Some(eligible) must only visit non-monomorphic minorants"); - if let Some((h15, _)) = family_entropy(genome_mask) { - builder.set(family_idx, h15 as f32); - } - })?; + scan_layer_families( + layer_dir, + n_parts, + n_genomes, + with_counts, + k, + &cache, + &Selection::Some(&eligible), + |family_idx, mask, genome_mask| { + debug_assert!( + mask.family_size() >= 2, + "Selection::Some(eligible) must only visit non-monomorphic minorants" + ); + if let Some((h15, _)) = family_entropy(genome_mask) { + builder.set(family_idx, h15 as f32); + } + }, + )?; builder.close().map_err(OKIError::Io)?; pb.inc(1); } diff --git a/src/obikphylo/src/siblings/sankoff_bundle.rs b/src/obikphylo/src/siblings/sankoff_bundle.rs index dabe3b9e..b360721b 100644 --- a/src/obikphylo/src/siblings/sankoff_bundle.rs +++ b/src/obikphylo/src/siblings/sankoff_bundle.rs @@ -29,18 +29,18 @@ use std::sync::Arc; use ndarray::Array2; -use obikpartitionner::KmerPartition; +use obikpartitionner::KmerPartitions; use obisys::progress_bar; -use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; +use obikindex::{OKIError, OKIResult}; -use super::alignment::{iupac_code, SnpAlignment}; +use super::alignment::{SnpAlignment, iupac_code}; use super::cache::PartitionCache; use super::cardinality::CardinalityTally; use super::distance::{BasePairTally, RawSnpDistanceOutput}; -use super::family_scan::{scan_layer_families, sibling_layer_dirs, Selection}; -use super::subsample::{compute_selections, EntropyBias}; +use super::family_scan::{Selection, scan_layer_families, sibling_layer_dirs}; +use super::subsample::{EntropyBias, compute_selections}; /// Every output the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` pipeline needs, /// computed together from one shared, possibly-subsampled/entropy-biased @@ -85,7 +85,7 @@ impl SankoffBundleExt for KmerIndex { let k = self.kmer_size(); let n_bits = n_parts.trailing_zeros() as usize; - let partition = KmerPartition::open_with_config( + let partition = KmerPartitions::open_with_config( self.root_path(), self.kmer_size(), self.minimizer_size(), @@ -108,25 +108,34 @@ impl SankoffBundleExt for KmerIndex { None => Selection::All, Some(set) => Selection::Some(set), }; - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |_family_idx, _mask, genome_mask| { - let single_form = |g: usize| -> Option { - let m = genome_mask[g]; - (m.count_ones() == 1).then(|| m.trailing_zeros() as u8) - }; - for i in 0..n_genomes { - let Some(bi) = single_form(i) else { continue }; - for j in (i + 1)..n_genomes { - let Some(bj) = single_form(j) else { continue }; - if bi == bj { - shared[[i, j]] += 1; - shared[[j, i]] += 1; - } else { - snp[[i, j]] += 1; - snp[[j, i]] += 1; + scan_layer_families( + layer_dir, + n_parts, + n_genomes, + with_counts, + k, + &cache, + &selection, + |_family_idx, _mask, genome_mask| { + let single_form = |g: usize| -> Option { + let m = genome_mask[g]; + (m.count_ones() == 1).then(|| m.trailing_zeros() as u8) + }; + for i in 0..n_genomes { + let Some(bi) = single_form(i) else { continue }; + for j in (i + 1)..n_genomes { + let Some(bj) = single_form(j) else { continue }; + if bi == bj { + shared[[i, j]] += 1; + shared[[j, i]] += 1; + } else { + snp[[i, j]] += 1; + snp[[j, i]] += 1; + } } } - } - })?; + }, + )?; pb.inc(1); } pb.finish_and_clear(); @@ -168,69 +177,83 @@ impl SankoffBundleExt for KmerIndex { None => Selection::All, Some(set) => Selection::Some(set), }; - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |_family_idx, mask, genome_mask| { - let variable = mask.family_size() >= 2; + scan_layer_families( + layer_dir, + n_parts, + n_genomes, + with_counts, + k, + &cache, + &selection, + |_family_idx, mask, genome_mask| { + let variable = mask.family_size() >= 2; - // base-pair tally — same pairwise single-form resolution as pass A. - let single_form = |g: usize| -> Option { - let m = genome_mask[g]; - (m.count_ones() == 1).then(|| m.trailing_zeros() as u8) - }; - for i in 0..n_genomes { - let Some(bi) = single_form(i) else { continue }; - for j in (i + 1)..n_genomes { - let Some(bj) = single_form(j) else { continue }; - if !included[[i, j]] { - continue; - } - if bi != bj { - bp_counts[bi as usize][bj as usize] += 1; - bp_counts[bj as usize][bi as usize] += 1; - } else if variable { - bp_same[bi as usize] += 1; - } - } - } - - // cardinality tally — restricted to variable families, see - // `CardinalityExt::cardinality_tally`'s docs for why. - if variable { + // base-pair tally — same pairwise single-form resolution as pass A. + let single_form = |g: usize| -> Option { + let m = genome_mask[g]; + (m.count_ones() == 1).then(|| m.trailing_zeros() as u8) + }; for i in 0..n_genomes { - let card_i = genome_mask[i].count_ones() as usize; + let Some(bi) = single_form(i) else { continue }; for j in (i + 1)..n_genomes { + let Some(bj) = single_form(j) else { continue }; if !included[[i, j]] { continue; } - let card_j = genome_mask[j].count_ones() as usize; - card_counts[card_i][card_j] += 1; - if card_i != card_j { - card_counts[card_j][card_i] += 1; + if bi != bj { + bp_counts[bi as usize][bj as usize] += 1; + bp_counts[bj as usize][bi as usize] += 1; + } else if variable { + bp_same[bi as usize] += 1; } } } - } - // pseudo-alignment — same variable-family gate as - // `snp_pseudo_alignment`. Every family the shared selection - // actually visits already satisfies this when the - // selection is `Selection::Some` (only non-monomorphic - // minorants are ever selected), but the explicit check - // still matters under `Selection::All` (no `--subsample`), - // which visits monomorphic minorants too. - if variable { - for (g, &m) in genome_mask.iter().enumerate() { - sequences[g].push(iupac_code(m)); + // cardinality tally — restricted to variable families, see + // `CardinalityExt::cardinality_tally`'s docs for why. + if variable { + 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; + card_counts[card_i][card_j] += 1; + if card_i != card_j { + card_counts[card_j][card_i] += 1; + } + } + } } - } - })?; + + // pseudo-alignment — same variable-family gate as + // `snp_pseudo_alignment`. Every family the shared selection + // actually visits already satisfies this when the + // selection is `Selection::Some` (only non-monomorphic + // minorants are ever selected), but the explicit check + // still matters under `Selection::All` (no `--subsample`), + // which visits monomorphic minorants too. + if variable { + for (g, &m) in genome_mask.iter().enumerate() { + sequences[g].push(iupac_code(m)); + } + } + }, + )?; pb.inc(1); } pb.finish_and_clear(); Ok(SankoffBundle { raw, - base_pair_tally: BasePairTally { counts: bp_counts, same: bp_same }, - cardinality_tally: CardinalityTally { counts: card_counts }, + base_pair_tally: BasePairTally { + counts: bp_counts, + same: bp_same, + }, + cardinality_tally: CardinalityTally { + counts: card_counts, + }, alignment: SnpAlignment { sequences }, }) } diff --git a/src/obikphylo/src/siblings/stats.rs b/src/obikphylo/src/siblings/stats.rs index 06cf4698..2a636187 100644 --- a/src/obikphylo/src/siblings/stats.rs +++ b/src/obikphylo/src/siblings/stats.rs @@ -2,16 +2,16 @@ use std::sync::Arc; use rayon::prelude::*; -use obikpartitionner::KmerPartition; +use obikpartitionner::KmerPartitions; use obisys::progress_bar; -use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; +use obikindex::{OKIError, OKIResult}; use super::ANNEX_FILE_NAME; use super::SiblingAnnex; use super::cache::PartitionCache; -use super::family_scan::{scan_layer_families, Selection}; +use super::family_scan::{Selection, scan_layer_families}; /// Distribution of family sizes (1-4), read back from an already-built /// annex (see [`super::build::SiblingAnnexBuildExt::build_sibling_annex`]) @@ -81,7 +81,9 @@ impl SiblingStatsExt for KmerIndex { |mut counts, layer_dir| -> OKIResult<[u64; 4]> { let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; for slot in 0..annex.len() { - let Some(mask) = annex.get(slot) else { continue }; + let Some(mask) = annex.get(slot) else { + continue; + }; if !mask.is_minorant() { continue; // family tallied once, at its minorant } @@ -112,7 +114,7 @@ impl SiblingStatsExt for KmerIndex { // 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( + let partition = KmerPartitions::open_with_config( self.root_path(), self.kmer_size(), self.minimizer_size(), @@ -131,18 +133,27 @@ impl SiblingStatsExt for KmerIndex { ..Default::default() }; for layer_dir in &layer_dirs { - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| { - // "Genome g represents this family" means g carries - // *any* of its members, not just the minorant's own — - // `genome_mask[g] != 0` is exactly that. - let s = mask.siblings() as usize; - stats.counts[s] += 1; - for (g, &m) in genome_mask.iter().enumerate() { - if m != 0 { - stats.per_genome[g][s] += 1; + scan_layer_families( + layer_dir, + n_parts, + n_genomes, + with_counts, + k, + &cache, + &Selection::All, + |_family_idx, mask, genome_mask| { + // "Genome g represents this family" means g carries + // *any* of its members, not just the minorant's own — + // `genome_mask[g] != 0` is exactly that. + let s = mask.siblings() as usize; + stats.counts[s] += 1; + for (g, &m) in genome_mask.iter().enumerate() { + if m != 0 { + stats.per_genome[g][s] += 1; + } } - } - })?; + }, + )?; pb.inc(1); } pb.finish_and_clear();