Bring optimised Jaccard code into main

Cargo.toml

@@ -1,12 +1,13 @@
 [package]
 name = "sketchlib"
-version = "0.3.0"
+version = "0.4.0"
 authors = [
     "John Lees <jlees@ebi.ac.uk>",
     "Nicholas Croucher <n.croucher@imperial.ac.uk>",
     "Johanna von Wachsmann <wachsmannj@ebi.ac.uk>",
     "Victor Rodriguez Bouza <vrbouza@ebi.ac.uk>",
-    "Joel Hellewell <joel@ebi.ac.uk",
+    "Matthew Russell <marussel@ebi.ac.uk>",
+    "Tommi Maklin",
 ]
 edition = "2021"
 description = "Genome and amino-acid sketching"
@@ -31,8 +32,10 @@ crate-type = ["cdylib", "lib"]
 [features]
 3di = ["dep:pyo3", "dep:pyo3-build-config"]
 
+
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 # data structures
+aligned-vec = "0.6.4"
 hashbrown = { version = "0.14", features = ["rayon", "serde"] }
 ahash = ">=0.8.7"
 roaring = { version = "0.10", features = ["serde"] }
@@ -63,6 +66,7 @@ num_cpus = "1.0"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 # data structures
+aligned-vec = "0.6.4"
 hashbrown = { version = "0.14", features = ["rayon", "serde"] }
 # ahash = ">=0.8.7"
 roaring = { version = "0.10", features = ["serde"] }

README.md

