corems.chroma_peak.calc.subset

  1# This file contains functions for subsetting dataframes that contain mass feature data.
  2# This is based on the deimos package, found here: https://github.com/pnnl/deimos/blob/master/deimos/subset.py with some modifications.
  3
  4import multiprocessing as mp
  5from functools import partial
  6
  7import numpy as np
  8import pandas as pd
  9
 10class MultiSamplePartitions:
 11    '''
 12    Generator object that will lazily build and return each partition constructed
 13    from multiple samples.
 14
 15    Attributes
 16    ----------
 17    features : :obj:`~pandas.DataFrame`
 18        Input feature coordinates and intensities.
 19    split_on : str
 20        Dimension to partition the data.
 21    size : int
 22        Target partition size.
 23    tol : float
 24        Largest allowed distance between unique `split_on` observations.
 25    n_partitions : int
 26        Number of partitions in the data.
 27
 28    '''
 29
 30    def __init__(self,
 31                 features, 
 32                 split_on: str = 'mz', 
 33                 size: int = 500, 
 34                 tol: float = 25E-6, 
 35                 relative: bool = False):
 36        '''
 37        Initialize :obj:`~deimos.subset.Partitions` instance.
 38
 39        Parameters
 40        ----------
 41        features : :obj:`~pandas.DataFrame`
 42            Input feature coordinates and intensities.
 43        split_on : str
 44            Dimension to partition the data.
 45        size : int
 46            Target partition size.
 47        tol : float
 48            Largest allowed distance between unique `split_on` observations.
 49
 50        '''
 51        if not isinstance(split_on, str):
 52            raise TypeError(f"Expected 'split_on' to be a string, got {type(split_on).__name__}")
 53        if not isinstance(size, int):
 54            raise TypeError(f"Expected 'size' to be an integer, got {type(size).__name__}")
 55        if not isinstance(tol, float):
 56            raise TypeError(f"Expected 'tol' to be a float, got {type(tol).__name__}")
 57        if not isinstance(relative, bool):
 58            raise TypeError(f"Expected 'relative' to be a boolean, got {type(relative).__name__}")
 59
 60        self.features = features
 61        self.split_on = split_on
 62        self.size = size
 63        self.tol = tol
 64        self.relative = relative
 65
 66        self._compute_splits()
 67
 68    def _compute_splits(self):
 69        '''
 70        Determines data splits for partitioning.
 71
 72        '''
 73
 74        self.counter = 0
 75
 76        idx = self.features.groupby(by=self.split_on).size().sort_index()
 77
 78        counts = idx.values
 79        idx = idx.index
 80
 81        if self.relative:
 82            dxs = np.diff(idx) / idx[:-1]
 83        else:
 84            dxs = np.diff(idx)
 85
 86        # if relative, convert tol to absolute
 87        bins = []
 88        current_count = counts[0]
 89        current_bin = [idx[0]]
 90        self._counts = []
 91
 92        for i, dx in zip(range(1, len(idx)), dxs):
 93            if (current_count + counts[i] <= self.size) or (dx <= self.tol):
 94                current_bin.append(idx[i])
 95                current_count += counts[i]
 96
 97            else:
 98                bins.append(np.array(current_bin))
 99                self._counts.append(current_count)
100
101                current_bin = [idx[i]]
102                current_count = counts[i]
103
104        # Add last unadded bin
105        bins.append(np.array(current_bin))
106        self._counts.append(current_count)
107
108        self.bounds = np.array([[x.min(), x.max()] for x in bins])
109
110        # Number of partitions in the data
111        self.n_partitions = len(bins)
112
113    def __iter__(self):
114        return self
115
116    def __next__(self):
117        if self.counter < len(self.bounds):
118            q = '({} >= {}) & ({} <= {})'.format(self.split_on,
119                                                 self.bounds[self.counter][0],
120                                                 self.split_on,
121                                                 self.bounds[self.counter][1])
122
123            subset = self.features.query(q)
124
125            self.counter += 1
126            if len(subset.index) > 1:
127                return subset
128            else:
129                return None
130
131        raise StopIteration
132
133    def map(self, func, processes=1, **kwargs):
134        '''
135        Maps `func` to each partition, then returns the combined result.
136
137        Parameters
138        ----------
139        func : function
140            Function to apply to partitions.
141        processes : int
142            Number of parallel processes. If less than 2, a serial mapping is
143            applied.
144        kwargs
145            Keyword arguments passed to `func`.
146
147        Returns
148        -------
149        :obj:`~pandas.DataFrame`
150            Combined result of `func` applied to partitions.
151
152        '''
153
154        # Serial
155        if processes < 2:
156            result = [func(x, **kwargs) for x in self]
157
158        # Parallel
159        else:
160            with mp.Pool(processes=processes) as p:
161                result = list(p.imap(partial(func, **kwargs), self))
162
163        # Add partition index
164        for i in range(len(result)):
165            if result[i] is not None:
166                result[i]['partition_idx'] = i
167
168        # Combine partitions
169        return pd.concat(result, ignore_index=True)
170
171def multi_sample_partition(features, split_on='mz', size=500, tol=25E-6, relative=True):
172    '''
173    Partitions data along a given dimension. For use with features across
174    multiple samples, e.g. in alignment.
175
176    Parameters
177    ----------
178    features : :obj:`~pandas.DataFrame`
179        Input feature coordinates and intensities.
180    split_on : str
181        Dimension to partition the data.
182    size : int
183        Target partition size.
184    tol : float
185        Largest allowed distance between unique `split_on` observations.
186    relative : bool
187        If `True`, the `tol` parameter is interpreted as a relative tolerance.
188
189    Returns
190    -------
191    :obj:`~deimos.subset.Partitions`
192        A generator object that will lazily build and return each partition.
193
194    '''
195
196    return MultiSamplePartitions(features, split_on, size, tol, relative)