From 20d893b5abeedd1ae4ff612f06693c55547b9e6c Mon Sep 17 00:00:00 2001 From: John Lees Date: Thu, 14 May 2026 17:36:02 +0100 Subject: [PATCH 1/2] Bring optimised Jaccard code into main --- Cargo.toml | 3 + src/distances/jaccard.rs | 160 +++++++++++++++--- src/lib.rs | 4 +- src/sketch/mod.rs | 102 ++++++++--- src/sketch/multisketch.rs | 9 +- src/sketch/sketch_datafile.rs | 48 +++++- tests/sketch.rs | 27 ++- .../test_results_correct/dists_knn_ani.stdout | 8 +- .../test_results_correct/dists_knn_ca.stdout | 8 +- .../dists_knn_jaccard.stdout | 8 +- .../test_results_correct/dists_subset.stdout | 4 +- tests/test_results_correct/inverted.skq | 2 +- .../inverted_precluster.stdout | 8 +- .../inverted_precluster_ani.stdout | 8 +- .../inverted_query_count.stdout | 8 +- .../inverted_sketch_info.stdout | 2 +- .../read_sketch_full_info.stdout | 4 +- 17 files changed, 326 insertions(+), 87 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 37c57c7..8b2cc6e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,9 @@ crate-type = ["cdylib", "lib"] [features] 3di = ["dep:pyo3", "dep:pyo3-build-config"] +[dependencies] +aligned-vec = "0.6.4" + [target.'cfg(not(target_arch = "wasm32"))'.dependencies] # data structures hashbrown = { version = "0.14", features = ["rayon", "serde"] } diff --git a/src/distances/jaccard.rs b/src/distances/jaccard.rs index 2ce8ae6..4cf9d6f 100644 --- a/src/distances/jaccard.rs +++ b/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( @@ -12,25 +12,11 @@ pub fn jaccard_index( completeness_cutoff: f64, ) -> 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)) - .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; + let samebits = jaccard_same_bits(sketch1, sketch2); + let raw_jaccard = samebits as f64 / unionsize; + let mut jaccard_index = random_match_correction(raw_jaccard); - 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; + log::trace!("samebits:{samebits} raw_jaccard:{raw_jaccard} jaccard:{jaccard_index}"); // Apply completeness correction if both completeness values are provided if let (Some(c1_val), Some(c2_val)) = (c1, c2) { @@ -44,6 +30,142 @@ pub fn jaccard_index( 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))] +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() +} + +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +unsafe fn jaccard_neon_unroll2_inner(a: &[u64], b: &[u64]) -> u32 { + use std::arch::aarch64::*; + + 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 +} + +#[inline(always)] +fn random_match_correction(jaccard: f64) -> f64 { + let bb = (1u32 << (BIN_BITS as u32)) as f64; + ((bb * jaccard - 1.0).max(0.0)) / (bb - 1.0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn random_match_correction_maps_expected_random_to_zero() { + let expected_random = 1.0 / ((1u32 << (BIN_BITS as u32)) as f64); + assert_eq!(random_match_correction(expected_random), 0.0); + assert_eq!(random_match_correction(1.0), 1.0); + } + + #[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); + } +} + /// Converts between Jaccard distance and ANI, using a Poisson model of mutations #[inline(always)] pub fn ani_pois(jaccard: f64, k: f64) -> f64 { diff --git a/src/lib.rs b/src/lib.rs index 4d59b65..0557ccb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,8 +9,8 @@ //! above these distances is unsupported and may lead to poor estimation of distances //! without clear warning. //! - Short k-mer lengths are likely to match at random, see [PopPUNK's docs](https://poppunk-docs.bacpop.org/sketching.html#choosing-the-right-k-mer-lengths) -//! for information on how to select good lengths. Note that this library does not support -//! random match correction. +//! for information on how to select good lengths. Jaccard distances include random match +//! correction for the retained sketch signature width. //! - ANI distance resolution is highly affected by sketch size at higher mismatch //! levels, so note that if you see lots of samples at around 80% they may be much lower than //! this. We recommend checking the Jaccard values in this case, if they are close to 0 diff --git a/src/sketch/mod.rs b/src/sketch/mod.rs index 8fa71d3..905c0fd 100644 --- a/src/sketch/mod.rs +++ b/src/sketch/mod.rs @@ -4,6 +4,7 @@ use std::fmt; #[cfg(not(target_arch = "wasm32"))] use std::sync::mpsc; +use aligned_vec::{AVec, ConstAlign}; #[cfg(not(target_arch = "wasm32"))] use indicatif::ParallelProgressIterator; #[cfg(not(target_arch = "wasm32"))] @@ -31,9 +32,28 @@ pub mod sketch_datafile; use self::sketch_datafile::SketchArrayWriter; /// Bin bits (lowest of 64-bits to keep) -pub const BBITS: u64 = 14; -/// Total width of all bins (used as sign % sign_mod) -pub const SIGN_MOD: u64 = (1 << 61) - 1; +pub const BIN_BITS: usize = 16; +#[doc(hidden)] +pub const BBITS: usize = BIN_BITS; +/// Byte alignment used for bit-packed sketch buffers. +pub const SKETCH_ALIGNMENT: usize = 64; + +/// Aligned storage for bit-packed sketch buffers. +pub type SketchVec = AVec>; + +pub(crate) fn aligned_sketch_vec() -> SketchVec { + SketchVec::new(SKETCH_ALIGNMENT) +} + +pub(crate) fn aligned_sketch_vec_with_capacity(capacity: usize) -> SketchVec { + SketchVec::with_capacity(SKETCH_ALIGNMENT, capacity) +} + +pub(crate) fn aligned_sketch_vec_filled(len: usize, value: u64) -> SketchVec { + let mut vec = aligned_sketch_vec_with_capacity(len); + vec.resize(len, value); + vec +} /// Get the number of elements in the sketch vectors for a given sketch size /// @@ -49,15 +69,15 @@ pub const SIGN_MOD: u64 = (1 << 61) - 1; pub fn num_bins(sketch_size: u64) -> (u64, u64, u64) { let sketchsize64 = sketch_size.div_ceil(u64::BITS as u64); let signs_size = sketchsize64 * (u64::BITS as u64); - let usigs_size = sketchsize64 * BBITS; + let usigs_size = sketchsize64 * (BIN_BITS as u64); (sketchsize64, signs_size, usigs_size) } /// A single sample's sketch -#[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq)] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] pub struct Sketch { - #[serde(skip)] - usigs: Vec, + #[serde(skip, default = "aligned_sketch_vec")] + usigs: SketchVec, name: String, index: Option, rc: bool, @@ -68,6 +88,22 @@ pub struct Sketch { non_acgt: usize, } +impl Default for Sketch { + fn default() -> Self { + Self { + usigs: aligned_sketch_vec(), + name: String::default(), + index: None, + rc: false, + reads: false, + seq_length: 0, + densified: false, + acgt: [0; 4], + non_acgt: 0, + } + } +} + // TODO: should this take hash_it and filter as input? impl Sketch { /// Sketch a sample from a hash generator over its k-mers, transposing @@ -81,7 +117,7 @@ impl Sketch { ) -> Self { let (_sketchsize64, num_bins, usigs_size) = num_bins(sketch_size); let flattened_size_u64 = usigs_size as usize * kmer_lengths.len(); - let mut usigs = Vec::with_capacity(flattened_size_u64); + let mut usigs = aligned_sketch_vec_with_capacity(flattened_size_u64); let mut read_filter = if seq_hashes.reads() { let mut filter = KmerFilter::new(min_count); @@ -98,11 +134,11 @@ impl Sketch { log::debug!("Running sketching at k={k}"); let (signs, k_densified) = Self::get_signs(seq_hashes, *k, &mut read_filter, num_bins); densified |= k_densified; - minhash_sum += (signs[0] as f64) / (SIGN_MOD as f64); + minhash_sum += (signs[0] as f64) / (u64::MAX as f64); // Transpose the bins and save to the sketch map log::debug!("Transposing bins"); - let mut kmer_usigs = vec![0; usigs_size as usize]; + let mut kmer_usigs = aligned_sketch_vec_filled(usigs_size as usize, 0); Self::fill_usigs(&mut kmer_usigs, &signs); usigs.append(&mut kmer_usigs); } @@ -143,9 +179,8 @@ impl Sketch { seq_hashes.set_k(kmer_size); // Calculate bin minima across all sequence - let bin_size: u64 = SIGN_MOD.div_ceil(num_bins); for hash in seq_hashes.iter() { - Self::bin_sign(&mut signs, hash % SIGN_MOD, bin_size, filter); + Self::bin_sign(&mut signs, hash, num_bins, filter); } // Densify let densified = Self::densify_bin(&mut signs); @@ -167,9 +202,8 @@ impl Sketch { seq_hashes.set_k(kmer_size); // Calculate bin minima across all sequence - let bin_size: u64 = SIGN_MOD.div_ceil(num_bins); for hash in seq_hashes.iter() { - Self::bin_sign(&mut signs, hash % SIGN_MOD, bin_size, filter); + Self::bin_sign(&mut signs, hash, num_bins, filter); } signs @@ -191,12 +225,12 @@ impl Sketch { } /// Take the (transposed) sketch, emptying it from the [`Sketch`] - pub fn get_usigs(&mut self) -> Vec { - std::mem::take(&mut self.usigs) + pub fn get_usigs(&mut self) -> SketchVec { + std::mem::replace(&mut self.usigs, aligned_sketch_vec()) } - fn bin_sign(signs: &mut [u64], sign: u64, binsize: u64, read_filter: &mut Option) { - let binidx = (sign / binsize) as usize; + fn bin_sign(signs: &mut [u64], sign: u64, num_bins: u64, read_filter: &mut Option) { + let binidx = Self::bin_index(sign, num_bins); // log::trace!("sign:{sign} idx:{binidx} curr_sign:{}", signs[binidx]); if let Some(filter) = read_filter { if sign < signs[binidx] && filter.filter(sign) == Ordering::Equal { @@ -207,6 +241,11 @@ impl Sketch { } } + #[inline(always)] + fn bin_index(sign: u64, num_bins: u64) -> usize { + (((sign as u128) * (num_bins as u128)) >> u64::BITS) as usize + } + #[inline(always)] fn bit_at_pos(x: u64, pos: u64) -> u64 { (x & (1_u64 << pos)) >> pos @@ -215,9 +254,9 @@ impl Sketch { fn fill_usigs(usigs: &mut [u64], signs: &[u64]) { for (sign_index, sign) in signs.iter().enumerate() { let leftshift = sign_index % (u64::BITS as usize); - for i in 0..BBITS { - let orval = Self::bit_at_pos(*sign, i) << leftshift; - usigs[sign_index / (u64::BITS as usize) * (BBITS as usize) + (i as usize)] |= orval; + for i in 0..BIN_BITS { + let orval = Self::bit_at_pos(*sign, i as u64) << leftshift; + usigs[sign_index / (u64::BITS as usize) * BIN_BITS + i] |= orval; } } } @@ -258,6 +297,25 @@ impl Sketch { } } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sketch_vec_is_64_byte_aligned_after_allocation() { + let mut sketch_vec = aligned_sketch_vec_with_capacity(16); + sketch_vec.resize(16, 0); + assert_eq!(sketch_vec.as_ptr() as usize % SKETCH_ALIGNMENT, 0); + } + + #[test] + fn bin_index_uses_full_u64_hash_space() { + assert_eq!(Sketch::bin_index(0, 64), 0); + assert_eq!(Sketch::bin_index(1_u64 << 63, 64), 32); + assert_eq!(Sketch::bin_index(u64::MAX, 64), 63); + } +} + impl fmt::Display for Sketch { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { writeln!( @@ -294,7 +352,7 @@ pub fn sketch_files( quiet: bool, ) -> Vec { let bin_stride = 1; - let kmer_stride = (sketch_size * BBITS) as usize; + let kmer_stride = sketch_size as usize * BIN_BITS; let sample_stride = kmer_stride * k.len(); #[cfg(feature = "3di")] diff --git a/src/sketch/multisketch.rs b/src/sketch/multisketch.rs index 510775e..6a14a65 100644 --- a/src/sketch/multisketch.rs +++ b/src/sketch/multisketch.rs @@ -11,10 +11,11 @@ use hashbrown::{HashMap, HashSet}; use serde::{Deserialize, Serialize}; use crate::hashing::HashType; +use crate::sketch::aligned_sketch_vec; use crate::sketch::num_bins; use crate::sketch::sketch_datafile::SketchArrayReader; use crate::sketch::sketch_datafile::SketchArrayWriter; -use crate::sketch::Sketch; +use crate::sketch::{Sketch, SketchVec}; use super::sketch_datafile::append_batch; @@ -33,8 +34,8 @@ pub struct MultiSketch { // NB: another way to do this is with the ouroboros crate, which allows this to reference self // But this requires manual impl for ser and deser, and does the same indirection as an index anyway so not worth it block_reindex: Option>, - #[serde(skip)] - sketch_bins: Vec, + #[serde(skip, default = "aligned_sketch_vec")] + sketch_bins: SketchVec, bin_stride: usize, kmer_stride: usize, sample_stride: usize, @@ -67,7 +68,7 @@ impl MultiSketch { sketch_metadata: mem::take(sketches), name_map, block_reindex: None, - sketch_bins: Vec::new(), + sketch_bins: aligned_sketch_vec(), bin_stride: 1, kmer_stride, sample_stride: kmer_stride * kmer_lengths.len(), diff --git a/src/sketch/sketch_datafile.rs b/src/sketch/sketch_datafile.rs index a71d4d5..9907f16 100644 --- a/src/sketch/sketch_datafile.rs +++ b/src/sketch/sketch_datafile.rs @@ -8,6 +8,8 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use anyhow::{Error, Result}; +use super::{aligned_sketch_vec_with_capacity, SketchVec}; + /// Write to an .skd or .skq file /// This writer is not parallel safe, serial only! #[derive(Debug)] @@ -29,7 +31,7 @@ pub fn append_batch( ) -> Result<(), Error> { if let Some(mmap) = input_reader.mmap_file() { for &sample_idx in sample_indices { - let mut sample_data = Vec::with_capacity(input_reader.sample_stride); + let mut sample_data = aligned_sketch_vec_with_capacity(input_reader.sample_stride); for bin_idx in 0..input_reader.sample_stride { let start_byte = (sample_idx * input_reader.sample_stride + bin_idx) * mem::size_of::(); @@ -156,11 +158,11 @@ impl SketchArrayReader { } /// Read all sketch data from an .skd into memory - pub fn read_all_from_skd(&mut self, total_number_bins: usize) -> Vec { + pub fn read_all_from_skd(&mut self, total_number_bins: usize) -> SketchVec { // Fixed-size buffer for u64 let mut buffer = [0u8; std::mem::size_of::()]; // Stream the whole file and convert to u64 vec - let mut flat_sketch_array: Vec = Vec::with_capacity(total_number_bins); + let mut flat_sketch_array = aligned_sketch_vec_with_capacity(total_number_bins); while let Ok(_read) = self.sketch_reader.read_exact(&mut buffer) { flat_sketch_array.push(u64::from_le_bytes(buffer)); } @@ -169,9 +171,9 @@ impl SketchArrayReader { /// Read selected samples from an .skd into memory, providing the selected /// samples with `sample_indices` - pub fn read_batch_from_skd(&self, sample_indices: &[usize], sample_stride: usize) -> Vec { - let mut flat_sketch_array: Vec = - Vec::with_capacity(sample_stride * sample_indices.len()); + pub fn read_batch_from_skd(&self, sample_indices: &[usize], sample_stride: usize) -> SketchVec { + let mut flat_sketch_array = + aligned_sketch_vec_with_capacity(sample_stride * sample_indices.len()); // TODO possible improvement would be to combine adjacent indices into ranges if let Some(mmap_array) = &self.mmap_file { for sample_idx in sample_indices { @@ -200,3 +202,37 @@ impl SketchArrayReader { Ok(mmap) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::sketch::SKETCH_ALIGNMENT; + use std::fs; + + #[test] + fn skd_reads_are_64_byte_aligned() { + let filename = std::env::temp_dir().join(format!( + "sketchlib_skd_alignment_{}.skd", + std::process::id() + )); + let filename = filename.to_str().expect("temporary path must be UTF-8"); + + { + let mut writer = SketchArrayWriter::new(filename, 1, 4, 4); + writer.write_sketch(&[1_u64, 2, 3, 4]); + writer.flush().expect("failed to flush test sketch"); + } + + let mut reader = SketchArrayReader::open(filename, false, 1, 4, 4); + let all_sketches = reader.read_all_from_skd(4); + assert_eq!(all_sketches.as_ptr() as usize % SKETCH_ALIGNMENT, 0); + + let reader = SketchArrayReader::open(filename, true, 1, 4, 4); + let batch_sketches = reader.read_batch_from_skd(&[0], 4); + assert_eq!(batch_sketches.as_ptr() as usize % SKETCH_ALIGNMENT, 0); + + drop(batch_sketches); + drop(reader); + fs::remove_file(filename).expect("failed to remove test sketch"); + } +} diff --git a/tests/sketch.rs b/tests/sketch.rs index b66c1e0..fb05115 100644 --- a/tests/sketch.rs +++ b/tests/sketch.rs @@ -162,17 +162,36 @@ mod tests { .assert() .success(); - // Check .skm the same + // Check legacy metadata and data still load, without requiring identical + // derived stride fields from old hash-space behavior. let new_sketch: MultiSketch = MultiSketch::load_metadata(&sandbox.file_string("new_db", TestDir::Output)) .expect("Failed to load new sketch"); - let expected_sketch = + let mut expected_sketch = MultiSketch::load_metadata(&sandbox.file_string("legacy_db", TestDir::Input)) .expect("Failed to load legacy sketch"); + expected_sketch.read_sketch_data(&sandbox.file_string("legacy_db", TestDir::Input)); + + assert!( + new_sketch.is_compatible_with(&expected_sketch), + "Legacy sketch should be compatible with newly generated sketch metadata" + ); assert_eq!( - new_sketch, expected_sketch, - "Sketch metadata does not match legacy and new version" + new_sketch.number_samples_loaded(), + expected_sketch.number_samples_loaded(), + "Legacy sketch sample count does not match new sketch" + ); + for idx in 0..new_sketch.number_samples_loaded() { + assert_eq!( + new_sketch.get_sample_name(idx), + expected_sketch.get_sample_name(idx), + "Legacy sketch sample name does not match new sketch" + ); + } + assert!( + !expected_sketch.get_sketch_slice(0, 0).is_empty(), + "Legacy sketch data did not load" ); } } diff --git a/tests/test_results_correct/dists_knn_ani.stdout b/tests/test_results_correct/dists_knn_ani.stdout index be4e0af..d1fda9d 100644 --- a/tests/test_results_correct/dists_knn_ani.stdout +++ b/tests/test_results_correct/dists_knn_ani.stdout @@ -1,4 +1,4 @@ -14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.9887168 -14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.9887168 -R6.fa.gz TIGR4.fa.gz 0.99009615 -TIGR4.fa.gz R6.fa.gz 0.99009615 +14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.9891178 +14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.9891178 +R6.fa.gz TIGR4.fa.gz 0.9904671 +TIGR4.fa.gz R6.fa.gz 0.9904671 diff --git a/tests/test_results_correct/dists_knn_ca.stdout b/tests/test_results_correct/dists_knn_ca.stdout index 7e77fde..c1a9fe4 100644 --- a/tests/test_results_correct/dists_knn_ca.stdout +++ b/tests/test_results_correct/dists_knn_ca.stdout @@ -1,4 +1,4 @@ -14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.008820182 0.2137754 -14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.008820182 0.2137754 -R6.fa.gz TIGR4.fa.gz 0.010040261 0.1535005 -TIGR4.fa.gz R6.fa.gz 0.010040261 0.1535005 +14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.011049058 0.17300157 +14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.011049058 0.17300157 +R6.fa.gz TIGR4.fa.gz 0.0099692345 0.15083496 +TIGR4.fa.gz R6.fa.gz 0.0099692345 0.15083496 diff --git a/tests/test_results_correct/dists_knn_jaccard.stdout b/tests/test_results_correct/dists_knn_jaccard.stdout index fa4f633..aac95d9 100644 --- a/tests/test_results_correct/dists_knn_jaccard.stdout +++ b/tests/test_results_correct/dists_knn_jaccard.stdout @@ -1,4 +1,4 @@ -14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.34842753 -14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.34842753 -R6.fa.gz TIGR4.fa.gz 0.31618232 -TIGR4.fa.gz R6.fa.gz 0.31618232 +14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.33927667 +14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.33927667 +R6.fa.gz TIGR4.fa.gz 0.3071305 +TIGR4.fa.gz R6.fa.gz 0.3071305 diff --git a/tests/test_results_correct/dists_subset.stdout b/tests/test_results_correct/dists_subset.stdout index 7b4ec99..9cdd0d6 100644 --- a/tests/test_results_correct/dists_subset.stdout +++ b/tests/test_results_correct/dists_subset.stdout @@ -1,3 +1,3 @@ -R6.fa.gz TIGR4.fa.gz 0.010040261 0.1535005 +R6.fa.gz TIGR4.fa.gz 0.0099692345 0.15083496 R6.fa.gz 14412_3#84.contigs_velvet.fa.gz 1 1 -TIGR4.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.061570305 0.9990812 +TIGR4.fa.gz 14412_3#84.contigs_velvet.fa.gz 1 1 diff --git a/tests/test_results_correct/inverted.skq b/tests/test_results_correct/inverted.skq index 5e0c52e..04e106c 100644 --- a/tests/test_results_correct/inverted.skq +++ b/tests/test_results_correct/inverted.skq @@ -1 +1 @@ -<vpʊ!i<3!!pO lDl>1n F~S>1n \ No newline at end of file +%+ҐXZ1;%+Xd5bt"R"^ư\dGtx>R"^ư\dG \ No newline at end of file diff --git a/tests/test_results_correct/inverted_precluster.stdout b/tests/test_results_correct/inverted_precluster.stdout index 879e2ba..4ee4da2 100644 --- a/tests/test_results_correct/inverted_precluster.stdout +++ b/tests/test_results_correct/inverted_precluster.stdout @@ -1,4 +1,4 @@ -14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.35546875 -14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.35546875 -R6.fa.gz TIGR4.fa.gz 0.29492188 -TIGR4.fa.gz R6.fa.gz 0.29492188 +14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.35352102 +14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.35352102 +R6.fa.gz TIGR4.fa.gz 0.32324713 +TIGR4.fa.gz R6.fa.gz 0.32324713 diff --git a/tests/test_results_correct/inverted_precluster_ani.stdout b/tests/test_results_correct/inverted_precluster_ani.stdout index d492dfa..8695191 100644 --- a/tests/test_results_correct/inverted_precluster_ani.stdout +++ b/tests/test_results_correct/inverted_precluster_ani.stdout @@ -1,4 +1,4 @@ -14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.98840284 -14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.98840284 -R6.fa.gz TIGR4.fa.gz 0.99095666 -TIGR4.fa.gz R6.fa.gz 0.99095666 +14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz 0.98849016 +14412_3#84.contigs_velvet.fa.gz 14412_3#82.contigs_velvet.fa.gz 0.98849016 +R6.fa.gz TIGR4.fa.gz 0.9898019 +TIGR4.fa.gz R6.fa.gz 0.9898019 diff --git a/tests/test_results_correct/inverted_query_count.stdout b/tests/test_results_correct/inverted_query_count.stdout index 7a10a4e..aba31d7 100644 --- a/tests/test_results_correct/inverted_query_count.stdout +++ b/tests/test_results_correct/inverted_query_count.stdout @@ -1,5 +1,5 @@ Query 14412_3#82.contigs_velvet.fa.gz 14412_3#84.contigs_velvet.fa.gz R6.fa.gz TIGR4.fa.gz -14412_3#82.contigs_velvet.fa.gz 10 5 0 0 -R6.fa.gz 0 0 10 7 -14412_3#84.contigs_velvet.fa.gz 5 10 0 0 -TIGR4.fa.gz 0 0 7 10 +14412_3#82.contigs_velvet.fa.gz 10 6 0 0 +R6.fa.gz 0 0 10 9 +14412_3#84.contigs_velvet.fa.gz 6 10 0 0 +TIGR4.fa.gz 0 0 9 10 diff --git a/tests/test_results_correct/inverted_sketch_info.stdout b/tests/test_results_correct/inverted_sketch_info.stdout index f31f8f7..a6722d8 100644 --- a/tests/test_results_correct/inverted_sketch_info.stdout +++ b/tests/test_results_correct/inverted_sketch_info.stdout @@ -7,4 +7,4 @@ rc=true inverted=true max_hashes_per_bin=2 min_hashes_per_bin=1 -avg_hashes_per_bin=1.381 +avg_hashes_per_bin=1.413 diff --git a/tests/test_results_correct/read_sketch_full_info.stdout b/tests/test_results_correct/read_sketch_full_info.stdout index fee0685..37cf35b 100644 --- a/tests/test_results_correct/read_sketch_full_info.stdout +++ b/tests/test_results_correct/read_sketch_full_info.stdout @@ -1,4 +1,4 @@ Name Sequence length Base frequencies Missing/ambig bases From reads Single strand Densified -test_1 1 [603, 330, 334, 603] 0 true false true -test_2 1 [617, 312, 314, 617] 0 true false true +test_1 91 [603, 330, 334, 603] 0 true false true +test_2 91 [617, 312, 314, 617] 0 true false true From 707c60f8369f85dc959bc4b067efca9e782397e3 Mon Sep 17 00:00:00 2001 From: John Lees Date: Fri, 22 May 2026 15:49:29 +0100 Subject: [PATCH 2/2] Some code review suggestions --- src/distances/jaccard.rs | 36 ++++++++++++++++- src/io.rs | 4 +- src/sketch/mod.rs | 22 ++++++---- tests/inverted.rs | 86 ++++++++++++++++++++++++++++++++++++++++ tests/sketch.rs | 1 + 5 files changed, 139 insertions(+), 10 deletions(-) diff --git a/src/distances/jaccard.rs b/src/distances/jaccard.rs index 4cf9d6f..131c30e 100644 --- a/src/distances/jaccard.rs +++ b/src/distances/jaccard.rs @@ -44,6 +44,7 @@ fn jaccard_same_bits(sketch1: &[u64], sketch2: &[u64]) -> u32 { } #[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); @@ -134,10 +135,41 @@ mod tests { use super::*; #[test] - fn random_match_correction_maps_expected_random_to_zero() { - let expected_random = 1.0 / ((1u32 << (BIN_BITS as u32)) as f64); + fn random_match_correction_handles_expected_random_and_bounds() { + let bb = (1u32 << (BIN_BITS as u32)) as f64; + let expected_random = 1.0 / bb; + let raw_above_random = 0.25; + let expected_corrected = (bb * raw_above_random - 1.0) / (bb - 1.0); + assert_eq!(random_match_correction(expected_random), 0.0); + assert_eq!(random_match_correction(expected_random / 2.0), 0.0); assert_eq!(random_match_correction(1.0), 1.0); + assert_eq!( + random_match_correction(raw_above_random), + expected_corrected + ); + } + + #[test] + fn jaccard_index_matches_expected_samebits_formula() { + let sketchsize64 = 1; + let unionsize = (u64::BITS as u64 * sketchsize64) as f64; + let bb = (1u32 << (BIN_BITS as u32)) as f64; + let expected_random_samebits = unionsize / bb; + let target_samebits = 42_u32; + let intersize = target_samebits as f64 - expected_random_samebits; + let previous_formula = intersize / (unionsize - expected_random_samebits); + + let sketch1 = [0_u64; BIN_BITS]; + let mut sketch2 = [0_u64; BIN_BITS]; + sketch2[0] = (1_u64 << (u64::BITS - target_samebits)) - 1; + + assert_eq!( + jaccard_same_bits_general(&sketch1, &sketch2), + target_samebits + ); + let current_formula = jaccard_index(&sketch1, &sketch2, sketchsize64, None, None, 0.0); + assert!((current_formula - previous_formula).abs() <= f64::EPSILON); } #[test] diff --git a/src/io.rs b/src/io.rs index 3be7af8..ad22b91 100644 --- a/src/io.rs +++ b/src/io.rs @@ -260,7 +260,9 @@ pub fn read_completeness_file( .filter_map(|line| { if let Some((genome_id, completeness_str)) = line.split_once('\t') { if let Ok(completeness) = completeness_str.parse::() { - sketches.get_sample_index(genome_id).map(|index| (index, completeness)) + sketches + .get_sample_index(genome_id) + .map(|index| (index, completeness)) } else { None } diff --git a/src/sketch/mod.rs b/src/sketch/mod.rs index f19fb75..acdc945 100644 --- a/src/sketch/mod.rs +++ b/src/sketch/mod.rs @@ -33,8 +33,6 @@ use self::sketch_datafile::SketchArrayWriter; /// Bin bits (lowest of 64-bits to keep) pub const BIN_BITS: usize = 16; -#[doc(hidden)] -pub const BBITS: usize = BIN_BITS; /// Byte alignment used for bit-packed sketch buffers. pub const SKETCH_ALIGNMENT: usize = 64; @@ -226,6 +224,7 @@ impl Sketch { /// Take the (transposed) sketch, emptying it from the [`Sketch`] pub fn get_usigs(&mut self) -> SketchVec { + // Move the old aligned buffer out and install an empty one; elements are not copied. std::mem::replace(&mut self.usigs, aligned_sketch_vec()) } @@ -252,11 +251,20 @@ impl Sketch { } fn fill_usigs(usigs: &mut [u64], signs: &[u64]) { - for (sign_index, sign) in signs.iter().enumerate() { - let leftshift = sign_index % (u64::BITS as usize); - for i in 0..BIN_BITS { - let orval = Self::bit_at_pos(*sign, i as u64) << leftshift; - usigs[sign_index / (u64::BITS as usize) * BIN_BITS + i] |= orval; + debug_assert_eq!(signs.len() % (u64::BITS as usize), 0); + debug_assert_eq!(usigs.len(), signs.len() / (u64::BITS as usize) * BIN_BITS); + + for (usig_chunk, sign_chunk) in usigs + .chunks_exact_mut(BIN_BITS) + .zip(signs.chunks_exact(u64::BITS as usize)) + { + for (bit_pos, usig) in usig_chunk.iter_mut().enumerate() { + *usig = sign_chunk + .iter() + .enumerate() + .fold(0, |bits, (sign_index, &sign)| { + bits | (Self::bit_at_pos(sign, bit_pos as u64) << sign_index) + }); } } } diff --git a/tests/inverted.rs b/tests/inverted.rs index 3c701dd..65cc6b1 100644 --- a/tests/inverted.rs +++ b/tests/inverted.rs @@ -6,10 +6,26 @@ use crate::common::*; #[cfg(test)] mod tests { + use sketchlib::sketch::{ + multisketch::MultiSketch, sketch_datafile::SketchArrayReader, BIN_BITS, + }; use snapbox::assert_data_eq; use super::*; + fn unpack_skd_bin(bit_planes: &[u64], bin_idx: usize) -> u16 { + let word_idx = bin_idx / (u64::BITS as usize); + let bit_offset = bin_idx % (u64::BITS as usize); + let plane_offset = word_idx * BIN_BITS; + + bit_planes[plane_offset..(plane_offset + BIN_BITS)] + .iter() + .enumerate() + .fold(0_u16, |bin_value, (bit_pos, &plane)| { + bin_value | (((plane >> bit_offset) & 1) as u16) << bit_pos + }) + } + #[test] fn inverted_sketch() { let sandbox = TestSetup::setup(); @@ -48,6 +64,76 @@ mod tests { .stdout_eq(sandbox.snapbox_file("inverted_sketch_full_info.stdout", TestDir::Correct)); } + #[test] + fn inverted_skq_matches_standard_skd_bins() { + let sandbox = TestSetup::setup(); + + sandbox.copy_input_file_to_wd("14412_3#82.contigs_velvet.fa.gz"); + sandbox.copy_input_file_to_wd("14412_3#84.contigs_velvet.fa.gz"); + sandbox.copy_input_file_to_wd("R6.fa.gz"); + sandbox.copy_input_file_to_wd("TIGR4.fa.gz"); + sandbox.copy_input_file_to_wd("rfile.txt"); + + Command::new(cmd::cargo_bin!("sketchlib")) + .current_dir(sandbox.get_wd()) + .arg("inverted") + .arg("build") + .arg("-o") + .arg("inverted") + .args(["-v", "-k", "21", "-s", "64"]) + .arg("-f") + .arg("rfile.txt") + .arg("--write-skq") + .assert() + .success(); + + Command::new(cmd::cargo_bin!("sketchlib")) + .current_dir(sandbox.get_wd()) + .arg("sketch") + .arg("-o") + .arg("standard") + .args(["-v", "--k-vals", "21", "-s", "64"]) + .arg("-f") + .arg("rfile.txt") + .assert() + .success(); + + let mut standard_sketches = + MultiSketch::load_metadata(&sandbox.file_string("standard", TestDir::Output)) + .expect("failed to load standard sketch metadata"); + standard_sketches.read_sketch_data(&sandbox.file_string("standard", TestDir::Output)); + + let sketch_size = standard_sketches.sketch_size as usize; + assert_eq!(sketch_size, 64); + assert_eq!(standard_sketches.kmer_lengths(), &[21]); + + let mut skq_reader = SketchArrayReader::open( + &sandbox.file_string("inverted.skq", TestDir::Output), + false, + 1, + 1, + sketch_size, + ); + let skq_bins = + skq_reader.read_all_from_skq(sketch_size * standard_sketches.number_samples_loaded()); + assert_eq!( + skq_bins.len(), + sketch_size * standard_sketches.number_samples_loaded() + ); + + for sample_idx in 0..standard_sketches.number_samples_loaded() { + let skd_bins = standard_sketches.get_sketch_slice(sample_idx, 0); + let skq_offset = sample_idx * sketch_size; + for bin_idx in 0..sketch_size { + assert_eq!( + skq_bins[skq_offset + bin_idx], + unpack_skd_bin(skd_bins, bin_idx), + "sample {sample_idx} bin {bin_idx} differs between .skq and .skd" + ); + } + } + } + #[test] fn inverted_sketch_with_metadata() { let sandbox = TestSetup::setup(); diff --git a/tests/sketch.rs b/tests/sketch.rs index fb05115..77e23c6 100644 --- a/tests/sketch.rs +++ b/tests/sketch.rs @@ -70,6 +70,7 @@ mod tests { .assert() .stdout_eq(sandbox.snapbox_file("read_sketch_info.stdout", TestDir::Correct)); + // FASTQ read length is estimated from minima normalized over the full u64 hash space. Command::new(cmd::cargo_bin!("sketchlib")) .current_dir(sandbox.get_wd()) .arg("info")