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