@@ -15,8 +15,11 @@ This is a reimplementation and extension of [pp-sketchlib](https://github.com/ba
 in the rust language. This version is optimised for larger sample numbers, particularly
 allowing subsets of samples to be compared.
 
-v0.2.0 is the first stable release. We intend to keep the file format unchanged after this
-point so sketch libraries will not need to be rebuilt.
+### News
+
+- **Data format changed in v0.4.0**. If you have an older database (see `sketchlib info <db_prefix> | grep "sketch_version"`)
+you should resketch samples with >v0.4. Otherwise you will need to use an older release available on crates.io.
+- v0.2.0 was the first stable release.
 
 ## Documentation
 

src/distances/jaccard.rs

@@ -1,6 +1,6 @@
 //! Implementation of Jaccard, core and accessory distance calculations
 use crate::sketch::multisketch::MultiSketch;
-use crate::sketch::BBITS;
+use crate::sketch::BIN_BITS;
 
 /// Returns the Jaccard index between two samples
 pub fn jaccard_index(
@@ -11,37 +11,180 @@
     c2: Option<f64>,
     completeness_cutoff: f64,
 ) -> f64 {
+    let samebits = jaccard_same_bits(sketch1, sketch2) as f64;
     let unionsize = (u64::BITS as u64 * sketchsize64) as f64;
-    let samebits: u32 = sketch1
-        .chunks_exact(BBITS as usize)
-        .zip(sketch2.chunks_exact(BBITS as usize))
+    // Correction for random matches
+    let expected_random = unionsize / (1u64 << BIN_BITS) as f64;
+    let mut jaccard_index =
+        ((samebits - expected_random) / (unionsize - expected_random)).clamp(0.0, 1.0);
+
+    log::trace!("samebits:{samebits} expected_random:{expected_random} jaccard:{jaccard_index}");
+
+    // Apply completeness correction if both completeness values are provided
+    if let (Some(c1_val), Some(c2_val)) = (c1, c2) {
+        if c1_val * c2_val >= completeness_cutoff {
+            // Cap the corrected Jaccard index at 1.0 to prevent negative distances
+            jaccard_index = completeness_correction(jaccard_index, c1_val, c2_val).min(1.0);
+        }
+    }
+
+    jaccard_index
+}
+
+#[inline(always)]
+fn jaccard_same_bits(sketch1: &[u64], sketch2: &[u64]) -> u32 {
+    #[cfg(target_arch = "aarch64")]
+    unsafe {
+        jaccard_neon_unroll2_inner(sketch1, sketch2)
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        jaccard_same_bits_general(sketch1, sketch2)
+    }
+}
+
+#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
+#[inline(always)]
+pub(crate) fn jaccard_same_bits_general(sketch1: &[u64], sketch2: &[u64]) -> u32 {
+    debug_assert_eq!(sketch1.len(), sketch2.len());
+    debug_assert_eq!(sketch1.len() % BIN_BITS, 0);
+    sketch1
+        .chunks_exact(BIN_BITS)
+        .zip(sketch2.chunks_exact(BIN_BITS))
         .map(|(chunk1, chunk2)| {
             let mut bits: u64 = !0;
             chunk1.iter().zip(chunk2.iter()).for_each(|(&s1, &s2)| {
                 bits &= !(s1 ^ s2);
             });
             bits.count_ones()
         })
-        .sum();
-    let maxnbits = sketchsize64 as u32 * u64::BITS;
-    let expected_samebits = maxnbits >> BBITS;
+        .sum()
+}
 
-    log::trace!("samebits:{samebits} expected_samebits:{expected_samebits} maxnbits:{maxnbits}");
-    let diff = samebits.saturating_sub(expected_samebits);
-    let intersize = (diff as f64 * maxnbits as f64) / (maxnbits - expected_samebits) as f64;
-    log::trace!("intersize:{intersize} unionsize:{unionsize}");
-    let mut jaccard_index = intersize / unionsize;
+#[cfg(target_arch = "aarch64")]
+#[target_feature(enable = "neon")]
+unsafe fn jaccard_neon_unroll2_inner(a: &[u64], b: &[u64]) -> u32 {
+    use std::arch::aarch64::*;
 
-    // Apply completeness correction if both completeness values are provided
-    if let (Some(c1_val), Some(c2_val)) = (c1, c2) {
-        if c1_val * c2_val >= completeness_cutoff {
-            jaccard_index = completeness_correction(jaccard_index, c1_val, c2_val);
-            // Cap the corrected Jaccard index at 1.0 to prevent negative distances
-            jaccard_index = jaccard_index.min(1.0);
+    debug_assert_eq!(a.len(), b.len());
+    debug_assert_eq!(a.len() % BIN_BITS, 0);
+
+    let chunk = BIN_BITS;
+    let n_chunks = a.len() / chunk;
+    let n_pairs = n_chunks / 2;
+    let mut total = 0u32;
+
+    #[inline(always)]
+    unsafe fn one_chunk(ap: *const u8, bp: *const u8) -> u8 {
+        let a0 = vld1q_u8(ap);
+        let a1 = vld1q_u8(ap.add(16));
+        let a2 = vld1q_u8(ap.add(32));
+        let a3 = vld1q_u8(ap.add(48));
+        let a4 = vld1q_u8(ap.add(64));
+        let a5 = vld1q_u8(ap.add(80));
+        let a6 = vld1q_u8(ap.add(96));
+        let a7 = vld1q_u8(ap.add(112));
+        let b0 = vld1q_u8(bp);
+        let b1 = vld1q_u8(bp.add(16));
+        let b2 = vld1q_u8(bp.add(32));
+        let b3 = vld1q_u8(bp.add(48));
+        let b4 = vld1q_u8(bp.add(64));
+        let b5 = vld1q_u8(bp.add(80));
+        let b6 = vld1q_u8(bp.add(96));
+        let b7 = vld1q_u8(bp.add(112));
+        let x0 = veorq_u8(a0, b0);
+        let x1 = veorq_u8(a1, b1);
+        let x2 = veorq_u8(a2, b2);
+        let x3 = veorq_u8(a3, b3);
+        let x4 = veorq_u8(a4, b4);
+        let x5 = veorq_u8(a5, b5);
+        let x6 = veorq_u8(a6, b6);
+        let x7 = veorq_u8(a7, b7);
+        let or01 = vorrq_u8(x0, x1);
+        let or23 = vorrq_u8(x2, x3);
+        let or45 = vorrq_u8(x4, x5);
+        let or67 = vorrq_u8(x6, x7);
+        let or_all = vorrq_u8(vorrq_u8(or01, or23), vorrq_u8(or45, or67));
+        let not_or = vmvnq_u8(or_all);
+        vaddv_u8(vcnt_u8(vand_u8(vget_low_u8(not_or), vget_high_u8(not_or))))
+    }
+
+    for i in 0..n_pairs {
+        let base = i * 2 * chunk;
+        total += one_chunk(a.as_ptr().add(base) as _, b.as_ptr().add(base) as _) as u32;
+        total += one_chunk(
+            a.as_ptr().add(base + chunk) as _,
+            b.as_ptr().add(base + chunk) as _,
+        ) as u32;
+    }
+    if n_chunks % 2 != 0 {
+        let base = (n_chunks - 1) * chunk;
+        total += one_chunk(a.as_ptr().add(base) as _, b.as_ptr().add(base) as _) as u32;
+    }
+    total
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn scalar_same_bits_counts_matching_bins() {
+        let sketch1 = [u64::MAX; BIN_BITS * 2];
+        let mut sketch2 = [u64::MAX; BIN_BITS * 2];
+        sketch2[BIN_BITS] = 0;
+
+        assert_eq!(jaccard_same_bits_general(&sketch1, &sketch1), 128);
+        assert_eq!(jaccard_same_bits_general(&sketch1, &sketch2), 64);
+    }
+
+    #[cfg(target_arch = "aarch64")]
+    #[test]
+    fn neon_same_bits_matches_scalar() {
+        let mut sketch1 = [0u64; BIN_BITS * 3];
+        let mut sketch2 = [0u64; BIN_BITS * 3];
+        for i in 0..sketch1.len() {
+            sketch1[i] = (i as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15);
+            sketch2[i] = sketch1[i] ^ ((i as u64) << (i % 17));
         }
+
+        let scalar = jaccard_same_bits_general(&sketch1, &sketch2);
+        let neon = unsafe { jaccard_neon_unroll2_inner(&sketch1, &sketch2) };
+        assert_eq!(neon, scalar);
     }
 
-    jaccard_index
+    #[test]
+    fn random_match_correction_applied_for_large_sketch_size() {
+        // sketchsize64 chosen so that unionsize = sketchsize64 * 64 exceeds
+        // 2^BIN_BITS, the regime the random-match correction is meant for.
+        let sketchsize64: u64 = 2048;
+        let unionsize = (u64::BITS as u64 * sketchsize64) as f64;
+        assert!(unionsize > (1u64 << BIN_BITS) as f64);
+
+        let n_chunks = sketchsize64 as usize;
+        let sketch1 = vec![u64::MAX; BIN_BITS * n_chunks];
+        let mut sketch2 = sketch1.clone();
+        // Mismatch exactly one chunk (one 64-bit "column") out of n_chunks.
+        for word in sketch2[..BIN_BITS].iter_mut() {
+            *word = 0;
+        }
+
+        let samebits = jaccard_same_bits_general(&sketch1, &sketch2) as f64;
+        let expected_random = unionsize / (1u64 << BIN_BITS) as f64;
+        let expected_jaccard =
+            ((samebits - expected_random) / (unionsize - expected_random)).clamp(0.0, 1.0);
+        // The correction should be non-trivial (i.e. not equal to the raw,
+        // uncorrected ratio) for this sketch size.
+        assert!((expected_jaccard - samebits / unionsize).abs() > f64::EPSILON);
+
+        let jaccard = jaccard_index(&sketch1, &sketch2, sketchsize64, None, None, 0.0);
+        assert_eq!(jaccard, expected_jaccard);
+
+        // Identical sketches should give a jaccard index of exactly 1.0.
+        let identical = jaccard_index(&sketch1, &sketch1, sketchsize64, None, None, 0.0);
+        assert_eq!(identical, 1.0);
+    }
 }
 
 /// Converts between Jaccard distance and ANI, using a Poisson model of mutations

src/distances/mod.rs

@@ -359,8 +359,7 @@ pub fn cross_dists_knn<'a>(
                     if ani {
                         heap.into_sorted_vec().iter().zip(row_dist_slice).for_each(
                             |(inverse_ani, output_ani)| {
-                                *output_ani =
-                                    SparseJaccard(inverse_ani.0, 1.0_f32 - inverse_ani.1);
+                                *output_ani = SparseJaccard(inverse_ani.0, 1.0_f32 - inverse_ani.1);
                             },
                         );
                     } else {

src/hashing/aahash_iterator.rs

@@ -79,11 +79,11 @@ impl AaHashIterator {
     }
 
     /// Create a new aaHash iterator from a fastX reader, at the set comparison level
-    pub fn new<I: Iterator<Item=(Vec<u8>, Option<Vec<u8>>)>>(
+    pub fn new<I: Iterator<Item = (Vec<u8>, Option<Vec<u8>>)>>(
         readers: &mut [I],
         file: &String,
         level: AaLevel,
-        concat_fasta: bool
+        concat_fasta: bool,
     ) -> Vec<Self> {
         let mut hash_vec = Vec::new();
 
@@ -124,7 +124,7 @@ impl AaHashIterator {
     }
 
     /// Create a new iterator from a 3di embedding file of a structure
-    pub fn from_3di_file<I: Iterator<Item=(Vec<u8>, Option<Vec<u8>>)>>(
+    pub fn from_3di_file<I: Iterator<Item = (Vec<u8>, Option<Vec<u8>>)>>(
         readers: &mut [I],
         file: &String,
     ) -> Vec<Self> {

src/hashing/nthash_iterator.rs

@@ -89,14 +89,13 @@ impl RollHash for NtHashIterator {
 impl NtHashIterator {
     #[cfg(not(target_arch = "wasm32"))]
     /// Creates a new ntHash iterator, by loading DNA sequences into memory
-    pub fn new<I: Iterator<Item=(Vec<u8>, Option<Vec<u8>>)>>(
+    pub fn new<I: Iterator<Item = (Vec<u8>, Option<Vec<u8>>)>>(
         files: &mut [I],
         k: usize,
         rc: bool,
         min_qual: u8,
         reads: bool,
     ) -> Vec<Self> {
-
         let mut seq = Vec::new();
         let mut offsets = Vec::new();
         let mut acgt = [0, 0, 0, 0];
@@ -192,15 +191,14 @@ impl NtHashIterator {
     }
 
     #[cfg(not(target_arch = "wasm32"))]
-    fn add_dna_seq<I: Iterator<Item=(Vec<u8>, Option<Vec<u8>>)>>(
+    fn add_dna_seq<I: Iterator<Item = (Vec<u8>, Option<Vec<u8>>)>>(
         records: &mut I,
         min_qual: u8,
         seq: &mut Vec<u8>,
         offsets: &mut Vec<usize>,
         acgt: &mut [usize; 4],
         non_acgt: &mut usize,
     ) {
-
         // Current byte
         let mut b = 0;
         // Nucleotide index within byte 0-3

src/inverted.rs

@@ -27,12 +27,13 @@ use super::hashing::{
 use crate::distances::distance_matrix::square_to_condensed;
 #[cfg(not(target_arch = "wasm32"))]
 use crate::io::InputFastx;
+use crate::sketch::multisketch::{parse_version, MIN_SKETCH_VERSION};
 #[cfg(target_arch = "wasm32")]
 use crate::sketch::Sketch;
 #[cfg(not(target_arch = "wasm32"))]
 use crate::sketch::{sketch_datafile::SketchArrayWriter, Sketch};
 use crate::utils::get_progress_bar;
-use anyhow::Error;
+use anyhow::{bail, Error};
 use std::fs::File;
 use std::io::{BufReader, BufWriter};
 
@@ -43,6 +44,28 @@ use wasm_bindgen_file_reader::WebSysFile;
 
 type InvSketches = (Vec<Vec<u16>>, Vec<String>);
 
+/// Fails if a loaded .ski index was built with a sketchlib version older
+/// than the current minimum sketch format version.
+fn check_version(sketch_version: &str) -> Result<(), Error> {
+    let version_ok =
+        parse_version(sketch_version).is_some_and(|version| version >= MIN_SKETCH_VERSION);
+    if !version_ok {
+        let found_version = if sketch_version.is_empty() {
+            "<unknown>"
+        } else {
+            sketch_version
+        };
+        log::error!(
+            "Inverted index was created with sketchlib v{found_version}, which is older than the minimum supported v{}.{}.{}. Please re-sketch with the current version.",
+            MIN_SKETCH_VERSION.0,
+            MIN_SKETCH_VERSION.1,
+            MIN_SKETCH_VERSION.2
+        );
+        bail!("Incompatible sketch file version");
+    }
+    Ok(())
+}
+
 /// An inverted index and associated metadata
 #[derive(Serialize, Deserialize, Default, Clone, PartialEq)]
 pub struct Inverted {
@@ -210,6 +233,7 @@ impl Inverted {
         let ski_file = BufReader::new(File::open(filename)?);
         let decompress_reader = snap::read::FrameDecoder::new(ski_file);
         let ski_obj: Self = rmp_serde::decode::from_read(decompress_reader)?;
+        check_version(&ski_obj.sketch_version)?;
         Ok(ski_obj)
     }
 
@@ -221,6 +245,7 @@ impl Inverted {
         let ski_file = BufReader::new(WebSysFile::new(file.clone()));
         let decompress_reader = snap::read::FrameDecoder::new(ski_file);
         let ski_obj: Self = rmp_serde::decode::from_read(decompress_reader)?;
+        check_version(&ski_obj.sketch_version)?;
         Ok(ski_obj)
     }