bug fix

Author: dagou

kr2r/src/args.rs

@@ -61,7 +61,7 @@ pub struct ClassifyArgs {
 
     /// File path for outputting normal Kraken output.
     #[clap(long = "output-dir", value_parser)]
-    pub kraken_output_dir: Option<PathBuf>,
+    pub output_dir: Option<PathBuf>,
 
     /// Enable paired-end processing.
     #[clap(short = 'P', long = "paired-end-processing", action)]

kr2r/src/bin/direct.rs

@@ -28,7 +28,7 @@ pub struct Args {
 
     /// File path for outputting normal Kraken output.
     #[clap(long = "output-dir", value_parser)]
-    pub kraken_output_dir: Option<PathBuf>,
+    pub output_dir: Option<PathBuf>,
 
     /// Enable paired-end processing.
     #[clap(short = 'P', long = "paired-end-processing", action)]
@@ -161,7 +161,7 @@ fn process_fastx_file<R>(
 where
     R: Reader,
 {
-    let mut writer: Box<dyn Write + Send> = match &args.kraken_output_dir {
+    let mut writer: Box<dyn Write + Send> = match &args.output_dir {
         Some(ref file_path) => {
             let filename = file_path.join(format!("output_{}.txt", file_index));
             let file = File::create(filename)?;
@@ -225,7 +225,7 @@ where
 
     let thread_sequences = seq_counter.load(Ordering::SeqCst);
     let thread_classified = classify_counter.load(Ordering::SeqCst);
-    if let Some(output) = &args.kraken_output_dir {
+    if let Some(output) = &args.output_dir {
         let filename = output.join(format!("output_{}.kreport2", file_index));
         report_kraken_style(
             filename,
@@ -248,7 +248,7 @@ fn process_files(
     chtable: &CHTable,
     taxonomy: &Taxonomy,
 ) -> Result<()> {
-    let (mut file_index, mut file_writer) = if let Some(out_dir) = &args.kraken_output_dir {
+    let (mut file_index, mut file_writer) = if let Some(out_dir) = &args.output_dir {
         let file_path = out_dir.join("sample_file.map");
         let file_writer = create_sample_file(&file_path);
         let file_index = get_lastest_file_index(&file_path)?;
@@ -295,7 +295,7 @@ fn process_files(
             total_seqs += thread_sequences;
             total_unclassified += thread_unclassified;
         }
-        if let Some(output) = &args.kraken_output_dir {
+        if let Some(output) = &args.output_dir {
             let filename = output.join("output.kreport2");
             report_kraken_style(
                 filename,

kr2r/src/bin/kun.rs

@@ -92,9 +92,10 @@ impl From<ClassifyArgs> for resolve::Args {
         Self {
             database: item.database,
             chunk_dir: item.chunk_dir,
+            num_threads: item.num_threads,
             confidence_threshold: item.confidence_threshold,
             minimum_hit_groups: item.minimum_hit_groups,
-            kraken_output_dir: item.kraken_output_dir,
+            output_dir: item.output_dir,
             report_kmer_data: item.report_kmer_data,
             report_zero_counts: item.report_zero_counts,
         }

kr2r/src/bin/resolve.rs

@@ -9,7 +9,7 @@ use kr2r::HitGroup;
 // use rayon::prelude::*;
 use seqkmer::{buffer_map_parallel, trim_pair_info, OptionPair};
 use std::collections::HashMap;
-use std::fs::File;
+use std::fs::{create_dir_all, File};
 use std::io::{self, BufRead, BufReader, BufWriter, Read, Result, Write};
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -63,7 +63,11 @@ pub struct Args {
 
     /// File path for outputting normal Kraken output.
     #[clap(long = "output-dir", value_parser)]
-    pub kraken_output_dir: Option<PathBuf>,
+    pub output_dir: Option<PathBuf>,
+
+    /// The number of threads to use.
+    #[clap(short = 'p', long = "num-threads", value_parser, default_value_t = num_cpus::get())]
+    pub num_threads: usize,
 
     // /// output file contains all unclassified sequence
     // #[clap(long, value_parser, default_value_t = false)]
@@ -128,7 +132,7 @@ fn process_batch<P: AsRef<Path>>(
 
         buffer_map_parallel(
             &hit_counts,
-            num_cpus::get(),
+            args.num_threads,
             |(k, rows)| {
                 if let Some(item) = id_map.get(&k) {
                     let mut rows = rows.to_owned();
@@ -168,7 +172,9 @@ fn process_batch<P: AsRef<Path>>(
             },
             |result| {
                 while let Some(Some(res)) = result.next() {
-                    writer.write_all(res.as_bytes()).unwrap();
+                    writer
+                        .write_all(res.as_bytes())
+                        .expect("write output content error");
                 }
             },
         )
@@ -194,6 +200,10 @@ pub fn run(args: Args) -> Result<()> {
     let mut total_seqs = 0;
     let mut total_unclassified = 0;
 
+    if let Some(output) = &args.output_dir {
+        create_dir_all(output)?;
+    }
+
     // 开始计时
     let start = Instant::now();
     println!("resolve start...");
@@ -202,7 +212,7 @@ pub fn run(args: Args) -> Result<()> {
         let sample_id_map = read_id_to_seq_map(&sample_id_files[i])?;
 
         let thread_sequences = sample_id_map.len();
-        let mut writer: Box<dyn Write + Send> = match &args.kraken_output_dir {
+        let mut writer: Box<dyn Write + Send> = match &args.output_dir {
             Some(ref file_path) => {
                 let filename = file_path.join(format!("output_{}.txt", i));
                 let file = File::create(filename)?;
@@ -235,7 +245,7 @@ pub fn run(args: Args) -> Result<()> {
                 .merge(&entry.value())
                 .unwrap();
         });
-        if let Some(output) = &args.kraken_output_dir {
+        if let Some(output) = &args.output_dir {
             let filename = output.join(format!("output_{}.kreport2", i));
             report_kraken_style(
                 filename,
@@ -252,7 +262,7 @@ pub fn run(args: Args) -> Result<()> {
         total_unclassified += thread_sequences - thread_classified;
     }
 
-    if let Some(output) = &args.kraken_output_dir {
+    if let Some(output) = &args.output_dir {
         if !sample_files.is_empty() {
             let min = &sample_files.keys().min().cloned().unwrap();
             let max = &sample_files.keys().max().cloned().unwrap();