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.
This commit is contained in:
Generated
+2
@@ -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",
|
||||
|
||||
@@ -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"] }
|
||||
|
||||
@@ -3,4 +3,5 @@ pub mod estimate;
|
||||
pub mod filter;
|
||||
pub mod index;
|
||||
pub mod merge;
|
||||
pub mod select;
|
||||
pub mod superkmer;
|
||||
|
||||
@@ -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<AggOpArg> 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: `<name>:<pred>` (repeatable; mutually exclusive with --aggregate-by)
|
||||
#[arg(long, value_name = "NAME:PRED", conflicts_with = "aggregate_by")]
|
||||
pub group: Vec<String>,
|
||||
|
||||
/// Per-group aggregation operator: `<name>:<op>` (repeatable)
|
||||
#[arg(long, value_name = "NAME:OP")]
|
||||
pub group_op: Vec<String>,
|
||||
|
||||
/// Auto-create one group per unique value of metadata key <KEY>
|
||||
#[arg(long, value_name = "KEY", conflicts_with = "group")]
|
||||
pub aggregate_by: Option<String>,
|
||||
|
||||
/// Aggregation operator for all auto-generated groups
|
||||
#[arg(long, value_name = "OP")]
|
||||
pub aggregate_op: Option<AggOpArg>,
|
||||
|
||||
/// Output columns in order: group names or genome labels, comma-separated
|
||||
#[arg(long, value_name = "COL,...", value_delimiter = ',')]
|
||||
pub select: Option<Vec<String>>,
|
||||
|
||||
/// 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 `<name>:<value>` 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 <name>:<value>, 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();
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
obikindex = { path = "../obikindex" }
|
||||
obikfilter = { path = "../obikfilter" }
|
||||
obikalgorithm = { path = "../obikalgorithm" }
|
||||
obicompactvec = { path = "../obicompactvec" }
|
||||
obisys = { path = "../obisys" }
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
//! Resolving CLI-level group/column definitions into a concrete
|
||||
//! `Vec<OutputCol>` — 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<AggOp>,
|
||||
/// 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<OutputCol>, bool), String> {
|
||||
let genomes = meta.genomes().map_err(|e| e.to_string())?;
|
||||
|
||||
// ── 1. group_indices: name → Vec<usize>, in definition order ──────────
|
||||
let mut group_order: Vec<String> = Vec::new();
|
||||
let mut group_indices: HashMap<String, Vec<usize>> = HashMap::new();
|
||||
|
||||
if let Some(key) = p.aggregate_by {
|
||||
let mut value_to_indices: BTreeMap<String, Vec<usize>> = 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<String, AggOp> = 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<String> = 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<OutputCol> = 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))
|
||||
}
|
||||
@@ -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};
|
||||
|
||||
@@ -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<Self, String> {
|
||||
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 ─────────────────────────────────────────────────────────────────
|
||||
|
||||
Reference in New Issue
Block a user