WIP: gather directly from fasta

pyproject.toml

@@ -2,6 +2,7 @@
 name = "sourmash_plugin_branchwater"
 description = "fast command-line extensions for sourmash"
 readme = "README.md"
+version = "0.9.12"
 requires-python = ">=3.10"
 classifiers = [
     "Programming Language :: Rust",

src/fastagather.rs

@@ -0,0 +1,138 @@
+use crate::utils::buildutils::BuildCollection;
+use anyhow::{bail, Result};
+
+use needletail::parse_fastx_file;
+use sourmash::selection::Selection;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use crate::utils::{
+    consume_query_by_gather, load_collection, load_sketches_above_threshold, write_prefetch,
+    ReportType,
+};
+
+#[allow(clippy::too_many_arguments)]
+pub fn fastagather(
+    query_filename: String,
+    against_filepath: String,
+    input_moltype: String,
+    threshold_bp: u64,
+    selection: &Selection,
+    prefetch_output: Option<String>,
+    gather_output: Option<String>,
+    allow_failed_sigpaths: bool,
+) -> Result<()> {
+    // to start, implement straightforward record --> sketch --> gather
+    // other ideas:
+    // - add full-file (lower resolution) prefetch first, to reduce search space
+    // - parallelize and/or batch records
+
+    // Build signature templates based on parsed parameters
+    let sig_template_result = BuildCollection::from_selection(selection);
+    let sig_template = match sig_template_result {
+        Ok(sig_template) => sig_template,
+        Err(e) => {
+            bail!("Failed to build template signatures: {}", e);
+        }
+    };
+
+    if sig_template.size() != 1 {
+        bail!("FASTAgather requires a single signature type for search.");
+    }
+
+    let input_moltype = input_moltype.to_ascii_lowercase();
+
+    let against_selection = selection;
+    // get scaled from selection here
+    let scaled = selection.scaled().unwrap(); // rm this unwrap?
+                                              //against_selection.set_scaled(scaled as u32);
+
+    // calculate the minimum number of hashes based on desired threshold
+    let threshold_hashes = {
+        let x = threshold_bp / scaled as u64;
+        if x > 0 {
+            x
+        } else {
+            1
+        }
+    };
+
+    // load collection to match against.
+    let against_collection = load_collection(
+        &against_filepath,
+        &against_selection,
+        ReportType::Against,
+        allow_failed_sigpaths,
+    )?;
+
+    let failed_records = AtomicUsize::new(0);
+    // open file and start iterating through sequences
+    // Open fasta file reader
+    let mut reader = match parse_fastx_file(query_filename.clone()) {
+        Ok(r) => r,
+        Err(err) => {
+            bail!("Error opening file {}: {:?}", query_filename, err);
+        }
+    };
+
+    // later: can we parallelize across records or sigs? Do we want to batch groups of records for improved gather efficiency?
+    while let Some(record_result) = reader.next() {
+        // clone sig_templates for use
+        let mut sigcoll = sig_template.clone();
+        match record_result {
+            Ok(record) => {
+                let query_name = std::str::from_utf8(record.id())
+                    .expect("record.id() contains invalid UTF-8")
+                    .to_string();
+                //let query_name = record.id().to_string();   //query_sig.name(); // this is actually just record.id --> so maybe don't get it from sig here?
+                if let Err(err) =
+                    sigcoll.build_singleton_sigs(record, &input_moltype, query_filename.clone())
+                {
+                    eprintln!(
+                        "Error building signatures from file: {}, {:?}",
+                        query_filename, err
+                    );
+                    failed_records.fetch_add(1, Ordering::SeqCst);
+                }
+                // in each iteration, this should just be a single signature made from the single record
+                for query_sig in sigcoll.sigs.iter() {
+                    let query_md5 = query_sig.md5sum();
+                    let query_mh = query_sig.minhash().expect("could not get minhash from sig");
+
+                    // now do prefetch/gather
+                    let prefetch_result = load_sketches_above_threshold(
+                        against_collection.clone(), // need to get rid of this clone
+                        &query_mh,
+                        threshold_hashes,
+                    )?;
+                    let matchlist = prefetch_result.0;
+                    let _skipped_paths = prefetch_result.1;
+                    let _failed_paths = prefetch_result.2;
+
+                    if prefetch_output.is_some() {
+                        write_prefetch(
+                            query_filename.clone(),
+                            query_name.clone(),
+                            query_md5,
+                            prefetch_output.clone(),
+                            &matchlist,
+                        )
+                        .ok();
+                    }
+
+                    consume_query_by_gather(
+                        query_name.clone(),
+                        query_filename.clone(),
+                        query_mh.clone(),
+                        scaled as u32,
+                        matchlist,
+                        threshold_hashes,
+                        gather_output.clone(),
+                    )
+                    .ok();
+                }
+            }
+            Err(err) => eprintln!("Error while processing record: {:?}", err),
+        }
+    }
+    Ok(())
+}

src/lib.rs

@@ -13,6 +13,7 @@ use crate::utils::build_selection;
 use crate::utils::is_revindex_database;
 mod check;
 mod cluster;
+mod fastagather;
 mod fastgather;
 mod fastmultigather;
 mod fastmultigather_rocksdb;
@@ -367,6 +368,41 @@ fn do_cluster(
     }
 }
 
