Playing with MQF

Cargo.toml

@@ -52,6 +52,7 @@ pdatastructs = { git = "https://github.com/luizirber/pdatastructs.rs", branch =
 itertools = "0.8.0"
 typed-builder = "0.3.0"
 csv = "1.0.7"
+tempfile = "3.1.0"
 
 [target.'cfg(all(target_arch = "wasm32", target_vendor="unknown"))'.dependencies.wasm-bindgen]
 version = "^0.2"
@@ -62,6 +63,9 @@ version = "0.1"
 path = "ocf"
 default-features = false
 
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies.mqf]
+version = "1.0.0"
+
 [dev-dependencies]
 proptest = "^0.8"
 criterion = "^0.2"

Makefile

@@ -46,4 +46,7 @@ last-tag:
 wasm:
 	wasm-pack build
 
+wasi:
+	cargo wasi build
+
 FORCE:

examples/generate_mqfs.rs

@@ -0,0 +1,81 @@
+use std::collections::HashMap;
+use std::fs::File;
+use std::io::BufReader;
+use std::path::PathBuf;
+use std::result::Result;
+
+use mqf::MQF;
+
+use sourmash::signature::Signature;
+use sourmash::sketch::Sketch;
+
+fn main() {
+    let mh_paths: HashMap<u8, String> = {
+        [
+            (6, "6d6e87e1154e95b279e5e7db414bc37b".into()),
+            (7, "60f7e23c24a8d94791cc7a8680c493f9".into()),
+            (8, "0107d767a345eff67ecdaed2ee5cd7ba".into()),
+            (9, "f71e78178af9e45e6f1d87a0c53c465c".into()),
+            (10, "f0c834bc306651d2b9321fb21d3e8d8f".into()),
+            (11, "4e94e60265e04f0763142e20b52c0da1".into()),
+            (12, "b59473c94ff2889eca5d7165936e64b3".into()),
+        ]
+        .into_iter()
+        .cloned()
+        .collect()
+    };
+
+    let mut filename = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+    filename.push("tests/test-data/.sbt.v5_mhmt/");
+
+    let mh_sigs: HashMap<u8, Signature> = mh_paths
+        .into_iter()
+        .map(|(k, v)| {
+            let mut mhpath = filename.clone();
+            mhpath.push(v);
+
+            let file = File::open(mhpath).unwrap();
+            let reader = BufReader::new(file);
+            let sig: Result<Vec<Signature>, _> = serde_json::from_reader(reader);
+            (k, sig.unwrap()[0].clone())
+        })
+        .collect();
+
+    let mut mqfs: HashMap<u8, MQF> = HashMap::default();
+
+    for i in (0..=5).rev() {
+        println!("Creating MQF {}", i);
+        let mut mqf = MQF::new(1, 18);
+
+        let left = i * 2 + 1;
+        for child in left..=left + 1 {
+            if mh_sigs.contains_key(&child) {
+                println!("Loading values from MH {}", child);
+                if let Sketch::MinHash(sig) = &mh_sigs[&child].signatures[0] {
+                    sig.mins()
+                        .iter()
+                        .map(|h| {
+                            dbg!(*h % u64::pow(2, 26));
+                            mqf.insert(*h % u64::pow(2, 26), 1)
+                            //mqf.insert(*h, 1)
+                        })
+                        .count();
+                };
+            } else if mqfs.contains_key(&child) {
+                let mut cmqf = mqfs.get_mut(&child).unwrap();
+                mqf.merge(&mut cmqf).expect("Error merging");
+            } else {
+                // TODO: shouldn't happen...
+                unimplemented!()
+            }
+        }
+
+        println!("Save MQF {} for later", i);
+        mqfs.insert(i, mqf);
+
+        println!("Saving MQFs to disk");
+        let mut internal = filename.clone();
+        internal.push(format!("internal.{}", i));
+        mqfs[&i].serialize(internal).unwrap();
+    }
+}

src/index/mod.rs

@@ -14,10 +14,10 @@ pub mod search;
 use std::path::Path;
 use std::rc::Rc;
 
-use serde_derive::{Deserialize, Serialize};
-
+use cfg_if::cfg_if;
 use failure::Error;
 use lazy_init::Lazy;
+use serde_derive::{Deserialize, Serialize};
 use typed_builder::TypedBuilder;
 
 use crate::index::sbt::{Node, SBT};
