mscore/timstof/
frame.rs

1use std::fmt;
2use std::collections::BTreeMap;
3use std::fmt::{Formatter};
4use bincode::{Decode, Encode};
5use itertools;
6use itertools::izip;
7use ordered_float::OrderedFloat;
8use rand::Rng;
9use serde::{Deserialize, Serialize};
10use crate::timstof::spectrum::TimsSpectrum;
11use crate::data::spectrum::{MsType, MzSpectrum, IndexedMzSpectrum, Vectorized, ToResolution};
12use crate::simulation::annotation::{PeakAnnotation, TimsFrameAnnotated};
13use crate::timstof::vec_utils::{filter_with_mask, find_sparse_local_maxima_mask};
14
15#[derive(Clone)]
16pub struct RawTimsFrame {
17    pub frame_id: i32,
18    pub retention_time: f64,
19    pub ms_type: MsType,
20    pub scan: Vec<u32>,
21    pub tof: Vec<u32>,
22    pub intensity: Vec<f64>,
23}
24
25impl RawTimsFrame {
26    pub fn smooth(mut self, window: u32) -> Self {
27        let mut smooth_intensities: Vec<f64> = self.intensity.clone();
28        for (current_index, current_tof) in self.tof.iter().enumerate()
29        {
30            let current_intensity: f64 = self.intensity[current_index];
31            for (_next_index, next_tof) in
32                self.tof[current_index + 1..].iter().enumerate()
33            {
34                let next_index: usize = _next_index + current_index + 1;
35                let next_intensity: f64 = self.intensity[next_index];
36                if (next_tof - current_tof) <= window {
37                    smooth_intensities[current_index] += next_intensity;
38                    smooth_intensities[next_index] += current_intensity;
39                } else {
40                    break;
41                }
42            }
43        }
44        self.intensity = smooth_intensities;
45
46        self
47    }
48    pub fn centroid(mut self, window: u32) -> Self {
49        let local_maxima: Vec<bool> = find_sparse_local_maxima_mask(
50            &self.tof,
51            &self.intensity,
52            window,
53        );
54        self.tof = filter_with_mask(&self.tof, &local_maxima);
55        self.intensity = filter_with_mask(&self.intensity, &local_maxima);
56        self.scan = filter_with_mask(&self.scan, &local_maxima);
57        self
58    }
59}
60
61#[derive(Clone, Debug, Default, Serialize, Deserialize, Encode, Decode)]
62pub struct ImsFrame {
63    pub retention_time: f64,
64    pub mobility: Vec<f64>,
65    pub mz: Vec<f64>,
66    pub intensity: Vec<f64>,
67}
68
69impl ImsFrame {
70    /// Creates a new `ImsFrame` instance.
71    ///
72    /// # Arguments
73    ///
74    /// * `retention_time` - The retention time in seconds.
75    /// * `mobility` - A vector of inverse ion mobilities.
76    /// * `mz` - A vector of m/z values.
77    /// * `intensity` - A vector of intensity values.
78    ///
79    /// # Examples
80    ///
81    /// ```
82    /// use mscore::timstof::frame::ImsFrame;
83    ///
84    /// let frame = ImsFrame::new(100.0, vec![0.1, 0.2], vec![100.5, 200.5], vec![50.0, 60.0]);
85    /// ```
86    pub fn new(retention_time: f64, mobility: Vec<f64>, mz: Vec<f64>, intensity: Vec<f64>) -> Self {
87        ImsFrame { retention_time, mobility, mz, intensity }
88    }
89}
90
91impl fmt::Display for ImsFrame {
92    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
93        write!(f, "ImsFrame(rt: {}, data points: {})", self.retention_time, self.mobility.len())
94    }
95}
96
97#[derive(Clone)]
98pub struct ImsFrameVectorized {
99    pub retention_time: f64,
100    pub mobility: Vec<f64>,
101    pub indices: Vec<i32>,
102    pub values: Vec<f64>,
103    pub resolution: i32,
104}
105
106#[derive(Clone, Debug, Serialize, Deserialize, Encode, Decode)]
107pub struct TimsFrame {
108    pub frame_id: i32,
109    pub ms_type: MsType,
110    pub scan: Vec<i32>,
111    pub tof: Vec<i32>,
112    pub ims_frame: ImsFrame,
113}
114
115impl Default for TimsFrame {
116    fn default() -> Self {
117        TimsFrame {
118            frame_id: 0, // Replace with a suitable default value
119            ms_type: MsType::Unknown,
120            scan: Vec::new(),
121            tof: Vec::new(),
122            ims_frame: ImsFrame::default(), // Uses the default implementation for `ImsFrame`
123        }
124    }
125}
126
127impl TimsFrame {
128    /// Creates a new `TimsFrame` instance.
129    ///
130    /// # Arguments
131    ///
132    /// * `frame_id` - index of frame in TDF raw file.
133    /// * `ms_type` - The type of frame.
134    /// * `retention_time` - The retention time in seconds.
135    /// * `scan` - A vector of scan IDs.
136    /// * `mobility` - A vector of inverse ion mobilities.
137    /// * `tof` - A vector of time-of-flight values.
138    /// * `mz` - A vector of m/z values.
139    /// * `intensity` - A vector of intensity values.
140    ///
141    /// # Examples
142    ///
143    /// ```
144    /// use mscore::data::spectrum::MsType;
145    /// use mscore::timstof::frame::TimsFrame;
146    /// use mscore::timstof::frame::ImsFrame;
147    ///
148    /// let frame = TimsFrame::new(1, MsType::Precursor, 100.0, vec![1, 2], vec![0.1, 0.2], vec![1000, 2000], vec![100.5, 200.5], vec![50.0, 60.0]);
149    /// ```
150    pub fn new(frame_id: i32, ms_type: MsType, retention_time: f64, scan: Vec<i32>, mobility: Vec<f64>, tof: Vec<i32>, mz: Vec<f64>, intensity: Vec<f64>) -> Self {
151        TimsFrame { frame_id, ms_type, scan, tof, ims_frame: ImsFrame { retention_time, mobility, mz, intensity } }
152    }
153
154    ///
155    /// Convert a given TimsFrame to an ImsFrame.
156    ///
157    /// # Examples
158    ///
159    /// ```
160    /// use mscore::data::spectrum::MsType;
161    /// use mscore::timstof::frame::TimsFrame;
162    ///
163    /// let frame = TimsFrame::new(1, MsType::Precursor, 100.0, vec![1, 2], vec![0.1, 0.2], vec![1000, 2000], vec![100.5, 200.5], vec![50.0, 60.0]);
164    /// let ims_spectrum = frame.get_ims_frame();
165    /// ```
166    pub fn get_ims_frame(&self) -> ImsFrame { self.ims_frame.clone() }
167
168    ///
169    /// Convert a given TimsFrame to a vector of TimsSpectrum.
170    ///
171    /// # Examples
172    ///
173    /// ```
174    /// use mscore::data::spectrum::MsType;
175    /// use mscore::timstof::frame::TimsFrame;
176    ///
177    /// let frame = TimsFrame::new(1, MsType::Precursor, 100.0, vec![1, 2], vec![0.1, 0.2], vec![1000, 2000], vec![100.5, 200.5], vec![50.0, 60.0]);
178    /// let tims_spectra = frame.to_tims_spectra();
179    /// ```
180    pub fn to_tims_spectra(&self) -> Vec<TimsSpectrum> {
181        // use a sorted map where scan is used as key
182        let mut spectra = BTreeMap::<i32, (f64, Vec<i32>, Vec<f64>, Vec<f64>)>::new();
183
184        // all indices and the intensity values are sorted by scan and stored in the map as a tuple (mobility, tof, mz, intensity)
185        for (scan, mobility, tof, mz, intensity) in itertools::multizip((
186            &self.scan,
187            &self.ims_frame.mobility,
188            &self.tof,
189            &self.ims_frame.mz,
190            &self.ims_frame.intensity)) {
191            let entry = spectra.entry(*scan).or_insert_with(|| (*mobility, Vec::new(), Vec::new(), Vec::new()));
192            entry.1.push(*tof);
193            entry.2.push(*mz);
194            entry.3.push(*intensity);
195        }
196
197        // convert the map to a vector of TimsSpectrum
198        let mut tims_spectra: Vec<TimsSpectrum> = Vec::new();
199
200        for (scan, (mobility, tof, mz, intensity)) in spectra {
201            let spectrum = IndexedMzSpectrum::new(tof, mz, intensity);
202            tims_spectra.push(TimsSpectrum::new(self.frame_id, scan, self.ims_frame.retention_time, mobility, self.ms_type.clone(), spectrum));
203        }
204
205        tims_spectra
206    }
207
208    ///
209    /// Filter a given TimsFrame by m/z, scan, and intensity.
210    ///
211    /// # Arguments
212    ///
213    /// * `mz_min` - The minimum m/z value.
214    /// * `mz_max` - The maximum m/z value.
215    /// * `scan_min` - The minimum scan value.
216    /// * `scan_max` - The maximum scan value.
217    /// *
218    /// * `intensity_min` - The minimum intensity value.
219    /// * `intensity_max` - The maximum intensity value.
220    ///
221    /// # Examples
222    ///
223    /// ```
224    /// use mscore::data::spectrum::MsType;
225    /// use mscore::timstof::frame::TimsFrame;
226    ///
227    /// let frame = TimsFrame::new(1, MsType::Precursor, 100.0, vec![1, 2], vec![0.1, 0.2], vec![1000, 2000], vec![100.5, 200.5], vec![50.0, 60.0]);
228    /// let filtered_frame = frame.filter_ranged(100.0, 200.0, 1, 2, 0.0, 1.6, 50.0, 60.0);
229    /// ```
230    pub fn filter_ranged(&self, mz_min: f64, mz_max: f64, scan_min: i32, scan_max: i32, inv_mob_min: f64, inv_mob_max: f64, intensity_min: f64, intensity_max: f64) -> TimsFrame {
231
232        let mut scan_vec = Vec::new();
233        let mut mobility_vec = Vec::new();
234        let mut tof_vec = Vec::new();
235        let mut mz_vec = Vec::new();
236        let mut intensity_vec = Vec::new();
237
238        for (mz, intensity, scan, mobility, tof) in itertools::multizip((&self.ims_frame.mz, &self.ims_frame.intensity, &self.scan, &self.ims_frame.mobility, &self.tof)) {
239            if mz >= &mz_min && mz <= &mz_max && scan >= &scan_min && scan <= &scan_max && mobility >= &inv_mob_min && mobility <= &inv_mob_max && intensity >= &intensity_min && intensity <= &intensity_max {
240                scan_vec.push(*scan);
241                mobility_vec.push(*mobility);
242                tof_vec.push(*tof);
243                mz_vec.push(*mz);
244                intensity_vec.push(*intensity);
245            }
246        }
247
248        TimsFrame::new(self.frame_id, self.ms_type.clone(), self.ims_frame.retention_time, scan_vec, mobility_vec, tof_vec, mz_vec, intensity_vec)
249    }
250
251    pub fn top_n(&self, n: usize) -> TimsFrame {
252        let mut indices: Vec<usize> = (0..self.ims_frame.intensity.len()).collect();
253        indices.sort_by(|a, b| self.ims_frame.intensity[*b].partial_cmp(&self.ims_frame.intensity[*a]).unwrap());
254        indices.truncate(n);
255
256        let scan = indices.iter().map(|&i| self.scan[i]).collect();
257        let mobility = indices.iter().map(|&i| self.ims_frame.mobility[i]).collect();
258        let tof = indices.iter().map(|&i| self.tof[i]).collect();
259        let mz = indices.iter().map(|&i| self.ims_frame.mz[i]).collect();
260        let intensity = indices.iter().map(|&i| self.ims_frame.intensity[i]).collect();
261
262        TimsFrame::new(self.frame_id, self.ms_type.clone(), self.ims_frame.retention_time, scan, mobility, tof, mz, intensity)
263    }
264
265    pub fn to_windows_indexed(&self, window_length: f64, overlapping: bool, min_peaks: usize, min_intensity: f64) -> (Vec<i32>, Vec<i32>, Vec<TimsSpectrum>) {
266        // split by scan (ion mobility)
267        let spectra = self.to_tims_spectra();
268
269        let windows: Vec<_> = spectra.iter().map(|spectrum|
270            spectrum.to_windows(window_length, overlapping, min_peaks, min_intensity))
271            .collect();
272
273        let mut scan_indices = Vec::new();
274
275        for tree in windows.iter() {
276            for (_, window) in tree {
277                scan_indices.push(window.scan)
278            }
279        }
280
281        let mut spectra = Vec::new();
282        let mut window_indices = Vec::new();
283
284        for window in windows {
285            for (i, spectrum) in window {
286                spectra.push(spectrum);
287                window_indices.push(i);
288            }
289        }
290
291        (scan_indices, window_indices, spectra)
292    }
293
294    pub fn to_windows(&self, window_length: f64, overlapping: bool, min_peaks: usize, min_intensity: f64) -> Vec<TimsSpectrum> {
295        let (_, _, widows) = self.to_windows_indexed(window_length, overlapping, min_peaks, min_intensity);
296        widows
297    }
298
299    pub fn from_windows(windows: Vec<TimsSpectrum>) -> TimsFrame {
300
301        let first_window = windows.first().unwrap();
302
303        let mut scan = Vec::new();
304        let mut tof = Vec::new();
305        let mut mzs = Vec::new();
306        let mut intensity = Vec::new();
307        let mut mobility = Vec::new();
308
309        for window in windows.iter() {
310            for (i, mz) in window.spectrum.mz_spectrum.mz.iter().enumerate() {
311                scan.push(window.scan);
312                mobility.push(window.mobility);
313                tof.push(window.spectrum.index[i]);
314                mzs.push(*mz);
315                intensity.push(window.spectrum.mz_spectrum.intensity[i]);
316            }
317        }
318
319        TimsFrame::new(first_window.frame_id, first_window.ms_type.clone(), first_window.retention_time, scan, mobility, tof, mzs, intensity)
320    }
321
322    pub fn from_tims_spectra(spectra: Vec<TimsSpectrum>) -> TimsFrame {
323
324        // Helper to quantize mz to an integer key
325        let quantize = |mz: f64| -> i64 {
326            (mz * 1_000_000.0).round() as i64
327        };
328
329        // Step 1: Get frame coordinates
330        let first_spec = spectra.first();
331        let frame_id = match first_spec {
332            Some(first_spec) => first_spec.frame_id,
333            _ => 1
334        };
335        
336        let ms_type = match first_spec {
337            Some(first_spec) => first_spec.ms_type.clone(),
338            _ => MsType::Unknown,
339        };
340        
341        let retention_time = match first_spec { 
342            Some(first_spec) => first_spec.retention_time,
343            _ => 0.0
344        };
345
346        let mut frame_map: BTreeMap<i32, (f64, BTreeMap<i64, (i32, f64)>)> = BTreeMap::new();
347        let mut capacity_count = 0;
348
349        // Step 2: Group by scan and unroll all spectra to a single vector per scan
350        for spectrum in &spectra {
351            let inv_mobility = spectrum.mobility;
352            let entry = frame_map.entry(spectrum.scan).or_insert_with(|| (inv_mobility, BTreeMap::new()));
353            for (i, mz) in spectrum.spectrum.mz_spectrum.mz.iter().enumerate() {
354                let tof = spectrum.spectrum.index[i];
355                let intensity = spectrum.spectrum.mz_spectrum.intensity[i];
356                entry.1.entry(quantize(*mz)).and_modify(|e| *e = (tof, e.1 + intensity)).or_insert((tof, intensity));
357                capacity_count += 1;
358            }
359        }
360
361        // Step 3: Unroll the map to vectors
362        let mut scan = Vec::with_capacity(capacity_count);
363        let mut mobility = Vec::with_capacity(capacity_count);
364        let mut tof = Vec::with_capacity(capacity_count);
365        let mut mzs = Vec::with_capacity(capacity_count);
366        let mut intensity = Vec::with_capacity(capacity_count);
367
368        for (scan_val, (mobility_val, mz_map)) in frame_map {
369            for (mz_val, (tof_val, intensity_val)) in mz_map {
370                scan.push(scan_val);
371                mobility.push(mobility_val);
372                tof.push(tof_val);
373                mzs.push(mz_val as f64 / 1_000_000.0);
374                intensity.push(intensity_val);
375            }
376        }
377
378        TimsFrame::new(frame_id, ms_type, retention_time, scan, mobility, tof, mzs, intensity)
379    }
380
381    pub fn to_dense_windows(&self, window_length: f64, overlapping: bool, min_peaks: usize, min_intensity: f64, resolution: i32) -> (Vec<f64>, Vec<f64>, Vec<f64>, Vec<i32>, Vec<i32>, usize, usize) {
382        let factor = (10.0f64).powi(resolution);
383        let num_colums = ((window_length * factor).round() + 1.0) as usize;
384
385        let (scans, window_indices, spectra) = self.to_windows_indexed(window_length, overlapping, min_peaks, min_intensity);
386        let mut mobilities = Vec::with_capacity(spectra.len());
387        let mut mzs = Vec::with_capacity(spectra.len());
388
389        // go over all spectra and fill mobilities and mzs
390        for spectrum in &spectra {
391            mobilities.push(spectrum.mobility);
392            mzs.push(spectrum.spectrum.mz_spectrum.mz.first().unwrap().clone());
393        }
394
395        let vectorized_spectra = spectra.iter().map(|spectrum| spectrum.vectorized(resolution)).collect::<Vec<_>>();
396
397        let mut flat_matrix: Vec<f64> = vec![0.0; spectra.len() * num_colums];
398
399        for (row_index, (window_index, spectrum)) in itertools::multizip((&window_indices, vectorized_spectra)).enumerate() {
400
401            let vectorized_window_index = match *window_index >= 0 {
402                true => (*window_index as f64 * window_length * factor).round() as i32,
403                false => (((-1.0 * (*window_index as f64)) * window_length - (0.5 * window_length)) * factor).round() as i32,
404            };
405
406            for (i, index) in spectrum.vector.mz_vector.indices.iter().enumerate() {
407                let zero_based_index = (index - vectorized_window_index) as usize;
408                let current_index = row_index * num_colums + zero_based_index;
409                flat_matrix[current_index] = spectrum.vector.mz_vector.values[i];
410            }
411
412        }
413        (flat_matrix, mobilities, mzs, scans, window_indices, spectra.len(), num_colums)
414    }
415
416
417    pub fn to_indexed_mz_spectrum(&self) -> IndexedMzSpectrum {
418        let mut grouped_data: BTreeMap<i32, Vec<(f64, f64)>> = BTreeMap::new();
419
420        // Group by 'tof' with 'mz' and 'intensity'
421        for (&tof, (&mz, &intensity)) in self.tof.iter().zip(self.ims_frame.mz.iter().zip(self.ims_frame.intensity.iter())) {
422            grouped_data.entry(tof).or_insert_with(Vec::new).push((mz, intensity));
423        }
424
425        let mut index = Vec::new();
426        let mut mz = Vec::new();
427        let mut intensity = Vec::new();
428
429        for (&tof_val, values) in &grouped_data {
430            let sum_intensity: f64 = values.iter().map(|&(_, i)| i).sum();
431            let avg_mz: f64 = values.iter().map(|&(m, _)| m).sum::<f64>() / values.len() as f64;
432
433            index.push(tof_val);
434            mz.push(avg_mz);
435            intensity.push(sum_intensity);
436        }
437
438        IndexedMzSpectrum {
439            index,
440            mz_spectrum: MzSpectrum { mz, intensity },
441        }
442    }
443
444    pub fn generate_random_sample(&self, take_probability: f64) -> TimsFrame {
445        assert!(take_probability >= 0.0 && take_probability <= 1.0);
446
447        let mut rng = rand::thread_rng();
448        let mut scan = Vec::new();
449        let mut mobility = Vec::new();
450        let mut tof = Vec::new();
451        let mut mz = Vec::new();
452        let mut intensity = Vec::new();
453
454        for (s, m, t, mz_val, i) in itertools::multizip((&self.scan, &self.ims_frame.mobility, &self.tof, &self.ims_frame.mz, &self.ims_frame.intensity)) {
455            if rng.gen::<f64>() <= take_probability {
456                scan.push(*s);
457                mobility.push(*m);
458                tof.push(*t);
459                mz.push(*mz_val);
460                intensity.push(*i);
461            }
462        }
463
464        TimsFrame::new(self.frame_id, self.ms_type.clone(), self.ims_frame.retention_time, scan, mobility, tof, mz, intensity)
465    }
466
467    pub fn to_noise_annotated_tims_frame(&self) -> TimsFrameAnnotated {
468        let mut annotations = Vec::with_capacity(self.ims_frame.mz.len());
469        let tof_values = self.tof.clone();
470        let mz_values = self.ims_frame.mz.clone();
471        let scan_values = self.scan.clone();
472        let inv_mobility_values = self.ims_frame.mobility.clone();
473        let intensity_values = self.ims_frame.intensity.clone();
474
475        for intensity in &intensity_values {
476            annotations.push(PeakAnnotation::new_random_noise(*intensity));
477        }
478
479        TimsFrameAnnotated::new(
480            self.frame_id,
481            self.ims_frame.retention_time,
482            self.ms_type.clone(),
483            tof_values.iter().map(|&x| x as u32).collect(),
484            mz_values,
485            scan_values.iter().map(|&x| x as u32).collect(),
486            inv_mobility_values,
487            intensity_values,
488            annotations,
489        )
490    }
491
492    pub fn get_inverse_mobility_along_scan_marginal(&self) -> f64 {
493        let mut marginal_map: BTreeMap<i32, (f64, f64)> = BTreeMap::new();
494        // go over all data points of scan, inv_mob and intensity
495        for (scan, inv_mob, intensity) in izip!(&self.scan, &self.ims_frame.mobility, &self.ims_frame.intensity) {
496            // create a key for the map
497            let key = *scan;
498            // get the entry from the map or insert a new one
499            let entry = marginal_map.entry(key).or_insert((0.0, 0.0));
500            // update the entry with the current intensity adding it to the existing intensity
501            entry.0 += *intensity;
502            // update the entry with the current inverse mobility, overwriting the existing value
503            entry.1 = *inv_mob;
504        }
505
506        // get the inverse mobility with the highest intensity
507        let (_, max_inv_mob) = marginal_map.iter().max_by(|a, b| a.1.0.partial_cmp(&b.1.0).unwrap_or(std::cmp::Ordering::Equal)).unwrap_or((&0, &(0.0, 0.0))).1;
508
509        *max_inv_mob
510    }
511
512    /// Calculate the weighted mean and variance of `inv_mob` values based on their intensities.
513    pub fn get_mobility_mean_and_variance(&self) -> (f64, f64) {
514        let mut mobility_map: BTreeMap<OrderedFloat<f64>, f64> = BTreeMap::new();
515
516        // Aggregate intensity values for each `inv_mob`
517        for (inv_mob, intensity) in izip!(&self.ims_frame.mobility, &self.ims_frame.intensity) {
518            let entry = mobility_map.entry(OrderedFloat(*inv_mob)).or_insert(0.0);
519            *entry += *intensity;
520        }
521
522        // Calculate weighted mean
523        let mut total_weight = 0.0;
524        let mut weighted_sum = 0.0;
525        for (&inv_mob, &intensity) in &mobility_map {
526            total_weight += intensity;
527            weighted_sum += inv_mob.into_inner() * intensity;
528        }
529        let mean = weighted_sum / total_weight;
530
531        // Calculate weighted variance
532        let mut weighted_squared_diff_sum = 0.0;
533        for (&inv_mob, &intensity) in &mobility_map {
534            let diff = inv_mob.into_inner() - mean;
535            weighted_squared_diff_sum += intensity * diff * diff;
536        }
537        let variance = weighted_squared_diff_sum / total_weight;
538
539        (mean, variance)
540    }
541
542    pub fn get_tims_spectrum(&self, scan_number: i32) -> Option<TimsSpectrum> {
543        let mut tof = Vec::new();
544        let mut mz = Vec::new();
545        let mut intensity = Vec::new();
546
547        for (s, t, m, i) in itertools::multizip((&self.scan, &self.tof, &self.ims_frame.mz, &self.ims_frame.intensity)) {
548            if *s == scan_number {
549                tof.push(*t);
550                mz.push(*m);
551                intensity.push(*i);
552            }
553        }
554
555        if mz.is_empty() {
556            return None;
557        }
558
559        let mobility = self.ims_frame.mobility.iter()
560            .zip(&self.scan)
561            .find(|(_, s)| **s == scan_number)
562            .map(|(m, _)| *m)?;
563
564        Some(TimsSpectrum {
565            frame_id: self.frame_id,
566            scan: scan_number,
567            retention_time: self.ims_frame.retention_time,
568            mobility,
569            ms_type: self.ms_type.clone(),
570            spectrum: IndexedMzSpectrum::new(tof, mz, intensity),
571        })
572    }
573
574    pub fn fold_along_scan_axis(self, fold_width: usize) -> TimsFrame {
575        // extract tims spectra from frame
576        let spectra = self.to_tims_spectra();
577
578        // create a new collection of merged spectra,where spectra are first grouped by the key they create when divided by fold_width
579        // and then merge them by addition
580        let mut merged_spectra: BTreeMap<i32, TimsSpectrum> = BTreeMap::new();
581        for spectrum in spectra {
582            let key = spectrum.scan / fold_width as i32;
583
584            // if the key already exists, merge the spectra
585            if let Some(existing_spectrum) = merged_spectra.get_mut(&key) {
586
587                let merged_spectrum = existing_spectrum.clone() + spectrum;
588                // update the existing spectrum with the merged one
589                *existing_spectrum = merged_spectrum;
590
591            } else {
592                // otherwise, insert the new spectrum
593                merged_spectra.insert(key, spectrum);
594            }
595        }
596
597        // convert the merged spectra back to a TimsFrame
598        TimsFrame::from_tims_spectra(merged_spectra.into_values().collect())
599    }
600}
601
602struct AggregateData {
603    intensity_sum: f64,
604    ion_mobility_sum: f64,
605    tof_sum: i64,
606    count: i32,
607}
608
609impl std::ops::Add for TimsFrame {
610    type Output = Self;
611    fn add(self, other: Self) -> TimsFrame {
612        let mut combined_map: BTreeMap<(i32, i64), AggregateData> = BTreeMap::new();
613
614        let quantize = |mz: f64| -> i64 {
615            (mz * 1_000_000.0).round() as i64
616        };
617
618        let add_to_map = |map: &mut BTreeMap<(i32, i64), AggregateData>, mz, ion_mobility, tof, scan, intensity| {
619            let key = (scan, quantize(mz));
620            let entry = map.entry(key).or_insert(AggregateData { intensity_sum: 0.0, ion_mobility_sum: 0.0, tof_sum: 0, count: 0 });
621            entry.intensity_sum += intensity;
622            entry.ion_mobility_sum += ion_mobility;
623            entry.tof_sum += tof as i64;
624            entry.count += 1;
625        };
626
627        for (mz, tof, ion_mobility, scan, intensity) in izip!(&self.ims_frame.mz, &self.tof, &self.ims_frame.mobility, &self.scan, &self.ims_frame.intensity) {
628            add_to_map(&mut combined_map, *mz, *ion_mobility, *tof, *scan, *intensity);
629        }
630
631        for (mz, tof, ion_mobility, scan, intensity) in izip!(&other.ims_frame.mz, &other.tof, &other.ims_frame.mobility, &other.scan, &other.ims_frame.intensity) {
632            add_to_map(&mut combined_map, *mz, *ion_mobility, *tof, *scan, *intensity);
633        }
634
635        let mut mz_combined = Vec::new();
636        let mut tof_combined = Vec::new();
637        let mut ion_mobility_combined = Vec::new();
638        let mut scan_combined = Vec::new();
639        let mut intensity_combined = Vec::new();
640
641        for ((scan, quantized_mz), data) in combined_map {
642            mz_combined.push(quantized_mz as f64 / 1_000_000.0);
643            tof_combined.push(data.tof_sum / data.count as i64);
644            ion_mobility_combined.push(data.ion_mobility_sum / data.count as f64);
645            scan_combined.push(scan);
646            intensity_combined.push(data.intensity_sum);
647        }
648
649        let frame = TimsFrame {
650            frame_id: self.frame_id,
651            ms_type: if self.ms_type == other.ms_type { self.ms_type.clone() } else { MsType::Unknown },
652            scan: scan_combined,
653            tof: tof_combined.iter().map(|&x| x as i32).collect(),
654            ims_frame: ImsFrame {
655                retention_time: self.ims_frame.retention_time,
656                mobility: ion_mobility_combined,
657                mz: mz_combined,
658                intensity: intensity_combined,
659            },
660        };
661
662        return frame;
663    }
664}
665
666impl fmt::Display for TimsFrame {
667    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
668
669        let (mz, i) = self.ims_frame.mz.iter()
670            .zip(&self.ims_frame.intensity)
671            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
672            .unwrap();
673
674        write!(f, "TimsFrame(id: {}, type: {}, rt: {}, data points: {}, max by intensity: (mz: {}, intensity: {}))",
675               self.frame_id, self.ms_type, self.ims_frame.retention_time, self.scan.len(), format!("{:.3}", mz), i)
676    }
677}
678
679impl Vectorized<TimsFrameVectorized> for TimsFrame {
680    fn vectorized(&self, resolution: i32) -> TimsFrameVectorized {
681        let binned_frame = self.to_resolution(resolution);
682        // Translate the m/z values into integer indices
683        let indices: Vec<i32> = binned_frame.ims_frame.mz.iter().map(|&mz| (mz * 10f64.powi(resolution)).round() as i32).collect();
684        // Create a vector of values
685        return TimsFrameVectorized {
686            frame_id: self.frame_id,
687            ms_type: self.ms_type.clone(),
688            scan: binned_frame.scan,
689            tof: binned_frame.tof,
690            ims_frame: ImsFrameVectorized {
691                retention_time: binned_frame.ims_frame.retention_time,
692                mobility: binned_frame.ims_frame.mobility,
693                indices,
694                values: binned_frame.ims_frame.intensity,
695                resolution,
696            },
697        };
698    }
699}
700
701///
702/// Convert a given TimsFrame to a vector of TimsSpectrum.
703///
704/// # Arguments
705///
706/// * `resolution` - The resolution to which the m/z values should be rounded.
707///
708/// # Examples
709///
710/// ```
711/// use mscore::data::spectrum::MsType;
712/// use mscore::timstof::frame::TimsFrame;
713/// use mscore::timstof::spectrum::TimsSpectrum;
714/// use mscore::data::spectrum::IndexedMzSpectrum;
715/// use mscore::data::spectrum::ToResolution;
716///
717/// let frame = TimsFrame::new(1, MsType::Precursor, 100.0, vec![1, 2], vec![0.1, 0.2], vec![1000, 2000], vec![100.5, 200.5], vec![50.0, 60.0]);
718/// let low_res_frame = frame.to_resolution(1);
719/// ```
720impl ToResolution for TimsFrame {
721    fn to_resolution(&self, resolution: i32) -> TimsFrame {
722        let factor = (10.0f64).powi(resolution);
723
724        // Using a tuple of (scan, mz_bin) as a key
725        // Value will store sum of intensities, sum of tofs, sum of mobilities and their count for averaging
726        let mut bin_map: BTreeMap<(i32, i32), (f64, f64, f64, i32)> = BTreeMap::new();
727
728        for i in 0..self.ims_frame.mz.len() {
729            let rounded_mz = (self.ims_frame.mz[i] * factor).round() as i32;
730            let scan_val = self.scan[i];
731            let intensity_val = self.ims_frame.intensity[i] as f64;
732            let tof_val = self.tof[i] as f64;
733            let mobility_val = self.ims_frame.mobility[i] as f64;
734
735            let entry = bin_map.entry((scan_val, rounded_mz)).or_insert((0.0, 0.0, 0.0, 0));
736            entry.0 += intensity_val;
737            entry.1 += tof_val;
738            entry.2 += mobility_val;
739            entry.3 += 1;
740        }
741
742        let mut new_mz = Vec::with_capacity(bin_map.len());
743        let mut new_scan = Vec::with_capacity(bin_map.len());
744        let mut new_intensity = Vec::with_capacity(bin_map.len());
745        let mut new_tof = Vec::with_capacity(bin_map.len());
746        let mut new_mobility = Vec::with_capacity(bin_map.len());
747
748        for ((scan, mz_bin), (intensity_sum, tof_sum, mobility_sum, count)) in bin_map {
749            new_mz.push(mz_bin as f64 / factor);
750            new_scan.push(scan);
751            new_intensity.push(intensity_sum);
752            new_tof.push((tof_sum / count as f64) as i32);
753            new_mobility.push(mobility_sum / count as f64);
754        }
755
756        TimsFrame {
757            frame_id: self.frame_id,
758            ms_type: self.ms_type.clone(),
759            scan: new_scan,
760            tof: new_tof,
761            ims_frame: ImsFrame {
762                retention_time: self.ims_frame.retention_time,
763                mobility: new_mobility,
764                mz: new_mz,
765                intensity: new_intensity,
766            },
767        }
768    }
769}
770
771#[derive(Clone)]
772pub struct TimsFrameVectorized {
773    pub frame_id: i32,
774    pub ms_type: MsType,
775    pub scan: Vec<i32>,
776    pub tof: Vec<i32>,
777    pub ims_frame: ImsFrameVectorized,
778}
779
780impl TimsFrameVectorized {
781    pub fn filter_ranged(&self, mz_min: f64, mz_max: f64, scan_min: i32, scan_max: i32, inv_mob_min: f64, inv_mob_max: f64, intensity_min: f64, intensity_max: f64) -> TimsFrameVectorized {
782        let mut scan_vec = Vec::new();
783        let mut mobility_vec = Vec::new();
784        let mut tof_vec = Vec::new();
785        let mut mz_vec = Vec::new();
786        let mut intensity_vec = Vec::new();
787        let mut indices_vec = Vec::new();
788
789        for (mz, intensity, scan, mobility, tof, index) in itertools::multizip((&self.ims_frame.values, &self.ims_frame.values, &self.scan, &self.ims_frame.mobility, &self.tof, &self.ims_frame.indices)) {
790            if mz >= &mz_min && mz <= &mz_max && scan >= &scan_min && scan <= &scan_max && mobility >= &inv_mob_min && mobility <= &inv_mob_max && intensity >= &intensity_min && intensity <= &intensity_max {
791                scan_vec.push(*scan);
792                mobility_vec.push(*mobility);
793                tof_vec.push(*tof);
794                mz_vec.push(*mz);
795                intensity_vec.push(*intensity);
796                indices_vec.push(*index);
797            }
798        }
799
800        TimsFrameVectorized {
801            frame_id: self.frame_id,
802            ms_type: self.ms_type.clone(),
803            scan: scan_vec,
804            tof: tof_vec,
805            ims_frame: ImsFrameVectorized {
806                retention_time: self.ims_frame.retention_time,
807                mobility: mobility_vec,
808                indices: indices_vec,
809                values: mz_vec,
810                resolution: self.ims_frame.resolution,
811            },
812        }
813    }
814}
815
816impl fmt::Display for TimsFrameVectorized {
817    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
818
819        let (mz, i) = self.ims_frame.values.iter()
820            .zip(&self.ims_frame.values)
821            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
822            .unwrap();
823
824        write!(f, "TimsFrame(id: {}, type: {}, rt: {}, data points: {}, max by intensity: (mz: {}, intensity: {}))",
825               self.frame_id, self.ms_type, self.ims_frame.retention_time, self.scan.len(), format!("{:.3}", mz), i)
826    }
827}