Merge pull request #37 from bacpop/i35_cli_update

CLI improvements

Author: johnlees

Cargo.toml

@@ -1,11 +1,12 @@
 [package]
 name = "sketchlib"
-version = "0.1.1"
+version = "0.1.2"
 authors = [
     "John Lees <[email protected]>",
     "Nicholas Croucher <[email protected]>",
     "Johanna von Wachsmann <[email protected]>",
     "Victor Rodriguez Bouza <[email protected]>",
+    "Joel Hellewell <[email protected]",
 ]
 edition = "2021"
 description = "Genome and amino-acid sketching"

README.md

@@ -22,9 +22,11 @@ With all options we typically recommend using `-v` to see all progress during th
 
 Using input fasta/fastq files, create a sketch database. Run `sketchlib sketch -h` to see the help.
 
-- List .fasta files on the command line, or use `-f` to provide a file(s). From file,
-these are one line per sample listing the name and fasta file, or name and two read files
-(fastq). Inputs can be gzipped or not, this is automatically detected.
+- List .fasta files on the command line, or use `-f` to provide a file(s). Inputs can be gzipped or not, this is automatically detected.
+From file, these are one line per sample listing:
+    - One column (fasta input): file name, which is also used as the sample name
+    - Two columns (fasta input): sample name and file name
+    - Three columns (fastq input): sample name and two read files
 - To set the k-mer size in the sketch database you can either give a list of sizes with `--k-vals`
 or a sequence `--k-seq` with start,stop,step. e.g. `--k-seq 17,29,4` would sketch at k=17, 21, 25 and 29.
 - Set the sketch size with `-s`. Typically 1000 is enough for species level resolution, 10000 for within-species/strain

src/cli.rs

@@ -1,5 +1,5 @@
 //! Command line interface, built using [`crate::clap` with `Derive`](https://docs.rs/clap/latest/clap/_derive/_tutorial/index.html)
-use clap::{ArgGroup, Parser, Subcommand};
+use clap::{ArgGroup, Args, Parser, Subcommand};
 
 use super::hashing::{AaLevel, HashType, DEFAULT_LEVEL};
 
