Merge pull request #44 from theGreatHerrLebert/david@mixture

David@mixture

imspy/imspy/dda.py

@@ -0,0 +1,37 @@
+import numpy as np
+import pandas as pd
+
+import imspy_connector as pims
+
+from imspy.spectrum import IndexedMzSpectrum
+
+
+class FragmentDDA:
+    def __init__(self, frame_id: int, precursor_id: int, selected_fragment: IndexedMzSpectrum):
+        self._fragment_ptr = pims.PyTimsFragmentDDA(frame_id, precursor_id, selected_fragment.get_spec_ptr())
+
+    @classmethod
+    def from_py_tims_fragment_dda(cls, fragment: pims.PyTimsFragmentDDA):
+        instance = cls.__new__(cls)
+        instance._fragment_ptr = fragment
+        return instance
+
+    @property
+    def frame_id(self) -> int:
+        return self._fragment_ptr.frame_id
+
+    @property
+    def precursor_id(self) -> int:
+        return self._fragment_ptr.precursor_id
+
+    @property
+    def selected_fragment(self) -> IndexedMzSpectrum:
+        return IndexedMzSpectrum.from_py_indexed_mz_spectrum(self._fragment_ptr.selected_fragment)
+
+    def __repr__(self):
+        return f"FragmentDDA(frame_id={self.frame_id}, precursor_id={self.precursor_id}, " \
+               f"selected_fragment={self.selected_fragment})"
+
+    def get_fragment_ptr(self):
+        return self._fragment_ptr
+

imspy/imspy/handle.py

@@ -10,6 +10,7 @@
 
 from abc import ABC
 
+from imspy.dda import FragmentDDA
 from imspy.frame import TimsFrame
 from imspy.slice import TimsSlice
 
@@ -21,19 +22,21 @@ def __init__(self, data_path: str):
         Args:
             data_path (str): Path to the data.
         """
+        self.__dataset = None
+        self.binary_path = None
+
         self.data_path = data_path
-        self.bp: List[str] = obb.get_so_paths()
         self.meta_data = self.__load_meta_data()
         self.precursor_frames = self.meta_data[self.meta_data["MsMsType"] == 0].Id.values.astype(np.int32)
         self.fragment_frames = self.meta_data[self.meta_data["MsMsType"] > 0].Id.values.astype(np.int32)
-        self.__handle = None
         self.__current_index = 1
 
         # Try to load the data with the first binary found
         appropriate_found = False
-        for so_path in self.bp:
+        for so_path in obb.get_so_paths():
             try:
-                self.__handle = pims.PyTimsDataHandle(self.data_path, so_path)
+                self.__dataset = pims.PyTimsDataset(self.data_path, so_path)
+                self.binary_path = so_path
                 appropriate_found = True
                 break
             except Exception:
@@ -48,7 +51,7 @@ def acquisition_mode(self) -> str:
         Returns:
             str: Acquisition mode.
         """
-        return self.__handle.get_acquisition_mode_as_string()
+        return self.__dataset.get_acquisition_mode_as_string()
 
     @property
     def acquisition_mode_numerical(self) -> int:
@@ -57,7 +60,7 @@ def acquisition_mode_numerical(self) -> int:
         Returns:
             int: Acquisition mode as a numerical value.
         """
-        return self.__handle.get_acquisition_mode()
+        return self.__dataset.get_acquisition_mode()
 
     @property
     def frame_count(self) -> int:
@@ -66,7 +69,7 @@ def frame_count(self) -> int:
         Returns:
             int: Number of frames.
         """
-        return self.__handle.frame_count
+        return self.__dataset.frame_count
 
     def __load_meta_data(self) -> pd.DataFrame:
         """Get the meta data.
@@ -85,7 +88,7 @@ def get_tims_frame(self, frame_id: int) -> TimsFrame:
         Returns:
             TimsFrame: TimsFrame.
         """
-        return TimsFrame.from_py_tims_frame(self.__handle.get_tims_frame(frame_id))
+        return TimsFrame.from_py_tims_frame(self.__dataset.get_frame(frame_id))
 
     def get_tims_slice(self, frame_ids: NDArray[np.int32]) -> TimsSlice:
         """Get a TimsFrame.
@@ -96,14 +99,14 @@ def get_tims_slice(self, frame_ids: NDArray[np.int32]) -> TimsSlice:
         Returns:
             TimsFrame: TimsFrame.
         """
-        return TimsSlice.from_py_tims_slice(self.__handle.get_tims_slice(frame_ids))
+        return TimsSlice.from_py_tims_slice(self.__dataset.get_slice(frame_ids))
 
     def __iter__(self):
         return self
 
     def __next__(self):
         if self.__current_index <= self.frame_count:
-            frame_ptr = self.__handle.get_tims_frame(self.__current_index)
+            frame_ptr = self.__dataset.get_frame(self.__current_index)
             self.__current_index += 1
             if frame_ptr is not None:
                 return TimsFrame.from_py_tims_frame(frame_ptr)
