From 0ecaa84da1331851df8afeb71ff80c86fc3ecb92 Mon Sep 17 00:00:00 2001 From: John Lees Date: Fri, 21 Aug 2026 18:57:21 +0100 Subject: [PATCH] Add some interface for wasm to deal with the breaking DB format changes --- src/distances/distance_matrix.rs | 5 +++ src/distances/jaccard.rs | 39 +++++++++++++++++--- src/distances/mod.rs | 30 ++++++++++++++-- src/sketch/mod.rs | 7 +++- src/sketch/multisketch.rs | 61 +++++++++++++++++++++++++++++++- 5 files changed, 134 insertions(+), 8 deletions(-) diff --git a/src/distances/distance_matrix.rs b/src/distances/distance_matrix.rs index 6c6e742..2249d50 100644 --- a/src/distances/distance_matrix.rs +++ b/src/distances/distance_matrix.rs @@ -312,6 +312,11 @@ impl<'a> SparseDistanceMatrix<'a> { pub fn dists_mut(&mut self) -> &mut DistVec { &mut self.distances } + + /// Reference to the completed sparse distance entries. + pub fn dists_as_ref(&self) -> &DistVec { + &self.distances + } } impl<'a> Distances<'a> for SparseDistanceMatrix<'a> { diff --git a/src/distances/jaccard.rs b/src/distances/jaccard.rs index 2ce8ae6..c52b45c 100644 --- a/src/distances/jaccard.rs +++ b/src/distances/jaccard.rs @@ -10,11 +10,23 @@ pub fn jaccard_index( c1: Option, c2: Option, completeness_cutoff: f64, +) -> f64 { + jaccard_index_generic::(sketch1, sketch2, sketchsize64, c1, c2, completeness_cutoff) +} + +/// Returns the Jaccard index using the selected sketch bin width. +pub fn jaccard_index_generic( + sketch1: &[u64], + sketch2: &[u64], + sketchsize64: u64, + c1: Option, + c2: Option, + 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)) + .chunks_exact(BITS as usize) + .zip(sketch2.chunks_exact(BITS as usize)) .map(|(chunk1, chunk2)| { let mut bits: u64 = !0; chunk1.iter().zip(chunk2.iter()).for_each(|(&s1, &s2)| { @@ -24,7 +36,7 @@ pub fn jaccard_index( }) .sum(); let maxnbits = sketchsize64 as u32 * u64::BITS; - let expected_samebits = maxnbits >> BBITS; + let expected_samebits = maxnbits >> BITS; log::trace!("samebits:{samebits} expected_samebits:{expected_samebits} maxnbits:{maxnbits}"); let diff = samebits.saturating_sub(expected_samebits); @@ -65,6 +77,25 @@ pub fn core_acc_dist( query_sketch_idx: usize, completeness_vec: Option<&Vec>, completeness_cutoff: f64, +) -> (f32, f32) { + core_acc_dist_generic::( + ref_sketches, + query_sketches, + ref_sketch_idx, + query_sketch_idx, + completeness_vec, + completeness_cutoff, + ) +} + +/// Calculates core/accessory distances using the selected sketch bin width. +pub fn core_acc_dist_generic( + ref_sketches: &MultiSketch, + query_sketches: &MultiSketch, + ref_sketch_idx: usize, + query_sketch_idx: usize, + completeness_vec: Option<&Vec>, + completeness_cutoff: f64, ) -> (f32, f32) { if ref_sketches.kmer_lengths().len() < 2 { panic!("Need at least two k-mer lengths to calculate core/accessory distances"); @@ -76,7 +107,7 @@ pub fn core_acc_dist( for (k_idx, k) in ref_sketches.kmer_lengths().iter().enumerate() { let c1 = completeness_vec.map(|cv| cv[ref_sketch_idx]); let c2 = completeness_vec.map(|cv| cv[query_sketch_idx]); - let y = jaccard_index( + let y = jaccard_index_generic::( ref_sketches.get_sketch_slice(ref_sketch_idx, k_idx), query_sketches.get_sketch_slice(query_sketch_idx, k_idx), ref_sketches.sketchsize64, diff --git a/src/distances/mod.rs b/src/distances/mod.rs index 02cd9bd..09deed6 100644 --- a/src/distances/mod.rs +++ b/src/distances/mod.rs @@ -10,6 +10,7 @@ use crate::cli::RetainUnmatched; use crate::get_progress_bar; use crate::inverted::Inverted; use crate::sketch::multisketch::MultiSketch; +use crate::sketch::BBITS; pub mod distance_matrix; use self::distance_matrix::*; @@ -138,6 +139,31 @@ pub fn self_dists_knn<'a>( quiet: bool, completeness_vec: Option<&Vec>, completeness_cutoff: f64, +) -> SparseDistanceMatrix<'a> { + self_dists_knn_generic::( + sketches, + n, + knn, + dist_type, + quiet, + completeness_vec, + completeness_cutoff, + ) +} + +/// Self kNN distances using a fixed sketch bin width. +/// +/// Mandrake calls this with [`crate::sketch::CURRENT_BBITS`] for current +/// databases; legacy 14-bit databases are intentionally outside the browser +/// contract. +pub fn self_dists_knn_generic<'a, const BITS: u64>( + sketches: &'a MultiSketch, + n: usize, + knn: usize, + dist_type: DistType, + quiet: bool, + completeness_vec: Option<&Vec>, + completeness_cutoff: f64, ) -> SparseDistanceMatrix<'a> { let mut sp_distances = SparseDistanceMatrix::new(sketches, knn, dist_type); let k_vals = sp_distances.k_vals(); @@ -162,7 +188,7 @@ pub fn self_dists_knn<'a>( // This uses Option::map to safely access the completeness value for each sample. let c1 = completeness_vec.map(|cv| cv[i]); let c2 = completeness_vec.map(|cv| cv[j]); - let dist = jaccard_index( + let dist = jaccard_index_generic::( i_sketch, sketches.get_sketch_slice(j, k_idx), sketches.sketchsize64, @@ -203,7 +229,7 @@ pub fn self_dists_knn<'a>( if i == j { continue; } - let dists = core_acc_dist( + let dists = core_acc_dist_generic::( sketches, sketches, i, diff --git a/src/sketch/mod.rs b/src/sketch/mod.rs index 0ac6c81..6458ace 100644 --- a/src/sketch/mod.rs +++ b/src/sketch/mod.rs @@ -30,8 +30,13 @@ pub mod sketch_datafile; #[cfg(not(target_family = "wasm"))] use self::sketch_datafile::SketchArrayWriter; -/// Bin bits (lowest of 64-bits to keep) +/// Bin bits (lowest of 64-bits to keep) used by the legacy sketch API. +/// +/// Current databases use [`CURRENT_BBITS`]; the generic distance functions +/// take the bin width explicitly so both formats can be read by the library. pub const BBITS: u64 = 14; +/// Bin width used by current (version 0.4+) sketch databases. +pub const CURRENT_BBITS: u64 = 16; /// Total width of all bins (used as sign % sign_mod) pub const SIGN_MOD: u64 = (1 << 61) - 1; diff --git a/src/sketch/multisketch.rs b/src/sketch/multisketch.rs index 510775e..0b51474 100644 --- a/src/sketch/multisketch.rs +++ b/src/sketch/multisketch.rs @@ -4,7 +4,7 @@ use anyhow::{bail, Error, Result}; use core::panic; use std::fmt; use std::fs::File; -use std::io::{BufReader, BufWriter}; +use std::io::{BufReader, BufWriter, Cursor}; use std::mem; use hashbrown::{HashMap, HashSet}; @@ -15,6 +15,7 @@ use crate::sketch::num_bins; use crate::sketch::sketch_datafile::SketchArrayReader; use crate::sketch::sketch_datafile::SketchArrayWriter; use crate::sketch::Sketch; +use crate::sketch::CURRENT_BBITS; use super::sketch_datafile::append_batch; @@ -102,6 +103,40 @@ impl MultiSketch { Ok(skm_obj) } + /// Loads metadata from an already materialised `.skm` file. + pub fn load_metadata_bytes(bytes: &[u8]) -> Result { + let skm_file = BufReader::new(Cursor::new(bytes)); + let decompress_reader = snap::read::FrameDecoder::new(skm_file); + let mut skm_obj: Self = ciborium::de::from_reader(decompress_reader)?; + if skm_obj.sketchsize64 == 0 { + skm_obj.sketchsize64 = skm_obj.sketch_size; + skm_obj.sketch_size *= 64; + } + Ok(skm_obj) + } + + /// Loads a paired `.skm`/`.skd` database from bytes. + pub fn load_bytes(metadata: &[u8], data: &[u8]) -> Result { + let mut sketches = Self::load_metadata_bytes(metadata)?; + sketches.read_sketch_data_bytes(data)?; + Ok(sketches) + } + + /// Returns whether this database predates the 16-bit current format. + pub fn is_legacy_format(&self) -> bool { + let mut version = self.sketch_version.split('.'); + let major = version.next().and_then(|part| part.parse::().ok()); + let minor = version.next().and_then(|part| part.parse::().ok()); + let current_version = + matches!((major, minor), (Some(major), Some(minor)) if major > 0 || minor >= 4); + let current_stride = self + .sketchsize64 + .checked_mul(CURRENT_BBITS) + .map(|stride| stride as usize) + .is_some_and(|stride| self.kmer_stride == stride); + !(current_version && current_stride) + } + /// Number of samples loaded from the .skm/.skd pub fn number_samples_loaded(&self) -> usize { match &self.block_reindex { @@ -183,6 +218,30 @@ impl MultiSketch { sketch_reader.read_all_from_skd(self.sample_stride * self.sketch_metadata.len()); } + /// Reads all sketch bins from an already materialised `.skd` file. + pub fn read_sketch_data_bytes(&mut self, bytes: &[u8]) -> Result<(), Error> { + if !bytes.len().is_multiple_of(std::mem::size_of::()) { + bail!("sketch data length is not a multiple of eight bytes"); + } + let expected = self + .sample_stride + .checked_mul(self.sketch_metadata.len()) + .ok_or_else(|| anyhow::anyhow!("sketch data size overflows memory limits"))?; + let values = bytes + .chunks_exact(8) + .map(|chunk| u64::from_le_bytes(chunk.try_into().expect("exact chunk size"))) + .collect::>(); + if values.len() != expected { + bail!( + "sketch data contains {} bins but metadata requires {}", + values.len(), + expected + ); + } + self.sketch_bins = values; + Ok(()) + } + /// Read a subset of the bins from an .skd file, in a memory efficient manner pub fn read_sketch_data_block(&mut self, file_prefix: &str, names: &[String]) { // Find the given names in the sketch metadata