OpenSWATH XIC parquet reader

docs/API.rst

@@ -90,6 +90,7 @@ Classes
 
    ResultsLoader
    MzMLDataLoader
+   OpenSwathXICParquetLoader
    SqMassLoader
    XICParquetDataLoader
    SpectralLibraryLoader
@@ -125,10 +126,12 @@ Classes
    SqMassDataAccess
    MzMLDataAccess
    OSWDataAccess
+   OpenSwathXICParquetAccess
    ResultsTSVDataAccess
    TransitionPQPDataAccess
    TransitionTSVDataAccess
    XICParquetDataAccess
+
 
 
 Abstract Classes

massdash/loaders/OpenSwathXICParquetLoader.py

@@ -0,0 +1,79 @@
+"""
+massdash/loaders/OpenSwathXICParquetLoader
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+"""
+
+
+from typing import List, Dict, Union
+from os.path import basename
+from pandas.core.api import DataFrame as DataFrame
+import pandas as pd
+
+# Loaders
+from .GenericChromatogramLoader import GenericChromatogramLoader
+from .ResultsLoader import ResultsLoader
+from .access import OpenSwathXICParquetAccess
+# Structs
+from ..structs import TransitionGroup, TransitionGroupCollection
+# Utils
+from massdash.util import LOGGER
+
+class OpenSwathXICParquetLoader(GenericChromatogramLoader):
+
+    ''' 
+    Class for loading Chromatograms and peak features from SqMass files and OSW files
+    Inherits from GenericChromatogramLoader
+    '''
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs) 
+        self.dataAccess = [OpenSwathXICParquetAccess(f) for f in self.dataFiles]
+
+    @ResultsLoader.cache_results
+    def loadTransitionGroupsDf(self, pep_id: str, charge: int) -> pd.DataFrame:
+        columns=['run_name', 'rt', 'intensity', 'annotation']
+        out = {}
+        for t in self.dataAccess:
+
+            df =  t.getChromatogramDfFromSequenceAndCharge(pep_id, charge)
+            # only add if there is data
+            if not df.empty:
+                out[t.runName] = df
+            else:
+                print(f"Warning: no data found for peptide in transition file {t.filename}")
+
+        if out == {}:
+            return pd.DataFrame(columns=columns)
+        else:
+            return pd.concat(out).reset_index().drop('level_1', axis=1).rename(columns=dict(level_0='run'))
+
+    @ResultsLoader.cache_results
+    def loadTransitionGroups(self, pep_id: str, charge: int, runNames: Union[None, str, List[str]] =None) -> Dict[str, TransitionGroupCollection]:
+        '''
+        Loads the transition group for a given peptide ID and charge across all files
+        Args:
+            pep_id (str): Peptide ID
+            charge (int): Charge
+            runNames (None | str | List[str]): Name of the run to extract the transition group from. If None, all runs are extracted. If str, only the specified run is extracted. If List[str], only the specified runs are extracted.
+        Returns:
+            Dict[str, TransitionGroup]: Dictionary of TransitionGroups, with keys as sqMass filenames
+        '''
+
+        out = TransitionGroupCollection()
+
+        if runNames is None:
+            for t in self.dataAccess:
+                out[t.runName] = t.getChromatogramsFromSequenceAndCharge(pep_id, charge)
+        elif isinstance(runNames, str):
+            t = self.dataAccess[self.runNames.index(runNames)]
+            out[runNames] = t.getChromatogramsFromSequenceAndCharge(pep_id, charge) 
+        elif isinstance(runNames, list):
+            out = TransitionGroupCollection()
+            for r in runNames:
+                for t in self.dataAccess:
+                    if t.runName == r:
+                        out[t.runName] = t.getChromatogramsFromSequenceAndCharge(pep_id, charge)
+        else:
+            raise ValueError("runName must be none, a string or list of strings")
+
+        return out

massdash/loaders/__init__.py

@@ -8,6 +8,7 @@
 from .GenericChromatogramLoader import GenericChromatogramLoader
 from .GenericSpectrumLoader import GenericSpectrumLoader
 from .MzMLDataLoader import MzMLDataLoader
+from .OpenSwathXICParquetLoader import OpenSwathXICParquetLoader
 from .ResultsLoader import ResultsLoader
 from .SpectralLibraryLoader import SpectralLibraryLoader
 from .SqMassLoader import SqMassLoader
@@ -17,6 +18,7 @@
             "GenericChromatogramLoader",
             "GenericSpectrumLoader",
             "MzMLDataLoader", 
+            "OpenSwathXICParquetLoader",
             "ResultsLoader",
             "SpectralLibraryLoader",
             "SqMassLoader",

massdash/loaders/access/OpenSwathXICParquetAccess.py

