From 96dfb5300b1ad176ff904f07722f6a69c6028f5a Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 26 Aug 2026 08:26:17 +0200 Subject: [PATCH] Add unitig extraction command to obikdump and obikmer2 Introduces a new unitig extraction feature that reads k-mers from an index, filters them by group metadata, and partitions the computation using rayon. The implementation constructs per-partition de Bruijn graphs, merges them, computes node degrees, and writes the resulting sequences as FASTA. A corresponding CLI command is added to obikmer2 to expose this functionality with configurable filtering and progress reporting. --- src/Cargo.lock | 2 + src/obikdump/Cargo.toml | 12 +++--- src/obikdump/src/lib.rs | 2 + src/obikdump/src/unitig.rs | 68 ++++++++++++++++++++++++++++++ src/obikmer2/src/cmd/mod.rs | 1 + src/obikmer2/src/cmd/unitig/mod.rs | 46 ++++++++++++++++++++ src/obikmer2/src/main.rs | 3 ++ 7 files changed, 129 insertions(+), 5 deletions(-) create mode 100644 src/obikdump/src/unitig.rs create mode 100644 src/obikmer2/src/cmd/unitig/mod.rs diff --git a/src/Cargo.lock b/src/Cargo.lock index b0a15fcd..9d76b1b7 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1523,6 +1523,8 @@ version = "0.1.0" name = "obikdump" version = "0.1.0" dependencies = [ + "obidebruinj", + "obifastwrite", "obikfilter", "obikidxcache", "obikindex", diff --git a/src/obikdump/Cargo.toml b/src/obikdump/Cargo.toml index 2d7ecc60..802ede05 100644 --- a/src/obikdump/Cargo.toml +++ b/src/obikdump/Cargo.toml @@ -4,8 +4,10 @@ version = "0.1.0" edition = "2024" [dependencies] -obikindex = { path = "../obikindex" } -obikfilter = { path = "../obikfilter" } -obikidxcache = { path = "../obikidxcache" } -obikseq = { path = "../obikseq" } -rayon = "1" +obikindex = { path = "../obikindex" } +obikfilter = { path = "../obikfilter" } +obikidxcache = { path = "../obikidxcache" } +obikseq = { path = "../obikseq" } +obidebruinj = { path = "../obidebruinj" } +obifastwrite = { path = "../obifastwrite" } +rayon = "1" diff --git a/src/obikdump/src/lib.rs b/src/obikdump/src/lib.rs index 577fd412..706fcc78 100644 --- a/src/obikdump/src/lib.rs +++ b/src/obikdump/src/lib.rs @@ -6,5 +6,7 @@ //! reverse), same pattern as `obikindexer`/`obikquery`. mod dump; +mod unitig; pub use dump::IndexDump; +pub use unitig::IndexUnitigs; diff --git a/src/obikdump/src/unitig.rs b/src/obikdump/src/unitig.rs new file mode 100644 index 00000000..b4a3c527 --- /dev/null +++ b/src/obikdump/src/unitig.rs @@ -0,0 +1,68 @@ +//! Assembling the full k-mer set of a *complete* source index into unitigs +//! and writing them as FASTA — the same `IndexCache`-per-partition, +//! `FilteredPartitionIter` read path as [`crate::IndexDump::dump`], feeding +//! an `obidebruinj::GraphDeBruijn` instead of a CSV row. `KmerIndex` is a +//! foreign type, so this is an extension trait rather than an inherent `impl`. + +use std::io::Write; + +use obidebruinj::GraphDeBruijn; +use obifastwrite::write_unitig; +use obikfilter::{FilteredPartitionIter, KmerFilter}; +use obikidxcache::index_cache::IndexCache; +use obikindex::{KmerIndex, OKIError, OKIResult}; +use rayon::prelude::*; + +pub trait IndexUnitigs { + /// Build the de Bruijn graph from every partition's filtered kmers, then + /// write the resulting unitigs as FASTA to `out`. Returns the number of + /// unitigs written. + fn write_unitigs( + &self, + out: &mut W, + filters: &[Box], + on_partition: F, + ) -> OKIResult; +} + +impl IndexUnitigs for KmerIndex { + fn write_unitigs( + &self, + out: &mut W, + filters: &[Box], + on_partition: F, + ) -> OKIResult { + let k = self.kmer_size(); + let n_genomes = self.meta().genomes().map_err(OKIError::Io)?.len().max(1); + let use_counts = self.meta().config.with_counts; + let n = self.n_partitions(); + + let g = (0..n) + .into_par_iter() + .try_fold(GraphDeBruijn::new, |mut local_g, i| -> OKIResult { + let cache = IndexCache::new(self, Some(vec![i])); + cache.iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, _row| { + local_g.push(kmer); + true + })?; + on_partition(); + Ok(local_g) + }) + .try_reduce(GraphDeBruijn::new, |mut a, b| { + a.merge(b); + Ok(a) + })?; + + g.compute_degrees_and_mark_starts(); + + let mut n_written = 0usize; + g.try_for_each_unitig(|unitig| { + write_unitig(unitig, k, 0, n_written, out)?; + n_written += 1; + Ok(()) + }) + .map_err(OKIError::Io)?; + + Ok(n_written) + } +} diff --git a/src/obikmer2/src/cmd/mod.rs b/src/obikmer2/src/cmd/mod.rs index 732ba742..58081aa7 100644 --- a/src/obikmer2/src/cmd/mod.rs +++ b/src/obikmer2/src/cmd/mod.rs @@ -7,3 +7,4 @@ pub mod merge; mod predicate; pub mod select; pub mod superkmer; +pub mod unitig; diff --git a/src/obikmer2/src/cmd/unitig/mod.rs b/src/obikmer2/src/cmd/unitig/mod.rs new file mode 100644 index 00000000..8a725b48 --- /dev/null +++ b/src/obikmer2/src/cmd/unitig/mod.rs @@ -0,0 +1,46 @@ +use std::io::{self, BufWriter}; +use std::path::PathBuf; + +use clap::Args; +use obikdump::IndexUnitigs; +use obikfilter::KmerFilter; +use obikindex::KmerIndex; +use obisys::progress_bar; +use tracing::info; + +use super::predicate::GroupFilterArgs; + +#[derive(Args)] +pub struct UnitigArgs { + /// Index directory + pub index: PathBuf, + + #[command(flatten)] + pub group_filter: GroupFilterArgs, +} + +pub fn run(args: UnitigArgs) { + let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + }); + + info!( + "unitig: building de Bruijn graph from {} partition(s) (k={})", + idx.n_partitions(), + idx.kmer_size(), + ); + + let filters: Vec> = vec![Box::new(args.group_filter.build_filter(&idx.meta()))]; + let pb = progress_bar("unitig", idx.n_partitions() as u64, "partitions"); + + let mut out = BufWriter::new(io::stdout()); + + let n = idx.write_unitigs(&mut out, &filters, || pb.inc(1)).unwrap_or_else(|e| { + eprintln!("unitig error: {e}"); + std::process::exit(1); + }); + pb.finish_and_clear(); + + info!("unitig: {n} unitig(s) written"); +} diff --git a/src/obikmer2/src/main.rs b/src/obikmer2/src/main.rs index 743ff416..a8153de1 100644 --- a/src/obikmer2/src/main.rs +++ b/src/obikmer2/src/main.rs @@ -25,6 +25,8 @@ enum Commands { Select(cmd::select::SelectArgs), /// Dump an index's kmers as a CSV table Dump(cmd::dump::DumpArgs), + /// Assemble an index's kmers into unitigs and write them as FASTA + Unitig(cmd::unitig::UnitigArgs), /// Estimate approximate-evidence false-positive rates for given parameters Estimate(cmd::estimate::EstimateArgs), /// Read/write genome metadata (CSV) on an already-built index @@ -47,6 +49,7 @@ fn main() { Commands::Filter(args) => cmd::filter::run(args), Commands::Select(args) => cmd::select::run(args), Commands::Dump(args) => cmd::dump::run(args), + Commands::Unitig(args) => cmd::unitig::run(args), Commands::Estimate(args) => cmd::estimate::run(args), Commands::Annotate(args) => cmd::annotate::run(args), }