MRG: improve downsampling behavior on KmerMinHash; fix RevIndex::gather bug around scaled.

include/sourmash.h

@@ -30,6 +30,10 @@ enum SourmashErrorCode {
   SOURMASH_ERROR_CODE_NON_EMPTY_MIN_HASH = 106,
   SOURMASH_ERROR_CODE_MISMATCH_NUM = 107,
   SOURMASH_ERROR_CODE_NEEDS_ABUNDANCE_TRACKING = 108,
+  SOURMASH_ERROR_CODE_CANNOT_UPSAMPLE_SCALED = 109,
+  SOURMASH_ERROR_CODE_NO_MIN_HASH_FOUND = 110,
+  SOURMASH_ERROR_CODE_EMPTY_SIGNATURE = 111,
+  SOURMASH_ERROR_CODE_MULTIPLE_SKETCHES_FOUND = 112,
   SOURMASH_ERROR_CODE_INVALID_DNA = 1101,
   SOURMASH_ERROR_CODE_INVALID_PROT = 1102,
   SOURMASH_ERROR_CODE_INVALID_CODON_LENGTH = 1103,

src/core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sourmash"
-version = "0.15.2"
+version = "0.16.0"
 authors = ["Luiz Irber <[email protected]>", "N. Tessa Pierce-Ward <[email protected]>"]
 description = "tools for comparing biological sequences with k-mer sketches"
 repository = "https://github.com/sourmash-bio/sourmash"

src/core/src/collection.rs

@@ -241,6 +241,7 @@ mod test {
     use crate::prelude::Select;
     use crate::selection::Selection;
     use crate::signature::Signature;
+    #[cfg(all(feature = "branchwater", not(target_arch = "wasm32")))]
     use crate::Result;
 
     #[test]

src/core/src/errors.rs

@@ -7,6 +7,9 @@ pub enum SourmashError {
     #[error("internal error: {message:?}")]
     Internal { message: String },
 
+    #[error("new scaled smaller than previous; cannot upsample")]
+    CannotUpsampleScaled,
+
     #[error("must have same num: {n1} != {n2}")]
     MismatchNum { n1: u32, n2: u32 },
 
@@ -28,6 +31,15 @@ pub enum SourmashError {
     #[error("sketch needs abundance for this operation")]
     NeedsAbundanceTracking,
 
+    #[error("Expected a MinHash sketch in this signature")]
+    NoMinHashFound,
+
+    #[error("Empty signature")]
+    EmptySignature,
+
+    #[error("Multiple sketches found, expected one")]
+    MultipleSketchesFound,
+
     #[error("Invalid hash function: {function:?}")]
     InvalidHashFunction { function: String },
 
@@ -104,6 +116,10 @@ pub enum SourmashErrorCode {
     NonEmptyMinHash = 1_06,
     MismatchNum = 1_07,
     NeedsAbundanceTracking = 1_08,
+    CannotUpsampleScaled = 1_09,
+    NoMinHashFound = 1_10,
+    EmptySignature = 1_11,
+    MultipleSketchesFound = 1_12,
     // Input sequence errors
     InvalidDNA = 11_01,
     InvalidProt = 11_02,
@@ -132,6 +148,7 @@ impl SourmashErrorCode {
         match error {
             SourmashError::Internal { .. } => SourmashErrorCode::Internal,
             SourmashError::Panic { .. } => SourmashErrorCode::Panic,
+            SourmashError::CannotUpsampleScaled { .. } => SourmashErrorCode::CannotUpsampleScaled,
             SourmashError::MismatchNum { .. } => SourmashErrorCode::MismatchNum,
             SourmashError::NeedsAbundanceTracking { .. } => {
                 SourmashErrorCode::NeedsAbundanceTracking
@@ -142,6 +159,9 @@ impl SourmashErrorCode {
             SourmashError::MismatchSeed => SourmashErrorCode::MismatchSeed,
             SourmashError::MismatchSignatureType => SourmashErrorCode::MismatchSignatureType,
             SourmashError::NonEmptyMinHash { .. } => SourmashErrorCode::NonEmptyMinHash,
+            SourmashError::NoMinHashFound => SourmashErrorCode::NoMinHashFound,
+            SourmashError::EmptySignature => SourmashErrorCode::EmptySignature,
+            SourmashError::MultipleSketchesFound => SourmashErrorCode::MultipleSketchesFound,
             SourmashError::InvalidDNA { .. } => SourmashErrorCode::InvalidDNA,
             SourmashError::InvalidProt { .. } => SourmashErrorCode::InvalidProt,
             SourmashError::InvalidCodonLength { .. } => SourmashErrorCode::InvalidCodonLength,

src/core/src/index/mod.rs

@@ -17,6 +17,7 @@ use getset::{CopyGetters, Getters, Setters};
 use log::trace;
 use serde::{Deserialize, Serialize};
 use stats::{median, stddev};
+use std::cmp::max;
 use typed_builder::TypedBuilder;
 
 use crate::ani_utils::{ani_ci_from_containment, ani_from_containment};
@@ -205,7 +206,6 @@ where
     }
 }
 
-// note all mh should be selected/downsampled prior to being passed in here.
 #[allow(clippy::too_many_arguments)]
 pub fn calculate_gather_stats(
     orig_query: &KmerMinHash,
@@ -220,10 +220,21 @@ pub fn calculate_gather_stats(
     confidence: Option<f64>,
 ) -> Result<GatherResult> {
     // get match_mh
-    let match_mh = match_sig.minhash().unwrap();
+    let match_mh = match_sig.minhash().expect("cannot retrieve sketch");
+
+    let max_scaled = max(match_mh.scaled(), query.scaled());
+    let query = query
+        .downsample_scaled(max_scaled)
+        .expect("cannot downsample query");
+    let match_mh = match_mh
+        .clone()
+        .downsample_scaled(max_scaled)
+        .expect("cannot downsample match");
 
     // calculate intersection
-    let isect = match_mh.intersection(&query)?;
+    let isect = match_mh
+        .intersection(&query)
+        .expect("could not do intersection");
     let isect_size = isect.0.len();
     trace!("isect_size: {}", isect_size);
     trace!("query.size: {}", query.size());

src/core/src/index/revindex/disk_revindex.rs

@@ -372,7 +372,6 @@ impl RevIndexOps for RevIndex {
         let mut query = KmerMinHashBTree::from(orig_query.clone());
         let mut sum_weighted_found = 0;
         let _selection = selection.unwrap_or_else(|| self.collection.selection());
-        let mut orig_query_ds = orig_query.clone();
         let total_weighted_hashes = orig_query.sum_abunds();
 
         // or set this with user --track-abundance?
@@ -393,29 +392,33 @@ impl RevIndexOps for RevIndex {
                 break;
             }
 
-            // this should downsample mh for us
             let match_sig = self.collection.sig_for_dataset(dataset_id)?;
 
             // get downsampled minhashes for comparison.
             let match_mh = match_sig.minhash().unwrap().clone();
-            query = query.downsample_scaled(match_mh.scaled())?;
-            orig_query_ds = orig_query_ds.downsample_scaled(match_mh.scaled())?;
+            let scaled = query.scaled();
+
+            let match_mh = match_mh
+                .downsample_scaled(scaled)
+                .expect("cannot downsample match");
 
             // just calculate essentials here
             let gather_result_rank = matches.len();
 
             let query_mh = KmerMinHash::from(query.clone());
 
             // grab the specific intersection:
-            let isect = match_mh.intersection(&query_mh)?;
+            let isect = match_mh
+                .intersection(&query_mh)
+                .expect("failed to intersect");
             let mut isect_mh = match_mh.clone();
             isect_mh.clear();
             isect_mh.add_many(&isect.0)?;
 
             // Calculate stats
             let gather_result = calculate_gather_stats(
-                &orig_query_ds,
-                KmerMinHash::from(query.clone()),
+                &orig_query,
+                query_mh,
                 match_sig,
                 match_size,
                 gather_result_rank,
@@ -424,7 +427,8 @@ impl RevIndexOps for RevIndex {
                 calc_abund_stats,
                 calc_ani_ci,
                 ani_confidence_interval_fraction,
-            )?;
+            )
+            .expect("could not calculate gather stats");
             // keep track of the sum weighted found
             sum_weighted_found = gather_result.sum_weighted_found();
             matches.push(gather_result);

src/core/src/index/revindex/mod.rs

@@ -903,14 +903,16 @@ mod test {
 
         let (counter, query_colors, hash_to_color) = index.prepare_gather_counters(&query);
 
-        let matches_external = index.gather(
-            counter,
-            query_colors,
-            hash_to_color,
-            0,
-            &query,
-            Some(selection.clone()),
-        )?;
+        let matches_external = index
+            .gather(
+                counter,
+                query_colors,
+                hash_to_color,
+                0,
+                &query,
+                Some(selection.clone()),
+            )
+            .expect("failed to gather!");
 
         {
             let mut index = index;

src/core/src/signature.rs

@@ -842,7 +842,7 @@
                 // TODO: also account for LargeMinHash
                 if let Sketch::MinHash(mh) = sketch {
                     if (mh.scaled() as u32) < sel_scaled {
-                        *sketch = Sketch::MinHash(mh.downsample_scaled(sel_scaled as u64)?);
+                        *sketch = Sketch::MinHash(mh.clone().downsample_scaled(sel_scaled as u64)?);
                     }
                 }
             }
@@ -887,6 +887,28 @@
     }
 }
 
+impl TryInto<KmerMinHash> for Signature {
+    type Error = Error;
+
+    fn try_into(self) -> Result<KmerMinHash, Error> {
+        match self.signatures.len() {
+            1 => self
+                .signatures
+                .into_iter()
+                .find_map(|sk| {
+                    if let Sketch::MinHash(mh) = sk {
+                        Some(mh)
+                    } else {
+                        None
+                    }
+                })
+                .ok_or_else(|| Error::NoMinHashFound),
+            0 => Err(Error::EmptySignature),
+            _ => Err(Error::MultipleSketchesFound),
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::fs::File;