@@ -31,6 +31,13 @@ use crate::sketch::Sketch;
 pub type MHBT = SBT<Node<Nodegraph>, Dataset<Signature>>;
 pub type UKHSTree = SBT<Node<FlatUKHS>, Dataset<Signature>>;
 
+cfg_if! {
+    if #[cfg(not(target_arch = "wasm32"))] {
+      use mqf::MQF;
+      pub type MHMT = SBT<Node<MQF>, Dataset<Signature>>;
+    }
+}
+
 pub trait Index {
     type Item;
 

src/index/sbt/mhmt.rs

@@ -0,0 +1,135 @@
+use std::io::{Read, Write};
+
+use failure::Error;
+use mqf::MQF;
+
+use crate::index::sbt::{FromFactory, Node, Update, SBT};
+use crate::index::storage::{ReadData, ReadDataError, ToWriter};
+use crate::index::{Comparable, Dataset};
+use crate::signature::{Signature, SigsTrait};
+use crate::sketch::Sketch;
+
+impl ToWriter for MQF {
+    fn to_writer<W>(&self, writer: &mut W) -> Result<(), Error>
+    where
+        W: Write,
+    {
+        // TODO: using tempfile for now, but ideally want to avoid that
+        let mut tmpfile = tempfile::NamedTempFile::new()?;
+        self.serialize(tmpfile.path()).unwrap(); // TODO: convert this to a proper error
+
+        let mut buffer = Vec::new();
+        tmpfile.read_to_end(&mut buffer)?;
+        writer.write_all(&buffer)?;
+
+        Ok(())
+    }
+}
+
+impl ReadData<MQF> for Node<MQF> {
+    fn data(&self) -> Result<&MQF, Error> {
+        if let Some(storage) = &self.storage {
+            Ok(self.data.get_or_create(|| {
+                let raw = storage.load(&self.filename).unwrap();
+
+                // TODO: using tempfile for now, but ideally want to avoid that
+                let mut tmpfile = tempfile::NamedTempFile::new().unwrap();
+                tmpfile.write_all(&mut &raw[..]).unwrap();
+
+                MQF::deserialize(tmpfile.path()).unwrap()
+            }))
+        } else if let Some(data) = self.data.get() {
+            Ok(data)
+        } else {
+            Err(ReadDataError::LoadError.into())
+        }
+    }
+}
+
+impl<L> FromFactory<Node<MQF>> for SBT<Node<MQF>, L> {
+    fn factory(&self, _name: &str) -> Result<Node<MQF>, Error> {
+        unimplemented!()
+    }
+}
+
+impl Update<Node<MQF>> for Node<MQF> {
+    fn update(&self, _other: &mut Node<MQF>) -> Result<(), Error> {
+        unimplemented!();
+    }
+}
+
+impl Update<Node<MQF>> for Dataset<Signature> {
+    fn update(&self, _other: &mut Node<MQF>) -> Result<(), Error> {
+        unimplemented!();
+    }
+}
+
+impl Comparable<Node<MQF>> for Node<MQF> {
+    fn similarity(&self, other: &Node<MQF>) -> f64 {
+        let ng: &MQF = self.data().unwrap();
+        let ong: &MQF = other.data().unwrap();
+        unimplemented!();
+        //ng.similarity(&ong)
+    }
+
+    fn containment(&self, other: &Node<MQF>) -> f64 {
+        let ng: &MQF = self.data().unwrap();
+        let ong: &MQF = other.data().unwrap();
+        unimplemented!();
+        //ng.containment(&ong)
+    }
+}
+
+impl Comparable<Dataset<Signature>> for Node<MQF> {
+    fn similarity(&self, other: &Dataset<Signature>) -> f64 {
+        let ng: &MQF = self.data().unwrap();
+        let oth: &Signature = other.data().unwrap();
+
+        // TODO: select the right signatures...
+        if let Sketch::MinHash(sig) = &oth.signatures[0] {
+            if sig.size() == 0 {
+                return 0.0;
+            }
+
+            let matches: usize = sig
+                .mins
+                .iter()
+                .filter(|h| dbg!(ng.count_key(**h % u64::pow(2, 26))) > 0)
+                //.filter(|h| dbg!(ng.count_key(**h)) > 0)
+                .count();
+
+            let min_n_below = self.metadata["min_n_below"] as f64;
+
+            // This overestimates the similarity, but better than truncating too
+            // soon and losing matches
+            matches as f64 / min_n_below
+        } else {
+            //TODO what if it is not a minhash?
+            unimplemented!()
+        }
+    }
+
+    fn containment(&self, other: &Dataset<Signature>) -> f64 {
+        let ng: &MQF = self.data().unwrap();
+        let oth: &Signature = other.data().unwrap();
+
+        // TODO: select the right signatures...
+        if let Sketch::MinHash(sig) = &oth.signatures[0] {
+            if sig.size() == 0 {
+                return 0.0;
+            }
+
+            let matches: usize = sig
+                .mins
+                .iter()
+                .filter(|h| ng.count_key(**h % u64::pow(2, 26)) > 0)
+                //.filter(|h| ng.count_key(**h) > 0)
+                .count();
+
+            matches as f64 / sig.size() as f64
+        } else {
+            //TODO what if it is not a minhash?
+            unimplemented!()
+        }
+    }
+}

src/index/sbt/mod.rs

@@ -1,6 +1,9 @@
 pub mod mhbt;
 pub mod ukhs;
 
+#[cfg(not(target_arch = "wasm32"))]
+pub mod mhmt;
+
 use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
@@ -812,6 +815,78 @@ mod test {
         tmpfile.seek(SeekFrom::Start(0)).unwrap();
     }
 
+    #[cfg(not(target_arch = "wasm32"))]
+    #[test]
+    fn load_mhmt() {
+        let mut filename = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+        filename.push("tests/test-data/v5_mhmt.sbt.json");
+
+        let mut sbt = crate::index::MHMT::from_path(filename).expect("Loading error");
+
+        let mut filename = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+        filename.push("tests/test-data/.sbt.v3/60f7e23c24a8d94791cc7a8680c493f9");
+
+        let mut reader = BufReader::new(File::open(filename).unwrap());
+        let sigs = Signature::load_signatures(&mut reader, 31, Some("DNA".into()), None).unwrap();
+        let sig_data = sigs[0].clone();
+
+        let data = Lazy::new();
+        data.get_or_create(|| sig_data);
+
+        let leaf = Dataset::builder()
+            .data(Rc::new(data))
+            .filename("")
+            .name("")
+            .metadata("")
+            .storage(None)
+            .build();
+
+        let results = sbt.find(search_minhashes, &leaf, 0.5).unwrap();
+        //assert_eq!(results.len(), 1);
+        println!("results: {:?}", results);
+        println!("leaf: {:?}", leaf);
+
+        let results = sbt.find(search_minhashes, &leaf, 0.1).unwrap();
+        assert_eq!(results.len(), 2);
+        println!("results: {:?}", results);
+        println!("leaf: {:?}", leaf);
+
+        let mut linear = LinearIndex::builder().storage(sbt.storage()).build();
+        for l in &sbt.leaves {
+            linear.insert(l.1).unwrap();
+        }
+
+        println!(
+            "linear leaves {:?} {:?}",
+            linear.datasets.len(),
+            linear.datasets
+        );
+
+        let results = linear.find(search_minhashes, &leaf, 0.5).unwrap();
+        assert_eq!(results.len(), 1);
+        println!("results: {:?}", results);
+        println!("leaf: {:?}", leaf);
+
+        let results = linear.find(search_minhashes, &leaf, 0.1).unwrap();
+        assert_eq!(results.len(), 2);
+        println!("results: {:?}", results);
+        println!("leaf: {:?}", leaf);
+
+        let results = linear
+            .find(search_minhashes_containment, &leaf, 0.5)
+            .unwrap();
+        assert_eq!(results.len(), 2);
+        println!("results: {:?}", results);
+        println!("leaf: {:?}", leaf);
+
+        let results = linear
+            .find(search_minhashes_containment, &leaf, 0.1)
+            .unwrap();
+        assert_eq!(results.len(), 4);
+        println!("results: {:?}", results);
+        println!("leaf: {:?}", leaf);
+    }
+
     #[test]
     fn load_sbt() {
         let mut filename = PathBuf::from(env!("CARGO_MANIFEST_DIR"));

src/sketch/minhash.rs

@@ -478,6 +478,10 @@ impl KmerMinHash {
     pub fn dayhoff(&self) -> bool {
         self.dayhoff
     }
+
+    pub fn mins(&self) -> Vec<u64> {
+        self.mins.clone()
+    }
 }
 
 impl SigsTrait for KmerMinHash {