diff --git a/src/Cargo.lock b/src/Cargo.lock index 33d3ddb3..06828b83 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1542,12 +1542,17 @@ name = "obikfilter" version = "0.1.0" dependencies = [ "obicompactvec", + "obidebruinj", + "obikalgorithm", "obikentropy", "obikidxcache", "obikindex", + "obikindexer", "obikseq", "obiskio", + "obisys", "obitaxonomy", + "tracing", ] [[package]] @@ -1662,6 +1667,7 @@ dependencies = [ "csv", "obifastwrite", "obikalgorithm", + "obikfilter", "obikindex", "obikindexer", "obikmerge", diff --git a/src/obicompactvec/src/lib.rs b/src/obicompactvec/src/lib.rs index e1a15a38..bb7782b2 100644 --- a/src/obicompactvec/src/lib.rs +++ b/src/obicompactvec/src/lib.rs @@ -8,6 +8,7 @@ mod format; mod rankselect; mod intmatrix; mod layer_meta; +mod matrix_builder; mod meta; mod mmap_file; mod reader; @@ -27,6 +28,7 @@ pub use builder::PersistentCompactIntVecBuilder; pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask}; pub use intmatrix::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, pack_compact_int_matrix}; pub use layer_meta::LayerMeta; +pub use matrix_builder::{ColBuilder, MatrixBuilder}; pub use reader::{PersistentCompactIntVec, Iter as CompactIntVecIter}; pub use sparse_intmatrix::{PersistentSparseCompactIntMatrix, PersistentSparseCompactIntMatrixBuilder, pack_sparse_compact_int_matrix}; pub use storage_kind::StorageKind; diff --git a/src/obicompactvec/src/matrix_builder.rs b/src/obicompactvec/src/matrix_builder.rs new file mode 100644 index 00000000..db7bb8ae --- /dev/null +++ b/src/obicompactvec/src/matrix_builder.rs @@ -0,0 +1,106 @@ +//! Bit/Int-dispatching wrappers around the persistent matrix/column +//! builders — pure row/column mechanics, no domain reference (genome, +//! k-mer, …). [`ColBuilder`] wraps one column; [`MatrixBuilder`] wraps a +//! whole matrix and hands out [`ColBuilder`]s. Shared by every consumer +//! that builds a destination matrix whose content kind (count vs +//! presence) is chosen at runtime: `obikmerge`, `obikselect`, `obikfilter`. + +use std::io; +use std::path::Path; + +use crate::{PersistentBitMatrixBuilder as BitMB, PersistentBitVecBuilder}; +use crate::{PersistentCompactIntMatrixBuilder as IntMB, PersistentCompactIntVecBuilder}; +use crate::{TempBitVec, TempCompactIntVec}; + +// ── ColBuilder ──────────────────────────────────────────────────────────────── + +pub enum ColBuilder { + Bit(PersistentBitVecBuilder), + Int(PersistentCompactIntVecBuilder), +} + +impl ColBuilder { + pub fn set_val(&mut self, slot: usize, value: u32) { + match self { + ColBuilder::Bit(b) => b.set(slot, value > 0), + ColBuilder::Int(b) => b.set(slot, value), + } + } + + pub fn close(self) -> io::Result<()> { + match self { + ColBuilder::Bit(b) => b.close(), + ColBuilder::Int(b) => b.close(), + } + } +} + +// ── MatrixBuilder ───────────────────────────────────────────────────────────── + +/// Wraps whichever matrix builder `presence` calls for, so callers never +/// have to know the on-disk column naming or the matrix `meta.json` schema +/// — both stay private to this crate. `resume` reopens a matrix directory +/// already closed by a previous builder session, continuing from its +/// current `n_cols` instead of starting a fresh matrix at 0. +pub enum MatrixBuilder { + Bit(BitMB), + Int(IntMB), +} + +impl MatrixBuilder { + pub fn new(presence: bool, n: usize, dir: &Path) -> io::Result { + Ok(if presence { + MatrixBuilder::Bit(BitMB::new(n, dir)?) + } else { + MatrixBuilder::Int(IntMB::new(n, dir)?) + }) + } + + pub fn resume(presence: bool, dir: &Path) -> io::Result { + Ok(if presence { + MatrixBuilder::Bit(BitMB::resume(dir)?) + } else { + MatrixBuilder::Int(IntMB::resume(dir)?) + }) + } + + /// Add a column with no data written (all-zero/false) — for columns + /// absent from a given source (e.g. destination-only genomes not + /// contributed by the source being merged in). + pub fn add_absent_col(&mut self) -> io::Result<()> { + match self { + MatrixBuilder::Bit(b) => b.add_col()?.close(), + MatrixBuilder::Int(b) => b.add_col()?.close(), + } + } + + pub fn add_col(&mut self) -> io::Result { + Ok(match self { + MatrixBuilder::Bit(b) => ColBuilder::Bit(b.add_col()?), + MatrixBuilder::Int(b) => ColBuilder::Int(b.add_col()?), + }) + } + + /// Add a column already computed as a bit vector. + pub fn add_col_from_bit(&mut self, src: &TempBitVec) -> io::Result<()> { + match self { + MatrixBuilder::Bit(b) => b.add_col_from(src), + MatrixBuilder::Int(b) => b.add_col_from_bit(src), + } + } + + /// Add a column already computed as an int vector. + pub fn add_col_from_int(&mut self, src: &TempCompactIntVec) -> io::Result<()> { + match self { + MatrixBuilder::Bit(b) => b.add_col_from_int(src), + MatrixBuilder::Int(b) => b.add_col_from(src), + } + } + + pub fn close(self) -> io::Result<()> { + match self { + MatrixBuilder::Bit(b) => b.close(), + MatrixBuilder::Int(b) => b.close(), + } + } +} diff --git a/src/obikfilter/Cargo.toml b/src/obikfilter/Cargo.toml index 7643cad7..21bc3a1f 100644 --- a/src/obikfilter/Cargo.toml +++ b/src/obikfilter/Cargo.toml @@ -11,3 +11,8 @@ obiskio = { path = "../obiskio" } obitaxonomy = { path = "../obitaxonomy" } obikentropy = { path = "../obikentropy" } obikidxcache = { path = "../obikidxcache" } +obikindexer = { path = "../obikindexer" } +obidebruinj = { path = "../obidebruinj" } +obikalgorithm = { path = "../obikalgorithm" } +obisys = { path = "../obisys" } +tracing = "0.1.44" diff --git a/src/obikfilter/src/filter_algo.rs b/src/obikfilter/src/filter_algo.rs new file mode 100644 index 00000000..2db2664f --- /dev/null +++ b/src/obikfilter/src/filter_algo.rs @@ -0,0 +1,141 @@ +//! Removing kmers from a source `obikindex::KmerIndex`, according to a list +//! of [`crate::KmerFilter`]s. Structured exactly like `obikselect::Select`/ +//! `obikmerge::Merge`: two-phase construction (`new` + setters), then +//! `obikalgorithm::Algorithm::run`, via `obikindex::IndexBuilder` +//! (`create_skeleton` + `finalize_indexed`) — no in-place variant, same as +//! `Select`/`Merge`. Unlike `Select`, every surviving layer's identity +//! (unitigs/MPHF/evidence) is rebuilt from scratch, not copied — see +//! [`crate::filter_partition`]. + +use std::io; +use std::path::PathBuf; + +use obikalgorithm::Algorithm; +use obikindex::{IndexBuilder, IndexState, KmerIndex}; +use obisys::{PartitionRunner, Progress, Reporter, Stage}; +use tracing::info; + +use crate::filter::KmerFilter; +use crate::filter_partition::filter_partition; + +pub struct Filter<'a> { + src: &'a KmerIndex, + output: PathBuf, + filters: &'a [Box], + presence: bool, + force: bool, + sparse: bool, + reporter: Reporter, + on_progress: Option>, +} + +impl<'a> Filter<'a> { + pub fn new(src: &'a KmerIndex, output: impl Into, filters: &'a [Box]) -> Self { + Self { + src, + output: output.into(), + filters, + presence: false, + force: false, + sparse: true, + reporter: Reporter::new(), + on_progress: None, + } + } + + /// Output presence/absence instead of counts, even if the source has + /// count data (default: `false` — keep the source's own content kind). + pub fn presence(mut self, v: bool) -> Self { + self.presence = v; + self + } + + /// Remove a pre-existing index at `output` before creating the + /// filtered one (default: `false`, fails if one already exists). + pub fn force(mut self, v: bool) -> Self { + self.force = v; + self + } + + /// Pack the output's presence matrices in the compact sparse format + /// rather than dense (default: `true`). + pub fn sparse(mut self, v: bool) -> Self { + self.sparse = v; + self + } + + /// Progress callback, called once per completed partition. + pub fn on_progress(mut self, cb: impl FnMut(Progress) + Send + 'a) -> Self { + self.on_progress = Some(Box::new(cb)); + self + } + + /// Per-stage timing report accumulated during `run` — call after `run` returns. + pub fn reporter(&self) -> &Reporter { + &self.reporter + } +} + +impl Algorithm for Filter<'_> { + type Output = KmerIndex; + + fn run(&mut self) -> obikalgorithm::Result { + let src = self.src; + let output = self.output.clone(); + + if src.state()? != IndexState::Indexed { + return Err(format!("{}: source index is not fully built", src.dir().display()).into()); + } + + let genomes = src.meta().genomes()?; + let n_genomes = genomes.len(); + let use_counts = !self.presence && src.meta().config.with_counts; + info!( + "filter: {} partition(s), {} genome(s), output={}", + src.n_partitions(), + n_genomes, + if use_counts { "count" } else { "presence" }, + ); + + KmerIndex::clear_output_for_create(&output, self.force)?; + + let mut config = src.meta().config.clone(); + config.with_counts = use_counts; + let block_bits = src.block_bits(); + let mode = src.evidence_mode().clone(); + + let t = Stage::start("filter"); + let dst = KmerIndex::create_skeleton(&output, config, genomes)?; + + let n_partitions = src.n_partitions(); + let order: Vec = (0..n_partitions).collect(); + let runner = PartitionRunner::new(); + let on_progress = &mut self.on_progress; + let mut done: u64 = 0; + let filters = self.filters; + + runner + .run( + &order, + |i| filter_partition(&dst, src, i, filters, use_counts, !use_counts, block_bits, &mode), + |_, _, _| { + done += 1; + if let Some(cb) = on_progress.as_mut() { + cb(Progress { + position: done, + total: Some(n_partitions as u64), + }); + } + }, + ) + .map_err(|e| io::Error::other(e.to_string()))?; + + self.reporter.push(t.stop()); + + let t = Stage::start("finalize"); + let dst = KmerIndex::finalize_indexed(&output, &mut self.reporter, self.sparse)?; + self.reporter.push(t.stop()); + + Ok(dst) + } +} diff --git a/src/obikfilter/src/filter_partition.rs b/src/obikfilter/src/filter_partition.rs new file mode 100644 index 00000000..6c2c1957 --- /dev/null +++ b/src/obikfilter/src/filter_partition.rs @@ -0,0 +1,77 @@ +//! Per-partition rebuild — the mechanics [`crate::filter_algo::Filter`] +//! runs once per partition. Unlike `obikselect::select_partition` (columns +//! only, kmer identity copied as-is), this removes kmers, so every layer's +//! identity (unitigs + MPHF + evidence) has to be rebuilt from scratch, not +//! copied — removing kmers changes MPHF slot assignment layer-wide. + +use std::collections::HashMap; + +use obidebruinj::GraphDeBruijn; +use obikidxcache::index_cache::IndexCache; +use obikindex::layer::{IndexMode, KmerLayer, TypedLayer}; +use obikindex::{KmerIndex, OKIError, OKIResult}; +use obikindexer::write_graph_as_unitigs; +use obikseq::CanonicalKmer; + +use crate::filter::KmerFilter; +use crate::partition_iter::FilteredPartitionIter; + +/// Rewrite partition `i` from `src` into `dst`, keeping only kmers that pass +/// every filter in `filters`. +/// +/// `use_counts` selects which values are read from `src` (count vs +/// presence-collapsed); `output_presence` selects the destination's own +/// matrix kind — independent, since `--presence` can force a presence +/// output from a count source. +pub(crate) fn filter_partition( + dst: &KmerIndex, + src: &KmerIndex, + i: usize, + filters: &[Box], + use_counts: bool, + output_presence: bool, + block_bits: u8, + mode: &IndexMode, +) -> OKIResult<()> { + let n_genomes = src.meta().genomes().map_err(OKIError::Io)?.len(); + + let cache = IndexCache::new(src, Some(vec![i])); + let dst_partition = dst.partition(i)?; + let n_layers = src.partition(i)?.n_layers(); + + // Bucket survivors by layer — bounded to one partition's worth of data, + // the same order of magnitude `IndexCache` already holds open for this + // partition. + let mut by_layer: HashMap>> = HashMap::new(); + cache.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |_part, layer, kmer, row| { + by_layer.entry(layer).or_default().insert(kmer, row); + true + })?; + + for l in 0..n_layers { + let survivors = by_layer.remove(&l).unwrap_or_default(); + + let dst_layer_dir = KmerLayer::new(&dst_partition, l) + .create() + .map_err(OKIError::Io)? + .dir() + .to_path_buf(); + + let mut g = GraphDeBruijn::new(); + for &kmer in survivors.keys() { + g.push(kmer); + } + write_graph_as_unitigs(g, &dst_layer_dir)?; + + TypedLayer::<()>::build_with_matrix( + &dst_layer_dir, + block_bits, + mode, + output_presence, + n_genomes, + |kmer| survivors.get(&kmer).cloned().expect("kmer pushed into the graph must be a survivor"), + )?; + } + + Ok(()) +} diff --git a/src/obikfilter/src/lib.rs b/src/obikfilter/src/lib.rs index 38a7b296..611ef45a 100644 --- a/src/obikfilter/src/lib.rs +++ b/src/obikfilter/src/lib.rs @@ -11,6 +11,8 @@ //! `Filter` algorithm (index-to-index rebuild). mod filter; +mod filter_algo; +mod filter_partition; mod partition_iter; mod predicate; @@ -18,5 +20,6 @@ pub use filter::{ GroupQuorumFilter, KmerFilter, MaxGenomeCount, MaxGenomeFraction, MaxTotalCount, MinComplexity, MinGenomeCount, MinGenomeFraction, MinTotalCount, passes_all, }; +pub use filter_algo::Filter; pub use partition_iter::FilteredPartitionIter; pub use predicate::{GenomeSelector, GroupFilterParams, MetaPred, Selection}; diff --git a/src/obikindex/src/index/common.rs b/src/obikindex/src/index/common.rs deleted file mode 100644 index f22fdc12..00000000 --- a/src/obikindex/src/index/common.rs +++ /dev/null @@ -1,25 +0,0 @@ -use crate::index::error::{OKIError, OKIResult}; -use obicompactvec::{PersistentBitVecBuilder, PersistentCompactIntVecBuilder}; - -// ── ColBuilder ──────────────────────────────────────────────────────────────── - -pub enum ColBuilder { - Bit(PersistentBitVecBuilder), - Int(PersistentCompactIntVecBuilder), -} - -impl ColBuilder { - pub fn set_val(&mut self, slot: usize, value: u32) { - match self { - ColBuilder::Bit(b) => b.set(slot, value > 0), - ColBuilder::Int(b) => b.set(slot, value), - } - } - - pub fn close(self) -> OKIResult<()> { - match self { - ColBuilder::Bit(b) => b.close().map_err(OKIError::Io), - ColBuilder::Int(b) => b.close().map_err(OKIError::Io), - } - } -} diff --git a/src/obikindex/src/index/mod.rs b/src/obikindex/src/index/mod.rs index 8a60f488..becde665 100644 --- a/src/obikindex/src/index/mod.rs +++ b/src/obikindex/src/index/mod.rs @@ -2,12 +2,10 @@ pub mod error; pub mod meta; pub mod state; mod builder; -mod common; mod kmer_index; pub use error::{OKIError, OKIResult}; pub use builder::IndexBuilder; -pub use common::ColBuilder; pub use kmer_index::KmerIndex; pub use meta::{GenomeInfo, IndexConfig, IndexMeta, META_FILENAME}; pub use state::IndexState; diff --git a/src/obikindex/src/layer/typed_layer.rs b/src/obikindex/src/layer/typed_layer.rs index 998294ca..800d3424 100644 --- a/src/obikindex/src/layer/typed_layer.rs +++ b/src/obikindex/src/layer/typed_layer.rs @@ -1,7 +1,7 @@ // use crate::layer::utils::layer_dir; use obicompactvec::{ - BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrix, - PersistentCompactIntMatrixBuilder, PersistentSparseBitMatrix, + BinaryMatrix, ColBuilder, MatrixBuilder, PersistentBitMatrix, PersistentBitMatrixBuilder, + PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, PersistentSparseBitMatrix, }; use obikseq::CanonicalKmer; use obiskio::{UnitigFileReader, UnitigFileWriter}; @@ -281,6 +281,51 @@ impl TypedLayer<()> { MphfLayer::build(out_dir, block_bits, mode, &mut |_, _| Ok(())) } + /// Build MPHF + evidence + a full multi-column matrix in one MPHF pass — + /// unlike `TypedLayer::::build`'s single + /// `count_of` column, this fills `n_cols` columns at once from `row_of`, + /// so the (expensive) MPHF construction only runs once regardless of how + /// many genome columns the destination has. Columns are written via + /// random-access, mmap-backed `ColBuilder::set_val(slot, _)` calls as + /// the MPHF reports each kmer's slot — no ordering assumption, no + /// buffering of the matrix in memory. + /// + /// `presence` picks the on-disk matrix kind (bit vs int), same meaning + /// as `obicompactvec::MatrixBuilder::new`. `row_of(kmer)` must return + /// exactly `n_cols` values, one per destination column, in column order. + pub fn build_with_matrix( + out_dir: &Path, + block_bits: u8, + mode: &IndexMode, + presence: bool, + n_cols: usize, + row_of: impl Fn(CanonicalKmer) -> Box<[u32]>, + ) -> OKIResult { + let data_dir = out_dir.join(if presence { PRESENCE_DIR } else { COUNTS_DIR }); + fs::create_dir_all(&data_dir).map_err(OKIError::Io)?; + let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers(); + + let mut mb = MatrixBuilder::new(presence, n, &data_dir).map_err(OKIError::Io)?; + let mut cols: Vec = (0..n_cols) + .map(|_| mb.add_col()) + .collect::>() + .map_err(OKIError::Io)?; + + let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| { + let row = row_of(kmer); + for (c, col) in cols.iter_mut().enumerate() { + col.set_val(slot, row[c]); + } + Ok(()) + })?; + + for col in cols { + col.close().map_err(OKIError::Io)?; + } + mb.close().map_err(OKIError::Io)?; + Ok(n_built) + } + /// Create a presence matrix for a set-membership layer (first merge). pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OKIResult<()> { let presence_dir = layer_dir.join(PRESENCE_DIR); diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 2ab9c8fe..cba763a7 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -17,6 +17,5 @@ pub use index::{ IndexBuilder, IndexConfig, IndexMeta, IndexState, KmerIndex, OKIError, OKIResult, META_FILENAME, - ColBuilder, }; pub use index::meta; diff --git a/src/obikmer2/Cargo.toml b/src/obikmer2/Cargo.toml index 695baaf0..971c6cc8 100644 --- a/src/obikmer2/Cargo.toml +++ b/src/obikmer2/Cargo.toml @@ -16,6 +16,7 @@ obikindex = { path = "../obikindex", default-features = false } obikindexer = { path = "../obikindexer" } obikalgorithm = { path = "../obikalgorithm" } obikmerge = { path = "../obikmerge" } +obikfilter = { path = "../obikfilter" } obifastwrite = { path = "../obifastwrite" } obiskbuilder = { path = "../obiskbuilder" } clap = { version = "4", features = ["derive"] } diff --git a/src/obikmer2/src/cmd/filter/mod.rs b/src/obikmer2/src/cmd/filter/mod.rs new file mode 100644 index 00000000..4da8cd6e --- /dev/null +++ b/src/obikmer2/src/cmd/filter/mod.rs @@ -0,0 +1,170 @@ +use std::path::PathBuf; + +use clap::Args; +use obikalgorithm::Algorithm; +use obikfilter::{ + Filter, GenomeSelector, GroupFilterParams, KmerFilter, MaxTotalCount, MinComplexity, + MinTotalCount, +}; +use obikindex::KmerIndex; +use obisys::{Progress, Reporter, Stage, progress_bar}; +use tracing::info; + +#[derive(Args)] +pub struct FilterArgs { + /// Source index directory + pub source: PathBuf, + + /// Output index directory + #[arg(short, long)] + pub output: PathBuf, + + /// Ingroup predicate (repeatable; AND). Forms: `key=v1|v2`, `key!=v`, `key~path`, `key!~path`, `*`/`all` + #[arg(long, value_name = "PRED")] + pub ingroup: Vec, + + /// Outgroup predicate (repeatable; OR). Forms: `key=v1|v2`, `key!=v`, `key~path`, `key!~path`, `*`/`all` + #[arg(long, value_name = "PRED")] + pub outgroup: Vec, + + /// Minimum number of ingroup genomes containing the k-mer + /// (negative: offset from group size, e.g. -1 = all but one) + #[arg(long, allow_hyphen_values = true)] + pub min_count: Option, + + /// Maximum number of ingroup genomes containing the k-mer + /// (negative: offset from group size, e.g. -1 = all but one) + #[arg(long, allow_hyphen_values = true)] + pub max_count: Option, + + /// Minimum fraction of ingroup genomes containing the k-mer [0.0-1.0] + /// (default 1.0 when --ingroup is set, 0.0 otherwise) + #[arg(long)] + pub min_frac: Option, + + /// Maximum fraction of ingroup genomes containing the k-mer [0.0-1.0] + #[arg(long)] + pub max_frac: Option, + + /// Minimum number of outgroup genomes containing the k-mer + /// (negative: offset from outgroup size, e.g. -1 = all but one) + #[arg(long, allow_hyphen_values = true)] + pub min_outgroup_count: Option, + + /// Maximum number of outgroup genomes containing the k-mer + /// (default 0 when --outgroup is set, no constraint otherwise; + /// negative: offset from outgroup size, e.g. -1 = all but one) + #[arg(long, allow_hyphen_values = true)] + pub max_outgroup_count: Option, + + /// Minimum fraction of outgroup genomes containing the k-mer [0.0-1.0] + #[arg(long)] + pub min_outgroup_frac: Option, + + /// Maximum fraction of outgroup genomes containing the k-mer [0.0-1.0] + #[arg(long)] + pub max_outgroup_frac: Option, + + /// Per-genome count threshold to consider a genome as "containing" the k-mer (default 0) + #[arg(long, default_value = "0")] + pub presence_threshold: u32, + + /// Minimum total count across all genomes (count index only) + #[arg(long)] + pub min_total_count: Option, + + /// Maximum total count across all genomes (count index only) + #[arg(long)] + pub max_total_count: Option, + + /// Minimum normalized entropy (complexity) to keep a k-mer + #[arg(long)] + pub min_complexity: Option, + + /// Maximum sub-word size for the complexity computation (only used when --min-complexity is set) + #[arg(long, default_value_t = 6)] + pub complexity_level_max: usize, + + /// Output as presence/absence instead of counts + #[arg(long)] + pub presence: bool, + + /// Pack the output's presence matrices in the dense format instead of the default sparse one + #[arg(long, default_value_t = false)] + pub dense: bool, + + /// Overwrite existing output directory + #[arg(short, long)] + pub force: bool, +} + +pub fn run(args: FilterArgs) { + let src = KmerIndex::open(&args.source).unwrap_or_else(|e| { + eprintln!("error opening source index: {e}"); + std::process::exit(1); + }); + + let selector = GenomeSelector::parse(&args.ingroup, &args.outgroup).unwrap_or_else(|e| { + eprintln!("error in --ingroup/--outgroup: {e}"); + std::process::exit(1); + }); + let group_filter = selector + .build_group_filter( + &src.meta(), + GroupFilterParams { + threshold: args.presence_threshold, + min_count: args.min_count, + max_count: args.max_count, + min_frac: args.min_frac, + max_frac: args.max_frac, + min_outgroup_count: args.min_outgroup_count, + max_outgroup_count: args.max_outgroup_count, + min_outgroup_frac: args.min_outgroup_frac, + max_outgroup_frac: args.max_outgroup_frac, + }, + ) + .unwrap_or_else(|e| { + eprintln!("error in filter parameters: {e}"); + std::process::exit(1); + }); + + let mut filters: Vec> = vec![Box::new(group_filter)]; + if let Some(v) = args.min_total_count { + filters.push(Box::new(MinTotalCount { total: v })); + } + if let Some(v) = args.max_total_count { + filters.push(Box::new(MaxTotalCount { total: v })); + } + if let Some(theta) = args.min_complexity { + filters.push(Box::new(MinComplexity { level_max: args.complexity_level_max, theta })); + } + + let n_genomes = src.meta().genomes().unwrap_or_else(|e| { + eprintln!("error reading index metadata: {e}"); + std::process::exit(1); + }).len(); + info!( + "filter: {} genome(s), source={}", + n_genomes, args.source.display() + ); + + let mut rep = Reporter::new(); + let t = Stage::start("filter"); + let pb = progress_bar("filter", src.n_partitions() as u64, "partitions"); + let mut alg = Filter::new(&src, &args.output, &filters) + .presence(args.presence) + .force(args.force) + .sparse(!args.dense) + .on_progress(|_: Progress| pb.inc(1)); + + let dst = alg.run().unwrap_or_else(|e| { + eprintln!("error filtering index: {e}"); + std::process::exit(1); + }); + pb.finish_and_clear(); + rep.push(t.stop()); + + info!("filtered index → {}", dst.dir().display()); + alg.reporter().print(); + rep.print(); +} diff --git a/src/obikmer2/src/cmd/mod.rs b/src/obikmer2/src/cmd/mod.rs index 0695669c..1d3f45c6 100644 --- a/src/obikmer2/src/cmd/mod.rs +++ b/src/obikmer2/src/cmd/mod.rs @@ -1,5 +1,6 @@ pub mod annotate; pub mod estimate; +pub mod filter; pub mod index; pub mod merge; pub mod superkmer; diff --git a/src/obikmer2/src/main.rs b/src/obikmer2/src/main.rs index fd4334ac..5fa7f9bc 100644 --- a/src/obikmer2/src/main.rs +++ b/src/obikmer2/src/main.rs @@ -19,6 +19,8 @@ enum Commands { Superkmer(cmd::superkmer::SuperkmerArgs), /// Merge multiple genome indexes into one Merge(cmd::merge::MergeArgs), + /// Filter kmers out of an index by genome metadata / abundance / complexity + Filter(cmd::filter::FilterArgs), /// Estimate approximate-evidence false-positive rates for given parameters Estimate(cmd::estimate::EstimateArgs), /// Read/write genome metadata (CSV) on an already-built index @@ -38,6 +40,7 @@ fn main() { Commands::Index(args) => cmd::index::run(args), Commands::Superkmer(args) => cmd::superkmer::run(args), Commands::Merge(args) => cmd::merge::run(args), + Commands::Filter(args) => cmd::filter::run(args), Commands::Estimate(args) => cmd::estimate::run(args), Commands::Annotate(args) => cmd::annotate::run(args), } diff --git a/src/obikmerge/src/partition_merge.rs b/src/obikmerge/src/partition_merge.rs index 55de21eb..6a60baba 100644 --- a/src/obikmerge/src/partition_merge.rs +++ b/src/obikmerge/src/partition_merge.rs @@ -6,8 +6,7 @@ //! instead. use std::fs; -use std::io; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use std::sync::{Arc, Mutex}; use obipipeline::{ @@ -16,66 +15,19 @@ use obipipeline::{ }; use tracing::debug; -use obicompactvec::{PersistentBitMatrixBuilder, PersistentCompactIntMatrixBuilder}; +use obicompactvec::{ColBuilder, MatrixBuilder}; use obikindex::layer::{IndexMode, KmerLayer, TypedLayer}; -use obikindex::{ColBuilder, KmerIndex, OKIError, OKIResult}; +use obikindex::{KmerIndex, OKIError, OKIResult}; use obikindexer::{build_graph, materialize_layer}; use obikseq::CanonicalKmer; use obiskio::UnitigFileReader; use crate::MergeMode; -// ── MatrixBuilder ───────────────────────────────────────────────────────────── -// -// Wraps whichever matrix builder `mode` calls for, so the merge pipeline never -// has to know the on-disk column naming (`col_NNNNNN.pbiv`/`.pciv`) or the -// matrix `meta.json` schema itself — both stay private to obicompactvec. -// `resume` reopens a matrix directory already closed by a previous builder -// session (an existing destination layer), continuing from its current -// `n_cols` instead of starting a fresh matrix at 0. - -enum MatrixBuilder { - Bit(PersistentBitMatrixBuilder), - Int(PersistentCompactIntMatrixBuilder), -} - -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)?), - }) - } - - 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)?), - }) - } - - /// Add a column with no data written (all-zero/false) — for genome - /// columns absent from this source (e.g. dst genomes in a new layer). - fn add_absent_col(&mut self) -> io::Result<()> { - match self { - MatrixBuilder::Bit(b) => b.add_col()?.close(), - MatrixBuilder::Int(b) => b.add_col()?.close(), - } - } - - fn add_col(&mut self) -> io::Result { - Ok(match self { - MatrixBuilder::Bit(b) => ColBuilder::Bit(b.add_col()?), - MatrixBuilder::Int(b) => ColBuilder::Int(b.add_col()?), - }) - } - - fn close(self) -> io::Result<()> { - match self { - MatrixBuilder::Bit(b) => b.close(), - MatrixBuilder::Int(b) => b.close(), - } - } +// `MatrixBuilder::new`/`resume` take a `presence: bool` — `MergeMode` maps +// to it directly (`Presence` → `true`, `Count` → `false`). +fn presence(mode: MergeMode) -> bool { + mode == MergeMode::Presence } #[cfg(test)] @@ -84,7 +36,7 @@ mod matrix_builder_tests { use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; - use super::{ColBuilder, MatrixBuilder, MergeMode}; + use super::{ColBuilder, MatrixBuilder}; /// Mirrors `merge_partition`'s "new layer" setup: absent (dst-genome) /// columns get no data, then source columns are filled — all through one @@ -93,7 +45,7 @@ mod matrix_builder_tests { fn new_layer_absent_then_source_columns_presence() { let dir = tempdir().unwrap(); let data_dir = dir.path().join("presence"); - let mut mb = MatrixBuilder::new(MergeMode::Presence, 3, &data_dir).unwrap(); + let mut mb = MatrixBuilder::new(true, 3, &data_dir).unwrap(); // Two absent (dst-genome) columns. mb.add_absent_col().unwrap(); @@ -123,7 +75,7 @@ mod matrix_builder_tests { fn new_layer_absent_then_source_columns_count() { let dir = tempdir().unwrap(); let data_dir = dir.path().join("counts"); - let mut mb = MatrixBuilder::new(MergeMode::Count, 2, &data_dir).unwrap(); + let mut mb = MatrixBuilder::new(false, 2, &data_dir).unwrap(); mb.add_absent_col().unwrap(); @@ -154,12 +106,12 @@ mod matrix_builder_tests { let data_dir = dir.path().join("presence"); // Previous merge: one dst-genome column already on disk. - let mut mb0 = MatrixBuilder::new(MergeMode::Presence, 3, &data_dir).unwrap(); + let mut mb0 = MatrixBuilder::new(true, 3, &data_dir).unwrap(); mb0.add_absent_col().unwrap(); mb0.close().unwrap(); // This merge: resume and append two more source columns. - let mut mb = MatrixBuilder::resume(MergeMode::Presence, &data_dir).unwrap(); + let mut mb = MatrixBuilder::resume(true, &data_dir).unwrap(); for vals in [[true, false, true], [false, true, false]] { let mut col = mb.add_col().unwrap(); match &mut col { @@ -368,7 +320,7 @@ pub(crate) fn merge_partition( MergeMode::Count => new_layer_dir.join("counts"), }; fs::create_dir_all(&data_dir)?; - let mut mb = MatrixBuilder::new(mode, n_new, &data_dir).map_err(OKIError::Io)?; + let mut mb = MatrixBuilder::new(presence(mode), n_new, &data_dir).map_err(OKIError::Io)?; for _ in 0..n_dst_genomes { mb.add_absent_col().map_err(OKIError::Io)?; } @@ -393,7 +345,7 @@ pub(crate) fn merge_partition( MergeMode::Presence => layer_dir.join("presence"), MergeMode::Count => layer_dir.join("counts"), }; - let mut mb = MatrixBuilder::resume(mode, &data_dir).map_err(OKIError::Io)?; + let mut mb = MatrixBuilder::resume(presence(mode), &data_dir).map_err(OKIError::Io)?; let cols = (0..n_src_total) .map(|_| mb.add_col().map_err(OKIError::Io)) .collect::>>()?; diff --git a/src/obikselect/src/select_layer.rs b/src/obikselect/src/select_layer.rs index cbd265ba..9c16fd5d 100644 --- a/src/obikselect/src/select_layer.rs +++ b/src/obikselect/src/select_layer.rs @@ -11,8 +11,8 @@ use std::io; use std::path::Path; use obicompactvec::{ - ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder, - PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, TempBitVec, TempCompactIntVec, + ColGroup, MatrixBuilder, MatrixGroupOps, PersistentBitMatrix, + PersistentCompactIntMatrix, TempBitVec, TempCompactIntVec, }; use obikindex::layer::{KmerLayer, LayerContent}; use obikindex::{KmerIndex, OKIError, OKIResult}; @@ -66,36 +66,15 @@ fn compute_group(mat: &dyn MatrixGroupOps, spec: &OutputCol, threshold: u32) -> }) } -// ── DstBuilder — destination matrix, generic over output content ───────────── +// ── AggResult → MatrixBuilder ───────────────────────────────────────────────── -enum DstBuilder { - Bit(PersistentBitMatrixBuilder), - Int(PersistentCompactIntMatrixBuilder), -} - -impl DstBuilder { - fn new(output_presence: bool, n: usize, dir: &Path) -> io::Result { - Ok(if output_presence { - DstBuilder::Bit(PersistentBitMatrixBuilder::new(n, dir)?) - } else { - DstBuilder::Int(PersistentCompactIntMatrixBuilder::new(n, dir)?) - }) - } - - fn add(&mut self, r: AggResult) -> io::Result<()> { - match (self, r) { - (DstBuilder::Bit(b), AggResult::Bit(v)) => b.add_col_from(&v), - (DstBuilder::Bit(b), AggResult::Int(v)) => b.add_col_from_int(&v), - (DstBuilder::Int(b), AggResult::Bit(v)) => b.add_col_from_bit(&v), - (DstBuilder::Int(b), AggResult::Int(v)) => b.add_col_from(&v), - } - } - - fn close(self) -> io::Result<()> { - match self { - DstBuilder::Bit(b) => b.close(), - DstBuilder::Int(b) => b.close(), - } +/// Add one already-aggregated column to `mb` — the only piece `MatrixBuilder` +/// (generic row/column mechanics, `obicompactvec`) can't know itself, since +/// `AggResult` is select's own aggregation-output type. +fn add_result(mb: &mut MatrixBuilder, r: AggResult) -> io::Result<()> { + match r { + AggResult::Bit(v) => mb.add_col_from_bit(&v), + AggResult::Int(v) => mb.add_col_from_int(&v), } } @@ -156,10 +135,10 @@ pub(crate) fn select_partition( let data_dir = dst_layer_dir.join(data_subdir); fs::create_dir_all(&data_dir).map_err(OKIError::Io)?; - let mut builder = DstBuilder::new(output_presence, n, &data_dir).map_err(OKIError::Io)?; + let mut builder = MatrixBuilder::new(output_presence, n, &data_dir).map_err(OKIError::Io)?; for spec in specs { let r = compute_group(group_mat.as_ref(), spec, threshold).map_err(OKIError::Io)?; - builder.add(r).map_err(OKIError::Io)?; + add_result(&mut builder, r).map_err(OKIError::Io)?; } builder.close().map_err(OKIError::Io)?; }