current state

Author: Johanna

src/cli.rs

@@ -147,6 +147,45 @@ pub enum Commands {
         #[arg(long, value_parser = valid_cpus, default_value_t = 1)]
         threads: usize,
     },
+    /// Uses an inverted index for faster and sparser pairwise comparisons
+    Inverted {
+        /// List of input FASTA files
+        #[arg(long, group = "input", num_args = 1.., value_delimiter = ',')]
+        seq_files: Option<Vec<String>>,
+
+        /// File listing input files (tab separated name, sequences)
+        #[arg(short, group = "input")]
+        file_list: Option<String>,
+
+        /// Output filename for the merged sketch
+        #[arg(required = true, short)]
+        output: String,
+
+        /// Ignore reverse complement (all contigs are oriented along same strand)
+        #[arg(long, default_value_t = DEFAULT_STRAND)]
+        single_strand: bool,
+
+        /// Minimum k-mer count (with reads)
+        #[arg(long, default_value_t = DEFAULT_MINCOUNT)]
+        min_count: u16,
+
+        /// Minimum k-mer quality (with reads)
+        #[arg(long, default_value_t = DEFAULT_MINQUAL)]
+        min_qual: u8,
+
+        /// Treat every sequence in an input file as a new sample (aa only)
+        // TODO: for now, could be extended to dna, but probably no need
+        #[arg(long, default_value_t = false)]
+        concat_fasta: bool,
+
+        /// Number of CPU threads
+        #[arg(long, value_parser = valid_cpus, default_value_t = 1)]
+        threads: usize,
+
+        /// aaHash 'level'
+        #[arg(long, value_enum, default_value_t = DEFAULT_LEVEL)]
+        level: AaLevel,
+    },
     /// Merge two sketch files (.skm and .skd pair)
     Merge {
         /// The first .skd (sketch data) file

src/hashing/mod.rs

@@ -125,4 +125,6 @@ pub trait RollHash: Iterator<Item = u64> {
     fn reads(&self) -> bool {
         false
     }
+
+    
 }

src/lib.rs

@@ -118,8 +118,9 @@ pub fn main() -> Result<(), Error> {
                 rc,
                 *min_count,
                 *min_qual,
+                false,
             );
-            let sketch_vec = MultiSketch::new(&mut sketches, sketch_size, &kmers, seq_type);
+            let sketch_vec = MultiSketch::new(&mut sketches, sketch_size, &kmers, seq_type, false);
             sketch_vec
                 .save_metadata(output)
                 .expect("Error saving metadata");
@@ -400,6 +401,106 @@ pub fn main() -> Result<(), Error> {
             utils::save_sketch_data(ref_db_name1, ref_db_name2, output)
         }
 
+        Commands::Inverted {
+            seq_files,
+            file_list,
+            output,
+            single_strand,
+            min_count,
+            min_qual,
+            concat_fasta,
+            threads,
+            level,
+        } => {
+            // An extra thread is needed for the writer. This doesn't 'overuse' CPU
+            check_threads(*threads + 1);
+            //get input files
+            log::info!("Getting input files");
+            let input_files: Vec<(String, String, Option<String>)> =
+                get_input_list(file_list, seq_files);
+            log::info!("Parsed {} samples in input list", input_files.len());
+            // read out sketching information needed to sketch the new files
+            let kmers: &[usize] = &[2];
+            let rc = !*single_strand;
+            let sketch_size = 3;
+            let seq_type = &HashType::DNA;
+            if *concat_fasta && matches!(*seq_type, HashType::DNA | HashType::PDB) {
+                panic!("--concat-fasta currently only supported with --seq-type aa");
+            }
+
+            log::info!(
+                "Running sketching: k:{:?}; sketch_size:{}; seq:{:?}; threads:{}",
+                kmers,
+                sketch_size * u64::BITS as u64,
+                seq_type,
+                threads,
+            );
+
+            let seq_type = if let HashType::AA(_) = seq_type {
+                HashType::AA(level.clone())
+            } else {
+                seq_type.clone()
+            };
+            // Before sketch_files
+            println!("Debug: Starting sketch_files with:");
+            println!("Debug: input_files: {:?}", input_files);
+            println!("Debug: kmers: {:?}", kmers);
+            println!("Debug: sketch_size: {}", sketch_size);
+            println!("Debug: seq_type: {:?}", seq_type);
+            println!("Debug: rc: {}", rc);
+            println!("Debug: min_count: {}", min_count);
+            println!("Debug: min_qual: {}", min_qual);
+
+            let mut mini_sketches = sketch_files(
+                output,
+                &input_files,
+                *concat_fasta,
+                &kmers,
+                sketch_size,
+                &seq_type,
+                rc,
+                *min_count,
+                *min_qual,
+                true,
+            );
+
+            println!("Debug: Sketches after sketch_files:");
+            for (i, sketch) in mini_sketches.iter().enumerate() {
+                println!("Debug: Sketch {}", i);
+                println!("Debug:   name: {}", sketch.name());
+                let usigs = sketch.get_usigs_debug();
+                println!("Debug:   usigs length: {}", usigs.len());
+                if !usigs.is_empty() {
+                    println!("Debug:   first few usigs: {:?}", &usigs[..usigs.len().min(5)]);
+                }
+            }
+
+            // After creating mini_sketches
+            println!("Debug: mini_sketches length: {}", mini_sketches.len());
+
+            // After creating sketch_vec
+            let sketch_vec = MultiSketch::new(&mut mini_sketches, sketch_size, &kmers, seq_type, true); 
+            println!("Debug: sketch_vec creation completed");
+            println!("Debug: sketch_vec.sketch_bins length: {}", MultiSketch::get_sketch_bins_len(&sketch_vec));
+            println!("Debug: sketch_vec contents: {:?}", sketch_vec);
+
+            // Before inverting index
+            println!("Debug: About to invert index");
+            let mini_inverted_index = MultiSketch::invert_index(&sketch_vec);
+
+                
+            // Write inverted index to a file
+            sketch_vec
+                .save_metadata(output)
+                .expect("Error saving metadata");
+            match MultiSketch::write_inverted_index_to_file(&mini_inverted_index, output) {
+                Ok(_) => println!("Successfully wrote inverted index to {}", output),
+                Err(e) => eprintln!("Error writing inverted index to file: {}", e),
+            }
+
+            Ok(())
+        }
+
         Commands::Append {
             db,
             seq_files,
@@ -461,13 +562,14 @@ pub fn main() -> Result<(), Error> {
                 rc,
                 *min_count,
                 *min_qual,
+                false,
             );
             let mut db2_metadata =
-                MultiSketch::new(&mut db2_sketches, sketch_size, kmers, seq_type);
+                MultiSketch::new(&mut db2_sketches, sketch_size, kmers, seq_type, false);
 
             // save skd data from db1 and from freshly sketched input files
             log::info!("Merging and saving sketch data to {}.skd", output);
-            // let mut output_file = File::create(format!("{}.skd", output))?;
+
             let mut output_file = OpenOptions::new()
                 .create(true)
                 .append(true)

src/multisketch.rs

@@ -1,6 +1,6 @@
 use anyhow::bail;
 use anyhow::Error;
-use anyhow::{Result, anyhow};
+use anyhow::{anyhow, Result};
 // use thiserror::Error;
 use core::panic;
 use std::fmt;
@@ -15,7 +15,11 @@ use crate::hashing::HashType;
 use crate::sketch::{Sketch, BBITS};
 use crate::sketch_datafile::SketchArrayFile;
 
+use rayon::prelude::*;
 use std::collections::HashSet;
+
+use std::io::Write;
+use std::path::Path;
 #[derive(Serialize, Deserialize)]
 pub struct MultiSketch {
     pub sketch_size: u64,
@@ -41,6 +45,7 @@ impl MultiSketch {
         sketch_size: u64,
         kmer_lengths: &[usize],
         hash_type: HashType,
+        inverted: bool,
     ) -> Self {
         let mut name_map = HashMap::with_capacity(sketches.len());
         for sketch in sketches.iter() {
@@ -63,6 +68,8 @@ impl MultiSketch {
         }
     }
 
+    
+
     /// Saves the metadata
     pub fn save_metadata(&self, file_prefix: &str) -> Result<(), Error> {
         let filename = format!("{}.skm", file_prefix);
@@ -205,10 +212,8 @@ impl MultiSketch {
         let mut removed_samples = Vec::new();
 
         for sketch in &self.sketch_metadata {
-            
             if !genome_ids_to_remove.contains(&(*sketch.name()).to_string()) {
                 new_sketch_metadata.push(sketch.clone());
-
             } else {
                 removed_samples.push(sketch.name());
             }
@@ -218,8 +223,11 @@ impl MultiSketch {
         let set2: HashSet<&str> = genome_ids_to_remove.iter().map(AsRef::as_ref).collect();
         let missing: Vec<&&str> = set2.difference(&set1).collect();
         if !missing.is_empty() {
-            bail!("The following samples have not been found in the database: {:?}", missing);
-        }        
+            bail!(
+                "The following samples have not been found in the database: {:?}",
+                missing
+            );
+        }
 
         self.sketch_metadata = new_sketch_metadata;
         self.save_metadata(output_file_name)?;
@@ -231,7 +239,7 @@ impl MultiSketch {
         input_prefix: &str,
         output_file: &str,
         genome_ids_to_remove: &[String],
-    ) -> anyhow::Result<()>  {
+    ) -> anyhow::Result<()> {
         let mut positions_to_remove = Vec::new();
         let mut missing_ids = Vec::new();
 
@@ -266,10 +274,85 @@ impl MultiSketch {
         Ok(())
     }
 
-    // This function is called when sketches are merged, not when they are
-    // first sketched (this is handled by sketch::sketch_files())
+    pub fn invert_index(sketches: &MultiSketch) -> HashMap<u64, HashSet<usize>> {
+        println!("Debug: Starting invert_index function");
+        println!("Debug: Sample stride: {}", sketches.sample_stride);
+        println!("Debug: Sketch bins length: {}", sketches.sketch_bins.len());
+    
+        // HashMap storing the inverted index
+        let mut inverted_index: HashMap<u64, HashSet<usize>> = HashMap::default();
+    
+        // Parallise the inversion for each sample
+        let local_indices: Vec<HashMap<u64, HashSet<usize>>> = sketches
+            .sketch_bins
+            .par_chunks(sketches.sample_stride)
+            .enumerate()
+            .map(|(genome_id, sample_hash)| {
+                println!("Debug: Processing genome_id: {}", genome_id);
+                println!("Debug: Sample hash length: {}", sample_hash.len());
+                
+                // Print all hashes for this genome
+                println!("Debug: Hashes for genome {}: {:?}", genome_id, sample_hash);
+    
+                let mut local_index: HashMap<u64, HashSet<usize>> = HashMap::default();
+                for &hash in sample_hash {
+                    local_index
+                        .entry(hash)
+                        .or_insert_with(HashSet::default)
+                        .insert(genome_id);
+                }
+                println!("Debug: Local index for genome {} size: {}", genome_id, local_index.len());
+                // Print the hash-to-genome mappings for this local index
+                println!("Debug: Local index contents for genome {}: {:?}", genome_id, local_index);
+                local_index
+            })
+            .collect();
+    
+        println!("Debug: Number of local indices: {}", local_indices.len());
+    
+        // Merge all local inverted indices into a global one
+        for (i, local) in local_indices.iter().enumerate() {
+            println!("Debug: Merging local index {}, size: {}", i, local.len());
+            for (hash, genome_set) in local {
+                inverted_index
+                    .entry(*hash)
+                    .or_insert_with(HashSet::default)
+                    .extend(genome_set);
+            }
+        }
+    
+        println!("Debug: Final inverted index size: {}", inverted_index.len());
+        println!("Debug: Inverted index contents: {:?}", inverted_index);
+    
+        inverted_index
+    }
+    // remove, only need this for debugging
+    pub fn get_sketch_bins_len(&self) -> usize {
+        self.sketch_bins.len()
+    }
+
+    pub fn write_inverted_index_to_file<P: AsRef<Path>>(
+        inverted_index: &HashMap<u64, HashSet<usize>>,
+        file_path: P,
+    ) -> std::io::Result<()> {
+        let mut file = File::create(file_path)?;
+    
+        for (hash, genome_set) in inverted_index {
+            writeln!(
+                file,
+                "{} {}",
+                hash,
+                genome_set.iter().map(|id| id.to_string()).collect::<Vec<_>>().join(" ")
+            )?;
+        }
+    
+        Ok(())
+    }
 }
 
+// This function is called when sketches are merged, not when they are
+// first sketched (this is handled by sketch::sketch_files())
+
 impl fmt::Debug for MultiSketch {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(