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<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 vectorized_spectra = spectra.iter().map(|spectrum| spectrum.vectorized(resolution)).collect::<Vec<_>>();
387
388        let mut flat_matrix: Vec<f64> = vec![0.0; spectra.len() * num_colums];
389
390        for (row_index, (window_index, spectrum)) in itertools::multizip((&window_indices, vectorized_spectra)).enumerate() {
391
392            let vectorized_window_index = match *window_index >= 0 {
393                true => (*window_index as f64 * window_length * factor).round() as i32,
394                false => (((-1.0 * (*window_index as f64)) * window_length - (0.5 * window_length)) * factor).round() as i32,
395            };
396
397            for (i, index) in spectrum.vector.mz_vector.indices.iter().enumerate() {
398                let zero_based_index = (index - vectorized_window_index) as usize;
399                let current_index = row_index * num_colums + zero_based_index;
400                flat_matrix[current_index] = spectrum.vector.mz_vector.values[i];
401            }
402
403        }
404        (flat_matrix, scans, window_indices, spectra.len(), num_colums)
405    }
406
407
408    pub fn to_indexed_mz_spectrum(&self) -> IndexedMzSpectrum {
409        let mut grouped_data: BTreeMap<i32, Vec<(f64, f64)>> = BTreeMap::new();
410
411        // Group by 'tof' with 'mz' and 'intensity'
412        for (&tof, (&mz, &intensity)) in self.tof.iter().zip(self.ims_frame.mz.iter().zip(self.ims_frame.intensity.iter())) {
413            grouped_data.entry(tof).or_insert_with(Vec::new).push((mz, intensity));
414        }
415
416        let mut index = Vec::new();
417        let mut mz = Vec::new();
418        let mut intensity = Vec::new();
419
420        for (&tof_val, values) in &grouped_data {
421            let sum_intensity: f64 = values.iter().map(|&(_, i)| i).sum();
422            let avg_mz: f64 = values.iter().map(|&(m, _)| m).sum::<f64>() / values.len() as f64;
423
424            index.push(tof_val);
425            mz.push(avg_mz);
426            intensity.push(sum_intensity);
427        }
428
429        IndexedMzSpectrum {
430            index,
431            mz_spectrum: MzSpectrum { mz, intensity },
432        }
433    }
434
435    pub fn generate_random_sample(&self, take_probability: f64) -> TimsFrame {
436        assert!(take_probability >= 0.0 && take_probability <= 1.0);
437
438        let mut rng = rand::thread_rng();
439        let mut scan = Vec::new();
440        let mut mobility = Vec::new();
441        let mut tof = Vec::new();
442        let mut mz = Vec::new();
443        let mut intensity = Vec::new();
444
445        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)) {
446            if rng.gen::<f64>() <= take_probability {
447                scan.push(*s);
448                mobility.push(*m);
449                tof.push(*t);
450                mz.push(*mz_val);
451                intensity.push(*i);
452            }
453        }
454
455        TimsFrame::new(self.frame_id, self.ms_type.clone(), self.ims_frame.retention_time, scan, mobility, tof, mz, intensity)
456    }
457
458    pub fn to_noise_annotated_tims_frame(&self) -> TimsFrameAnnotated {
459        let mut annotations = Vec::with_capacity(self.ims_frame.mz.len());
460        let tof_values = self.tof.clone();
461        let mz_values = self.ims_frame.mz.clone();
462        let scan_values = self.scan.clone();
463        let inv_mobility_values = self.ims_frame.mobility.clone();
464        let intensity_values = self.ims_frame.intensity.clone();
465
466        for intensity in &intensity_values {
467            annotations.push(PeakAnnotation::new_random_noise(*intensity));
468        }
469
470        TimsFrameAnnotated::new(
471            self.frame_id,
472            self.ims_frame.retention_time,
473            self.ms_type.clone(),
474            tof_values.iter().map(|&x| x as u32).collect(),
475            mz_values,
476            scan_values.iter().map(|&x| x as u32).collect(),
477            inv_mobility_values,
478            intensity_values,
479            annotations,
480        )
481    }
482
483    pub fn get_inverse_mobility_along_scan_marginal(&self) -> f64 {
484        let mut marginal_map: BTreeMap<i32, (f64, f64)> = BTreeMap::new();
485        // go over all data points of scan, inv_mob and intensity
486        for (scan, inv_mob, intensity) in izip!(&self.scan, &self.ims_frame.mobility, &self.ims_frame.intensity) {
487            // create a key for the map
488            let key = *scan;
489            // get the entry from the map or insert a new one
490            let entry = marginal_map.entry(key).or_insert((0.0, 0.0));
491            // update the entry with the current intensity adding it to the existing intensity
492            entry.0 += *intensity;
493            // update the entry with the current inverse mobility, overwriting the existing value
494            entry.1 = *inv_mob;
495        }
496
497        // get the inverse mobility with the highest intensity
498        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;
499
500        *max_inv_mob
501    }
502
503    /// Calculate the weighted mean and variance of `inv_mob` values based on their intensities.
504    pub fn get_mobility_mean_and_variance(&self) -> (f64, f64) {
505        let mut mobility_map: BTreeMap<OrderedFloat<f64>, f64> = BTreeMap::new();
506
507        // Aggregate intensity values for each `inv_mob`
508        for (inv_mob, intensity) in izip!(&self.ims_frame.mobility, &self.ims_frame.intensity) {
509            let entry = mobility_map.entry(OrderedFloat(*inv_mob)).or_insert(0.0);
510            *entry += *intensity;
511        }
512
513        // Calculate weighted mean
514        let mut total_weight = 0.0;
515        let mut weighted_sum = 0.0;
516        for (&inv_mob, &intensity) in &mobility_map {
517            total_weight += intensity;
518            weighted_sum += inv_mob.into_inner() * intensity;
519        }
520        let mean = weighted_sum / total_weight;
521
522        // Calculate weighted variance
523        let mut weighted_squared_diff_sum = 0.0;
524        for (&inv_mob, &intensity) in &mobility_map {
525            let diff = inv_mob.into_inner() - mean;
526            weighted_squared_diff_sum += intensity * diff * diff;
527        }
528        let variance = weighted_squared_diff_sum / total_weight;
529
530        (mean, variance)
531    }
532
533    pub fn get_tims_spectrum(&self, scan_number: i32) -> Option<TimsSpectrum> {
534        let mut tof = Vec::new();
535        let mut mz = Vec::new();
536        let mut intensity = Vec::new();
537
538        for (s, t, m, i) in itertools::multizip((&self.scan, &self.tof, &self.ims_frame.mz, &self.ims_frame.intensity)) {
539            if *s == scan_number {
540                tof.push(*t);
541                mz.push(*m);
542                intensity.push(*i);
543            }
544        }
545
546        if mz.is_empty() {
547            return None;
548        }
549
550        let mobility = self.ims_frame.mobility.iter()
551            .zip(&self.scan)
552            .find(|(_, s)| **s == scan_number)
553            .map(|(m, _)| *m)?;
554
555        Some(TimsSpectrum {
556            frame_id: self.frame_id,
557            scan: scan_number,
558            retention_time: self.ims_frame.retention_time,
559            mobility,
560            ms_type: self.ms_type.clone(),
561            spectrum: IndexedMzSpectrum::new(tof, mz, intensity),
562        })
563    }
564}
565
566struct AggregateData {
567    intensity_sum: f64,
568    ion_mobility_sum: f64,
569    tof_sum: i64,
570    count: i32,
571}
572
573impl std::ops::Add for TimsFrame {
574    type Output = Self;
575    fn add(self, other: Self) -> TimsFrame {
576        let mut combined_map: BTreeMap<(i32, i64), AggregateData> = BTreeMap::new();
577
578        let quantize = |mz: f64| -> i64 {
579            (mz * 1_000_000.0).round() as i64
580        };
581
582        let add_to_map = |map: &mut BTreeMap<(i32, i64), AggregateData>, mz, ion_mobility, tof, scan, intensity| {
583            let key = (scan, quantize(mz));
584            let entry = map.entry(key).or_insert(AggregateData { intensity_sum: 0.0, ion_mobility_sum: 0.0, tof_sum: 0, count: 0 });
585            entry.intensity_sum += intensity;
586            entry.ion_mobility_sum += ion_mobility;
587            entry.tof_sum += tof as i64;
588            entry.count += 1;
589        };
590
591        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) {
592            add_to_map(&mut combined_map, *mz, *ion_mobility, *tof, *scan, *intensity);
593        }
594
595        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) {
596            add_to_map(&mut combined_map, *mz, *ion_mobility, *tof, *scan, *intensity);
597        }
598
599        let mut mz_combined = Vec::new();
600        let mut tof_combined = Vec::new();
601        let mut ion_mobility_combined = Vec::new();
602        let mut scan_combined = Vec::new();
603        let mut intensity_combined = Vec::new();
604
605        for ((scan, quantized_mz), data) in combined_map {
606            mz_combined.push(quantized_mz as f64 / 1_000_000.0);
607            tof_combined.push(data.tof_sum / data.count as i64);
608            ion_mobility_combined.push(data.ion_mobility_sum / data.count as f64);
609            scan_combined.push(scan);
610            intensity_combined.push(data.intensity_sum);
611        }
612
613        let frame = TimsFrame {
614            frame_id: self.frame_id,
615            ms_type: if self.ms_type == other.ms_type { self.ms_type.clone() } else { MsType::Unknown },
616            scan: scan_combined,
617            tof: tof_combined.iter().map(|&x| x as i32).collect(),
618            ims_frame: ImsFrame {
619                retention_time: self.ims_frame.retention_time,
620                mobility: ion_mobility_combined,
621                mz: mz_combined,
622                intensity: intensity_combined,
623            },
624        };
625
626        return frame;
627    }
628}
629
630impl fmt::Display for TimsFrame {
631    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
632
633        let (mz, i) = self.ims_frame.mz.iter()
634            .zip(&self.ims_frame.intensity)
635            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
636            .unwrap();
637
638        write!(f, "TimsFrame(id: {}, type: {}, rt: {}, data points: {}, max by intensity: (mz: {}, intensity: {}))",
639               self.frame_id, self.ms_type, self.ims_frame.retention_time, self.scan.len(), format!("{:.3}", mz), i)
640    }
641}
642
643impl Vectorized<TimsFrameVectorized> for TimsFrame {
644    fn vectorized(&self, resolution: i32) -> TimsFrameVectorized {
645        let binned_frame = self.to_resolution(resolution);
646        // Translate the m/z values into integer indices
647        let indices: Vec<i32> = binned_frame.ims_frame.mz.iter().map(|&mz| (mz * 10f64.powi(resolution)).round() as i32).collect();
648        // Create a vector of values
649        return TimsFrameVectorized {
650            frame_id: self.frame_id,
651            ms_type: self.ms_type.clone(),
652            scan: binned_frame.scan,
653            tof: binned_frame.tof,
654            ims_frame: ImsFrameVectorized {
655                retention_time: binned_frame.ims_frame.retention_time,
656                mobility: binned_frame.ims_frame.mobility,
657                indices,
658                values: binned_frame.ims_frame.intensity,
659                resolution,
660            },
661        };
662    }
663}
664
665///
666/// Convert a given TimsFrame to a vector of TimsSpectrum.
667///
668/// # Arguments
669///
670/// * `resolution` - The resolution to which the m/z values should be rounded.
671///
672/// # Examples
673///
674/// ```
675/// use mscore::data::spectrum::MsType;
676/// use mscore::timstof::frame::TimsFrame;
677/// use mscore::timstof::spectrum::TimsSpectrum;
678/// use mscore::data::spectrum::IndexedMzSpectrum;
679/// use mscore::data::spectrum::ToResolution;
680///
681/// 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]);
682/// let low_res_frame = frame.to_resolution(1);
683/// ```
684impl ToResolution for TimsFrame {
685    fn to_resolution(&self, resolution: i32) -> TimsFrame {
686        let factor = (10.0f64).powi(resolution);
687
688        // Using a tuple of (scan, mz_bin) as a key
689        // Value will store sum of intensities, sum of tofs, sum of mobilities and their count for averaging
690        let mut bin_map: BTreeMap<(i32, i32), (f64, f64, f64, i32)> = BTreeMap::new();
691
692        for i in 0..self.ims_frame.mz.len() {
693            let rounded_mz = (self.ims_frame.mz[i] * factor).round() as i32;
694            let scan_val = self.scan[i];
695            let intensity_val = self.ims_frame.intensity[i] as f64;
696            let tof_val = self.tof[i] as f64;
697            let mobility_val = self.ims_frame.mobility[i] as f64;
698
699            let entry = bin_map.entry((scan_val, rounded_mz)).or_insert((0.0, 0.0, 0.0, 0));
700            entry.0 += intensity_val;
701            entry.1 += tof_val;
702            entry.2 += mobility_val;
703            entry.3 += 1;
704        }
705
706        let mut new_mz = Vec::with_capacity(bin_map.len());
707        let mut new_scan = Vec::with_capacity(bin_map.len());
708        let mut new_intensity = Vec::with_capacity(bin_map.len());
709        let mut new_tof = Vec::with_capacity(bin_map.len());
710        let mut new_mobility = Vec::with_capacity(bin_map.len());
711
712        for ((scan, mz_bin), (intensity_sum, tof_sum, mobility_sum, count)) in bin_map {
713            new_mz.push(mz_bin as f64 / factor);
714            new_scan.push(scan);
715            new_intensity.push(intensity_sum);
716            new_tof.push((tof_sum / count as f64) as i32);
717            new_mobility.push(mobility_sum / count as f64);
718        }
719
720        TimsFrame {
721            frame_id: self.frame_id,
722            ms_type: self.ms_type.clone(),
723            scan: new_scan,
724            tof: new_tof,
725            ims_frame: ImsFrame {
726                retention_time: self.ims_frame.retention_time,
727                mobility: new_mobility,
728                mz: new_mz,
729                intensity: new_intensity,
730            },
731        }
732    }
733}
734
735#[derive(Clone)]
736pub struct TimsFrameVectorized {
737    pub frame_id: i32,
738    pub ms_type: MsType,
739    pub scan: Vec<i32>,
740    pub tof: Vec<i32>,
741    pub ims_frame: ImsFrameVectorized,
742}
743
744impl TimsFrameVectorized {
745    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 {
746        let mut scan_vec = Vec::new();
747        let mut mobility_vec = Vec::new();
748        let mut tof_vec = Vec::new();
749        let mut mz_vec = Vec::new();
750        let mut intensity_vec = Vec::new();
751        let mut indices_vec = Vec::new();
752
753        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)) {
754            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 {
755                scan_vec.push(*scan);
756                mobility_vec.push(*mobility);
757                tof_vec.push(*tof);
758                mz_vec.push(*mz);
759                intensity_vec.push(*intensity);
760                indices_vec.push(*index);
761            }
762        }
763
764        TimsFrameVectorized {
765            frame_id: self.frame_id,
766            ms_type: self.ms_type.clone(),
767            scan: scan_vec,
768            tof: tof_vec,
769            ims_frame: ImsFrameVectorized {
770                retention_time: self.ims_frame.retention_time,
771                mobility: mobility_vec,
772                indices: indices_vec,
773                values: mz_vec,
774                resolution: self.ims_frame.resolution,
775            },
776        }
777    }
778}
779
780impl fmt::Display for TimsFrameVectorized {
781    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
782
783        let (mz, i) = self.ims_frame.values.iter()
784            .zip(&self.ims_frame.values)
785            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
786            .unwrap();
787
788        write!(f, "TimsFrame(id: {}, type: {}, rt: {}, data points: {}, max by intensity: (mz: {}, intensity: {}))",
789               self.frame_id, self.ms_type, self.ims_frame.retention_time, self.scan.len(), format!("{:.3}", mz), i)
790    }
791}