David@dl predictors

imspy/imspy/core/frame.py

@@ -37,6 +37,17 @@ def __init__(self, frame_id: int, ms_type: int, retention_time: float, scan: NDA
 
         self.__frame_ptr = pims.PyTimsFrame(frame_id, ms_type, retention_time, scan, mobility, tof, mz, intensity)
 
+    def __add__(self, other: 'TimsFrame') -> 'TimsFrame':
+        """Add two TimsFrames together.
+
+        Args:
+            other (TimsFrame): TimsFrame to add.
+
+        Returns:
+            TimsFrame: Sum of the two TimsFrames.
+        """
+        return TimsFrame.from_py_tims_frame(self.__frame_ptr + other.__frame_ptr)
+
     @classmethod
     def from_py_tims_frame(cls, frame: pims.PyTimsFrame):
         """Create a TimsFrame from a PyTimsFrame.

imspy_connector/src/py_tims_frame.rs

@@ -204,6 +204,11 @@ impl PyTimsFrame {
 
         Ok(tuple.into())
     }
+
+    fn __add__(&self, other: PyTimsFrame) -> PyTimsFrame {
+        let result = self.inner.clone() + other.inner.clone();
+        return PyTimsFrame { inner: result }
+    }
 }
 
 #[pyclass]

mscore/src/tims_frame.rs

@@ -2,6 +2,7 @@ use std::fmt;
 use std::collections::BTreeMap;
 use std::fmt::{Formatter};
 use itertools;
+use itertools::izip;
 
 use crate::mz_spectrum::{MsType, MzSpectrum, IndexedMzSpectrum, TimsSpectrum};
 
@@ -307,7 +308,68 @@ impl TimsFrame {
             mz_spectrum: MzSpectrum { mz, intensity },
         }
     }
+}
+
+struct AggregateData {
+    intensity_sum: f64,
+    ion_mobility_sum: f64,
+    tof_sum: i32,
+    count: i32,
+}
+
+impl std::ops::Add for TimsFrame {
+    type Output = Self;
+    fn add(self, other: Self) -> TimsFrame {
+        let mut combined_map: BTreeMap<(i64, i32), AggregateData> = BTreeMap::new();
+
+        let quantize = |mz: f64| -> i64 {
+            (mz * 1_000_000.0).round() as i64
+        };
+
+        let add_to_map = |map: &mut BTreeMap<(i64, i32), AggregateData>, mz, ion_mobility, tof, scan, intensity| {
+            let key = (quantize(mz), scan);
+            let entry = map.entry(key).or_insert(AggregateData { intensity_sum: 0.0, ion_mobility_sum: 0.0, tof_sum: 0, count: 0 });
+            entry.intensity_sum += intensity;
+            entry.ion_mobility_sum += ion_mobility;
+            entry.tof_sum += tof;
+            entry.count += 1;
+        };
+
+        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) {
+            add_to_map(&mut combined_map, *mz, *ion_mobility, *tof, *scan, *intensity);
+        }
 
+        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) {
+            add_to_map(&mut combined_map, *mz, *ion_mobility, *tof, *scan, *intensity);
+        }
+
+        let mut mz_combined = Vec::new();
+        let mut tof_combined = Vec::new();
+        let mut ion_mobility_combined = Vec::new();
+        let mut scan_combined = Vec::new();
+        let mut intensity_combined = Vec::new();
+
+        for ((quantized_mz, scan), data) in combined_map {
+            mz_combined.push(quantized_mz as f64 / 1_000_000.0);
+            tof_combined.push(data.tof_sum / data.count);
+            ion_mobility_combined.push(data.ion_mobility_sum / data.count as f64);
+            scan_combined.push(scan);
+            intensity_combined.push(data.intensity_sum);
+        }
+
+        TimsFrame {
+            frame_id: self.frame_id + other.frame_id,
+            ms_type: if self.ms_type == other.ms_type { self.ms_type.clone() } else { MsType::Unknown },
+            scan: scan_combined,
+            tof: tof_combined,
+            ims_frame: ImsFrame {
+                retention_time: (self.ims_frame.retention_time + other.ims_frame.retention_time) / 2.0,
+                mobility: ion_mobility_combined,
+                mz: mz_combined,
+                intensity: intensity_combined,
+            },
+        }
+    }
 }
 
 impl fmt::Display for TimsFrame {