feat: database splitting

FEAT: added option to iterate over fasta chunks

FEAT: Added cloning to Search and QuantSettings

CHORE: made reordering of target_decoys a seperate function outside of digest

CHORE: implemented build db from peptide list

FIX: reorder_peptide function return

FIX: ambiguous target/decoy peptides are now always target

FEAT: added fasta chunking params to db

CHORE: refactored process_chunk from runner

FEAT: added defaultto IndexedDatabase

FEAT: added quick_score option to quickly filter peptides

FEAT: added parsing of the prefilter fasta by chunk parameters

FEAT: added extra low memory option

FEAT: using heap to retain best scores in low_mem mode

FIX: merge issues

FEAT: auto calculate prefilter chunk size

chore: linearize history via rebase

Author: sander-willems-bruker

crates/sage-cli/src/input.rs

@@ -10,7 +10,7 @@ use sage_core::{
 };
 use serde::{Deserialize, Serialize};
 
-#[derive(Serialize)]
+#[derive(Serialize, Clone)]
 /// Actual search parameters - may include overrides or default values not set by user
 pub struct Search {
     pub version: String,
@@ -148,7 +148,7 @@ pub struct QuantOptions {
     pub lfq_options: Option<LfqOptions>,
 }
 
-#[derive(Serialize, Default)]
+#[derive(Serialize, Default, Clone)]
 pub struct QuantSettings {
     pub tmt: Option<Isobaric>,
     pub tmt_settings: TmtSettings,

crates/sage-cli/src/main.rs

@@ -1,10 +1,6 @@
 use clap::{value_parser, Arg, Command, ValueHint};
 use input::Input;
-
-mod input;
-mod output;
-mod runner;
-mod telemetry;
+use runner::Runner;
 
 fn main() -> anyhow::Result<()> {
     env_logger::Builder::default()
@@ -107,9 +103,11 @@ fn main() -> anyhow::Result<()> {
 
     let input = Input::from_arguments(matches)?;
 
-    let runner_ = input.build().and_then(runner::Runner::new)?;
+    let runner = input
+        .build()
+        .and_then(|parameters| Runner::new(parameters, parallel))?;
 
-    let tel = runner_.run(parallel, parquet)?;
+    let tel = runner.run(parallel, parquet)?;
 
     if send_telemetry {
         tel.send();

crates/sage-cli/src/runner.rs

@@ -6,15 +6,18 @@ use csv::ByteRecord;
 use log::info;
 use rayon::prelude::*;
 use sage_cloudpath::CloudPath;
-use sage_core::database::IndexedDatabase;
+use sage_core::database::{IndexedDatabase, Parameters, PeptideIx};
+use sage_core::fasta::Fasta;
 use sage_core::ion_series::Kind;
 use sage_core::lfq::{Peak, PrecursorId};
 use sage_core::mass::Tolerance;
+use sage_core::peptide::Peptide;
 use sage_core::scoring::Fragments;
 use sage_core::scoring::{Feature, Scorer};
 use sage_core::spectrum::{ProcessedSpectrum, SpectrumProcessor};
 use sage_core::tmt::TmtQuant;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::time::Instant;
 
 pub struct Runner {
@@ -24,7 +27,8 @@ pub struct Runner {
 }
 
 impl Runner {
-    pub fn new(parameters: Search) -> anyhow::Result<Self> {
+    pub fn new(parameters: Search, parallel: usize) -> anyhow::Result<Self> {
+        let mut parameters = parameters.clone();
         let start = Instant::now();
         let fasta = sage_cloudpath::util::read_fasta(
             &parameters.database.fasta,
@@ -38,7 +42,32 @@ impl Runner {
             )
         })?;
 
-        let database = parameters.database.clone().build(fasta);
+        let database = match parameters.database.prefilter {
+            false => parameters.database.clone().build(fasta),
+            true => {
+                parameters
+                    .database
+                    .auto_calculate_prefilter_chunk_size(&fasta);
+                if parameters.database.prefilter_chunk_size >= fasta.targets.len() {
+                    parameters.database.clone().build(fasta)
+                } else {
+                    info!(
+                        "using {} db chunks of size {}",
+                        (fasta.targets.len() + parameters.database.prefilter_chunk_size - 1)
+                            / parameters.database.prefilter_chunk_size,
+                        parameters.database.prefilter_chunk_size,
+                    );
+                    let mini_runner = Self {
+                        database: IndexedDatabase::default(),
+                        parameters: parameters.clone(),
+                        start,
+                    };
+                    let peptides = mini_runner.prefilter_peptides(parallel, fasta);
+                    parameters.database.clone().build_from_peptides(peptides)
+                }
+            }
+        };
+
         info!(
             "generated {} fragments, {} peptides in {}ms",
             database.fragments.len(),
@@ -52,6 +81,108 @@ impl Runner {
         })
     }
 
+    pub fn prefilter_peptides(self, parallel: usize, fasta: Fasta) -> Vec<Peptide> {
+        let spectra: Option<Vec<ProcessedSpectrum>> =
+            match parallel >= self.parameters.mzml_paths.len() {
+                true => Some(self.read_processed_spectra(&self.parameters.mzml_paths, 0, 0)),
+                false => None,
+            };
+        let mut all_peptides: Vec<Peptide> = fasta
+            .iter_chunks(self.parameters.database.prefilter_chunk_size)
+            .enumerate()
+            .flat_map(|(chunk_id, fasta_chunk)| {
+                let start = Instant::now();
+                info!("pre-filtering fasta chunk {}", chunk_id,);
+                let db = &self.parameters.database.clone().build(fasta_chunk);
+                info!(
+                    "generated {} fragments, {} peptides in {}ms",
+                    db.fragments.len(),
+                    db.peptides.len(),
+                    (Instant::now() - start).as_millis()
+                );
+                let scorer = Scorer {
+                    db,
+                    precursor_tol: self.parameters.precursor_tol,
+                    fragment_tol: self.parameters.fragment_tol,
+                    min_matched_peaks: self.parameters.min_matched_peaks,
+                    min_isotope_err: self.parameters.isotope_errors.0,
+                    max_isotope_err: self.parameters.isotope_errors.1,
+                    min_precursor_charge: self.parameters.precursor_charge.0,
+                    max_precursor_charge: self.parameters.precursor_charge.1,
+                    override_precursor_charge: self.parameters.override_precursor_charge,
+                    max_fragment_charge: self.parameters.max_fragment_charge,
+                    chimera: self.parameters.chimera,
+                    report_psms: self.parameters.report_psms + 1,
+                    wide_window: self.parameters.wide_window,
+                    annotate_matches: self.parameters.annotate_matches,
+                    score_type: self.parameters.score_type,
+                };
+                let peptide_idxs: HashSet<PeptideIx> = match &spectra {
+                    Some(spectra) => self.peptide_filter_processed_spectra(&scorer, spectra),
+                    None => self
+                        .parameters
+                        .mzml_paths
+                        .chunks(parallel)
+                        .enumerate()
+                        .flat_map(|(chunk_idx, chunk)| {
+                            let spectra_chunk =
+                                self.read_processed_spectra(chunk, chunk_idx, parallel);
+                            self.peptide_filter_processed_spectra(&scorer, &spectra_chunk)
+                        })
+                        .collect(),
+                }
+                .into_iter()
+                .collect();
+                let peptides: Vec<Peptide> = peptide_idxs
+                    .into_iter()
+                    .map(|idx| db[idx].clone())
+                    .collect();
+                info!(
+                    "found {} pre-filtered peptides for fasta chunk {}",
+                    peptides.len(),
+                    chunk_id,
+                );
+                peptides
+            })
+            .collect();
+        Parameters::reorder_peptides(&mut all_peptides);
+        all_peptides
+    }
+
+    fn peptide_filter_processed_spectra(
+        &self,
+        scorer: &Scorer,
+        spectra: &Vec<ProcessedSpectrum>,
+    ) -> Vec<PeptideIx> {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        let counter = AtomicUsize::new(0);
+        let start = Instant::now();
+
+        let peptide_idxs: Vec<_> = spectra
+            .par_iter()
+            .filter(|spec| spec.peaks.len() >= self.parameters.min_peaks && spec.level == 2)
+            .map(|x| {
+                let prev = counter.fetch_add(1, Ordering::Relaxed);
+                if prev > 0 && prev % 10_000 == 0 {
+                    let duration = Instant::now().duration_since(start).as_millis() as usize;
+
+                    let rate = prev * 1000 / (duration + 1);
+                    log::trace!("- searched {} spectra ({} spectra/s)", prev, rate);
+                }
+                x
+            })
+            .flat_map(|spec| {
+                scorer.quick_score(spec, self.parameters.database.prefilter_low_memory)
+            })
+            .collect();
+
+        let duration = Instant::now().duration_since(start).as_millis() as usize;
+        let prev = counter.load(Ordering::Relaxed);
+        let rate = prev * 1000 / (duration + 1);
+        log::info!("- search:  {:8} ms ({} spectra/s)", duration, rate);
+        peptide_idxs
+    }
+
     fn spectrum_fdr(&self, features: &mut [Feature]) -> usize {
         if sage_core::ml::linear_discriminant::score_psms(features, self.parameters.precursor_tol)
             .is_none()
@@ -76,8 +207,8 @@ impl Runner {
     fn search_processed_spectra(
         &self,
         scorer: &Scorer,
-        spectra: Vec<ProcessedSpectrum>,
-    ) -> SageResults {
+        spectra: &Vec<ProcessedSpectrum>,
+    ) -> Vec<Feature> {
         use std::sync::atomic::{AtomicUsize, Ordering};
         let counter = AtomicUsize::new(0);
         let start = Instant::now();
@@ -102,7 +233,14 @@ impl Runner {
         let prev = counter.load(Ordering::Relaxed);
         let rate = prev * 1000 / (duration + 1);
         log::info!("- search:  {:8} ms ({} spectra/s)", duration, rate);
+        features
+    }
 
+    fn complete_features(
+        &self,
+        spectra: Vec<ProcessedSpectrum>,
+        features: Vec<Feature>,
+    ) -> SageResults {
         let quant = self
             .parameters
             .quant
@@ -132,6 +270,17 @@ impl Runner {
         chunk_idx: usize,
         batch_size: usize,
     ) -> SageResults {
+        let spectra = self.read_processed_spectra(chunk, chunk_idx, batch_size);
+        let features = self.search_processed_spectra(scorer, &spectra);
+        self.complete_features(spectra, features)
+    }
+
+    fn read_processed_spectra(
+        &self,
+        chunk: &[String],
+        chunk_idx: usize,
+        batch_size: usize,
+    ) -> Vec<ProcessedSpectrum> {
         // Read all of the spectra at once - this can help prevent memory over-consumption issues
         info!(
             "processing files {} .. {} ",
@@ -190,7 +339,7 @@ impl Runner {
         let io_time = Instant::now() - start;
         info!("- file IO: {:8} ms", io_time.as_millis());
 
-        self.search_processed_spectra(scorer, spectra)
+        spectra
     }
 
     pub fn batch_files(&self, scorer: &Scorer, batch_size: usize) -> SageResults {

crates/sage/src/database.rs

@@ -84,6 +84,12 @@ pub struct Builder {
     pub generate_decoys: Option<bool>,
     /// Path to fasta database
     pub fasta: Option<String>,
+    /// Number of sequences to handle simultaneously when pre-filtering the db
+    pub prefilter_chunk_size: Option<usize>,
+    /// Pre-filter the database to minimize memory usage
+    pub prefilter: Option<bool>,
+    /// Pre-filter the database with a minimal amount of memory at the cost of speed
+    pub prefilter_low_memory: Option<bool>,
 }
 
 impl Builder {
@@ -102,6 +108,9 @@ impl Builder {
             max_variable_mods: self.max_variable_mods.map(|x| x.max(1)).unwrap_or(2),
             generate_decoys: self.generate_decoys.unwrap_or(true),
             fasta: self.fasta.expect("A fasta file must be provided!"),
+            prefilter_chunk_size: self.prefilter_chunk_size.unwrap_or(0),
+            prefilter: self.prefilter.unwrap_or(false),
+            prefilter_low_memory: self.prefilter_low_memory.unwrap_or(true),
         }
     }
 
@@ -124,9 +133,32 @@ pub struct Parameters {
     pub decoy_tag: String,
     pub generate_decoys: bool,
     pub fasta: String,
+    pub prefilter_chunk_size: usize,
+    pub prefilter: bool,
+    pub prefilter_low_memory: bool,
 }
 
 impl Parameters {
+    pub fn auto_calculate_prefilter_chunk_size(&mut self, fasta: &Fasta) {
+        const MAX_PEPS_PER_CHUNK: usize = 10_000_000;
+        self.prefilter_chunk_size = match self.prefilter_chunk_size {
+            0 => {
+                let enzyme = self.enzyme.clone().into();
+                let total_unmodified_pep_count: usize = fasta.digest(&enzyme).len();
+                let mod_count_estimate =
+                    (self.variable_mods.len() + 1) * (1 << self.max_variable_mods);
+                let chunk_count =
+                    mod_count_estimate * total_unmodified_pep_count / MAX_PEPS_PER_CHUNK;
+                if chunk_count == 0 {
+                    fasta.targets.len()
+                } else {
+                    fasta.targets.len() / chunk_count
+                }
+            }
+            x => x,
+        };
+    }
+
     pub fn digest(&self, fasta: &Fasta) -> Vec<Peptide> {
         log::trace!("digesting fasta");
         let enzyme = self.enzyme.clone().into();
@@ -172,6 +204,12 @@ impl Parameters {
             })
             .collect::<Vec<_>>();
 
+        Self::reorder_peptides(&mut target_decoys);
+
+        target_decoys
+    }
+
+    pub fn reorder_peptides(target_decoys: &mut Vec<Peptide>) {
         log::trace!("sorting and deduplicating peptides");
 
         // This is equivalent to a stable sort
@@ -187,6 +225,9 @@ impl Parameters {
                 && remove.cterm == keep.cterm
             {
                 keep.proteins.extend(remove.proteins.iter().cloned());
+                // When merging peptides from different Fastas,
+                // decoys in one fasta might be targets in another
+                keep.decoy &= remove.decoy;
                 true
             } else {
                 false
@@ -196,13 +237,15 @@ impl Parameters {
         target_decoys
             .par_iter_mut()
             .for_each(|peptide| peptide.proteins.sort_unstable());
-
-        target_decoys
     }
 
     // pub fn build(self) -> Result<IndexedDatabase, Box<dyn std::error::Error + Send + Sync + 'static>> {
     pub fn build(self, fasta: Fasta) -> IndexedDatabase {
         let target_decoys = self.digest(&fasta);
+        self.build_from_peptides(target_decoys)
+    }
+
+    pub fn build_from_peptides(self, target_decoys: Vec<Peptide>) -> IndexedDatabase {
         log::trace!("generating fragments");
 
         // Finally, perform in silico digest for our target sequences
@@ -321,6 +364,7 @@ pub struct Theoretical {
     pub fragment_mz: f32,
 }
 
+#[derive(Default)]
 pub struct IndexedDatabase {
     pub peptides: Vec<Peptide>,
     pub fragments: Vec<Theoretical>,
@@ -598,6 +642,9 @@ mod test {
             decoy_tag: "rev_".into(),
             generate_decoys: false,
             fasta: "none".into(),
+            prefilter: false,
+            prefilter_chunk_size: 0,
+            prefilter_low_memory: true,
         };
 
         let peptides = params.digest(&fasta);

crates/sage/src/fasta.rs

@@ -77,4 +77,14 @@ impl Fasta {
             })
             .collect()
     }
+
+    pub fn iter_chunks(&self, chunk_size: usize) -> impl Iterator<Item = Self> + '_ {
+        self.targets
+            .chunks(chunk_size)
+            .map(move |target_chunk| Self {
+                targets: target_chunk.to_vec(),
+                decoy_tag: self.decoy_tag.clone(),
+                generate_decoys: self.generate_decoys,
+            })
+    }
 }