diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5d88683..218c66d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,4 +21,4 @@ jobs: - run: rustup update ${{ matrix.toolchain }} && rustup default ${{ matrix.toolchain }} - run: cargo build --verbose - run: cargo test --verbose - - run: cargo build --verbose -F 3di + - run: cargo build --verbose -F python diff --git a/Cargo.toml b/Cargo.toml index 2f680aa..a67a9e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,6 @@ include = [ "/LICENSE", "/NOTICE", "/src", - "/python_mini3di", "/tests" ] keywords = ["bioinformatics", "genomics", "sequencing", "k-mer", "sketch"] @@ -29,7 +28,7 @@ categories = ["command-line-utilities", "science"] crate-type = ["cdylib", "lib"] [features] -3di = ["dep:pyo3", "dep:pyo3-build-config"] +python = ["dep:pyo3"] [target.'cfg(not(target_arch = "wasm32"))'.dependencies] # data structures @@ -48,7 +47,7 @@ memmap2 = "0.9" arrayref = "0.3" num-traits = "0.2" # ffi -pyo3 = { version = "0.21", features = ["auto-initialize"], optional = true } +pyo3 = { version = "0.21", features = ["abi3-py310"], optional = true } # logging log = "0.4" simple_logger = { version = "4", features = ["stderr"] } @@ -101,9 +100,6 @@ js-sys = {version = "0.3.51" } -[build-dependencies] -pyo3-build-config = {version = "0.22", optional = true} - [target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies] # testing snapbox = "0.6.21" diff --git a/maturin/pyproject.toml b/maturin/pyproject.toml new file mode 100644 index 0000000..d27f7fa --- /dev/null +++ b/maturin/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["maturin>=1.7,<2"] +build-backend = "maturin" + +[project] +name = "sketchlib" +version = "0.3.0" +description = "Python bindings for sketchlib.rust" +requires-python = ">=3.10" +license = "Apache-2.0" +authors = [ + { name = "John Lees", email = "jlees@ebi.ac.uk" }, +] + +[tool.maturin] +manifest-path = "../Cargo.toml" +features = ["python"] +module-name = "sketchlib" diff --git a/python_mini3di/3di_convert.py b/python_mini3di/3di_convert.py deleted file mode 100644 index c56cdcb..0000000 --- a/python_mini3di/3di_convert.py +++ /dev/null @@ -1,49 +0,0 @@ -import os, sys -import mini3di -from Bio.PDB import PDBParser -from warnings import warn - - -encoder = mini3di.Encoder() -parser = PDBParser(QUIET=True) -#parser = PDBParser() - -def pdb_to_3di(struct_name: str, filename: str) -> str: - struct = parser.get_structure(struct_name, filename) - - #print("filename: " + filename + " structname: " + struct_name) - - structure_string = "" - for chain in struct.get_chains(): - try: - states = encoder.encode_chain(chain) - sequence = encoder.build_sequence(states) - structure_string += f"{sequence}," - except IndexError: - warn("Not able to code into 3Di chain {} from protein ID {}".format(chain.__repr__(), struct_name), RuntimeWarning) - continue - - structure_string = structure_string.removesuffix(",") - - #print(structure_string) - - return structure_string - - -def main(): - thel = sys.argv - if len(thel) < 2: - raise RuntimeError("FATAL: no input list of files provided!") - - if not os.path.isfile(thel[1]): - raise RuntimeError("FATAL: not provided a valid file!!!") - - with open(thel[1], "r") as inf: - for line in inf: - tmpl = line.strip().split("\t") - pdb_to_3di(tmpl[0], tmpl[1]) - - -if __name__ == '__main__': - main() - diff --git a/src/api.rs b/src/api.rs new file mode 100644 index 0000000..28bcb2e --- /dev/null +++ b/src/api.rs @@ -0,0 +1,636 @@ +//! Reusable API functions shared by the CLI and Python bindings. + +use std::fmt; +use std::io::Write; +use std::sync::mpsc; + +use anyhow::{bail, Result}; +use hashbrown::{HashMap, HashSet}; +use indicatif::ParallelProgressIterator; +use rayon::prelude::*; + +use crate::cli::{ + check_and_set_threads, InvertedQueryType, RetainUnmatched, DEFAULT_MINCOUNT, DEFAULT_MINQUAL, +}; +use crate::distances::{ + cross_dists_all, cross_dists_knn, self_dists_all, self_dists_knn, self_dists_knn_precluster, + set_k, +}; +use crate::hashing::{AaLevel, HashType, DEFAULT_LEVEL}; +use crate::inverted::Inverted; +use crate::io::{ + get_input_list, parse_metadata_info, read_completeness_file, read_subset_names, + reorder_input_files, set_ostream, +}; +use crate::sketch::multisketch::MultiSketch; +use crate::sketch::sketch_datafile::SketchArrayReader; +use crate::sketch::{num_bins, sketch_files}; +use crate::utils::{get_progress_bar, strip_sketch_extension}; + +/// File paths produced by sketching. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SketchPaths { + /// Metadata file path. + pub skm: String, + /// Sketch data file path. + pub skd: String, +} + +/// File paths produced by inverted index construction. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InvertedPaths { + /// Inverted index file path. + pub ski: String, + /// Optional query sketch file path. + pub skq: Option, +} + +/// Basic database metadata. +#[derive(Clone, Debug, PartialEq)] +pub struct DbInfo { + /// Human-readable metadata in the same format as the CLI `info` command. + pub text: String, + /// Whether this is an inverted index. + pub inverted: bool, + /// Sketch size. + pub sketch_size: usize, + /// Number of samples. + pub n_samples: usize, + /// K-mer sizes. + pub kmers: Vec, +} + +/// Options for [`sketch`]. +#[derive(Clone, Debug)] +pub struct SketchOptions { + /// Input sequence files, mutually exclusive with `file_list`. + pub seq_files: Option>, + /// File listing input samples. + pub file_list: Option, + /// Output file prefix. + pub output: String, + /// K-mer sizes. + pub kmers: Vec, + /// Requested sketch size. + pub sketch_size: u64, + /// Sequence type. + pub seq_type: HashType, + /// aaHash level. + pub level: AaLevel, + /// Treat each FASTA record as a sample. + pub concat_fasta: bool, + /// Ignore reverse complement. + pub single_strand: bool, + /// Minimum k-mer count. + pub min_count: u16, + /// Minimum base quality. + pub min_qual: u8, + /// Number of worker threads. + pub threads: usize, + /// Suppress progress output. + pub quiet: bool, +} + +impl SketchOptions { + /// Construct DNA sketching options for a file-list input. + pub fn dna_file_list( + file_list: String, + output: String, + sketch_size: u64, + kmer_length: usize, + threads: usize, + ) -> Self { + Self { + seq_files: None, + file_list: Some(file_list), + output, + kmers: vec![kmer_length], + sketch_size, + seq_type: HashType::DNA, + level: DEFAULT_LEVEL, + concat_fasta: false, + single_strand: false, + min_count: DEFAULT_MINCOUNT, + min_qual: DEFAULT_MINQUAL, + threads, + quiet: true, + } + } +} + +/// Options for [`dist`]. +#[derive(Clone, Debug)] +pub struct DistOptions { + /// Reference sketch database prefix or file. + pub ref_db: String, + /// Optional query sketch database prefix or file. + pub query_db: Option, + /// Output distance filename. If omitted, writes to stdout. + pub output: Option, + /// Optional sparse nearest-neighbor count. + pub knn: Option, + /// Optional subset file for reference samples. + pub subset: Option, + /// Optional single k-mer length. + pub kmer: Option, + /// Return ANI rather than Jaccard distances. + pub ani: bool, + /// Number of worker threads. + pub threads: usize, + /// Optional reference completeness file. + pub ref_completeness_file: Option, + /// Optional query completeness file. + pub query_completeness_file: Option, + /// Completeness correction cutoff. + pub completeness_cutoff: f64, + /// Suppress progress output. + pub quiet: bool, +} + +/// Options for [`inverted_build`]. +#[derive(Clone, Debug)] +pub struct InvertedBuildOptions { + /// Input sequence files, mutually exclusive with `file_list`. + pub seq_files: Option>, + /// File listing input samples. + pub file_list: Option, + /// Output file prefix. + pub output: String, + /// Whether to write an `.skq` file. + pub write_skq: bool, + /// Optional species/label file. + pub species_names: Option, + /// Optional metadata file. + pub metadata: Option, + /// Inverted sketch size. + pub sketch_size: u64, + /// K-mer length. + pub kmer_length: usize, + /// Ignore reverse complement. + pub single_strand: bool, + /// Minimum k-mer count. + pub min_count: u16, + /// Minimum base quality. + pub min_qual: u8, + /// Number of worker threads. + pub threads: usize, + /// Suppress progress output. + pub quiet: bool, +} + +/// Options for [`inverted_precluster`]. +#[derive(Clone, Debug)] +pub struct InvertedPreclusterOptions { + /// Inverted index file. + pub ski: String, + /// Standard sketch database prefix or file. + pub skd: String, + /// Output distance filename. If omitted, writes to stdout. + pub output: Option, + /// Sparse nearest-neighbor count. + pub knn: usize, + /// Return ANI rather than Jaccard distances. + pub ani: bool, + /// Number of worker threads. + pub threads: usize, + /// Optional reference completeness file. + pub ref_completeness_file: Option, + /// Completeness correction cutoff. + pub completeness_cutoff: f64, + /// How to retain unmatched samples. + pub retain_unmatched: Option, + /// Suppress progress output. + pub quiet: bool, +} + +/// Create a sketch database. +pub fn sketch(options: SketchOptions) -> Result { + if options.concat_fasta && matches!(options.seq_type, HashType::DNA) { + bail!("--concat-fasta currently only supported with --seq-type aa"); + } + check_and_set_threads(options.threads + 1); + + let input_files = get_input_list(&options.file_list, &options.seq_files); + let rc = !options.single_strand; + let seq_type = if matches!(options.seq_type, HashType::AA(_)) { + HashType::AA(options.level) + } else { + options.seq_type + }; + let (_, sketch_bins, _) = num_bins(options.sketch_size); + let mut sketches = sketch_files( + &options.output, + &input_files, + options.concat_fasta, + &options.kmers, + sketch_bins, + &seq_type, + rc, + options.min_count, + options.min_qual, + options.quiet, + ); + let sketch_vec = MultiSketch::new(&mut sketches, sketch_bins, &options.kmers, seq_type); + sketch_vec.save_metadata(&options.output)?; + + Ok(SketchPaths { + skm: format!("{}.skm", options.output), + skd: format!("{}.skd", options.output), + }) +} + +/// Calculate distances from one or two sketch databases. +pub fn dist(options: DistOptions) -> Result> { + check_and_set_threads(options.threads); + let mut output_file = set_ostream(&options.output); + write_distances(&options, &mut output_file)?; + Ok(options.output) +} + +fn write_distances(options: &DistOptions, output_file: &mut impl Write) -> Result<()> { + let ref_db_name = strip_sketch_extension(&options.ref_db); + let mut references = MultiSketch::load_metadata(ref_db_name)?; + + if let Some(subset_file) = &options.subset { + let subset_names = read_subset_names(subset_file); + references.read_sketch_data_block(ref_db_name, &subset_names); + } else { + references.read_sketch_data(ref_db_name); + } + let n = references.number_samples_loaded(); + let ref_completeness_vec = if let Some(file_path) = &options.ref_completeness_file { + Some(read_completeness_file(file_path, &references)?) + } else { + None + }; + + let dist_type = set_k(&references, options.kmer, options.ani)?; + let queries = if let Some(query_db_name) = &options.query_db { + let query_db_name = strip_sketch_extension(query_db_name); + let mut queries = MultiSketch::load_metadata(query_db_name)?; + queries.read_sketch_data(query_db_name); + Some(queries) + } else { + None + }; + + match queries { + None => match options.knn { + None => { + let distances = self_dists_all( + &references, + n, + dist_type, + options.quiet, + ref_completeness_vec.as_ref(), + options.completeness_cutoff, + ); + write!(output_file, "{distances}")?; + } + Some(mut nn) => { + if nn >= n { + log::warn!("knn={nn} is higher than number of samples={n}"); + nn = n - 1; + } + let distances = self_dists_knn( + &references, + n, + nn, + dist_type, + options.quiet, + ref_completeness_vec.as_ref(), + options.completeness_cutoff, + ); + write!(output_file, "{distances}")?; + } + }, + Some(query_db) => { + let query_completeness_vec = if let Some(file_path) = &options.query_completeness_file { + Some(read_completeness_file(file_path, &query_db)?) + } else { + None + }; + let n_query = query_db.number_samples_loaded(); + match options.knn { + Some(mut nn) => { + if nn > n { + log::warn!("knn={nn} is higher than number of reference samples={n}"); + nn = n; + } + let distances = cross_dists_knn( + &references, + &query_db, + n, + n_query, + nn, + dist_type, + options.quiet, + ref_completeness_vec.as_ref(), + query_completeness_vec.as_ref(), + options.completeness_cutoff, + ); + write!(output_file, "{distances}")?; + } + None => { + let distances = cross_dists_all( + &references, + &query_db, + n, + n_query, + dist_type, + options.quiet, + ref_completeness_vec.as_ref(), + query_completeness_vec.as_ref(), + options.completeness_cutoff, + ); + write!(output_file, "{distances}")?; + } + } + } + } + Ok(()) +} + +/// Build an inverted index. +pub fn inverted_build(options: InvertedBuildOptions) -> Result { + check_and_set_threads(options.threads + 1); + let input_files = get_input_list(&options.file_list, &options.seq_files); + let mut different_samples = HashSet::new(); + for input_file in input_files.iter() { + different_samples.insert(input_file.0.clone()); + } + + let (file_order, map_names_labels) = if let Some(species_name_file) = &options.species_names { + reorder_input_files(&input_files, species_name_file) + } else { + default_file_order(&input_files) + }; + + let species_labels_vec = map_names_labels.map(|labels| { + let mut tmpvec = vec!["".to_string(); different_samples.len()]; + file_order + .iter() + .zip(&input_files) + .for_each(|(idx, (name, _))| { + tmpvec[*idx] = labels.get(name).unwrap_or(&"".to_owned()).clone(); + }); + tmpvec + }); + + let metadata_vec = if let Some(metadata_file) = &options.metadata { + let tmpdict = parse_metadata_info(metadata_file); + let mut tmpvec = vec!["".to_string(); different_samples.len()]; + file_order + .iter() + .zip(&input_files) + .for_each(|(idx, (name, _))| tmpvec[*idx] = tmpdict[name].clone()); + Some(tmpvec) + } else { + None + }; + + let skq_file = options.write_skq.then(|| format!("{}.skq", options.output)); + let rc = !options.single_strand; + let inverted = Inverted::new( + &input_files, + skq_file.clone(), + &file_order, + options.kmer_length, + options.sketch_size, + &HashType::DNA, + rc, + options.min_count, + options.min_qual, + options.quiet, + &metadata_vec, + &species_labels_vec, + ); + inverted.save(&options.output)?; + + Ok(InvertedPaths { + ski: format!("{}.ski", options.output), + skq: skq_file, + }) +} + +fn default_file_order( + input_files: &[(String, Vec)], +) -> (Vec, Option>) { + let tmp_name_set = input_files + .iter() + .map(|x| x.0.clone()) + .collect::>(); + if tmp_name_set.len() == input_files.len() { + ((0..input_files.len()).collect(), None) + } else { + let mut tmpoutvec = vec![0; input_files.len()]; + let mut tmpmap = HashMap::new(); + for (i, name) in tmp_name_set.iter().enumerate() { + tmpmap.insert(name.clone(), i); + } + for i in 0..tmpoutvec.len() { + tmpoutvec[i] = tmpmap[&input_files[i].0]; + } + (tmpoutvec, None) + } +} + +/// Run inverted-index preclustering distances. +pub fn inverted_precluster(options: InvertedPreclusterOptions) -> Result> { + check_and_set_threads(options.threads); + let mut output_file = set_ostream(&options.output); + write_inverted_precluster(&options, &mut output_file)?; + Ok(options.output) +} + +fn write_inverted_precluster( + options: &InvertedPreclusterOptions, + output_file: &mut impl Write, +) -> Result<()> { + let input_prefix = strip_sketch_extension(&options.ski); + let inverted_index = Inverted::load(input_prefix)?; + let skq_filename = format!("{input_prefix}.skq"); + let (mmap, bin_stride, kmer_stride, sample_stride) = + (false, 1, 1, inverted_index.sketch_size()); + let mut skq_reader = + SketchArrayReader::open(&skq_filename, mmap, bin_stride, kmer_stride, sample_stride); + let skq_bins = skq_reader.read_all_from_skq(sample_stride * inverted_index.sketch_size()); + + let ref_db_name = strip_sketch_extension(&options.skd); + let mut references = MultiSketch::load_metadata(ref_db_name)?; + references.read_sketch_data(ref_db_name); + let n = references.number_samples_loaded(); + let mut knn = options.knn; + if knn >= n { + log::warn!("knn={knn} is higher than number of samples={n}"); + knn = n - 1; + } + + let kmer = inverted_index.kmer(); + let dist_type = set_k(&references, Some(kmer), options.ani)?; + let ref_completeness_vec = if let Some(file_path) = &options.ref_completeness_file { + Some(read_completeness_file(file_path, &references)?) + } else { + None + }; + let distances = self_dists_knn_precluster( + &references, + &inverted_index, + &skq_bins, + skq_reader.sample_stride, + n, + knn, + dist_type, + options.quiet, + ref_completeness_vec.as_ref(), + options.completeness_cutoff, + &options.retain_unmatched, + ); + write!(output_file, "{distances}")?; + Ok(()) +} + +/// Compute query-vs-reference sparse distances. +pub fn query_dist( + reference_skm: String, + query_skm: String, + output: String, + kmer_length: usize, + knn: usize, + threads: usize, + ref_completeness_file: Option, + query_completeness_file: Option, + completeness_cutoff: f64, + quiet: bool, +) -> Result { + dist(DistOptions { + ref_db: reference_skm, + query_db: Some(query_skm), + output: Some(output.clone()), + knn: Some(knn), + subset: None, + kmer: Some(kmer_length), + ani: true, + threads, + ref_completeness_file, + query_completeness_file, + completeness_cutoff, + quiet, + })?; + Ok(output) +} + +/// Query an inverted index and write the result table. +pub fn inverted_query( + ski: &str, + seq_files: &Option>, + file_list: &Option, + output: &Option, + query_type: &InvertedQueryType, + min_count: u16, + min_qual: u8, + threads: usize, + quiet: bool, +) -> Result<()> { + let mut output_file = set_ostream(output); + let inverted_index = Inverted::load(strip_sketch_extension(ski))?; + let input_files = get_input_list(file_list, seq_files); + check_and_set_threads(threads + 1); + let (queries, query_names) = + inverted_index.sketch_queries(&input_files, min_count, min_qual, quiet); + + write!(output_file, "Query")?; + if *query_type == InvertedQueryType::MatchCount { + for name in inverted_index.sample_names() { + write!(output_file, "\t{name}")?; + } + writeln!(output_file)?; + } else { + writeln!(output_file, "\tMatches")?; + } + + let (tx, rx) = mpsc::channel(); + let progress_bar = get_progress_bar(queries.len(), false, quiet); + rayon::scope(|s| { + s.spawn(|_| { + queries + .par_iter() + .progress_with(progress_bar) + .zip(query_names) + .map(|(q, q_name)| match query_type { + InvertedQueryType::MatchCount => { + (q_name, inverted_index.query_against_inverted_index(q)) + } + InvertedQueryType::AllBins => (q_name, inverted_index.all_shared_bins(q)), + InvertedQueryType::AnyBins => (q_name, inverted_index.any_shared_bins(q)), + }) + .for_each_with(tx, |tx, dists| { + let _ = tx.send(dists); + }); + }); + }); + for (q_name, dist) in rx { + write!(output_file, "{q_name}")?; + if *query_type == InvertedQueryType::MatchCount { + for distance in dist { + write!(output_file, "\t{distance}")?; + } + } else if !dist.is_empty() { + write!( + output_file, + "\t{}", + inverted_index.sample_at(dist[0] as usize) + )?; + for r_name in dist + .iter() + .skip(1) + .map(|idx| inverted_index.sample_at(*idx as usize)) + { + write!(output_file, ",{r_name}")?; + } + } + writeln!(output_file)?; + } + Ok(()) +} + +/// Return database information. +pub fn db_info(db_file: &str, sample_info: bool) -> Result { + if db_file.ends_with(".ski") { + let ski_file = strip_sketch_extension(db_file); + let index = Inverted::load(ski_file)?; + let text = if sample_info { + format!("{index}") + } else { + format!("{index:?}") + }; + Ok(DbInfo { + text, + inverted: true, + sketch_size: index.sketch_size(), + n_samples: index.n_samples(), + kmers: vec![index.kmer()], + }) + } else { + let ref_db_name = strip_sketch_extension(db_file); + let sketches = MultiSketch::load_metadata(ref_db_name)?; + let text = if sample_info { + format!("{sketches}") + } else { + format!("{sketches:?}") + }; + Ok(DbInfo { + text, + inverted: false, + sketch_size: sketches.sketch_size as usize, + n_samples: sketches.number_samples_loaded(), + kmers: sketches.kmer_lengths().to_vec(), + }) + } +} + +impl fmt::Display for DbInfo { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.text) + } +} diff --git a/src/cli.rs b/src/cli.rs index 4570a7e..7c9bb0e 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -79,10 +79,13 @@ pub fn check_and_set_threads(threads: usize) { } else { log::info!("Using {threads} threads"); } - rayon::ThreadPoolBuilder::new() + if rayon::ThreadPoolBuilder::new() .num_threads(threads) .build_global() - .unwrap(); + .is_err() + { + log::debug!("Rayon global thread pool was already initialized"); + } } /// Options that apply to all subcommands @@ -140,11 +143,6 @@ pub enum Commands { #[arg(long, default_value_t = false)] concat_fasta: bool, - /// Input files are .pdb, convert them to 3Di first - #[cfg(feature = "3di")] - #[arg(long, default_value_t = false)] - convert_pdb: bool, - /// Output prefix #[arg(short)] output: String, diff --git a/src/hashing/aahash_iterator.rs b/src/hashing/aahash_iterator.rs index 0651d0e..a92dc5e 100644 --- a/src/hashing/aahash_iterator.rs +++ b/src/hashing/aahash_iterator.rs @@ -123,18 +123,6 @@ impl AaHashIterator { hash_vec } - /// Create a new iterator from a 3di embedding file of a structure - pub fn from_3di_file(files: &[String]) -> Vec { - Self::new(files, AaLevel::Level1, false) - } - - /// Create a new iterator from a 3di embedding string of a structure - pub fn from_3di_string(sequence: String) -> Vec { - let mut hash_it = Self::default(AaLevel::Level1); - hash_it.seq = sequence.into_bytes(); - vec![hash_it] - } - fn new_iterator( mut start: usize, level: &AaLevel, diff --git a/src/hashing/mod.rs b/src/hashing/mod.rs index e4cb14a..8f185c8 100644 --- a/src/hashing/mod.rs +++ b/src/hashing/mod.rs @@ -34,45 +34,18 @@ pub enum HashType { DNA, /// Amino acids at set [`AaLevel`] AA(AaLevel), - /// Structures using 3di embedding - PDB, } -// TODO: for PDB need to use 3Di sequences. These are from an embedding -// With foldseek, can run: -// foldseek createdb 5uak.pdb 5Uak_DB -// foldseek lndb 5uak_DB 5uak_DB_ss_h -// foldseek convert2fasta 5uak_DB_ss_h 5uak.fasta -// The Cpp code looks like a bit of a pain to build in: -// https://github.com/steineggerlab/foldseek/blob/master/lib/3di/structureto3di.cpp -// This python port is an alternative: -// https://github.com/althonos/mini3di -// (this seems to give different results to the above) -// Seems pretty easy to run python from within rust: -// https://pyo3.rs/v0.21.2/python-from-rust/calling-existing-code -// Or a language model: -// https://github.com/mheinzinger/ProstT5 - // NB: this is needed because ValueEnum (for clap) only works with non-unit types // So here set a default for the level and set it properly later (in lib.rs) impl clap::ValueEnum for HashType { fn value_variants<'a>() -> &'a [Self] { - &[ - HashType::DNA, - HashType::AA(DEFAULT_LEVEL), - // HashType::AA(AaLevel::Level2), - // HashType::AA(AaLevel::Level3), - HashType::PDB, - ] + &[HashType::DNA, HashType::AA(DEFAULT_LEVEL)] } fn to_possible_value<'a>(&self) -> ::std::option::Option { match self { Self::DNA => Some(clap::builder::PossibleValue::new("dna")), Self::AA(_) => Some(clap::builder::PossibleValue::new("aa")), - // Self::AA(AaLevel::Level1) => Some(clap::builder::PossibleValue::new("aa_1")), - // Self::AA(AaLevel::Level2) => Some(clap::builder::PossibleValue::new("aa_2")), - // Self::AA(AaLevel::Level3) => Some(clap::builder::PossibleValue::new("aa_3")), - Self::PDB => Some(clap::builder::PossibleValue::new("pdb")), } } } diff --git a/src/lib.rs b/src/lib.rs index 755d1f8..aaba7c6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,5 @@ -//! Fast distance calculations between biological sequences (DNA, AA or structures -//! via the 3di alphabet). Distances are based on bindash approximations of the Jaccard +//! Fast distance calculations between biological sequences (DNA or AA). +//! Distances are based on bindash approximations of the Jaccard //! distance, with the [PopPUNK method](https://poppunk.bacpop.org/index.html) to calculate core and accessory distances. nthash/aahash //! are used for hash functions to create the sketches. //! @@ -42,8 +42,6 @@ //! resolution and 100000-1000000 for SNP level resolution. //! - To sketch amino acid sequences use `--seq-type aa --concat-fasta` if you have the typical case //! of each fasta file being a multifasta with many aa sequences. Each one will then be its own sample. -//! - You can also sketch structures with .pdb input, see 'Enabling PDB->3Di' below. This is experimental. -//! //! ### Distances //! //! To compute internal all-vs-all core and accessory distances use: @@ -128,24 +126,9 @@ //! - `append` sketches new input samples, and adds them to an existing database. //! - `delete` removes samples from a sketch database. //! -//! ## Enabling PDB->3Di -//! conda doesn't work, so make sure it is deactivated -//! ```bash -//! export PYO3_PYTHON=python3 -//! python3 -m venv 3di_venv -//! source 3di_venv/bin/activate -//! python3 -m pip install numpy biopython mini3di -//! cargo run -F 3di -//! export PYTHONPATH=${PYTHONPATH}:$(realpath ./)/3di_venv/lib/python3.12/site-packages -//! ``` - #![warn(missing_docs)] #![allow(clippy::too_many_arguments)] -#[cfg(not(target_arch = "wasm32"))] -use std::io::Write; -#[cfg(not(target_arch = "wasm32"))] -use std::sync::mpsc; #[cfg(not(target_arch = "wasm32"))] use std::time::Instant; @@ -154,11 +137,10 @@ extern crate arrayref; extern crate num_cpus; use anyhow::Error; #[cfg(not(target_arch = "wasm32"))] -use indicatif::ParallelProgressIterator; -#[cfg(not(target_arch = "wasm32"))] -use rayon::prelude::*; - +pub mod api; pub mod cli; +#[cfg(all(not(target_arch = "wasm32"), feature = "python"))] +pub mod python; #[cfg(not(target_arch = "wasm32"))] use crate::cli::*; #[cfg(target_arch = "wasm32")] @@ -167,38 +149,26 @@ use crate::cli::{InvertedQueryType, DEFAULT_MINCOUNT, DEFAULT_MINQUAL}; #[cfg(not(target_arch = "wasm32"))] use crate::hashing::HashType; -#[cfg(not(target_arch = "wasm32"))] -use hashbrown::{HashMap, HashSet}; - pub mod sketch; #[cfg(not(target_arch = "wasm32"))] use crate::sketch::multisketch::MultiSketch; #[cfg(not(target_arch = "wasm32"))] -use crate::sketch::sketch_datafile::SketchArrayReader; -#[cfg(not(target_arch = "wasm32"))] -use crate::sketch::{num_bins, sketch_files}; +use crate::sketch::sketch_files; pub mod inverted; use crate::inverted::Inverted; pub mod distances; -#[cfg(not(target_arch = "wasm32"))] -use crate::distances::*; pub mod io; #[cfg(not(target_arch = "wasm32"))] -use crate::io::{ - get_input_list, parse_kmers, parse_metadata_info, read_completeness_file, read_subset_names, - reorder_input_files, set_ostream, -}; - -pub mod structures; +use crate::io::{get_input_list, parse_kmers}; pub mod hashing; pub mod utils; -use crate::utils::get_progress_bar; +pub use crate::utils::get_progress_bar; #[cfg(not(target_arch = "wasm32"))] use crate::utils::strip_sketch_extension; @@ -243,8 +213,6 @@ pub fn main() -> Result<(), Error> { seq_files, file_list, concat_fasta, - #[cfg(feature = "3di")] - convert_pdb, output, kmers, sketch_size, @@ -255,49 +223,21 @@ pub fn main() -> Result<(), Error> { min_qual, threads, } => { - if *concat_fasta && matches!(*seq_type, HashType::DNA | HashType::PDB) { - panic!("--concat-fasta currently only supported with --seq-type aa"); - } - - // An extra thread is needed for the writer. This doesn't 'overuse' CPU - check_and_set_threads(*threads + 1); - - // Read input - 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(kmers); - // Build, merge - let rc = !*single_strand; - // Set aa level - let seq_type = if let HashType::AA(_) = seq_type { - HashType::AA(level.clone()) - } else { - seq_type.clone() - }; - - let (_, sketch_bins, _) = num_bins(*sketch_size); - log::info!( - "Running sketching: k:{kmers:?}; sketch_size:{sketch_bins}; seq:{seq_type:?}; threads:{threads}" - ); - let mut sketches = sketch_files( - output, - &input_files, - *concat_fasta, - #[cfg(feature = "3di")] - *convert_pdb, - &kmers, - sketch_bins, - &seq_type, - rc, - *min_count, - *min_qual, - args.quiet, - ); - let sketch_vec = MultiSketch::new(&mut sketches, sketch_bins, &kmers, seq_type); - sketch_vec - .save_metadata(output) - .expect("Error saving metadata"); + api::sketch(api::SketchOptions { + seq_files: seq_files.clone(), + file_list: file_list.clone(), + output: output.clone(), + kmers: parse_kmers(kmers), + sketch_size: *sketch_size, + seq_type: seq_type.clone(), + level: level.clone(), + concat_fasta: *concat_fasta, + single_strand: *single_strand, + min_count: *min_count, + min_qual: *min_qual, + threads: *threads, + quiet: args.quiet, + })?; Ok(()) } Commands::Dist { @@ -313,142 +253,20 @@ pub fn main() -> Result<(), Error> { query_completeness_file, completeness_cutoff, } => { - check_and_set_threads(*threads); - - let mut output_file = set_ostream(output); - - let ref_db_name = utils::strip_sketch_extension(ref_db); - - let mut references = MultiSketch::load_metadata(ref_db_name) - .unwrap_or_else(|_| panic!("Could not read sketch metadata from {ref_db}.skm")); - - log::info!("Loading sketch data from {ref_db_name}.skd"); - if let Some(subset_file) = subset { - let subset_names = read_subset_names(subset_file); - references.read_sketch_data_block(ref_db_name, &subset_names); - } else { - references.read_sketch_data(ref_db_name); - } - log::info!("Read reference sketches:\n{references:?}"); - let n = references.number_samples_loaded(); - let ref_completeness_vec: Option> = - if let Some(file_path) = ref_completeness_file { - Some(read_completeness_file(file_path, &references)?) - } else { - None - }; - - let dist_type = set_k(&references, *kmer, *ani).unwrap_or_else(|e| { - panic!("Error setting k size: {e}"); - }); - - // Read queries if supplied. Note no subsetting here - let queries = if let Some(query_db_name) = query_db { - let mut queries = MultiSketch::load_metadata(query_db_name).unwrap_or_else(|_| { - panic!("Could not read sketch metadata from {query_db_name}.skm") - }); - log::info!("Loading query sketch data from {query_db_name}.skd"); - queries.read_sketch_data(query_db_name); - log::info!("Read query sketches:\n{queries:?}"); - Some(queries) - } else { - None - }; - - match queries { - None => { - // Ref v ref functions - match knn { - None => { - // Self mode (dense) - log::info!("Calculating all ref vs ref distances"); - let distances = self_dists_all( - &references, - n, - dist_type, - args.quiet, - ref_completeness_vec.as_ref(), - *completeness_cutoff, - ); - log::info!("Writing out in long matrix form"); - write!(output_file, "{distances}") - .expect("Error writing output distances"); - } - Some(mut nn) => { - // Self mode (sparse): a genome cannot be its own neighbour - if nn >= n { - log::warn!("knn={nn} is higher than number of samples={n}"); - nn = n - 1; - } - log::info!("Calculating sparse ref vs ref distances with {nn} nearest neighbours"); - let distances = self_dists_knn( - &references, - n, - nn, - dist_type, - args.quiet, - ref_completeness_vec.as_ref(), - *completeness_cutoff, - ); - - log::info!("Writing out in sparse matrix form"); - write!(output_file, "{distances}") - .expect("Error writing output distances"); - } - } - } - Some(query_db) => { - let query_completeness_vec: Option> = - if let Some(file_path) = query_completeness_file { - Some(read_completeness_file(file_path, &query_db)?) - } else { - None - }; - let n_query = query_db.number_samples_loaded(); - match knn { - Some(mut nn) => { - // Cross-query mode: query genomes never overlap ref genomes, so knn=n is valid - if nn > n { - log::warn!("knn={nn} is higher than number of reference samples={n}"); - nn = n; - } - // Cross-query mode (sparse kNN) - log::info!("Calculating sparse ref vs query distances with {nn} nearest neighbours"); - let distances = cross_dists_knn( - &references, - &query_db, - n, - n_query, - nn, - dist_type, - args.quiet, - ref_completeness_vec.as_ref(), - query_completeness_vec.as_ref(), - *completeness_cutoff, - ); - log::info!("Writing out in sparse matrix form"); - write!(output_file, "{distances}").expect("Error writing output distances"); - } - None => { - // Cross-query mode (dense, all pairs) - log::info!("Calculating all ref vs query distances"); - let distances = cross_dists_all( - &references, - &query_db, - n, - n_query, - dist_type, - args.quiet, - ref_completeness_vec.as_ref(), - query_completeness_vec.as_ref(), - *completeness_cutoff, - ); - log::info!("Writing out in long matrix form"); - write!(output_file, "{distances}").expect("Error writing output distances"); - } - } - } - } + api::dist(api::DistOptions { + ref_db: ref_db.clone(), + query_db: query_db.clone(), + output: output.clone(), + knn: *knn, + subset: subset.clone(), + kmer: *kmer, + ani: *ani, + threads: *threads, + ref_completeness_file: ref_completeness_file.clone(), + query_completeness_file: query_completeness_file.clone(), + completeness_cutoff: *completeness_cutoff, + quiet: args.quiet, + })?; Ok(()) } Commands::Merge { db1, db2, output } => { @@ -496,101 +314,21 @@ pub fn main() -> Result<(), Error> { sketch_size, kmer_length, } => { - // An extra thread is needed for the writer - check_and_set_threads(*threads + 1); - - // Get input files - log::info!("Getting input files"); - let input_files: Vec<(String, Vec)> = get_input_list(file_list, seq_files); - log::info!("Parsed {} samples in input list", input_files.len()); - - let mut differentsamples: HashSet = HashSet::new(); - - for i in input_files.iter() { - differentsamples.insert(i.0.clone()); - } - - // Reordering by species, or default - let (file_order, map_names_labels) = if let Some(species_name_file) = species_names - { - reorder_input_files(&input_files, species_name_file) - } else { - // Check first if there are repeated samples - - let tmpnamesset = input_files - .iter() - .map(|x| x.0.clone()) - .collect::>(); - if tmpnamesset.len() == input_files.len() { - ((0..input_files.len()).collect(), None) - } else { - let mut tmpoutvec: Vec = vec![0; input_files.len()]; - let mut tmpmap: HashMap = HashMap::new(); - - for (i, name) in tmpnamesset.iter().enumerate() { - tmpmap.insert(name.clone(), i); - } - for i in 0..tmpoutvec.len() { - tmpoutvec[i] = tmpmap[&input_files[i].0]; - } - - (tmpoutvec, None) - } - }; - - // If species labels were provided, create the list of them - let species_labels_vec = if let Some(themaplabels) = map_names_labels { - let mut tmpvec: Vec = vec!["".to_string(); differentsamples.len()]; - file_order - .iter() - .zip(&input_files) - .for_each(|(idx, (name, _))| { - // log::info!("{:?} {:?}", name, idx); - tmpvec[*idx] = themaplabels.get(name).unwrap_or(&"".to_owned()).clone(); - }); - Some(tmpvec) - } else { - None - }; - - // Parse metadata, if any - let metadata_vec; - if let Some(metadata_file) = metadata { - let tmpdict = parse_metadata_info(metadata_file); - let mut tmpvec: Vec = vec!["".to_string(); differentsamples.len()]; - file_order - .iter() - .zip(&input_files) - .for_each(|(idx, (name, _))| tmpvec[*idx] = tmpdict[name].clone()); - metadata_vec = Some(tmpvec); - } else { - metadata_vec = None; - }; - - let skq_file = if *write_skq { - Some(format!("{output}.skq")) - } else { - None - }; - - let rc = !*single_strand; - let seq_type = &HashType::DNA; - let inverted = Inverted::new( - &input_files, - skq_file, - &file_order, - *kmer_length, - *sketch_size, // unconstrained, equals the number of bins here, doesn't need to be a multiple of 64 - seq_type, - rc, - *min_count, - *min_qual, - args.quiet, - &metadata_vec, - &species_labels_vec, - ); - inverted.save(output)?; - log::info!("Index info:\n{inverted:?}"); + api::inverted_build(api::InvertedBuildOptions { + seq_files: seq_files.clone(), + file_list: file_list.clone(), + output: output.clone(), + write_skq: *write_skq, + species_names: species_names.clone(), + metadata: metadata.clone(), + sketch_size: *sketch_size, + kmer_length: *kmer_length, + single_strand: *single_strand, + min_count: *min_count, + min_qual: *min_qual, + threads: *threads, + quiet: args.quiet, + })?; Ok(()) } InvertedCommands::Query { @@ -603,80 +341,10 @@ pub fn main() -> Result<(), Error> { min_qual, threads, } => { - let mut output_file = set_ostream(output); - let inverted_index = Inverted::load(strip_sketch_extension(ski))?; - log::info!("Read inverted index:\n{inverted_index:?}"); - - // Get input files - log::info!("Getting input queries"); - let input_files: Vec<(String, Vec)> = get_input_list(file_list, seq_files); - log::info!("Parsed {} samples in input query list", input_files.len()); - - log::info!("Sketching input queries"); - check_and_set_threads(*threads + 1); // Writer thread - let (queries, query_names) = - inverted_index.sketch_queries(&input_files, *min_count, *min_qual, args.quiet); - - log::info!("Running queries in mode: {query_type}"); - // Header - write!(output_file, "Query")?; - if *query_type == InvertedQueryType::MatchCount { - for name in inverted_index.sample_names() { - write!(output_file, "\t{name}")?; - } - writeln!(output_file)?; - } else { - writeln!(output_file, "\tMatches")?; - } - - // Query loop (parallelised) - let (tx, rx) = mpsc::channel(); - let percent = false; - let progress_bar = get_progress_bar(queries.len(), percent, args.quiet); - rayon::scope(|s| { - s.spawn(|_| { - queries - .par_iter() - .progress_with(progress_bar) - .zip(query_names) - .map(|(q, q_name)| match query_type { - InvertedQueryType::MatchCount => { - (q_name, inverted_index.query_against_inverted_index(q)) - } - InvertedQueryType::AllBins => { - (q_name, inverted_index.all_shared_bins(q)) - } - InvertedQueryType::AnyBins => { - (q_name, inverted_index.any_shared_bins(q)) - } - }) - .for_each_with(tx, |tx, dists| { - let _ = tx.send(dists); - }); - }); - }); - for (q_name, dist) in rx { - write!(output_file, "{q_name}")?; - if *query_type == InvertedQueryType::MatchCount { - for distance in dist { - write!(output_file, "\t{distance}")?; - } - } else if !dist.is_empty() { - write!( - output_file, - "\t{}", - inverted_index.sample_at(dist[0] as usize) - )?; - for r_name in dist - .iter() - .skip(1) - .map(|idx| inverted_index.sample_at(*idx as usize)) - { - write!(output_file, ",{r_name}")?; - } - } - writeln!(output_file)?; - } + api::inverted_query( + ski, seq_files, file_list, output, query_type, *min_count, *min_qual, *threads, + args.quiet, + )?; Ok(()) } InvertedCommands::Precluster { @@ -684,15 +352,13 @@ pub fn main() -> Result<(), Error> { skd, count, output, - mut knn, + knn, ani, threads, ref_completeness_file, completeness_cutoff, retain_unmatched, } => { - check_and_set_threads(*threads); - // Load the inverted index let input_prefix = strip_sketch_extension(ski); let inverted_index = Inverted::load(input_prefix)?; @@ -710,80 +376,18 @@ pub fn main() -> Result<(), Error> { / 2 ); } else if let Some(ref_db_input) = skd { - let mut output_file = set_ostream(output); - - // Open the .skq - let skq_filename = &format!("{input_prefix}.skq"); - log::info!("Loading queries from {skq_filename}"); - let (mmap, bin_stride, kmer_stride, sample_stride) = - (false, 1, 1, inverted_index.sketch_size()); - let mut skq_reader = SketchArrayReader::open( - skq_filename, - mmap, - bin_stride, - kmer_stride, - sample_stride, - ); - let skq_bins = - skq_reader.read_all_from_skq(sample_stride * inverted_index.sketch_size()); - - // Load the .skd/.skm - let ref_db_name = utils::strip_sketch_extension(ref_db_input); - let mut references = - MultiSketch::load_metadata(ref_db_name).unwrap_or_else(|_| { - panic!("Could not read sketch metadata from {ref_db_name}.skm") - }); - log::info!("Loading sketch data from {ref_db_name}.skd"); - references.read_sketch_data(ref_db_name); - log::info!("Read reference sketches:\n{references:?}"); - let n = references.number_samples_loaded(); - if knn >= n { - log::warn!("knn={knn} is higher than number of samples={n}"); - knn = n - 1; - } - - // Check that k-mer exists in the .skd, and find its index - let kmer = inverted_index.kmer(); - // This panics if k not found. Maybe more graceful error if this happens - let dist_type = set_k(&references, Some(kmer), *ani).unwrap_or_else(|e| { - panic!("K-mer size {kmer} used for .ski not found in .skd: {e}"); - }); - - let ref_completeness_vec: Option> = - if let Some(file_path) = ref_completeness_file { - Some(read_completeness_file(file_path, &references)?) - } else { - None - }; - // Run the distances with both indexes - log::info!( - "Calculating sparse ref vs ref distances with {knn} nearest neighbours" - ); - log::info!( - "Preclustering with k={} and s={}", - kmer, - inverted_index.sketch_size() - ); - if let Some(ref mode) = retain_unmatched { - log::info!("Retain unmatched mode: {mode}"); - } - let distances = self_dists_knn_precluster( - &references, - &inverted_index, - &skq_bins, - skq_reader.sample_stride, - n, - knn, - dist_type, - args.quiet, - ref_completeness_vec.as_ref(), - *completeness_cutoff, - retain_unmatched, - ); - - // Write the results - log::info!("Writing out in sparse matrix form"); - write!(output_file, "{distances}").expect("Error writing output distances"); + api::inverted_precluster(api::InvertedPreclusterOptions { + ski: ski.clone(), + skd: ref_db_input.clone(), + output: output.clone(), + knn: *knn, + ani: *ani, + threads: *threads, + ref_completeness_file: ref_completeness_file.clone(), + completeness_cutoff: *completeness_cutoff, + retain_unmatched: retain_unmatched.clone(), + quiet: args.quiet, + })?; } Ok(()) @@ -822,7 +426,7 @@ pub fn main() -> Result<(), Error> { let rc = !*single_strand; let sketch_size = db_metadata.sketch_size; let seq_type = db_metadata.get_hash_type(); - if *concat_fasta && matches!(*seq_type, HashType::DNA | HashType::PDB) { + if *concat_fasta && matches!(*seq_type, HashType::DNA) { panic!("--concat-fasta currently only supported with --seq-type aa"); } @@ -844,8 +448,6 @@ pub fn main() -> Result<(), Error> { output, &input_files, *concat_fasta, - #[cfg(feature = "3di")] - false, kmers, sketch_size, &seq_type, @@ -911,37 +513,7 @@ pub fn main() -> Result<(), Error> { skm_file, sample_info, } => { - if skm_file.ends_with(".ski") { - let ski_file = &skm_file[0..skm_file.len() - 4]; - let index = Inverted::load(ski_file).unwrap_or_else(|err| { - println!("Read error: {err}"); - panic!("Could not read inverted index from {ski_file}.ski") - }); - if *sample_info { - log::info!("Printing sample info"); - println!("{index}"); - } else { - log::info!("Printing inverted index info"); - println!("{index:?}"); - } - } else { - let ref_db_name = if skm_file.ends_with(".skm") || skm_file.ends_with(".skd") { - &skm_file[0..skm_file.len() - 4] - } else { - skm_file.as_str() - }; - let sketches = MultiSketch::load_metadata(ref_db_name).unwrap_or_else(|_| { - panic!("Could not read sketch metadata from {ref_db_name}.skm") - }); - if *sample_info { - log::info!("Printing sample info"); - println!("{sketches}"); - } else { - log::info!("Printing database info"); - println!("{sketches:?}"); - } - } - + println!("{}", api::db_info(skm_file, *sample_info)?); print_success = false; // Turn the final message off Ok(()) } diff --git a/src/python.rs b/src/python.rs new file mode 100644 index 0000000..bc14fdc --- /dev/null +++ b/src/python.rs @@ -0,0 +1,293 @@ +//! Python bindings for sketchlib. + +use std::collections::HashMap; + +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +use crate::api::{ + self, DistOptions, InvertedBuildOptions, InvertedPreclusterOptions, SketchOptions, +}; +use crate::cli::{RetainUnmatched, DEFAULT_MINCOUNT, DEFAULT_MINQUAL}; +use crate::hashing::{HashType, DEFAULT_LEVEL}; + +fn py_runtime_error(err: anyhow::Error) -> PyErr { + PyRuntimeError::new_err(err.to_string()) +} + +fn parse_retain_unmatched(value: Option<&str>) -> PyResult> { + match value { + None => Ok(None), + Some("singleton") => Ok(Some(RetainUnmatched::Singleton)), + Some("bruteforce") => Ok(Some(RetainUnmatched::Bruteforce)), + Some(other) => Err(PyValueError::new_err(format!( + "retain_unmatched must be 'singleton', 'bruteforce', or None, got {other:?}" + ))), + } +} + +/// Create a sketch database and return `(skm_path, skd_path)`. +#[pyfunction] +#[pyo3(signature = ( + output_prefix, + kmer_length, + file_list=None, + seq_files=None, + sketch_size=1000, + threads=1, + min_count=DEFAULT_MINCOUNT, + min_qual=DEFAULT_MINQUAL, + single_strand=false, + quiet=true +))] +fn sketch( + output_prefix: String, + kmer_length: usize, + file_list: Option, + seq_files: Option>, + sketch_size: u64, + threads: usize, + min_count: u16, + min_qual: u8, + single_strand: bool, + quiet: bool, +) -> PyResult<(String, String)> { + let paths = api::sketch(SketchOptions { + seq_files, + file_list, + output: output_prefix, + kmers: vec![kmer_length], + sketch_size, + seq_type: HashType::DNA, + level: DEFAULT_LEVEL, + concat_fasta: false, + single_strand, + min_count, + min_qual, + threads, + quiet, + }) + .map_err(py_runtime_error)?; + Ok((paths.skm, paths.skd)) +} + +/// Calculate distances and return the output path, or `None` when writing to stdout. +#[pyfunction] +#[pyo3(signature = ( + ref_db, + output=None, + query_db=None, + knn=None, + subset=None, + kmer=None, + ani=false, + threads=1, + ref_completeness_file=None, + query_completeness_file=None, + completeness_cutoff=0.64, + quiet=true +))] +fn dist( + ref_db: String, + output: Option, + query_db: Option, + knn: Option, + subset: Option, + kmer: Option, + ani: bool, + threads: usize, + ref_completeness_file: Option, + query_completeness_file: Option, + completeness_cutoff: f64, + quiet: bool, +) -> PyResult> { + api::dist(DistOptions { + ref_db, + query_db, + output, + knn, + subset, + kmer, + ani, + threads, + ref_completeness_file, + query_completeness_file, + completeness_cutoff, + quiet, + }) + .map_err(py_runtime_error) +} + +/// Build an inverted index and return `(ski_path, skq_path_or_none)`. +#[pyfunction] +#[pyo3(signature = ( + output_prefix, + kmer_length, + file_list=None, + seq_files=None, + sketch_size=1000, + threads=1, + write_skq=true, + species_names=None, + metadata=None, + min_count=DEFAULT_MINCOUNT, + min_qual=DEFAULT_MINQUAL, + single_strand=false, + quiet=true +))] +fn inverted_build( + output_prefix: String, + kmer_length: usize, + file_list: Option, + seq_files: Option>, + sketch_size: u64, + threads: usize, + write_skq: bool, + species_names: Option, + metadata: Option, + min_count: u16, + min_qual: u8, + single_strand: bool, + quiet: bool, +) -> PyResult<(String, Option)> { + let paths = api::inverted_build(InvertedBuildOptions { + seq_files, + file_list, + output: output_prefix, + write_skq, + species_names, + metadata, + sketch_size, + kmer_length, + single_strand, + min_count, + min_qual, + threads, + quiet, + }) + .map_err(py_runtime_error)?; + Ok((paths.ski, paths.skq)) +} + +/// Run inverted-index preclustering and return the output path, or `None` for stdout. +#[pyfunction] +#[pyo3(signature = ( + ski, + skd, + output=None, + knn=50, + ani=false, + threads=1, + ref_completeness_file=None, + completeness_cutoff=0.64, + retain_unmatched=None, + quiet=true +))] +fn inverted_precluster( + ski: String, + skd: String, + output: Option, + knn: usize, + ani: bool, + threads: usize, + ref_completeness_file: Option, + completeness_cutoff: f64, + retain_unmatched: Option<&str>, + quiet: bool, +) -> PyResult> { + api::inverted_precluster(InvertedPreclusterOptions { + ski, + skd, + output, + knn, + ani, + threads, + ref_completeness_file, + completeness_cutoff, + retain_unmatched: parse_retain_unmatched(retain_unmatched)?, + quiet, + }) + .map_err(py_runtime_error) +} + +/// Compute query-vs-reference distances and return the output path. +#[pyfunction] +#[pyo3(signature = ( + reference_skm, + query_skm, + output, + kmer_length, + knn, + threads=1, + ref_completeness_file=None, + query_completeness_file=None, + completeness_cutoff=0.64, + quiet=true +))] +fn query_dist( + reference_skm: String, + query_skm: String, + output: String, + kmer_length: usize, + knn: usize, + threads: usize, + ref_completeness_file: Option, + query_completeness_file: Option, + completeness_cutoff: f64, + quiet: bool, +) -> PyResult { + api::query_dist( + reference_skm, + query_skm, + output, + kmer_length, + knn, + threads, + ref_completeness_file, + query_completeness_file, + completeness_cutoff, + quiet, + ) + .map_err(py_runtime_error) +} + +/// Return database metadata as a dict. +#[pyfunction] +#[pyo3(signature = (db_file, sample_info=false))] +fn db_info(py: Python<'_>, db_file: String, sample_info: bool) -> PyResult { + let info = api::db_info(&db_file, sample_info).map_err(py_runtime_error)?; + let dict = PyDict::new_bound(py); + dict.set_item("text", info.text)?; + dict.set_item("inverted", info.inverted)?; + dict.set_item("sketch_size", info.sketch_size)?; + dict.set_item("n_samples", info.n_samples)?; + dict.set_item("kmers", info.kmers)?; + Ok(dict.into()) +} + +/// Return available binding functions. +#[pyfunction] +fn functions() -> HashMap<&'static str, &'static str> { + HashMap::from([ + ("sketch", "create .skm/.skd sketch files"), + ("dist", "calculate distances from sketch files"), + ("inverted_build", "create .ski/.skq inverted index files"), + ("inverted_precluster", "calculate preclustered distances"), + ("query_dist", "calculate query-vs-reference distances"), + ("db_info", "read sketch database metadata"), + ]) +} + +/// Python module entrypoint. +#[pymodule] +fn sketchlib(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(sketch, m)?)?; + m.add_function(wrap_pyfunction!(dist, m)?)?; + m.add_function(wrap_pyfunction!(inverted_build, m)?)?; + m.add_function(wrap_pyfunction!(inverted_precluster, m)?)?; + m.add_function(wrap_pyfunction!(query_dist, m)?)?; + m.add_function(wrap_pyfunction!(db_info, m)?)?; + m.add_function(wrap_pyfunction!(functions, m)?)?; + Ok(()) +} diff --git a/src/sketch/mod.rs b/src/sketch/mod.rs index e40dfe6..ea9c5a8 100644 --- a/src/sketch/mod.rs +++ b/src/sketch/mod.rs @@ -19,8 +19,6 @@ use super::hashing::{nthash_iterator::NtHashIterator, HashType}; use crate::hashing::aahash_iterator::AaHashIterator; #[cfg(not(target_arch = "wasm32"))] use crate::io::InputFastx; -#[cfg(feature = "3di")] -use crate::structures::pdb_to_3di; #[cfg(not(target_arch = "wasm32"))] use crate::utils::get_progress_bar; @@ -284,7 +282,6 @@ pub fn sketch_files( output_prefix: &str, input_files: &[InputFastx], concat_fasta: bool, - #[cfg(feature = "3di")] convert_pdb: bool, k: &[usize], sketch_size: u64, seq_type: &HashType, @@ -297,18 +294,6 @@ pub fn sketch_files( let kmer_stride = (sketch_size * BBITS) as usize; let sample_stride = kmer_stride * k.len(); - #[cfg(feature = "3di")] - let struct_strings = if convert_pdb { - log::info!("Converting PDB files into 3Di representations"); - Some(pdb_to_3di(input_files).expect("Error converting to 3Di")) - } else { - None - }; - #[cfg(not(feature = "3di"))] - let struct_strings: Option> = None; - - log::trace!("{struct_strings:?}"); - // Open output file let data_filename = format!("{output_prefix}.skd"); let mut serial_writer = @@ -326,8 +311,7 @@ pub fn sketch_files( input_files .par_iter() .progress_with(progress_bar) - .enumerate() - .map(|(idx, (name, fastxvec))| { + .map(|(name, fastxvec)| { // Read in sequence and set up rolling hash by alphabet type let mut hash_its: Vec> = match seq_type { HashType::DNA => NtHashIterator::new(fastxvec, k[0], rc, min_qual) @@ -340,20 +324,6 @@ pub fn sketch_files( .map(|it| Box::new(it) as Box) .collect() } - HashType::PDB => { - if let Some(di) = &struct_strings { - log::trace!("Length of string: {}", di.len()); - AaHashIterator::from_3di_string(di[idx].clone()) // TODO: clone is not ideal - .into_iter() - .map(|it| Box::new(it) as Box) - .collect() - } else { - AaHashIterator::from_3di_file(fastxvec) - .into_iter() - .map(|it| Box::new(it) as Box) - .collect() - } - } }; hash_its diff --git a/src/structures.rs b/src/structures.rs deleted file mode 100644 index c5f87f2..0000000 --- a/src/structures.rs +++ /dev/null @@ -1,65 +0,0 @@ -//! Support for .pdb files and the 3di alphabet -#[cfg(feature = "3di")] -use crate::io::InputFastx; -#[cfg(feature = "3di")] -use anyhow::Error; - -#[cfg(feature = "3di")] -use indicatif::{ProgressIterator, ProgressStyle}; -#[cfg(feature = "3di")] -use pyo3::prelude::*; - -#[cfg(feature = "3di")] -/// Uses python library to convert pdb files into 1D 3di representation -pub fn pdb_to_3di(input_files: &[InputFastx]) -> Result, Error> { - pyo3::prepare_freethreaded_python(); - let py_file = include_str!(concat!( - env!("CARGO_MANIFEST_DIR"), - "/python_mini3di/3di_convert.py" - )); - - let mut struct_strings = Vec::with_capacity(input_files.len()); - let bar_style = - ProgressStyle::with_template("{human_pos}/{human_len} {bar:80.cyan/blue} eta:{eta}") - .unwrap(); - Python::with_gil(|py| { - let converter_fun: Py = PyModule::from_code_bound(py, py_file, "", "") - .unwrap() - .getattr("pdb_to_3di") - .unwrap() - .into(); - // println!("n f1 f2: {:?} {:?} {:?}", input_files[0], input_files[1], input_files[2]); - - for ftup in input_files.iter().progress_with_style(bar_style) { - // println!("{:?}", ftup); - if ftup.1.len() == 1 { - let struct_string: Py = converter_fun - .call1(py, (ftup.0.clone(), ftup.1[0].clone())) - .unwrap(); - struct_strings.push(struct_string.extract::(py).unwrap()); - } else { - let mut struct_string = "".to_owned(); - for file in ftup.1.iter() { - let struct_string_tmp: Py = - converter_fun.call1(py, (ftup.0.clone(), file)).unwrap(); - struct_string.push_str(","); - struct_string - .push_str(struct_string_tmp.extract::(py).unwrap().as_str()); - } - struct_strings.push(struct_string); - } - } - }); - - // println!("Finished!"); - - Ok(struct_strings) -} - -// This shouldn't be needed, but I am putting it here just in case -/* -#[cfg(not(feature = "3di"))] -pub fn pdb_to_3di(_input_files: &[InputFastx]) -> Result, Error> { - unimplemented!("This executable was not compiled with the 3di feature") -} -*/ diff --git a/tests/api.rs b/tests/api.rs new file mode 100644 index 0000000..bdec152 --- /dev/null +++ b/tests/api.rs @@ -0,0 +1,64 @@ +use std::path::Path; + +use sketchlib::api::{self, DistOptions, SketchOptions}; +use sketchlib::cli::{DEFAULT_MINCOUNT, DEFAULT_MINQUAL}; +use sketchlib::hashing::{HashType, DEFAULT_LEVEL}; + +pub mod common; +use crate::common::*; + +#[test] +fn api_sketch_dist_and_repeated_thread_setup() { + let sandbox = TestSetup::setup(); + let rfile = sandbox.create_named_rfile( + "api_file_list.txt", + &["R6.fa.gz", "TIGR4.fa.gz", "14412_3#82.contigs_velvet.fa.gz"], + ); + let rfile = sandbox.file_string(&rfile, TestDir::Output); + + for output_name in ["api_db1", "api_db2"] { + let output = sandbox.file_string(output_name, TestDir::Output); + let paths = api::sketch(SketchOptions { + seq_files: None, + file_list: Some(rfile.clone()), + output, + kmers: vec![21], + sketch_size: 1000, + seq_type: HashType::DNA, + level: DEFAULT_LEVEL, + concat_fasta: false, + single_strand: false, + min_count: DEFAULT_MINCOUNT, + min_qual: DEFAULT_MINQUAL, + threads: 2, + quiet: true, + }) + .expect("API sketch should succeed"); + assert!(Path::new(&paths.skm).exists()); + assert!(Path::new(&paths.skd).exists()); + } + + let db_prefix = sandbox.file_string("api_db1", TestDir::Output); + let info = api::db_info(&db_prefix, false).expect("API db_info should succeed"); + assert!(!info.inverted); + assert_eq!(info.kmers, vec![21]); + assert_eq!(info.n_samples, 3); + + let dist_path = sandbox.file_string("api.dists", TestDir::Output); + api::dist(DistOptions { + ref_db: db_prefix, + query_db: None, + output: Some(dist_path.clone()), + knn: Some(1), + subset: None, + kmer: Some(21), + ani: true, + threads: 2, + ref_completeness_file: None, + query_completeness_file: None, + completeness_cutoff: 0.64, + quiet: true, + }) + .expect("API dist should succeed"); + assert!(Path::new(&dist_path).exists()); +}