diff --git a/src/Cargo.lock b/src/Cargo.lock index eeb66157..c08ff779 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1654,12 +1654,14 @@ name = "obikmer2" version = "1.2.2" dependencies = [ "clap", + "obifastwrite", "obikalgorithm", "obikindex", "obikindexer", "obikseq", "obipipeline", "obiread", + "obiskbuilder", "obisys", "tracing", "tracing-subscriber", diff --git a/src/obikmer2/Cargo.toml b/src/obikmer2/Cargo.toml index f332c200..8937e705 100644 --- a/src/obikmer2/Cargo.toml +++ b/src/obikmer2/Cargo.toml @@ -15,6 +15,8 @@ obisys = { path = "../obisys" } obikindex = { path = "../obikindex", default-features = false } obikindexer = { path = "../obikindexer" } obikalgorithm = { path = "../obikalgorithm" } +obifastwrite = { path = "../obifastwrite" } +obiskbuilder = { path = "../obiskbuilder" } clap = { version = "4", features = ["derive"] } tracing = "0.1.44" tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } diff --git a/src/obikmer2/src/cmd/mod.rs b/src/obikmer2/src/cmd/mod.rs index 33edc959..2547bc25 100644 --- a/src/obikmer2/src/cmd/mod.rs +++ b/src/obikmer2/src/cmd/mod.rs @@ -1 +1,2 @@ pub mod index; +pub mod superkmer; diff --git a/src/obikmer2/src/cmd/superkmer/mod.rs b/src/obikmer2/src/cmd/superkmer/mod.rs new file mode 100644 index 00000000..707651f9 --- /dev/null +++ b/src/obikmer2/src/cmd/superkmer/mod.rs @@ -0,0 +1,72 @@ +use std::io::{self, BufWriter, Write}; +use std::path::PathBuf; + +use clap::Args; +use obifastwrite::write_scatter; +use obikseq::{RoutableSuperKmer, set_k, set_m}; + +use obipipeline::{Throttled, throttle}; + +use crate::cli::{CommonArgs, PipelineData, partitions_to_bits}; + +#[derive(Args)] +pub struct SuperkmerArgs { + #[command(flatten)] + pub common: CommonArgs, +} + +// ── Stage functions ─────────────────────────────────────────────────────────── + +fn write_batch( + batch: Vec, + out: &mut BufWriter, + partition_bits: usize, + k: usize, + m: usize, +) -> io::Result<()> { + let partition_mask = (1u64 << partition_bits) - 1; + for rsk in batch { + let minimizer = *rsk.minimizer(); + let partition = (minimizer.seq_hash() & partition_mask) as usize; + write_scatter(rsk.superkmer(), out, k, m, partition, minimizer)?; + } + Ok(()) +} + +// ── Entry point ─────────────────────────────────────────────────────────────── + +pub fn run(args: SuperkmerArgs) { + args.common.validate(); + + let k = args.common.kmer_size; + let m = args.common.minimizer_size; + let theta = args.common.theta; + let level_max = args.common.level_max; + let partition_bits = partitions_to_bits(args.common.partitions); + let n_workers = args.common.threads.max(1); + let max_open = args.common.effective_max_open(); + + set_k(k); + set_m(m); + + let path_source = throttle(args.common.seqfile_paths(), max_open); + + let pipe = obipipeline::make_pipe! { + PipelineData : Throttled => Vec, + ||? { + let k = k; + move |pw: Throttled| { + let path_str = pw.item.to_str().unwrap_or("").to_owned(); + let _guard = pw.guard; + obiread::open_nuc_stream(&path_str, k) + } + } : Path => NucPage, + | { move |page| obiskbuilder::build_superkmers_page(page, k, level_max, theta) } : NucPage => Batch, + }; + + let mut out = BufWriter::new(io::stdout()); + for batch in pipe.apply(path_source, n_workers, 1) { + write_batch(batch, &mut out, partition_bits, k, m).expect("write error"); + } + out.flush().expect("flush error"); +} diff --git a/src/obikmer2/src/main.rs b/src/obikmer2/src/main.rs index a490edd1..9ec8afb8 100644 --- a/src/obikmer2/src/main.rs +++ b/src/obikmer2/src/main.rs @@ -15,6 +15,8 @@ struct Cli { enum Commands { /// Build the complete genome index (scatter → dereplicate → count → layered MPHF) Index(cmd::index::IndexArgs), + /// Extract super-k-mers from input sequences and scatter them by partition + Superkmer(cmd::superkmer::SuperkmerArgs), } fn main() { @@ -28,5 +30,6 @@ fn main() { let cli = Cli::parse(); match cli.command { Commands::Index(args) => cmd::index::run(args), + Commands::Superkmer(args) => cmd::superkmer::run(args), } }