From 737966899fdf1c9cc30e617e952cac3a88e7c837 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 26 Aug 2026 08:20:58 +0200 Subject: [PATCH] Add obikselect crate for k-mer selection and aggregation Introduces a new `obikselect` module to handle k-mer selection and column aggregation. Adds a `select` CLI command in `obikmer2` that supports group predicates, aggregate operators, and output column filtering. Implements operator parsing with case-insensitive matching and default rules, along with a centralized module for resolving group specifications from metadata or genome filters. Updates dependencies to include the new crate and its prerequisites. --- src/Cargo.lock | 2 + src/obikmer2/Cargo.toml | 1 + src/obikmer2/src/cmd/mod.rs | 1 + src/obikmer2/src/cmd/select/mod.rs | 144 +++++++++++++++++++++++++++++ src/obikmer2/src/main.rs | 3 + src/obikselect/Cargo.toml | 1 + src/obikselect/src/group_specs.rs | 116 +++++++++++++++++++++++ src/obikselect/src/lib.rs | 2 + src/obikselect/src/select_layer.rs | 21 +++++ 9 files changed, 291 insertions(+) create mode 100644 src/obikmer2/src/cmd/select/mod.rs create mode 100644 src/obikselect/src/group_specs.rs diff --git a/src/Cargo.lock b/src/Cargo.lock index 06828b83..a1cad46c 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1671,6 +1671,7 @@ dependencies = [ "obikindex", "obikindexer", "obikmerge", + "obikselect", "obikseq", "obipipeline", "obiread", @@ -1758,6 +1759,7 @@ version = "0.1.0" dependencies = [ "obicompactvec", "obikalgorithm", + "obikfilter", "obikindex", "obisys", "tracing", diff --git a/src/obikmer2/Cargo.toml b/src/obikmer2/Cargo.toml index 971c6cc8..0d849d15 100644 --- a/src/obikmer2/Cargo.toml +++ b/src/obikmer2/Cargo.toml @@ -17,6 +17,7 @@ obikindexer = { path = "../obikindexer" } obikalgorithm = { path = "../obikalgorithm" } obikmerge = { path = "../obikmerge" } obikfilter = { path = "../obikfilter" } +obikselect = { path = "../obikselect" } obifastwrite = { path = "../obifastwrite" } obiskbuilder = { path = "../obiskbuilder" } clap = { version = "4", features = ["derive"] } diff --git a/src/obikmer2/src/cmd/mod.rs b/src/obikmer2/src/cmd/mod.rs index 1d3f45c6..9b94a285 100644 --- a/src/obikmer2/src/cmd/mod.rs +++ b/src/obikmer2/src/cmd/mod.rs @@ -3,4 +3,5 @@ pub mod estimate; pub mod filter; pub mod index; pub mod merge; +pub mod select; pub mod superkmer; diff --git a/src/obikmer2/src/cmd/select/mod.rs b/src/obikmer2/src/cmd/select/mod.rs new file mode 100644 index 00000000..28539b4c --- /dev/null +++ b/src/obikmer2/src/cmd/select/mod.rs @@ -0,0 +1,144 @@ +use std::path::PathBuf; + +use clap::{Args, ValueEnum}; +use obikalgorithm::Algorithm; +use obikindex::KmerIndex; +use obikselect::{AggOp, ColumnSpecParams, Select, build_output_cols}; +use obisys::{Progress, Reporter, Stage, progress_bar}; +use tracing::info; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum AggOpArg { + Any, + All, + None, + Sum, + Min, + Max, +} + +impl From for AggOp { + fn from(a: AggOpArg) -> Self { + match a { + AggOpArg::Any => AggOp::Any, + AggOpArg::All => AggOp::All, + AggOpArg::None => AggOp::None, + AggOpArg::Sum => AggOp::Sum, + AggOpArg::Min => AggOp::Min, + AggOpArg::Max => AggOp::Max, + } + } +} + +#[derive(Args)] +pub struct SelectArgs { + /// Source index directory + pub source: PathBuf, + + /// Output index directory + #[arg(short, long)] + pub output: PathBuf, + + /// Define a named group: `:` (repeatable; mutually exclusive with --aggregate-by) + #[arg(long, value_name = "NAME:PRED", conflicts_with = "aggregate_by")] + pub group: Vec, + + /// Per-group aggregation operator: `:` (repeatable) + #[arg(long, value_name = "NAME:OP")] + pub group_op: Vec, + + /// Auto-create one group per unique value of metadata key + #[arg(long, value_name = "KEY", conflicts_with = "group")] + pub aggregate_by: Option, + + /// Aggregation operator for all auto-generated groups + #[arg(long, value_name = "OP")] + pub aggregate_op: Option, + + /// Output columns in order: group names or genome labels, comma-separated + #[arg(long, value_name = "COL,...", value_delimiter = ',')] + pub select: Option>, + + /// Minimum count to consider a genome as "carrying" the k-mer (logical ops only) + #[arg(long, default_value = "0")] + pub presence_threshold: u32, + + /// 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, +} + +/// Split a repeatable `:` argument. Exits on malformed input. +fn parse_name_value(s: &str, flag: &str) -> (String, String) { + match s.find(':') { + Some(pos) => (s[..pos].trim().to_string(), s[pos + 1..].to_string()), + None => { + eprintln!("error in {flag}: expected :, got: {s}"); + std::process::exit(1); + } + } +} + +pub fn run(args: SelectArgs) { + let src = KmerIndex::open(&args.source).unwrap_or_else(|e| { + eprintln!("error opening source index: {e}"); + std::process::exit(1); + }); + + let group_preds: Vec<(String, String)> = + args.group.iter().map(|s| parse_name_value(s, "--group")).collect(); + let group_ops: Vec<(String, String)> = + args.group_op.iter().map(|s| parse_name_value(s, "--group-op")).collect(); + + let src_is_count = src.meta().config.with_counts; + let (specs, output_presence) = build_output_cols( + &src.meta(), + ColumnSpecParams { + group_preds: &group_preds, + aggregate_by: args.aggregate_by.as_deref(), + group_ops: &group_ops, + aggregate_op: args.aggregate_op.map(AggOp::from), + select: args.select.as_deref(), + src_is_count, + }, + ) + .unwrap_or_else(|e| { + eprintln!("error building output columns: {e}"); + std::process::exit(1); + }); + + let n_genomes = src.meta().genomes().unwrap_or_else(|e| { + eprintln!("error reading index metadata: {e}"); + std::process::exit(1); + }).len(); + info!( + "select: {} genome(s) → {} output column(s), output={}", + n_genomes, + specs.len(), + if output_presence { "presence" } else { "count" }, + ); + + let mut rep = Reporter::new(); + let t = Stage::start("select"); + let pb = progress_bar("select", src.n_partitions() as u64, "partitions"); + let mut alg = Select::new(&src, &args.output, &specs, output_presence) + .threshold(args.presence_threshold) + .force(args.force) + .sparse(!args.dense) + .on_progress(|_: Progress| pb.inc(1)); + + let dst = alg.run().unwrap_or_else(|e| { + eprintln!("select error: {e}"); + std::process::exit(1); + }); + pb.finish_and_clear(); + rep.push(t.stop()); + + info!("selected index → {}", dst.dir().display()); + alg.reporter().print(); + rep.print(); +} diff --git a/src/obikmer2/src/main.rs b/src/obikmer2/src/main.rs index 5fa7f9bc..27f96d53 100644 --- a/src/obikmer2/src/main.rs +++ b/src/obikmer2/src/main.rs @@ -21,6 +21,8 @@ enum Commands { Merge(cmd::merge::MergeArgs), /// Filter kmers out of an index by genome metadata / abundance / complexity Filter(cmd::filter::FilterArgs), + /// Project/aggregate genome columns into a new index + Select(cmd::select::SelectArgs), /// Estimate approximate-evidence false-positive rates for given parameters Estimate(cmd::estimate::EstimateArgs), /// Read/write genome metadata (CSV) on an already-built index @@ -41,6 +43,7 @@ fn main() { Commands::Superkmer(args) => cmd::superkmer::run(args), Commands::Merge(args) => cmd::merge::run(args), Commands::Filter(args) => cmd::filter::run(args), + Commands::Select(args) => cmd::select::run(args), Commands::Estimate(args) => cmd::estimate::run(args), Commands::Annotate(args) => cmd::annotate::run(args), } diff --git a/src/obikselect/Cargo.toml b/src/obikselect/Cargo.toml index d5f50cdc..907e3239 100644 --- a/src/obikselect/Cargo.toml +++ b/src/obikselect/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" [dependencies] obikindex = { path = "../obikindex" } +obikfilter = { path = "../obikfilter" } obikalgorithm = { path = "../obikalgorithm" } obicompactvec = { path = "../obicompactvec" } obisys = { path = "../obisys" } diff --git a/src/obikselect/src/group_specs.rs b/src/obikselect/src/group_specs.rs new file mode 100644 index 00000000..71b2db32 --- /dev/null +++ b/src/obikselect/src/group_specs.rs @@ -0,0 +1,116 @@ +//! Resolving CLI-level group/column definitions into a concrete +//! `Vec` — the logic `obikmer2`'s `select` command delegates to, +//! so the command itself stays argument parsing + the `Select` call. Named +//! groups defined by a metadata predicate reuse `obikfilter::GenomeSelector` +//! (the same predicate parser `filter`'s `--ingroup`/`--outgroup` use) +//! instead of a separate ad hoc resolver. + +use std::collections::{BTreeMap, HashMap}; + +use obikfilter::GenomeSelector; +use obikindex::IndexMeta; + +use crate::select_layer::{AggOp, OutputCol}; + +pub struct ColumnSpecParams<'a> { + /// Named groups defined by a metadata predicate: `(name, predicate)`. + /// Mutually exclusive with `aggregate_by` (enforced at the CLI layer). + pub group_preds: &'a [(String, String)], + /// One group per unique value of this metadata key, instead of `group_preds`. + pub aggregate_by: Option<&'a str>, + /// Per-group operator override: `(name, op string)`. + pub group_ops: &'a [(String, String)], + /// Operator applied to every group without its own override — already + /// resolved (e.g. from a `clap::ValueEnum`), unlike `group_ops`' + /// free-form strings (embedded in a combined `name:op` argument). + pub aggregate_op: Option, + /// Explicit output column order — group names or genome labels. Defaults + /// to group definition order, or every source genome if there are no groups. + pub select: Option<&'a [String]>, + pub src_is_count: bool, +} + +/// Resolve `p` against `meta` into an ordered list of output columns. +/// +/// Returns `(specs, output_presence)`. +pub fn build_output_cols(meta: &IndexMeta, p: ColumnSpecParams) -> Result<(Vec, bool), String> { + let genomes = meta.genomes().map_err(|e| e.to_string())?; + + // ── 1. group_indices: name → Vec, in definition order ────────── + let mut group_order: Vec = Vec::new(); + let mut group_indices: HashMap> = HashMap::new(); + + if let Some(key) = p.aggregate_by { + let mut value_to_indices: BTreeMap> = BTreeMap::new(); + for (i, g) in genomes.iter().enumerate() { + if let Some(v) = g.meta.get(key) { + value_to_indices.entry(v.clone()).or_default().push(i); + } + } + for (v, idxs) in value_to_indices { + group_order.push(v.clone()); + group_indices.insert(v, idxs); + } + } else { + for (name, pred) in p.group_preds { + let idxs = GenomeSelector::parse(std::slice::from_ref(pred), &[])? + .run(meta)? + .ingroup_idx; + if !group_indices.contains_key(name) { + group_order.push(name.clone()); + } + group_indices.insert(name.clone(), idxs); + } + } + + // ── 2. per-group operators ────────────────────────────────────────── + let global_op = p.aggregate_op; + let mut group_op: HashMap = HashMap::new(); + for (name, op_str) in p.group_ops { + if !group_indices.contains_key(name) { + return Err(format!("--group-op references undefined group: {name}")); + } + group_op.insert(name.clone(), AggOp::parse(op_str)?); + } + + // ── 3. genome label → index, for pass-through columns ─────────────── + let label_to_idx: HashMap<&str, usize> = genomes.iter().enumerate() + .map(|(i, g)| (g.label.as_str(), i)) + .collect(); + + // ── 4. output column order ─────────────────────────────────────────── + let col_names: Vec = if let Some(sel) = p.select { + sel.to_vec() + } else if !group_order.is_empty() { + group_order + } else { + genomes.iter().map(|g| g.label.clone()).collect() + }; + + // ── 5. build OutputCol list ─────────────────────────────────────────── + let mut specs: Vec = Vec::with_capacity(col_names.len()); + for name in &col_names { + if let Some(idxs) = group_indices.get(name) { + let op = group_op.get(name).copied() + .or(global_op) + .unwrap_or_else(|| AggOp::default_for(p.src_is_count)); + specs.push(OutputCol { label: name.clone(), indices: idxs.clone(), op }); + } else if let Some(&idx) = label_to_idx.get(name.as_str()) { + specs.push(OutputCol { + label: name.clone(), + indices: vec![idx], + op: AggOp::default_for(p.src_is_count), + }); + } else { + return Err(format!("--select: unknown column '{name}' (not a group name or genome label)")); + } + } + if specs.is_empty() { + return Err("select: no output columns defined".to_string()); + } + + // ── 6. output content kind ──────────────────────────────────────────── + let output_presence = !p.src_is_count || specs.iter().all(|s| s.op.is_logical()); + + Ok((specs, output_presence)) +} diff --git a/src/obikselect/src/lib.rs b/src/obikselect/src/lib.rs index 78be8dc5..537b1a72 100644 --- a/src/obikselect/src/lib.rs +++ b/src/obikselect/src/lib.rs @@ -7,8 +7,10 @@ //! `obikmerge`: [`Select`] implements `obikalgorithm::Algorithm` (two-phase //! `new` + setters, then `run`). +mod group_specs; mod select; mod select_layer; +pub use group_specs::{ColumnSpecParams, build_output_cols}; pub use select::Select; pub use select_layer::{AggOp, OutputCol}; diff --git a/src/obikselect/src/select_layer.rs b/src/obikselect/src/select_layer.rs index 9c16fd5d..8b0a9700 100644 --- a/src/obikselect/src/select_layer.rs +++ b/src/obikselect/src/select_layer.rs @@ -33,6 +33,27 @@ impl AggOp { pub fn is_logical(self) -> bool { matches!(self, AggOp::Any | AggOp::All | AggOp::None) } + + /// Parse an aggregation operator name (`any`/`all`/`none`/`sum`/`min`/`max`, + /// case-insensitive) — the CLI-facing string form, e.g. `--group-op`/ + /// `--aggregate-op`. + pub fn parse(s: &str) -> Result { + match s.to_lowercase().as_str() { + "any" => Ok(AggOp::Any), + "all" => Ok(AggOp::All), + "none" => Ok(AggOp::None), + "sum" => Ok(AggOp::Sum), + "min" => Ok(AggOp::Min), + "max" => Ok(AggOp::Max), + other => Err(format!("unknown aggregation operator: {other}; valid: any, all, none, sum, min, max")), + } + } + + /// Default operator for a source of the given content kind: `Sum` for + /// counts, `Any` for presence. + pub fn default_for(src_is_count: bool) -> Self { + if src_is_count { AggOp::Sum } else { AggOp::Any } + } } // ── OutputCol ─────────────────────────────────────────────────────────────────