@@ -40,14 +40,30 @@ pub fn check_threads(threads: usize) {
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
 #[command(propagate_version = true)]
-pub struct Args {
+pub struct MainArgs {
     #[doc(hidden)]
     #[command(subcommand)]
     pub command: Commands,
 
     /// Show progress messages
     #[arg(short, long, global = true)]
     pub verbose: bool,
+
+    /// Don't show any messages
+    #[arg(long, global = true)]
+    pub quiet: bool,
+}
+
+#[derive(Args)]
+#[group(required = true, multiple = false)]
+pub struct Kmers {
+    /// K-mer list (comma separated k-mer values to sketch at)
+    #[arg(short, long, required = true, value_delimiter = ',')]
+    pub k_vals: Option<Vec<usize>>,
+
+    /// K-mer linear sequence (start,end,step)
+    #[arg(long, required = true, value_delimiter = ',')]
+    pub k_seq: Option<Vec<usize>>,
 }
 
 /// Subcommands and their specific options
@@ -61,10 +77,10 @@ pub enum Commands {
     /// Create sketches from input data
     Sketch {
         /// List of input FASTA files
-        #[arg(long, group = "input", num_args = 1.., value_delimiter = ',')]
+        #[arg(group = "input")]
         seq_files: Option<Vec<String>>,
 
-        /// File listing input files (tab separated name, sequences)
+        /// File listing input files (tab separated name, sequences, see README)
         #[arg(short, group = "input")]
         file_list: Option<String>,
 
@@ -82,13 +98,8 @@ pub enum Commands {
         #[arg(short)]
         output: String,
 
-        /// K-mer list
-        #[arg(short, long, group = "kmer", required = true, num_args = 1.., value_delimiter = ',')]
-        k_vals: Option<Vec<usize>>,
-
-        /// K-mer sequence: start end step
-        #[arg(long, group = "kmer", required = true, num_args = 3)]
-        k_seq: Option<Vec<usize>>,
+        #[command(flatten)]
+        kmers: Kmers,
 
         /// Sketch size
         #[arg(short, long, default_value_t = DEFAULT_SKETCHSIZE)]
@@ -173,7 +184,7 @@ pub enum Commands {
         db: String,
 
         /// List of input FASTA files
-        #[arg(long, group = "input", num_args = 1.., value_delimiter = ',')]
+        #[arg(group = "input")]
         seq_files: Option<Vec<String>>,
 
         /// File listing input files (tab separated name, sequences)
@@ -236,7 +247,7 @@ pub enum Commands {
     },
 }
 
-/// Function to parse command line args into [`Args`] struct
-pub fn cli_args() -> Args {
-    Args::parse()
+/// Function to parse command line args into [`MainArgs`] struct
+pub fn cli_args() -> MainArgs {
+    MainArgs::parse()
 }

src/hashing/aahash_iterator.rs

@@ -34,7 +34,10 @@ impl RollHash for AaHashIterator {
             self.fh = new_it.0;
             self.index = new_it.1;
         } else {
-            panic!("K-mer larger than smallest valid sequence, which is:\n{}", std::str::from_utf8(&self.seq).unwrap());
+            panic!(
+                "K-mer larger than smallest valid sequence, which is:\n{}",
+                std::str::from_utf8(&self.seq).unwrap()
+            );
         }
     }
 

src/io.rs

@@ -1,4 +1,6 @@
 //! Functions to read input fasta/fastq files
+use crate::cli::Kmers;
+
 use std::fs::File;
 use std::io::{stdout, BufRead, BufReader, BufWriter, Write};
 use std::path::Path;
@@ -10,9 +12,9 @@ pub type InputFastx = (String, String, Option<String>);
 pub fn read_input_fastas(seq_files: &[String]) -> Vec<InputFastx> {
     let mut input_files = Vec::new();
     // matches the file name (no extension) in a full path
-    let re_path = Regex::new(r"^.+/(.+)\.(?i:fa|fasta|fastq|fastq\.gz)$").unwrap();
+    let re_path = Regex::new(r"^.+/(.+\.?i:fa|fasta|fastq|fastq\.gz)$").unwrap();
     // matches the file name (no extension) with no path
-    let re_name = Regex::new(r"^(.+)\.(?i:fa|fasta|fastq|fastq\.gz)$").unwrap();
+    let re_name = Regex::new(r"^(.+\.?i:fa|fasta|fastq|fastq\.gz)$").unwrap();
     for file in seq_files {
         let caps = re_path.captures(file).or(re_name.captures(file));
         let name = match caps {
@@ -24,14 +26,14 @@ pub fn read_input_fastas(seq_files: &[String]) -> Vec<InputFastx> {
     input_files
 }
 
-pub fn parse_kmers(k_list: &Option<Vec<usize>>, k_seq: &Option<Vec<usize>>) -> Vec<usize> {
-    if k_list.is_some() && k_seq.is_some() {
+pub fn parse_kmers(k: &Kmers) -> Vec<usize> {
+    if k.k_vals.is_some() && k.k_seq.is_some() {
         panic!("Only one of --k-vals or --k-seq should be specified");
     }
 
-    let mut kmers = if let Some(k) = k_list {
+    let mut kmers = if let Some(k) = &k.k_vals {
         k.clone().to_vec()
-    } else if let Some(k) = k_seq {
+    } else if let Some(k) = &k.k_seq {
         (k[0]..=k[1]).step_by(k[2]).collect()
     } else {
         panic!("Must specify --k-vals or --k-seq");
@@ -79,18 +81,22 @@ pub fn get_input_list(
             for line in f.lines() {
                 let line = line.expect("Unable to read line in file_list");
                 let fields: Vec<&str> = line.split_whitespace().collect();
-                // Should be 2 entries for fasta, 3 for fastq
-                let second_file = match fields.len() {
-                    0..=1 => {
-                        panic!("Unable to parse line in file_list")
-                    }
-                    2 => None,
-                    3 => Some(fields[2].to_string()),
+                // 1 entry: fasta with name = file
+                // 2 entries: fasta name, file
+                // 3 entries: fastq name, file1, file2
+                let parsed_input = match fields.len() {
+                    1 => ((fields[0].to_string()), fields[0].to_string(), None),
+                    2 => ((fields[0].to_string()), fields[1].to_string(), None),
+                    3 => (
+                        (fields[0].to_string()),
+                        fields[0].to_string(),
+                        Some(fields[2].to_string()),
+                    ),
                     _ => {
                         panic!("Unable to parse line in file_list")
                     }
                 };
-                input_files.push((fields[0].to_string(), fields[1].to_string(), second_file));
+                input_files.push(parsed_input);
             }
             input_files
         }

src/lib.rs

@@ -16,7 +16,7 @@ use std::time::Instant;
 extern crate arrayref;
 extern crate num_cpus;
 use anyhow::Error;
-use indicatif::{ParallelProgressIterator, ProgressStyle};
+use indicatif::ParallelProgressIterator;
 use rayon::prelude::*;
 
 pub mod cli;
@@ -45,6 +45,8 @@ pub mod bloom_filter;
 pub mod hashing;
 
 pub mod utils;
+use crate::utils::get_progress_bar;
+
 use std::fs::{File, OpenOptions};
 use std::io::copy;
 
@@ -59,7 +61,9 @@ pub const CHUNK_SIZE: usize = 1000;
 #[doc(hidden)]
 pub fn main() -> Result<(), Error> {
     let args = cli_args();
-    if args.verbose {
+    if args.quiet {
+        simple_logger::init_with_level(log::Level::Error).unwrap();
+    } else if args.verbose {
         simple_logger::init_with_level(log::Level::Info).unwrap();
         // simple_logger::init_with_level(log::Level::Trace).unwrap();
     } else {
@@ -76,8 +80,7 @@ pub fn main() -> Result<(), Error> {
             #[cfg(feature = "3di")]
             convert_pdb,
             output,
-            k_vals,
-            k_seq,
+            kmers,
             mut sketch_size,
             seq_type,
             level,
@@ -97,7 +100,7 @@ pub fn main() -> Result<(), Error> {
             log::info!("Getting input files");
             let input_files = get_input_list(file_list, seq_files);
             log::info!("Parsed {} samples in input list", input_files.len());
-            let kmers = parse_kmers(k_vals, k_seq);
+            let kmers = parse_kmers(kmers);
             // Build, merge
             let rc = !*single_strand;
             // Set expected sketchsize
@@ -128,6 +131,7 @@ pub fn main() -> Result<(), Error> {
                 rc,
                 *min_count,
                 *min_qual,
+                args.quiet,
             );
             let sketch_vec = MultiSketch::new(&mut sketches, sketch_size, &kmers, seq_type);
             sketch_vec
@@ -196,8 +200,8 @@ pub fn main() -> Result<(), Error> {
             };
             log::info!("{dist_type}");
 
-            let bar_style =
-                ProgressStyle::with_template("{percent}% {bar:80.cyan/blue} eta:{eta}").unwrap();
+            let percent = true; // In progress bar, don't show total number as huge
+
             // TODO: possible improvement would be to load sketch slices when i, j change
             // This would require a change to core_acc where multiple k-mer lengths are loaded at once
             // Overall this would be nicer I think (not sure about speed)
@@ -210,10 +214,11 @@ pub fn main() -> Result<(), Error> {
                             log::info!("Calculating all ref vs ref distances");
                             let mut distances = DistanceMatrix::new(&references, None, dist_type);
                             let par_chunk = CHUNK_SIZE * distances.n_dist_cols();
+                            let progress_bar = get_progress_bar(par_chunk, percent, args.quiet);
                             distances
                                 .dists_mut()
                                 .par_chunks_mut(par_chunk)
-                                .progress_with_style(bar_style)
+                                .progress_with(progress_bar)
                                 .enumerate()
                                 .for_each(|(chunk_idx, dist_slice)| {
                                     // Get first i, j index for the chunk
@@ -228,7 +233,11 @@ pub fn main() -> Result<(), Error> {
                                                 references.get_sketch_slice(j, k),
                                                 references.sketch_size,
                                             );
-                                            dist = if *ani { ani_pois(dist, k_f32) } else { 1.0_f32 - dist };
+                                            dist = if *ani {
+                                                ani_pois(dist, k_f32)
+                                            } else {
+                                                1.0_f32 - dist
+                                            };
                                             dist_slice[dist_idx] = dist;
                                         } else {
                                             let dist =
@@ -259,13 +268,14 @@ pub fn main() -> Result<(), Error> {
                             log::info!("Calculating sparse ref vs ref distances with {nn} nearest neighbours");
                             let mut sp_distances =
                                 SparseDistanceMatrix::new(&references, nn, dist_type);
+                            let progress_bar = get_progress_bar(nn, percent, args.quiet);
                             // TODO is it possible to add a template to the trait so this code is only written once? Maybe not
                             match sp_distances.dists_mut() {
                                 DistVec::Jaccard(distances) => {
                                     let k = k_idx.unwrap();
                                     distances
                                         .par_chunks_mut(nn)
-                                        .progress_with_style(bar_style)
+                                        .progress_with(progress_bar)
                                         .enumerate()
                                         .for_each(|(i, row_dist_slice)| {
                                             let mut heap = BinaryHeap::with_capacity(nn);
@@ -279,8 +289,11 @@ pub fn main() -> Result<(), Error> {
                                                     references.get_sketch_slice(j, k),
                                                     references.sketch_size,
                                                 );
-                                                dist =
-                                                    if *ani { ani_pois(dist, k_f32) } else { 1.0_f32 - dist };
+                                                dist = if *ani {
+                                                    ani_pois(dist, k_f32)
+                                                } else {
+                                                    1.0_f32 - dist
+                                                };
                                                 let dist_item = SparseJaccard(j, dist);
                                                 if heap.len() < nn
                                                     || dist_item < *heap.peek().unwrap()
@@ -299,7 +312,7 @@ pub fn main() -> Result<(), Error> {
                                 DistVec::CoreAcc(distances) => {
                                     distances
                                         .par_chunks_mut(nn)
-                                        .progress_with_style(bar_style)
+                                        .progress_with(progress_bar)
                                         .enumerate()
                                         .for_each(|(i, row_dist_slice)| {
                                             let mut heap = BinaryHeap::with_capacity(nn);
@@ -339,10 +352,11 @@ pub fn main() -> Result<(), Error> {
                         DistanceMatrix::new(&references, Some(&query_db), dist_type);
                     let par_chunk = CHUNK_SIZE * distances.n_dist_cols();
                     let nq = query_db.number_samples_loaded();
+                    let progress_bar = get_progress_bar(par_chunk, percent, args.quiet);
                     distances
                         .dists_mut()
                         .par_chunks_mut(par_chunk)
-                        .progress_with_style(bar_style)
+                        .progress_with(progress_bar)
                         .enumerate()
                         .for_each(|(chunk_idx, dist_slice)| {
                             // Get first i, j index for the chunk
@@ -355,7 +369,11 @@ pub fn main() -> Result<(), Error> {
                                         query_db.get_sketch_slice(j, k),
                                         references.sketch_size,
                                     );
-                                    dist = if *ani { ani_pois(dist, k_f32) } else { 1.0_f32 - dist };
+                                    dist = if *ani {
+                                        ani_pois(dist, k_f32)
+                                    } else {
+                                        1.0_f32 - dist
+                                    };
                                     dist_slice[dist_idx] = dist;
                                 } else {
                                     let dist = core_acc_dist(&references, &query_db, i, j);
@@ -474,6 +492,7 @@ pub fn main() -> Result<(), Error> {
                 rc,
                 *min_count,
                 *min_qual,
+                args.quiet,
             );
             let mut db2_metadata =
                 MultiSketch::new(&mut db2_sketches, sketch_size, kmers, seq_type);
@@ -554,7 +573,7 @@ pub fn main() -> Result<(), Error> {
     let end = Instant::now();
 
     log::info!("Complete");
-    if print_success {
+    if print_success && !args.quiet {
         eprintln!(
             "🧬🖋️ sketchlib done in {}s",
             end.duration_since(start).as_secs()