allow checkm2 running in cluster mode

Author: AroneyS

src/cluster_argument_parsing.rs

@@ -11,6 +11,7 @@ use crate::skani::SkaniClusterer;
 use crate::skani::SkaniPreclusterer;
 use crate::ClusterDistanceFinder;
 use crate::PreclusterDistanceFinder;
+use crate::QualityFinder;
 use crate::SortedPairGenomeDistanceCache;
 use bird_tool_utils::clap_utils::*;
 use bird_tool_utils::clap_utils::{default_roff, monospace_roff};
@@ -117,6 +118,8 @@ pub struct GalahClustererCommandDefinition {
     pub dereplication_ani_argument: String,
     pub dereplication_prethreshold_ani_argument: String,
     pub dereplication_quality_formula_argument: String,
+    pub dereplication_run_checkm2_argument: String,
+    pub dereplication_checkm2_db_path_argument: String,
     pub dereplication_precluster_method_argument: String,
     pub dereplication_cluster_method_argument: String,
     pub dereplication_aligned_fraction_argument: String,
@@ -140,6 +143,8 @@ lazy_static! {
             dereplication_ani_argument: "ani".to_string(),
             dereplication_prethreshold_ani_argument: "precluster-ani".to_string(),
             dereplication_quality_formula_argument: "quality-formula".to_string(),
+            dereplication_run_checkm2_argument: "run-checkm2".to_string(),
+            dereplication_checkm2_db_path_argument: "checkm2-db-path".to_string(),
             dereplication_precluster_method_argument: "precluster-method".to_string(),
             dereplication_cluster_method_argument: "cluster-method".to_string(),
             dereplication_aligned_fraction_argument: "min-aligned-fraction".to_string(),
@@ -284,6 +289,15 @@ pub fn add_dereplication_filtering_parameters_to_section(section: Section) -> Se
             "Ignore genomes with more contamination than \
         this percentage. [default: not set]",
         ))
+        .flag(Flag::new().long("--run-checkm2").help(
+            "Run CheckM2 to generate quality scoring used for clustering. Requires \
+                --checkm2-db-path or CHECKM2DB env variable to be set.",
+        ))
+        .option(
+            Opt::new("DB_PATH").long("--checkm2-db-path").help(
+                "Path to CheckM2 database (required for running CheckM2) \
+                [default: from CHECKM2DB environment variable]",
+        ))
 }
 
 pub fn add_dereplication_clustering_parameters_to_section(
@@ -838,9 +852,10 @@ pub fn filter_genomes_through_checkm<'a>(
     match clap_matches.contains_id("checkm-tab-table")
         || clap_matches.contains_id("genome-info")
         || clap_matches.contains_id("checkm2-quality-report")
+        || clap_matches.contains_id("run-checkm2")
     {
         false => {
-            warn!("Since CheckM input is missing, genomes are not being ordered by quality. Instead the order of their input is being used");
+            warn!("Since CheckM input has not been provided and CheckM2 has been disabled, genomes are not being ordered by quality. Instead the order of their input is being used");
             Ok(genome_fasta_files.iter().map(|s| &**s).collect())
         }
         true => {
@@ -882,6 +897,27 @@ pub fn filter_genomes_through_checkm<'a>(
                     )
                     .expect("Error parsing genomeInfo file"),
                 }
+            } else if clap_matches.contains_id("run-checkm2") {
+                // Run CheckM2 as in analyse
+                let db_path = clap_matches.get_one::<String>("checkm2-db-path")
+                    .map(|s| s.to_string())
+                    .or_else(|| std::env::var("CHECKM2DB").ok())
+                    .expect("CheckM2 database path must be provided via --checkm2-db-path or CHECKM2DB env var");
+                use crate::checkm2::CheckM2Analyser;
+                let tmpdir = tempfile::tempdir().expect("Failed to create tempdir for CheckM2");
+                let tmp_path = tmpdir.path();
+                let mut analyser = CheckM2Analyser::new(db_path);
+                analyser.prepare_comp_cont(genome_fasta_files, 1, tmp_path);
+                CheckMResultEnum::CheckM2Result {
+                    result: {
+                        let quality_report_path =
+                            tmp_path.join("checkm2").join("quality_report.tsv");
+                        checkm::CheckM2QualityReport::read_file_path(
+                            quality_report_path.to_str().unwrap(),
+                        )
+                        .unwrap()
+                    },
+                }
             } else {
                 panic!("Programming error");
             };
@@ -1568,6 +1604,13 @@ pub fn add_cluster_subcommand(app: clap::Command) -> clap::Command {
                 "Parks2020_reduced",
                 "dRep"])
             .default_value(crate::DEFAULT_QUALITY_FORMULA))
