FEAT: DIAGrouper

alpharaw/dia/normal_dia.py

@@ -0,0 +1,55 @@
+import pandas as pd
+
+import typing
+
+from alpharaw.ms_data_base import MSData_Base
+
+from alpharaw.utils.df_processing import remove_unused_peaks
+from alpharaw.utils.timstof import convert_to_alphatims
+from alphatims.bruker import TimsTOF
+
+class NormalDIAGrouper():
+    def __init__(self, ms_data: MSData_Base):
+        self.ms_data = ms_data
+        self.ms_data.spectrum_df[
+            "dia_group"
+        ] = ms_data.spectrum_df.precursor_mz.astype(int)
+
+    def get_grouped_ms_data(self, 
+        dia_group:int=-1, 
+        return_alpharaw_data: bool=True,
+        return_alphatims_data: bool=True,
+    )->typing.Union[MSData_Base, TimsTOF, typing.Tuple[MSData_Base, TimsTOF]]:
+        """ Get compressed MS data for isolation window `dia_group`.
+
+        Args:
+            dia_group (int, optional): The DIA group, -1 means ms1. Defaults to -1.
+            return_alphatims_data (bool, optional): If return `MSData_Base`. Defaults to True
+            return_alphatims_data (bool, optional): If return alphatims object. Defaults to True.
+
+        Returns:
+            MSData_Base: Compressed MS data, if `return_alpharaw_data==True`
+            TimsTOF: Alphatims object for the window, if `return_alphatims_data==True`
+        """
+
+        spec_df = self.ms_data.spectrum_df.query(f"dia_group == {dia_group}")
+
+        if return_alphatims_data:
+            ms_data, ms_tims = convert_to_alphatims(
+                spec_df, self.ms_data.peak_df, dda=False
+            )
+            if return_alpharaw_data:
+                return ms_data, ms_tims
+            else:
+                return ms_tims
+        else:
+            ms_data = MSData_Base()
+
+            spec_df, peak_df = remove_unused_peaks(
+                spec_df, self.ms_data.peak_df
+            )
+
+            ms_data.spectrum_df = spec_df
+            ms_data.peak_df = peak_df
+            return ms_data
+

alpharaw/utils/timstof.py

@@ -0,0 +1,37 @@
+import pandas as pd
+import typing
+
+from alpharaw.ms_data_base import MSData_Base
+from alphatims.bruker import TimsTOF
+from alpharaw.wrappers.alphatims_wrapper import AlphaTimsWrapper
+from alpharaw.utils.df_processing import remove_unused_peaks
+
+def convert_to_alphatims(
+    spec_df:pd.DataFrame, 
+    peak_df:pd.DataFrame,
+    dda:bool = False,
+)->typing.Tuple[MSData_Base,TimsTOF]:
+    """ 
+    Convert any spectrum dataframe or sliced spectrum dataframe
+    and its peak dataframe into AlphaTims' TimsTOF object (AlphaTimsWrapper).
+
+    Args:
+        spec_df (pd.DataFrame): 
+            spectrum dataframe or sliced spectrum dataframe in AlphaRaw's format.
+        peak_df (pd.DataFrame): 
+            peak dataframe in AlphaRaw's format by removing unused peaks in spec_df.
+        dda (bool):
+            if dda data.
+
+    Returns:
+        MSData_Base: AlphaRaw object
+        TimsTOF: AlphaTims' TimsTOF object (AlphaTimsWrapper).
+    """
+    spec_df, peak_df = remove_unused_peaks(
+        spec_df, peak_df
+    )
+    ms_data = MSData_Base()
+    ms_data.spectrum_df = spec_df
+    ms_data.peak_df = peak_df
+
+    return ms_data, AlphaTimsWrapper(ms_data, dda=dda)

alpharaw/viz/xic_plot.py

@@ -11,45 +11,10 @@
     TimsTOF, 
 )
 
-from alpharaw.ms_data_base import MSData_Base
-from alpharaw.utils.df_processing import remove_unused_peaks
-from alpharaw.wrappers.alphatims_wrapper import AlphaTimsWrapper
-
 from .plot_utils import (
     plot_line_fast
 )
 
-def convert_to_timstof(
-    spec_df:pd.DataFrame, 
-    peak_df:pd.DataFrame,
-)->typing.Tuple[MSData_Base,TimsTOF]:
-    """ 
-    Convert any spectrum dataframe or sliced spectrum dataframe
-    and its peak dataframe into AlphaTims' TimsTOF object (AlphaTimsWrapper).
-
-    Args:
-        spec_df (pd.DataFrame): 
-            spectrum dataframe or sliced spectrum dataframe in AlphaRaw's format
-        peak_df (pd.DataFrame): 
-            peak dataframe in AlphaRaw's format
-        remove_used_peaks (bool, optional): 
-            For sliced spectrum dataframe, this will remove peaks in `peak_df` 
-            that are not in `spec_df`.
-            Defaults to False.
-
-    Returns:
-        MSData_Base: AlphaRaw object
-        TimsTOF: AlphaTims' TimsTOF object (AlphaTimsWrapper).
-    """
-    spec_df, peak_df = remove_unused_peaks(
-        spec_df, peak_df
-    )
-    ms_data = MSData_Base()
-    ms_data.spectrum_df = spec_df
-    ms_data.peak_df = peak_df
-
-    return ms_data, AlphaTimsWrapper(ms_data, dda=False)
-
 class XIC_Plot():
     # hovermode = "x" | "y" | "closest" | False | "x unified" | "y unified"
     hovermode = 'closest'