1+ """
2+ massdash/loaders/OpenSwathXICParquetLoader
3+ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
4+ """
5+
6+
7+ from typing import List , Dict , Union
8+ from os .path import basename
9+ from pandas .core .api import DataFrame as DataFrame
10+ import pandas as pd
11+
12+ # Loaders
13+ from .GenericChromatogramLoader import GenericChromatogramLoader
14+ from .ResultsLoader import ResultsLoader
15+ from .access import OpenSwathXICParquetAccess
16+ # Structs
17+ from ..structs import TransitionGroup , TransitionGroupCollection
18+ # Utils
19+ from ..util import LOGGER
20+
21+ class OpenSwathXICParquetLoader (GenericChromatogramLoader ):
22+
23+ '''
24+ Class for loading Chromatograms and peak features from XIC (PyProphet) parquet files and a results files.
25+ Inherits from GenericChromatogramLoader
26+ '''
27+
28+ def __init__ (self , ** kwargs ):
29+ super ().__init__ (** kwargs )
30+ self .dataAccess = [OpenSwathXICParquetAccess (f ) for f in self .dataFiles ]
31+
32+ @ResultsLoader .cache_results
33+ def loadTransitionGroupsDf (self , pep_id : str , charge : int ) -> pd .DataFrame :
34+ columns = ['run_name' , 'rt' , 'intensity' , 'annotation' ]
35+ out = {}
36+ for t in self .dataAccess :
37+
38+ df = t .getChromatogramDfFromSequenceAndCharge (pep_id , charge )
39+ # only add if there is data
40+ if not df .empty :
41+ out [t .runName ] = df
42+ else :
43+ LOGGER .warning (f"Warning: no data found for peptide in transition file { t .filename } )
44+
45+ if out == {}:
46+ return pd .DataFrame (columns = columns )
47+ else :
48+ return pd .concat (out ).reset_index ().drop ('level_1' , axis = 1 ).rename (columns = dict (level_0 = 'run' ))
49+
50+ @ResultsLoader .cache_results
51+ def loadTransitionGroups (self , pep_id : str , charge : int , runNames : Union [None , str , List [str ]] = None ) -> Dict [str , TransitionGroupCollection ]:
52+ '''
53+ Loads the transition group for a given peptide ID and charge across all files
54+ Args:
55+ pep_id (str): Peptide ID
56+ charge (int): Charge
57+ 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.
58+ Returns:
59+ Dict[str, TransitionGroup]: Dictionary of TransitionGroups, with keys as sqMass filenames
60+ '''
61+
62+ out = TransitionGroupCollection ()
63+
64+ def _assembleTransitionGroup (t ):
65+ chroms = t .getChromatogramsFromSequenceAndCharge (pep_id , charge )
66+ precursorChroms = [i for i in chroms if 'precursor' in i .label .lower ()]
67+ transitionChroms = [i for i in chroms if 'precursor' not in i .label .lower ()]
68+ if len (precursorChroms ) == 0 and len (transitionChroms ) == 0 : # do not create a transition group if there are no chromatograms
69+ return None
70+ else :
71+ return TransitionGroup (precursorChroms , transitionChroms , pep_id , charge )
72+
73+ if runNames is None :
74+ for t in self .dataAccess :
75+ out [t .runName ] = _assembleTransitionGroup (t )
76+ elif isinstance (runNames , str ):
77+ t = self .dataAccess [self .runNames .index (runNames )]
78+ out [runNames ] = _assembleTransitionGroup (t )
79+ elif isinstance (runNames , list ):
80+ out = TransitionGroupCollection ()
81+ for r in runNames :
82+ for t in self .dataAccess :
83+ if t .runName == r :
84+ out [t .runName ] = _assembleTransitionGroup (t )
85+ else :
86+ raise ValueError ("runName must be none, a string or list of strings" )
87+
88+ # if there are no chromatograms, return none
89+ if all ([i is None for i in out .values ()]):
90+ LOGGER .warning (f"No chromatograms found for peptide { pep_id } { charge } )
91+ return None
92+ else :
93+ return out
0 commit comments