+        .arg(Arg::new(&*GALAH_COMMAND_DEFINITION.dereplication_run_checkm2_argument)
+            .long("run-checkm2")
+            .help("Run CheckM2 for genome quality scoring during clustering")
+            .action(clap::ArgAction::SetTrue))
+        .arg(Arg::new(&*GALAH_COMMAND_DEFINITION.dereplication_checkm2_db_path_argument)
+            .long("checkm2-db-path")
+            .help("Path to CheckM2 database. If not specified, will use $CHECKM2_DB_PATH environment variable if set."))
         .arg(Arg::new(&*GALAH_COMMAND_DEFINITION.dereplication_prethreshold_ani_argument)
             .long("precluster-ani")
             .value_parser(clap::value_parser!(f32))

tests/test_cmdline.rs

@@ -3,6 +3,10 @@ extern crate assert_cli;
 #[cfg(test)]
 mod tests {
     use assert_cli::Assert;
+    use std::env;
+    use std::fs;
+    use std::path::Path;
+    use tempfile::tempdir;
 
     #[test]
     fn test_completeness_4contamination_quality_score() {
@@ -986,4 +990,98 @@ mod tests {
                 tests/data/abisko4/73.20110600_S2D.10.fna	tests/data/abisko4/73.20110600_S2D.10.fna\n")
                 .unwrap();
     }
+
+    #[test]
+    #[ignore]
+    fn test_cluster_real_checkm2() {
+        let checkm2_db_path = std::env::var("CHECKM2DB")
+            .expect("CHECKM2DB environment variable must be set to run this test");
+        println!("Using CheckM2 database at {}", checkm2_db_path);
+        Assert::main_binary()
+            .with_args(&[
+                "cluster",
+                "--quality-formula",
+                "completeness-4contamination",
+                "--run-checkm2",
+                "--genome-fasta-files",
+                "tests/data/set1/1mbp.fna",
+                "tests/data/set1/500kb.fna",
+                "tests/data/abisko4/73.20120800_S1D.21.fna",
+                "tests/data/abisko4/73.20110800_S2M.16.fna",
+                "--precluster-method",
+                "finch",
+                "--output-cluster-definition",
+                "/dev/stdout",
+            ])
+            .succeeds()
+            .stdout()
+            .is("\
+            tests/data/abisko4/73.20110800_S2M.16.fna\ttests/data/abisko4/73.20110800_S2M.16.fna\n\
+            tests/data/abisko4/73.20110800_S2M.16.fna\ttests/data/abisko4/73.20120800_S1D.21.fna\n\
+            tests/data/set1/500kb.fna\ttests/data/set1/500kb.fna\n\
+            tests/data/set1/500kb.fna\ttests/data/set1/1mbp.fna\n")
+            .unwrap();
+    }
+
+    fn setup_mock_bin(dir: &Path, genomes: &[(String, f64, f64)]) {
+        // CheckM2 writes to quality_report.tsv within the directory specified by -o
+        let mut checkm2_script = String::from("#!/bin/bash\n");
+        checkm2_script.push_str("out=\"\"\n");
+        checkm2_script.push_str("while [[ $# -gt 0 ]]; do\n");
+        checkm2_script.push_str("  case $1 in\n");
+        checkm2_script.push_str("    -o) out=$2; shift 2;;\n");
+        checkm2_script.push_str("    *) shift;;\n");
+        checkm2_script.push_str("  esac\n");
+        checkm2_script.push_str("done\n");
+        checkm2_script.push_str("mkdir -p \"$out\"\n");
+        checkm2_script.push_str("echo -e 'Name\tCompleteness\tContamination\tCompleteness_Model_Used\tTranslation_Table_Used\tCoding_Density\tContig_N50\tAverage_Gene_Length\tGenome_Size\tGC_Content\tTotal_Coding_Sequences\tTotal_Contigs\tMax_Contig_Length\tAdditional_Notes' > \"$out/quality_report.tsv\"\n");
+        for (genome, completeness, contamination) in genomes {
+            checkm2_script.push_str(&format!(
+                "echo -e '{genome}\t{completeness}\t{contamination}\tGradient Boost (General Model)\t11\t0.885\t5745\t235.3609865470852\t355151\t0.33\t446\t75\t24150\tNone' >> \"$out/quality_report.tsv\"\n"
+            ));
+        }
+        let checkm2 = dir.join("checkm2");
+        fs::write(&checkm2, checkm2_script).unwrap();
+        let _ = std::process::Command::new("chmod")
+            .arg("+x")
+            .arg(&checkm2)
+            .status();
+    }
+
+    #[test]
+    fn test_cluster_mock_checkm2() {
+        let tmpdir = tempdir().unwrap();
+        setup_mock_bin(
+            tmpdir.path(),
+            &[
+                (String::from("500kb"), 90.0, 5.0),
+                (String::from("1mbp"), 95.0, 2.0),
+            ],
+        );
+        let path = env::var("PATH").unwrap();
+        let new_path = format!("{}:{}", tmpdir.path().display(), path);
+
+        // Provide genomes in reverse order to their quality
+        Assert::main_binary()
+            .with_env(&[("PATH", new_path)])
+            .with_args(&[
+                "cluster",
+                "--quality-formula",
+                "completeness-4contamination",
+                "--run-checkm2",
+                "--genome-fasta-files",
+                "tests/data/set1/500kb.fna",
+                "tests/data/set1/1mbp.fna",
+                "--output-cluster-definition",
+                "/dev/stdout",
+                "--checkm2-db-path",
+                "/tmp/mockdb",
+            ])
+            .succeeds()
+            .stdout()
+            .is("\
+            tests/data/set1/1mbp.fna\ttests/data/set1/1mbp.fna\n\
+            tests/data/set1/1mbp.fna\ttests/data/set1/500kb.fna\n")
+            .unwrap();
+    }
 }