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

@@ -1523,6 +1523,8 @@ version = "0.1.0"
 name = "obikdump"
 version = "0.1.0"
 dependencies = [
+ "obidebruinj",
+ "obifastwrite",
  "obikfilter",
  "obikidxcache",
  "obikindex",

src/obikdump/Cargo.toml

@@ -4,8 +4,10 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-obikindex   = { path = "../obikindex" }
+obikindex     = { path = "../obikindex" }
-obikfilter  = { path = "../obikfilter" }
+obikfilter    = { path = "../obikfilter" }
-obikidxcache = { path = "../obikidxcache" }
+obikidxcache  = { path = "../obikidxcache" }
-obikseq     = { path = "../obikseq" }
+obikseq       = { path = "../obikseq" }
-rayon       = "1"
+obidebruinj   = { path = "../obidebruinj" }
+obifastwrite  = { path = "../obifastwrite" }
+rayon         = "1"

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;

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<W: Write + Send, F: Fn() + Send + Sync>(
+        &self,
+        out: &mut W,
+        filters: &[Box<dyn KmerFilter>],
+        on_partition: F,
+    ) -> OKIResult<usize>;
+}
+
+impl IndexUnitigs for KmerIndex {
+    fn write_unitigs<W: Write + Send, F: Fn() + Send + Sync>(
+        &self,
+        out: &mut W,
+        filters: &[Box<dyn KmerFilter>],
+        on_partition: F,
+    ) -> OKIResult<usize> {
+        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<GraphDeBruijn> {
+                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)
+    }
+}

src/obikmer2/src/cmd/mod.rs

@@ -7,3 +7,4 @@ pub mod merge;
 mod predicate;
 pub mod select;
 pub mod superkmer;
+pub mod unitig;

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<Box<dyn KmerFilter>> = 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");
+}

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),
     }