@@ -120,6 +123,11 @@ def __getitem__(self, index):
 
 
 class TimsDatasetDDA(TimsDataset):
+
+    def __init__(self, data_path: str):
+        super().__init__(data_path=data_path)
+        self.__dataset = pims.PyTimsDatasetDDA(self.data_path, self.binary_path)
+
     @property
     def selected_precursors(self):
         """Get precursors selected for fragmentation.
@@ -139,8 +147,24 @@ def pasef_meta_data(self):
         return pd.read_sql_query("SELECT * from PasefFrameMsMsInfo",
                                  sqlite3.connect(self.data_path + "/analysis.tdf"))
 
+    def get_pasef_fragments(self):
+        """Get PASEF fragments.
+
+        Args:
+            num_threads (int, optional): Number of threads. Defaults to 4.
+
+        Returns:
+            List[FragmentDDA]: List of PASEF fragments.
+        """
+        pasef_fragments = self.__dataset.get_pasef_fragments(1)
+        return [FragmentDDA.from_py_tims_fragment_dda(fragment) for fragment in pasef_fragments]
+
 
 class TimsDatasetDIA(TimsDataset):
+    def __init__(self, data_path: str):
+        super().__init__(data_path=data_path)
+        self.__dataset = pims.PyTimsDatasetDIA(self.data_path, self.binary_path)
+
     @property
     def pasef_meta_data(self):
         """Get PASEF meta data for DIA.

imspy/imspy/spectrum.py

@@ -90,6 +90,14 @@ def df(self) -> pd.DataFrame:
 
         return pd.DataFrame({'index': self.index, 'mz': self.mz, 'intensity': self.intensity})
 
+    def get_spec_ptr(self) -> pims.PyIndexedMzSpectrum:
+        """Get the spec_ptr.
+
+        Returns:
+            pims.PyIndexedMzSpectrum: spec_ptr.
+        """
+        return self.__spec_ptr
+
     def __repr__(self):
         return f"IndexedMzSpectrum(num_peaks={len(self.index)})"
 

imspy_connector/src/lib.rs

@@ -1,18 +1,21 @@
-mod py_handle;
+mod py_dataset;
 mod py_mz_spectrum;
 mod py_tims_frame;
 mod py_tims_slice;
-mod py_timstof_dda;
+mod py_dda;
+mod py_dia;
 
 use pyo3::prelude::*;
-use crate::py_handle::PyTimsDataHandle;
+use crate::py_dataset::PyTimsDataset;
 use crate::py_mz_spectrum::{PyMzSpectrum, PyIndexedMzSpectrum, PyTimsSpectrum, PyMzSpectrumVectorized};
 use crate::py_tims_frame::{PyTimsFrame, PyTimsFrameVectorized, PyRawTimsFrame};
 use crate::py_tims_slice::{PyTimsPlane, PyTimsSlice, PyTimsSliceVectorized};