@@ -0,0 +1,104 @@
+"""
+massdash/loaders/access/OpenSwathXICParquetAccess.py
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+"""
+
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+from typing import List
+from collections import OrderedDict
+import pyopenms as po
+import sqlite3
+import pandas as pd
+import base64
+import struct
+import zlib
+from pathlib import Path
+import pyarrow.dataset as ds
+
+# Structs
+from ...structs.Chromatogram import Chromatogram
+# Utils
+
+class OpenSwathXICParquetAccess:
+
+    def __init__(self, filename):
+        self.filename = filename
+        self.runName = str(Path(filename).stem)
+        self.parquet = ds.dataset(filename)
+
+    def getChromatogramsFromSequenceAndCharge(self, sequence: str, charge: int):
+        """
+        Get chromatograms for a given peptide sequence and charge
+        """
+        df = self.parquet.scanner(
+            columns=['RT_DATA', 'INTENSITY_DATA', 'RT_COMPRESSION', 'INTENSITY_COMPRESSION', 'TRANSITION_ORDINAL', 'TRANSITION_TYPE', 'PRODUCT_CHARGE', 'NATIVE_ID'],
+            filter=( 
+                (ds.field("MODIFIED_SEQUENCE") == sequence) &
+                (ds.field("PRECURSOR_CHARGE") == charge)
+            )
+        ).to_table().to_pandas()
+
+        # Create an ANNOTATION column, is the annotation for transitions and the native ID for precursors
+        mask = ~df['PRODUCT_CHARGE'].isnull()
+        df.loc[mask, 'ANNOTATION'] = (df.loc[mask, 'TRANSITION_TYPE'] +
+                                       df.loc[mask, 'TRANSITION_ORDINAL'].astype(int).astype(str) + 
+                                       '^' + 
+                                       df.loc[mask, 'PRODUCT_CHARGE'].astype(int).astype(str))
+        df.loc[~mask, 'ANNOTATION'] = df.loc[~mask, 'NATIVE_ID']
+
+        chroms = []
+        for _, row in df.iterrows():
+            rt_data = OpenSwathXICParquetAccess._decodeArray(row['RT_DATA'], row['RT_COMPRESSION'])
+            intensity_data = OpenSwathXICParquetAccess._decodeArray(row['INTENSITY_DATA'], row['INTENSITY_COMPRESSION'])
+            chroms.append(Chromatogram(rt_data, intensity_data, row['ANNOTATION']))
+
+        return chroms
+
+    @staticmethod
+    def _decodeArray(data, compr):
+        numpress_config = po.NumpressConfig()
+        result = []
+        if compr == 0:
+            return data
+        if compr == 1:
+            tmp = zlib.decompress(data)
+            return struct.unpack("<%sd" % (len(tmp) // 8), tmp)
+        elif compr == 5:
+            tmp = bytearray(zlib.decompress(data))
+            if len(tmp) > 0:
+                numpress_config.setCompression('linear')
+                po.MSNumpressCoder().decodeNP(base64.b64encode(tmp), result, False, numpress_config)
+                return result
+            else:
+                return [0]
+        elif compr == 6:
+            tmp = bytearray( zlib.decompress(data) )
+            if len(tmp) > 0:
+                numpress_config.setCompression('slof')
+                po.MSNumpressCoder().decodeNP(base64.b64encode(tmp), result, False, numpress_config)
+                return result
+            else:
+                return [0]
+        else:
+            raise Exception(f"Compression type {compr} not supported")
+
+    def getChromatogramDfFromSequenceAndCharge(self, sequence: str, charge: int) -> pd.DataFrame:
+        '''
+        Get chromatogram data as a dataframe
+        '''
+        chroms = self.getChromatogramsFromSequenceAndCharge(sequence, charge)
+        chroms_df = []
+        for c in chroms:
+            chroms_df.append(c.toPandasDf())
+
+        if len(chroms_df) == 0:
+            return pd.DataFrame(columns=['rt', 'intensity', 'annotation'])
+        else:
+            return pd.concat(chroms_df)
+
+    def __str__(self):
+        return f"<OpenSwathXICParquetAccess(filename={self.filename})>"
+
+    def __repr__(self):
+        return f"OpenSwathXICParquetAccess(filename={self.filename})"

massdash/loaders/access/__init__.py

@@ -8,6 +8,7 @@
 from .GenericResultsAccess import GenericResultsAccess
 from .MzMLDataAccess import MzMLDataAccess
 from .OSWDataAccess import OSWDataAccess
+from .OpenSwathXICParquetAccess import OpenSwathXICParquetAccess
 from .ResultsTSVDataAccess import ResultsTSVDataAccess
 from .SqMassDataAccess import SqMassDataAccess
 from .TransitionPQPDataAccess import TransitionPQPDataAccess
@@ -17,6 +18,7 @@
 __all__ = [ "GenericResultsAccess",
             "MzMLDataAccess",
             "OSWDataAccess",
+            "OpenSwathXICParquetAccess",
             "ResultsTSVDataAccess",
             "SqMassDataAccess",
             "TransitionPQPDataAccess",