diff --git a/src/Cargo.lock b/src/Cargo.lock index e7554bb6..63fec439 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1654,6 +1654,7 @@ name = "obikmer2" version = "1.2.2" dependencies = [ "clap", + "csv", "obifastwrite", "obikalgorithm", "obikindex", diff --git a/src/obikindex/examples/compare_sparse.rs b/src/obikindex/examples/compare_sparse.rs index a772b82b..16766b67 100644 --- a/src/obikindex/examples/compare_sparse.rs +++ b/src/obikindex/examples/compare_sparse.rs @@ -25,25 +25,25 @@ fn main() -> anyhow::Result<()> { let mut first_mismatch = None; for part in 0..n_parts { - let index_dir_sparse = sparse.index_dir(part); - let index_dir_dense = dense.index_dir(part); - - if !index_dir_sparse.exists() || !index_dir_dense.exists() { + let Ok(part_sparse) = sparse.partition(part) else { continue; - } + }; + let Ok(part_dense) = dense.partition(part) else { + continue; + }; for layer in 0..n_layers { - let layer_dir_sparse = sparse.layer_dir(part, layer); - let layer_dir_dense = dense.layer_dir(part, layer); - - if !layer_dir_sparse.exists() || !layer_dir_dense.exists() { + let Ok(layer_sparse) = part_sparse.layer(layer) else { continue; - } + }; + let Ok(layer_dense) = part_dense.layer(layer) else { + continue; + }; - let presence_sparse = layer_dir_sparse.join("presence"); - let presence_dense = layer_dir_dense.join("presence"); + let presence_sparse = layer_sparse.presence_dir(); + let layer_dir_dense = layer_dense.dir().to_path_buf(); - if !presence_sparse.exists() || !presence_dense.exists() { + if !presence_sparse.exists() || !layer_dense.presence_dir().exists() { continue; } diff --git a/src/obikindex/src/layer/typed_layer.rs b/src/obikindex/src/layer/typed_layer.rs index 5df93beb..befb3d80 100644 --- a/src/obikindex/src/layer/typed_layer.rs +++ b/src/obikindex/src/layer/typed_layer.rs @@ -150,7 +150,7 @@ impl HasStorageKind for PersistentBitMatrix { // ── Structures ──────────────────────────────────────────────────────────────── pub struct TypedLayer { - mphf: MphfLayer, + kmerstore: MphfLayer, data: D, } @@ -165,11 +165,14 @@ impl TypedLayer { pub fn open(path: &Path) -> OKIResult { let mphf = MphfLayer::open(path)?; let data = D::open(path)?; - Ok(Self { mphf, data }) + Ok(Self { + kmerstore: mphf, + data, + }) } pub fn query(&self, kmer: CanonicalKmer) -> Option> { - self.mphf.find(kmer).map(|slot| Hit { + self.kmerstore.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot), }) @@ -181,37 +184,37 @@ impl TypedLayer { /// sweep), so a plain `query` reading — and discarding — a full row /// per lookup would be wasted work. pub fn find_slot(&self, kmer: CanonicalKmer) -> Option { - self.mphf.find(kmer) + self.kmerstore.find(kmer) } pub fn n(&self) -> usize { - self.mphf.n() + self.kmerstore.n() } /// Raw MPHF lookup: kmer → slot, no membership check. pub fn hash(&self, kmer: CanonicalKmer) -> usize { - self.mphf.hash(kmer) + self.kmerstore.hash(kmer) } /// Batch raw MPHF lookup: kmers → slots, no membership check. pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec { - self.mphf.hash_batch(kmers) + self.kmerstore.hash_batch(kmers) } /// Iterate over all canonical kmers in the layer, in deterministic order. pub fn iter_kmers(&self) -> crate::layer::mphf_layer::KmerIter { - self.mphf.iter_kmers() + self.kmerstore.iter_kmers() } /// Iterate over all canonical kmers, each paired with its zero-based /// sequence index in `unitigs.bin`. pub fn enumerate_kmers(&self) -> std::iter::Enumerate { - self.mphf.enumerate_kmers() + self.kmerstore.enumerate_kmers() } /// Iterate over the layer's canonical kmers in batches of `n`. pub fn iter_kmers_batch(&self, n: usize) -> crate::layer::mphf_layer::KmerBatchIter { - self.mphf.iter_kmers_batch(n) + self.kmerstore.iter_kmers_batch(n) } /// Iterate over batches, each paired with the zero-based index of the @@ -220,13 +223,13 @@ impl TypedLayer { &self, n: usize, ) -> impl Iterator)> + Send + 'static { - self.mphf.enumerate_kmers_batch(n) + self.kmerstore.enumerate_kmers_batch(n) } /// Iterate over the layer's unitigs (whole sequences, not decomposed /// into k-mers). pub fn iter_unitigs(&self) -> crate::layer::mphf_layer::UnitigIter { - self.mphf.iter_unitigs() + self.kmerstore.iter_unitigs() } pub fn unitig_writer(out_dir: &Path) -> OKIResult { @@ -251,7 +254,7 @@ impl TypedLayer { /// already in memory (`MphfLayer`'s own `LayerEvidence`), no disk /// access. Available regardless of `D`, unlike `content`/`storage_kind`. pub fn evidence_kind(&self) -> crate::layer::mphf_layer::EvidenceKind { - self.mphf.evidence_kind() + self.kmerstore.evidence_kind() } } diff --git a/src/obikmer2/Cargo.toml b/src/obikmer2/Cargo.toml index b08c0eab..695baaf0 100644 --- a/src/obikmer2/Cargo.toml +++ b/src/obikmer2/Cargo.toml @@ -19,6 +19,7 @@ obikmerge = { path = "../obikmerge" } obifastwrite = { path = "../obifastwrite" } obiskbuilder = { path = "../obiskbuilder" } clap = { version = "4", features = ["derive"] } +csv = "1" tracing = "0.1.44" tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } diff --git a/src/obikmer2/src/cmd/annotate/mod.rs b/src/obikmer2/src/cmd/annotate/mod.rs new file mode 100644 index 00000000..a4018dde --- /dev/null +++ b/src/obikmer2/src/cmd/annotate/mod.rs @@ -0,0 +1,184 @@ +use std::collections::HashSet; +use std::io::{self, BufWriter, Write}; +use std::path::PathBuf; + +use clap::Args; +use obikindex::KmerIndex; +use tracing::info; + +#[derive(Args)] +pub struct AnnotateArgs { + /// Index directory to annotate (modified in-place) + pub index: PathBuf, + + /// CSV file with genome metadata (must contain an id column) + #[arg(long)] + pub csv: Option, + + /// CSV field separator + #[arg(long, default_value = ",")] + pub sep: char, + + /// Name of the column that contains genome labels + #[arg(long, default_value = "id")] + pub id_col: String, + + /// Value that means "delete / absent" (removes existing key if present) + #[arg(long, default_value = "NA")] + pub na_value: String, + + /// Do not overwrite existing metadata keys + #[arg(long)] + pub no_overwrite: bool, + + /// Dump all genome metadata as CSV (stdout) instead of reading a CSV + #[arg(long)] + pub dump: bool, +} + +pub fn run(args: AnnotateArgs) { + if args.dump { + run_dump(&args); + } else { + run_annotate(&args); + } +} + +fn run_dump(args: &AnnotateArgs) { + let idx = open_index(&args.index); + + let genomes = idx.meta().genomes().unwrap_or_else(|e| { + eprintln!("error reading index metadata: {e}"); + std::process::exit(1); + }); + let genomes = &genomes; + + // Collect all keys in stable order (sorted for determinism) + let mut key_set: HashSet = HashSet::new(); + for g in genomes { + for k in g.meta.keys() { + key_set.insert(k.clone()); + } + } + let mut keys: Vec = key_set.into_iter().collect(); + keys.sort(); + + let stdout = io::stdout(); + let mut out = BufWriter::new(stdout.lock()); + + // Header + write!(out, "id").unwrap(); + for k in &keys { + write!(out, "{}{k}", args.sep).unwrap(); + } + writeln!(out).unwrap(); + + // Rows + for g in genomes { + write!(out, "{}", g.label).unwrap(); + for k in &keys { + let v = g.meta.get(k).map(|s| s.as_str()).unwrap_or("NA"); + write!(out, "{}{v}", args.sep).unwrap(); + } + writeln!(out).unwrap(); + } +} + +fn run_annotate(args: &AnnotateArgs) { + let csv_path = match &args.csv { + Some(p) => p.clone(), + None => { + eprintln!("error: --csv is required unless --dump is used"); + std::process::exit(1); + } + }; + + let idx = open_index(&args.index); + + let mut genomes = idx.meta().genomes().unwrap_or_else(|e| { + eprintln!("error reading index metadata: {e}"); + std::process::exit(1); + }); + + // Build a label → genome index position map + let label_to_pos: std::collections::HashMap = genomes + .iter() + .enumerate() + .map(|(i, g)| (g.label.clone(), i)) + .collect(); + + let sep = args.sep as u8; + let mut rdr = csv::ReaderBuilder::new() + .delimiter(sep) + .from_path(&csv_path) + .unwrap_or_else(|e| { + eprintln!("error opening {}: {e}", csv_path.display()); + std::process::exit(1); + }); + + let headers = rdr + .headers() + .unwrap_or_else(|e| { + eprintln!("error reading CSV headers: {e}"); + std::process::exit(1); + }) + .clone(); + + let id_col_idx = headers.iter().position(|h| h == args.id_col).unwrap_or_else(|| { + eprintln!("error: id column '{}' not found in CSV", args.id_col); + std::process::exit(1); + }); + + let meta_cols: Vec<(usize, String)> = headers + .iter() + .enumerate() + .filter(|(i, _)| *i != id_col_idx) + .map(|(i, h)| (i, h.to_string())) + .collect(); + + let mut updated = 0usize; + let mut skipped = 0usize; + + for result in rdr.records() { + let record = result.unwrap_or_else(|e| { + eprintln!("error reading CSV record: {e}"); + std::process::exit(1); + }); + + let label = record.get(id_col_idx).unwrap_or("").to_string(); + let pos = match label_to_pos.get(&label) { + Some(&p) => p, + None => { + skipped += 1; + continue; + } + }; + + let genome = &mut genomes[pos]; + for (col_idx, key) in &meta_cols { + let val = record.get(*col_idx).unwrap_or(""); + if val == args.na_value { + genome.meta.remove(key); + } else if args.no_overwrite && genome.meta.contains_key(key) { + // skip + } else { + genome.meta.insert(key.clone(), val.to_string()); + } + } + updated += 1; + } + + idx.meta().set_genomes(genomes).unwrap_or_else(|e| { + eprintln!("error writing index metadata: {e}"); + std::process::exit(1); + }); + + info!("annotated {updated} genome(s), skipped {skipped} CSV row(s) with unknown label"); +} + +fn open_index(path: &PathBuf) -> KmerIndex { + KmerIndex::open(path).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + }) +} diff --git a/src/obikmer2/src/cmd/mod.rs b/src/obikmer2/src/cmd/mod.rs index 8419a2d8..0695669c 100644 --- a/src/obikmer2/src/cmd/mod.rs +++ b/src/obikmer2/src/cmd/mod.rs @@ -1,3 +1,4 @@ +pub mod annotate; pub mod estimate; pub mod index; pub mod merge; diff --git a/src/obikmer2/src/main.rs b/src/obikmer2/src/main.rs index 7cc80788..fd4334ac 100644 --- a/src/obikmer2/src/main.rs +++ b/src/obikmer2/src/main.rs @@ -21,6 +21,8 @@ enum Commands { Merge(cmd::merge::MergeArgs), /// Estimate approximate-evidence false-positive rates for given parameters Estimate(cmd::estimate::EstimateArgs), + /// Read/write genome metadata (CSV) on an already-built index + Annotate(cmd::annotate::AnnotateArgs), } fn main() { @@ -37,5 +39,6 @@ fn main() { Commands::Superkmer(args) => cmd::superkmer::run(args), Commands::Merge(args) => cmd::merge::run(args), Commands::Estimate(args) => cmd::estimate::run(args), + Commands::Annotate(args) => cmd::annotate::run(args), } }