+use crate::py_dda::{PyTimsDatasetDDA, PyTimsFragmentDDA};
 
 #[pymodule]
 fn imspy_connector(_py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_class::<PyTimsDataHandle>()?;
+    m.add_class::<PyTimsDataset>()?;
+    m.add_class::<PyTimsDatasetDDA>()?;
     m.add_class::<PyMzSpectrum>()?;
     m.add_class::<PyMzSpectrumVectorized>()?;
     m.add_class::<PyIndexedMzSpectrum>()?;
@@ -23,5 +26,6 @@ fn imspy_connector(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyTimsSlice>()?;
     m.add_class::<PyTimsSliceVectorized>()?;
     m.add_class::<PyTimsPlane>()?;
+    m.add_class::<PyTimsFragmentDDA>()?;
     Ok(())
 }

imspy_connector/src/py_dataset.rs

@@ -0,0 +1,44 @@
+use pyo3::prelude::*;
+use rustdf::data::dataset::TimsDataset;
+use rustdf::data::handle::TimsData;
+
+use crate::py_tims_frame::{PyTimsFrame};
+use crate::py_tims_slice::PyTimsSlice;
+
+#[pyclass]
+pub struct PyTimsDataset {
+    inner: TimsDataset,
+}
+
+#[pymethods]
+impl PyTimsDataset {
+    #[new]
+    pub fn new(data_path: &str, bruker_lib_path: &str) -> Self {
+        let dataset = TimsDataset::new(bruker_lib_path, data_path);
+        PyTimsDataset { inner: dataset }
+    }
+
+    pub fn get_frame(&self, frame_id: u32) -> PyTimsFrame {
+        PyTimsFrame { inner: self.inner.get_frame(frame_id) }
+    }
+
+    pub fn get_slice(&self, frame_ids: Vec<u32>) -> PyTimsSlice {
+        PyTimsSlice { inner: self.inner.get_slice(frame_ids) }
+    }
+
+    pub fn get_aquisition_mode(&self) -> String {
+        self.inner.get_aquisition_mode().to_string()
+    }
+
+    pub fn get_frame_count(&self) -> i32 {
+        self.inner.get_frame_count()
+    }
+
+    pub fn get_data_path(&self) -> &str {
+        self.inner.get_data_path()
+    }
+
+    pub fn get_bruker_lib_path(&self) -> &str {
+        self.inner.get_bruker_lib_path()
+    }
+}

imspy_connector/src/py_dda.rs

@@ -0,0 +1,78 @@
+use pyo3::prelude::*;
+
+use rustdf::data::dda::{PASEFDDAFragment, TimsDatasetDDA};
+use rustdf::data::handle::TimsData;
+use crate::py_mz_spectrum::PyIndexedMzSpectrum;
+use crate::py_tims_frame::PyTimsFrame;
+use crate::py_tims_slice::PyTimsSlice;
+
+#[pyclass]
+pub struct PyTimsDatasetDDA {
+    inner: TimsDatasetDDA,
+}
+
+#[pymethods]
+impl PyTimsDatasetDDA {
+    #[new]
+    pub fn new(data_path: &str, bruker_lib_path: &str) -> Self {
+        let dataset = TimsDatasetDDA::new(bruker_lib_path, data_path);
+        PyTimsDatasetDDA { inner: dataset }
+    }
+    pub fn get_frame(&self, frame_id: u32) -> PyTimsFrame {
+        PyTimsFrame { inner: self.inner.get_frame(frame_id) }
+    }
+
+    pub fn get_slice(&self, frame_ids: Vec<u32>) -> PyTimsSlice {
+        PyTimsSlice { inner: self.inner.get_slice(frame_ids) }
+    }
+
+    pub fn get_aquisition_mode(&self) -> String {
+        self.inner.get_aquisition_mode().to_string()
+    }
+
+    pub fn get_frame_count(&self) -> i32 {
+        self.inner.get_frame_count()
+    }
+
+    pub fn get_data_path(&self) -> &str {
+        self.inner.get_data_path()
+    }
+
+    pub fn get_bruker_lib_path(&self) -> &str {
+        self.inner.get_bruker_lib_path()
+    }
+
+    pub fn get_pasef_fragments(&self, num_threads: usize) -> Vec<PyTimsFragmentDDA> {
+        let pasef_fragments = self.inner.get_pasef_fragments(num_threads);
+        pasef_fragments.iter().map(|pasef_fragment| PyTimsFragmentDDA { inner: pasef_fragment.clone() }).collect()
+    }
+}
+
+#[pyclass]
+pub struct PyTimsFragmentDDA {
+    inner: PASEFDDAFragment,
+}
+
+#[pymethods]
+impl PyTimsFragmentDDA {
+    #[new]
+    pub fn new(frame_id: u32, precursor_id: u32, selected_fragment: &PyIndexedMzSpectrum) -> PyResult<Self> {
+
+        let pasef_fragment = PASEFDDAFragment {
+            frame_id,
+            precursor_id,
+            selected_fragment: selected_fragment.inner.clone(),
+        };
+
+        Ok(PyTimsFragmentDDA { inner: pasef_fragment })
+    }
+
+    #[getter]
+    pub fn frame_id(&self) -> u32 { self.inner.frame_id }
+
+    #[getter]
+    pub fn precursor_id(&self) -> u32 { self.inner.precursor_id }
+
+    #[getter]
+    pub fn selected_fragment(&self) -> PyIndexedMzSpectrum { PyIndexedMzSpectrum { inner: self.inner.selected_fragment.clone() } }
+}

imspy_connector/src/py_dia.rs

@@ -0,0 +1,43 @@
+use pyo3::prelude::*;
+
+use rustdf::data::dia::TimsDatasetDIA;
+use rustdf::data::handle::TimsData;
+use crate::py_tims_frame::PyTimsFrame;
+use crate::py_tims_slice::PyTimsSlice;
+
+#[pyclass]
+pub struct PyTimsDatasetDIA {
+    inner: TimsDatasetDIA,
+}
+
+#[pymethods]
+impl PyTimsDatasetDIA {
+    #[new]
+    pub fn new(data_path: &str, bruker_lib_path: &str) -> Self {
+        let dataset = TimsDatasetDIA::new(bruker_lib_path, data_path);
+        PyTimsDatasetDIA { inner: dataset }
+    }
+    pub fn get_frame(&self, frame_id: u32) -> PyTimsFrame {
+        PyTimsFrame { inner: self.inner.get_frame(frame_id) }
+    }
+
+    pub fn get_slice(&self, frame_ids: Vec<u32>) -> PyTimsSlice {
+        PyTimsSlice { inner: self.inner.get_slice(frame_ids) }
+    }
+
+    pub fn get_aquisition_mode(&self) -> String {
+        self.inner.get_aquisition_mode().to_string()
+    }
+
+    pub fn get_frame_count(&self) -> i32 {
+        self.inner.get_frame_count()
+    }
+
+    pub fn get_data_path(&self) -> &str {
+        self.inner.get_data_path()
+    }
+
+    pub fn get_bruker_lib_path(&self) -> &str {
+        self.inner.get_bruker_lib_path()
+    }
+}