+#[pyfunction]
+#[allow(clippy::too_many_arguments)]
+#[pyo3(signature = (query_filename, against_filepath, input_moltype, threshold_bp, ksize, scaled, moltype, output_path_prefetch=None, output_path_gather=None))]
+fn do_fastagather(
+    query_filename: String,
+    against_filepath: String,
+    input_moltype: String,
+    threshold_bp: u64,
+    ksize: u8,
+    scaled: Option<u32>,
+    moltype: String,
+    output_path_prefetch: Option<String>,
+    output_path_gather: Option<String>,
+) -> anyhow::Result<u8> {
+    let selection = build_selection(ksize, scaled, &moltype);
+    let allow_failed_sigpaths = true;
+
+    match fastagather::fastagather(
+        query_filename,
+        against_filepath,
+        input_moltype,
+        threshold_bp,
+        &selection,
+        output_path_prefetch,
+        output_path_gather,
+        allow_failed_sigpaths,
+    ) {
+        Ok(_) => Ok(0),
+        Err(e) => {
+            eprintln!("Error: {e}");
+            Ok(1)
+        }
+    }
+}
+
 /// Module interface for the `sourmash_plugin_branchwater` extension module.
 
 #[pymodule]
@@ -382,5 +418,6 @@ fn sourmash_plugin_branchwater(_py: Python, m: &Bound<'_, PyModule>) -> PyResult
     m.add_function(wrap_pyfunction!(do_pairwise, m)?)?;
     m.add_function(wrap_pyfunction!(do_cluster, m)?)?;
     m.add_function(wrap_pyfunction!(do_singlesketch, m)?)?;
+    m.add_function(wrap_pyfunction!(do_fastagather, m)?)?;
     Ok(())
 }

src/python/sourmash_plugin_branchwater/__init__.py

@@ -822,3 +822,94 @@ def main(self, args):
             notify(f"...clustering is done! results in '{args.output}'")
             notify(f"                       cluster counts in '{args.cluster_sizes}'")
         return status
+
+
+class Branchwater_Fastagather(CommandLinePlugin):
+    command = "fastagather"
+    description = "massively parallel gather directly from FASTA"
+
+    def __init__(self, p):
+        super().__init__(p)
+        p.add_argument("query_fa", help="FASTA file")
+        p.add_argument("against_paths", help="database file of sketches")
+        p.add_argument(
+            "-I",
+            "--input-moltype",
+            "--input-molecule-type",
+            choices=["DNA", "dna", "protein"],
+            default="DNA",
+            help="molecule type of input sequence (DNA or protein)",
+        )
+        p.add_argument(
+            "-o",
+            "--output-gather",
+            required=True,
+            help="save gather output (minimum metagenome cover) to this file",
+        )
+        p.add_argument(
+            "--output-prefetch", help="save prefetch output (all overlaps) to this file"
+        )
+        p.add_argument(
+            "-t",
+            "--threshold-bp",
+            default=4000,
+            type=float,
+            help="threshold in estimated base pairs, for reporting matches (default: 4kb)",
+        )
+        p.add_argument(
+            "-k",
+            "--ksize",
+            default=31,
+            type=int,
+            help="k-mer size at which to do comparisons (default: 31)",
+        )
+        p.add_argument(
+            "-s",
+            "--scaled",
+            default=1000,
+            type=int,
+            help="scaled factor at which to do comparisons (default: 1000)",
+        )
+        p.add_argument(
+            "-m",
+            "--moltype",
+            default="DNA",
+            choices=["DNA", "protein", "dayhoff", "hp"],
+            help="molecule type for search (DNA, protein, dayhoff, or hp; default DNA)",
+        )
+        p.add_argument(
+            "-c",
+            "--cores",
+            default=0,
+            type=int,
+            help="number of cores to use (default is all available)",
+        )
+
+    def main(self, args):
+        print_version()
+        notify(
+            f"ksize: {args.ksize} / scaled: {args.scaled} / moltype: {args.moltype} / threshold bp: {args.threshold_bp}"
+        )
+
+        num_threads = set_thread_pool(args.cores)
+
+        notify(
+            f"gathering all sketches in '{args.query_sig}' against '{args.against_paths}' using {num_threads} threads"
+        )
+        super().main(args)
+        status = sourmash_plugin_branchwater.do_fastagather(
+            args.query_fa,
+            args.against_paths,
+            args.input_moltype,
+            int(args.threshold_bp),
+            args.ksize,
+            args.scaled,
+            args.moltype,
+            args.output_gather,
+            args.output_prefetch,
+        )
+        if status == 0:
+            notify(f"...fastgather is done! gather results in '{args.output_gather}'")
+            if args.output_prefetch:
+                notify(f"prefetch results in '{args.output_prefetch}'")
+        return status

src/python/tests/sourmash_tst_utils.py

@@ -73,7 +73,6 @@ def _runscript(scriptname):
     smash_cli()
     return 0
 
-
 ScriptResults = collections.namedtuple("ScriptResults", ["status", "out", "err"])