#3: proper fingerprints writing with fixed size, #5: fingerprints similarity search

Author: mkviatkovskii

molecules.smi

@@ -1,2 +1,2 @@
-C
-CC
+CC(C)CC1=CC=C(C=C1)C(C)C(=O)O Ibuprofen
+CC(C1=CC2=C(C=C1)C=C(C=C2)OC)C(=O)O Naproxen

src/ringo/molecule/model/molecule.rs

@@ -3,7 +3,7 @@ use crate::ringo::molecule::model::atom::Atom;
 use crate::ringo::molecule::model::bond::Bond;
 use crate::ringo::molecule::model::element::atomic_weight;
 use crate::ringo::molecule::smiles::reader::molecule::parse_molecule;
-use crate::ringo::ringo::fingerprint::Fingerprint;
+use crate::ringo::ringo::fingerprint::{Fingerprint, FINGERPRINT_SIZE};
 use fixedbitset::FixedBitSet;
 use petgraph::stable_graph::{EdgeIndex, NodeIndex, StableGraph};
 use petgraph::visit::EdgeRef;
@@ -89,7 +89,7 @@ impl Molecule {
 
     // TODO: move to Descriptors crate
     pub fn ecfp(&self, radius: usize, fp_length: usize) -> Fingerprint {
-        let mut fp = FixedBitSet::new();
+        let mut fp = FixedBitSet::with_capacity(FINGERPRINT_SIZE);
 
         for node in self.graph.node_indices() {
             ecfp_recursive(

src/ringo/ringo/fingerprint.rs

@@ -41,9 +41,11 @@ mod tests {
         fp.0.set(1, true);
         fp.0.set(17, true);
 
-        let encoded = bincode::encode_to_vec(&fp, bincode::config::standard()).unwrap();
+        let mut buf = vec![0u8; FINGERPRINT_SIZE / 8];
+        let encoded = bincode::encode_into_slice(&fp, buf.as_mut_slice(), bincode::config::standard()).unwrap();
+
         let decoded: Fingerprint =
-            bincode::decode_from_slice(&encoded, bincode::config::standard())
+            bincode::decode_from_slice(&buf, bincode::config::standard())
                 .unwrap()
                 .0;
         assert_eq!(decoded.0.ones().collect::<Vec<usize>>(), vec![1, 17]);

src/ringo/ringo/index.rs

@@ -1,23 +1,37 @@
 use crate::ringo::molecule::smiles::reader::molecule::parse_molecule;
 use crate::ringo::ringo::index_item::IndexItem;
 use std::fs::File;
-use std::io::BufRead;
+use std::io::{BufRead, BufWriter, Write};
+use bincode::{encode_into_slice};
+use crate::ringo::ringo::fingerprint::FINGERPRINT_SIZE;
+
+#[cfg(windows)]
+const LINE_ENDING_LENGTH: usize = 2;
+#[cfg(not(windows))]
+const LINE_ENDING_LENGTH: usize = 1;
 
 fn index(smiles_file: &str) {
     // open file for reading
     let fi = File::open(smiles_file).expect("Could not open file");
 
     // open binary file for index
     let mut offset = 0;
-    // let mut fo = File::create(smiles_file.to_owned() + ".fp");
+    let mut fo = File::create(smiles_file.to_owned() + ".fp");
+    let mut writer = BufWriter::new(fo.unwrap());
+
     for line in std::io::BufReader::new(fi).lines() {
         let line = line.unwrap();
         let molecule = parse_molecule(&line).unwrap().1;
-        IndexItem {
+        let index_item = IndexItem {
             position: offset,
             fingerprint: molecule.ecfp(2, 512),
         };
-        offset += line.len() + 1;
+        offset += line.len() + LINE_ENDING_LENGTH;
+
+        let mut buf = vec![0u8; FINGERPRINT_SIZE / 8 + 8];
+
+        encode_into_slice(&index_item, buf.as_mut_slice(), bincode::config::standard()).unwrap();
+        writer.write(&buf).unwrap();
     }
 }
 

src/ringo/ringo/search.rs

@@ -1 +1,64 @@
+use std::fs::File;
+use std::io::{BufRead, BufReader, Read, Seek};
+use crate::ringo::math::similarity::tanimoto::tanimoto_bitset;
+use crate::ringo::molecule::smiles::reader::molecule::parse_molecule;
+use crate::ringo::ringo::index_item::IndexItem;
 
+pub struct SearchResult {
+    pub line: String,
+    pub similarity: f32
+}
+
+pub fn similarity_search(smiles_file: &str, query: &str, min_similarity: f32, limit: usize) -> Vec<SearchResult> {
+    let query = parse_molecule(query).unwrap().1;
+    let query_fp = query.ecfp(2, 512);
+
+    // smiles file
+    let fis = File::open(&smiles_file).expect("Could not open file");
+    let mut reader = BufReader::new(fis);
+
+    //fingerprings file
+    let fif = File::open(smiles_file.to_owned() + ".fp").expect("Could not open file");
+    let file_len = fif.metadata().unwrap().len();
+    let index_item_size = 72u8;
+    let index_count = file_len / index_item_size as u64;
+    let mut buf_reader = BufReader::new(fif);
+
+    let mut results = Vec::new();
+    
+    for i in 0..index_count {
+        // read index item from file
+        let mut buf = vec![0u8; index_item_size as usize];
+        buf_reader.read_exact(&mut buf).unwrap();
+
+        // decode index item
+        let index_item: IndexItem = bincode::decode_from_slice(&buf, bincode::config::standard()).unwrap().0;
+
+        // calculate similarity
+        let similarity = tanimoto_bitset(&index_item.fingerprint.0, &query_fp.0);
+
+        // print similarity if it is greater than min_similarity
+        if similarity >= min_similarity {
+            let position  = index_item.position;
+            reader.seek(std::io::SeekFrom::Start(position as u64)).unwrap();
+
+            let mut line = String::new();
+            reader.read_line(&mut line).unwrap();
+            // println!("{i} {similarity} {position} {line}");
+            results.push(SearchResult {
+                line: line,
+                similarity: similarity
+            });
+        }
+    }
+    
+    results
+}
+
+#[test]
+fn test_similarity_search() {
+    let results = similarity_search("molecules.smi", "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", 0.7, 100);
+    assert_eq!(results.len(), 1);
+    let results = similarity_search("molecules.smi", "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", 0.5, 100);
+    assert_eq!(results.len(), 2);
+}