(wip) added span-run and nsplit strategies

src/io/readers.rs

@@ -4,6 +4,7 @@ mod metadata_reader;
 mod precursor_reader;
 mod quad_settings_reader;
 mod spectrum_reader;
+mod tdf_utils;
 
 pub use frame_reader::*;
 pub use metadata_reader::*;

src/io/readers/precursor_reader/tdf/dia.rs

@@ -1,5 +1,6 @@
 use std::path::{Path, PathBuf};
 
+use crate::io::readers::tdf_utils::expand_quadrupole_settings;
 use crate::{
     domain_converters::{
         ConvertableDomain, Frame2RtConverter, Scan2ImConverter,
@@ -34,23 +35,8 @@ impl DIATDFPrecursorReader {
             SqlWindowGroup::from_sql_reader(&tdf_sql_reader).unwrap();
         let quadrupole_settings =
             QuadrupoleSettingsReader::new(tdf_sql_reader.get_path());
-        let mut expanded_quadrupole_settings: Vec<QuadrupoleSettings> = vec![];
-        for window_group in window_groups {
-            let window = window_group.window_group;
-            let frame = window_group.frame;
-            let group = &quadrupole_settings[window as usize - 1];
-            for sub_window in 0..group.isolation_mz.len() {
-                let sub_quad_settings = QuadrupoleSettings {
-                    index: frame,
-                    scan_starts: vec![group.scan_starts[sub_window]],
-                    scan_ends: vec![group.scan_ends[sub_window]],
-                    isolation_mz: vec![group.isolation_mz[sub_window]],
-                    isolation_width: vec![group.isolation_width[sub_window]],
-                    collision_energy: vec![group.collision_energy[sub_window]],
-                };
-                expanded_quadrupole_settings.push(sub_quad_settings)
-            }
-        }
+        let expanded_quadrupole_settings =
+            expand_quadrupole_settings(&window_groups, &quadrupole_settings);
         Self {
             path: path.as_ref().to_path_buf(),
             expanded_quadrupole_settings,

src/io/readers/spectrum_reader/tdf.rs

@@ -69,7 +69,11 @@ impl SpectrumReaderTrait for TDFSpectrumReader {
     }
 
     fn len(&self) -> usize {
-        self.precursor_reader.len()
+        debug_assert_eq!(
+            self.precursor_reader.len(),
+            self.raw_spectrum_reader.len()
+        );
+        self.raw_spectrum_reader.len()
     }
 
     fn get_path(&self) -> PathBuf {

src/io/readers/spectrum_reader/tdf/dda.rs

@@ -96,4 +96,8 @@ impl RawSpectrumReaderTrait for DDARawSpectrumReader {
         };
         raw_spectrum
     }
+
+    fn len(&self) -> usize {
+        self.offsets.len() - 1
+    }
 }

src/io/readers/spectrum_reader/tdf/dia.rs

@@ -1,3 +1,4 @@
+use crate::io::readers::tdf_utils::expand_quadrupole_settings;
 use crate::{
     io::readers::{
         file_readers::sql_reader::{
@@ -19,27 +20,12 @@ pub struct DIARawSpectrumReader {
 
 impl DIARawSpectrumReader {
     pub fn new(tdf_sql_reader: &SqlReader, frame_reader: FrameReader) -> Self {
-        let window_groups =
+        let window_groups: Vec<SqlWindowGroup> =
             SqlWindowGroup::from_sql_reader(&tdf_sql_reader).unwrap();
         let quadrupole_settings =
             QuadrupoleSettingsReader::new(&tdf_sql_reader.get_path());
-        let mut expanded_quadrupole_settings: Vec<QuadrupoleSettings> = vec![];
-        for window_group in window_groups {
-            let window = window_group.window_group;
-            let frame = window_group.frame;
-            let group = &quadrupole_settings[window as usize - 1];
-            for sub_window in 0..group.isolation_mz.len() {
-                let sub_quad_settings = QuadrupoleSettings {
-                    index: frame,
-                    scan_starts: vec![group.scan_starts[sub_window]],
-                    scan_ends: vec![group.scan_ends[sub_window]],
-                    isolation_mz: vec![group.isolation_mz[sub_window]],
-                    isolation_width: vec![group.isolation_width[sub_window]],
-                    collision_energy: vec![group.collision_energy[sub_window]],
-                };
-                expanded_quadrupole_settings.push(sub_quad_settings)
-            }
-        }
+        let expanded_quadrupole_settings =
+            expand_quadrupole_settings(&window_groups, &quadrupole_settings);
         Self {
             expanded_quadrupole_settings,
             frame_reader,
@@ -50,6 +36,15 @@ impl DIARawSpectrumReader {
 impl RawSpectrumReaderTrait for DIARawSpectrumReader {
     fn get(&self, index: usize) -> RawSpectrum {
         let quad_settings = &self.expanded_quadrupole_settings[index];
+        if index < 10 {
+            println!("{}", index);
+            println!("{:?}", quad_settings);
+        }
+        if index > (self.expanded_quadrupole_settings.len() - 10) {
+            println!("{}", index);
+            println!("{:?}", quad_settings);
+        }
+
         let collision_energy = quad_settings.collision_energy[0];
         let isolation_mz = quad_settings.isolation_mz[0];
         let isolation_width = quad_settings.isolation_width[0];
@@ -75,4 +70,8 @@ impl RawSpectrumReaderTrait for DIARawSpectrumReader {
         };
         raw_spectrum
     }
+
+    fn len(&self) -> usize {
+        self.expanded_quadrupole_settings.len()
+    }
 }

src/io/readers/spectrum_reader/tdf/raw_spectra.rs

@@ -110,8 +110,13 @@ impl RawSpectrumReader {
     pub fn get(&self, index: usize) -> RawSpectrum {
         self.raw_spectrum_reader.get(index)
     }
+
+    pub fn len(&self) -> usize {
+        self.raw_spectrum_reader.len()
+    }
 }
 
 pub trait RawSpectrumReaderTrait: Sync {
     fn get(&self, index: usize) -> RawSpectrum;
+    fn len(&self) -> usize;
 }

src/io/readers/tdf_utils.rs

@@ -0,0 +1,124 @@
+use crate::io::readers::file_readers::sql_reader::frame_groups::SqlWindowGroup;
+use crate::ms_data::QuadrupoleSettings;
+
+type SpanStep = (usize, usize);
+
+enum QuadWindowExpansionStrategy {
+    None,
+    Even(usize),
+    Uniform(SpanStep),
+}
+
+fn scan_range_subsplit(
+    start: usize,
+    end: usize,
+    strategy: &QuadWindowExpansionStrategy,
+) -> Vec<(usize, usize)> {
+    let out = match strategy {
+        QuadWindowExpansionStrategy::None => {
+            vec![(start, end)]
+        },
+        QuadWindowExpansionStrategy::Even(num_splits) => {
+            let sub_subwindow_width = (end - start) / (num_splits + 1);
+            let mut out = Vec::new();
+            for sub_subwindow in 0..num_splits.clone() {
+                let sub_subwindow_scan_start =
+                    start + (sub_subwindow_width * sub_subwindow);
+                let sub_subwindow_scan_end =
+                    start + (sub_subwindow_width * (sub_subwindow + 2));
+
+                out.push((sub_subwindow_scan_start, sub_subwindow_scan_end))
+            }
+            out
+        },
+        QuadWindowExpansionStrategy::Uniform((span, step)) => {
+            let mut curr_start = start.clone();
+            let mut curr_end = start + span;
+            let mut out = Vec::new();
+            while curr_end < end {
+                out.push((curr_start, curr_end));
+                curr_start += step;
+                curr_end += step;
+            }
+            out
+        },
+    };
+    out
+}
+
+pub fn expand_quadrupole_settings(
+    window_groups: &[SqlWindowGroup],
+    quadrupole_settings: &[QuadrupoleSettings],
+) -> Vec<QuadrupoleSettings> {
+    // Read the 'NUM_SUB_SUB_SPLITS' from env variables ... default to 1
+    // (for now)
+
+    let splits = match std::env::var("NUM_SUB_SUB_SPLITS") {
+        Ok(s) => match s.parse::<usize>() {
+            Ok(n) => {
+                println!("Number of splits: {} from env", n);
+                QuadWindowExpansionStrategy::Even(n)
+            },
+            Err(_) => {
+                println!("Invalid number of splits: {}", s);
+                QuadWindowExpansionStrategy::None
+            },
+        },
+        Err(_) => match std::env::var("SUB_SPLITS_SPAN") {
+            Ok(s) => match s.parse::<usize>() {
+                Ok(n) => {
+                    println!("Number of scans per split: {} from env", n);
+                    QuadWindowExpansionStrategy::Uniform((n, n / 2))
+                },
+                Err(_) => {
+                    println!("Invalid number of splits: {}", s);
+                    QuadWindowExpansionStrategy::None
+                },
+            },
+            Err(_) => QuadWindowExpansionStrategy::None,
+        },
+    };
+
+    let mut expanded_quadrupole_settings: Vec<QuadrupoleSettings> = vec![];
+    for window_group in window_groups {
+        let window = window_group.window_group;
+        let frame = window_group.frame;
+        let group = &quadrupole_settings[window as usize - 1];
+        for sub_window in 0..group.isolation_mz.len() {
+            let subwindow_scan_start = group.scan_starts[sub_window];
+            let subwindow_scan_end = group.scan_ends[sub_window];
+            for (sws, swe) in scan_range_subsplit(
+                subwindow_scan_start,
+                subwindow_scan_end,
+                &splits,
+            ) {
+                assert!(
+                    sws >= subwindow_scan_start,
+                    "{} >= {} not true",
+                    sws,
+                    subwindow_scan_start
+                );
+                assert!(
+                    swe <= subwindow_scan_end,
+                    "{} <= {} not true",
+                    swe,
+                    subwindow_scan_end
+                );
+                let sub_quad_settings = QuadrupoleSettings {
+                    index: frame,
+                    scan_starts: vec![sws],
+                    scan_ends: vec![swe],
+                    isolation_mz: vec![group.isolation_mz[sub_window]],
+                    isolation_width: vec![group.isolation_width[sub_window]],
+                    collision_energy: vec![group.collision_energy[sub_window]],
+                };
+                expanded_quadrupole_settings.push(sub_quad_settings)
+            }
+        }
+    }
+    println!(
+        "Number of expanded quad settings {}",
+        expanded_quadrupole_settings.len()
+    );
+    expanded_quadrupole_settings
+}

src/utils/vec_utils.rs

@@ -12,8 +12,8 @@ pub fn group_and_sum<T: Ord + Copy, U: std::ops::Add<Output = U> + Copy>(
         return (vec![], vec![]);
     }
     let order: Vec<usize> = argsort(&groups);
-    let mut new_groups: Vec<T> = vec![];
-    let mut new_values: Vec<U> = vec![];
+    let mut new_groups: Vec<T> = Vec::with_capacity(order.len());
+    let mut new_values: Vec<U> = Vec::with_capacity(order.len());
     let mut current_group: T = groups[order[0]];
     let mut current_value: U = values[order[0]];
     for &index in &order[1..] {