Merge pull request #58 from smoia/cristina

smoia · web-flow · commit 2a3dbfbac2e7 · 2023-07-07T20:16:54.000-04:00
A set of major changes that are getting merged at once (check PR text)
diff --git a/docs/user_guide/editing.rst b/docs/user_guide/editing.rst
@@ -41,9 +41,9 @@ removed. We can do that easily with :py:func:`~.operations.edit_physio`:
 
 This function will open up an interactive viewer, which supports scrolling
 through the time series (with the scroll wheel), rejection of noisy segments of
-data (left click + drag, red highlight), and deleting peaks / troughs that were
+data (left click + drag, blue highlight), and deleting peaks / troughs that were
 erroneously detected and shouldn't be considered at all (right click + drag,
-blue highlight):
+red highlight):
 
 .. image:: physio_edit.gif
 
diff --git a/peakdet/editor.py b/peakdet/editor.py
@@ -1,7 +1,5 @@
 # -*- coding: utf-8 -*-
-"""
-Functions and class for performing interactive editing of physiological data
-"""
+"""Functions and class for performing interactive editing of physiological data."""
 
 import functools
 import numpy as np
@@ -12,7 +10,7 @@
 
 class _PhysioEditor():
     """
-    Class for editing physiological data
+    Class for editing physiological data.
 
     Parameters
     ----------
@@ -25,30 +23,49 @@ def __init__(self, data):
         self.data = utils.check_physio(data, copy=True)
         fs = 1 if data.fs is None else data.fs
         self.time = np.arange(0, len(data.data) / fs, 1 / fs)
+        # Read if there is support data
+        self.suppdata = data.suppdata
 
         # we need to create these variables in case someone doesn't "quit"
         # the plot appropriately (i.e., clicks X instead of pressing ctrl+q)
-        self.deleted, self.rejected = set(), set()
+        self.deleted, self.rejected, self.included = set(), set(), set()
+
+        # make main plot objects depending on supplementary data
+        if self.suppdata is None:
+            self.fig, self._ax = plt.subplots(nrows=1, ncols=1,
+                                              tight_layout=True, sharex=True)
+        else:
+            self.fig, self._ax = plt.subplots(nrows=2, ncols=1,
+                                              tight_layout=True, sharex=True,
+                                              gridspec_kw={'height_ratios': [3, 2]})
 
-        # make main plot objects
-        self.fig, self.ax = plt.subplots(nrows=1, ncols=1, tight_layout=True)
         self.fig.canvas.mpl_connect('scroll_event', self.on_wheel)
         self.fig.canvas.mpl_connect('key_press_event', self.on_key)
 
-        # two selectors for rejection (left mouse) and deletion (right mouse)
-        reject = functools.partial(self.on_remove, reject=True)
-        delete = functools.partial(self.on_remove, reject=False)
-        self.span1 = SpanSelector(self.ax, reject, 'horizontal',
+        # Set axis handler
+        self.ax = self._ax if self.suppdata is None else self._ax[0]
+
+        # three selectors for:
+        #    1. rejection (central mouse),
+        #    2. addition (right mouse), and
+        #    3. deletion (left mouse)
+        delete = functools.partial(self.on_edit, method='delete')
+        reject = functools.partial(self.on_edit, method='reject')
+        insert = functools.partial(self.on_edit, method='insert')
+        self.span2 = SpanSelector(self.ax, delete, 'horizontal',
                                   button=1, useblit=True,
                                   rectprops=dict(facecolor='red', alpha=0.3))
-        self.span2 = SpanSelector(self.ax, delete, 'horizontal',
-                                  button=3, useblit=True,
+        self.span1 = SpanSelector(self.ax, reject, 'horizontal',
+                                  button=2, useblit=True,
                                   rectprops=dict(facecolor='blue', alpha=0.3))
+        self.span3 = SpanSelector(self.ax, insert, 'horizontal',
+                                  button=3, useblit=True,
+                                  rectprops=dict(facecolor='green', alpha=0.3))
 
         self.plot_signals(False)
 
     def plot_signals(self, plot=True):
-        """ Clears axes and plots data / peaks / troughs """
+        """Clear axes and plots data / peaks / troughs."""
         # don't reset x-/y-axis zooms on replot
         if plot:
             xlim, ylim = self.ax.get_xlim(), self.ax.get_ylim()
@@ -62,50 +79,79 @@ def plot_signals(self, plot=True):
                      self.data[self.data.peaks], '.r',
                      self.time[self.data.troughs],
                      self.data[self.data.troughs], '.g')
+
+        if self.suppdata is not None:
+            self._ax[1].plot(self.time, self.suppdata, 'k', linewidth=0.7)
+            self._ax[1].set_ylim(-.5, .5)
+
         self.ax.set(xlim=xlim, ylim=ylim, yticklabels='')
         self.fig.canvas.draw()
 
     def on_wheel(self, event):
-        """ Moves axis on wheel scroll """
+        """Move axis on wheel scroll."""
         (xlo, xhi), move = self.ax.get_xlim(), event.step * -10
         self.ax.set_xlim(xlo + move, xhi + move)
         self.fig.canvas.draw()
 
     def quit(self):
-        """ Quits editor """
+        """Quit editor."""
         plt.close(self.fig)
 
     def on_key(self, event):
-        """ Undoes last span select or quits peak editor """
+        """Undo last span select or quits peak editor."""
         # accept both control or Mac command key as selector
         if event.key in ['ctrl+z', 'super+d']:
             self.undo()
         elif event.key in ['ctrl+q', 'super+d']:
             self.quit()
 
-    def on_remove(self, xmin, xmax, *, reject):
-        """ Removes specified peaks by either rejection / deletion """
+    def on_edit(self, xmin, xmax, *, method):
+        """
+        Edit peaks by rejection, deletion, or insert.
+
+        Removes specified peaks by either rejection / deletion, OR
+        Include one peak by finding the max in the selection.
+
+        method accepts 'insert', 'reject', 'delete'
+        """
+        if method not in ['insert', 'reject', 'delete']:
+            raise ValueError(f'Action "{method}" not supported.')
+
         tmin, tmax = np.searchsorted(self.time, (xmin, xmax))
         pmin, pmax = np.searchsorted(self.data.peaks, (tmin, tmax))
-        bad = np.arange(pmin, pmax, dtype=int)
 
-        if len(bad) == 0:
-            return
+        if method == 'insert':
+            tmp = np.argmax(self.data.data[tmin:tmax]) if tmin != tmax else 0
+            newpeak = tmin + tmp
+            if newpeak == tmin:
+                self.plot_signals()
+                return
+        else:
+            bad = np.arange(pmin, pmax, dtype=int)
+            if len(bad) == 0:
+                self.plot_signals()
+                return
 
-        if reject:
+        if method == 'reject':
             rej, fcn = self.rejected, operations.reject_peaks
-        else:
+        elif method == 'delete':
             rej, fcn = self.deleted, operations.delete_peaks
 
-        # store edits in local history
-        rej.update(self.data.peaks[bad].tolist())
-        self.data = fcn(self.data, self.data.peaks[bad])
+        # store edits in local history & call function
+        if method == 'insert':
+            self.included.add(newpeak)
+            self.data = operations.add_peaks(self.data, newpeak)
+        else:
+            rej.update(self.data.peaks[bad].tolist())
+            self.data = fcn(self.data, self.data.peaks[bad])
+
         self.plot_signals()
 
     def undo(self):
-        """ Resets last span select peak removal """
+        """Reset last span select peak removal."""
         # check if last history entry was a manual reject / delete
-        if self.data._history[-1][0] not in ['reject_peaks', 'delete_peaks']:
+        relevant = ['reject_peaks', 'delete_peaks', 'add_peaks']
+        if self.data._history[-1][0] not in relevant:
             return
 
         # pop off last edit and delete
@@ -123,7 +169,12 @@ def undo(self):
                 peaks['remove']
             )
             self.deleted.difference_update(peaks['remove'])
-
+        elif func == 'add_peaks':
+            self.data._metadata['peaks'] = np.delete(
+                self.data._metadata['peaks'],
+                np.searchsorted(self.data._metadata['peaks'], peaks['add']),
+            )
+            self.included.remove(peaks['add'])
         self.data._metadata['troughs'] = utils.check_troughs(self.data,
                                                              self.data.peaks,
                                                              self.data.troughs)
diff --git a/peakdet/operations.py b/peakdet/operations.py
@@ -92,7 +92,11 @@ def interpolate_physio(data, target_fs, *, kind='cubic'):
 
     # interpolate data and generate new Physio object
     interp = interpolate.interp1d(t_orig, data, kind=kind)(t_new)
-    interp = utils.new_physio_like(data, interp, fs=target_fs)
+    if data.suppdata is None:
+        suppinterp = None
+    else:
+        suppinterp = interpolate.interp1d(t_orig, data.suppdata, kind=kind)(t_new)
+    interp = utils.new_physio_like(data, interp, fs=target_fs, suppdata=suppinterp)
 
     return interp
 
@@ -182,6 +186,52 @@ def reject_peaks(data, remove):
     return data
 
 
+@utils.make_operation()
+def add_peaks(data, add):
+    """
+    Add `newpeak` to add them in `data`
+
+    Parameters
+    ----------
+    data : Physio_like
+    add : int
+
+    Returns
+    -------
+    data : Physio_like
+    """
+
+    data = utils.check_physio(data, ensure_fs=False, copy=True)
+    idx = np.searchsorted(data._metadata['peaks'], add)
+    data._metadata['peaks'] = np.insert(data._metadata['peaks'], idx, add)
+    data._metadata['troughs'] = utils.check_troughs(data, data.peaks)
+
+    return data
+
+
+@utils.make_operation()
+def add_peaks(data, add):
+    """
+    Add `newpeak` to add them in `data`
+
+    Parameters
+    ----------
+    data : Physio_like
+    add : int
+
+    Returns
+    -------
+    data : Physio_like
+    """
+
+    data = utils.check_physio(data, ensure_fs=False, copy=True)
+    idx = np.searchsorted(data._metadata['peaks'], add)
+    data._metadata['peaks'] = np.insert(data._metadata['peaks'], idx, add)
+    data._metadata['troughs'] = utils.check_troughs(data, data.peaks)
+
+    return data
+
+
 def edit_physio(data):
     """
     Opens interactive plot with `data` to permit manual editing of time series
@@ -206,13 +256,14 @@ def edit_physio(data):
     # perform manual editing
     edits = editor._PhysioEditor(data)
     plt.show(block=True)
-    delete, reject = sorted(edits.deleted), sorted(edits.rejected)
 
     # replay editing on original provided data object
-    if reject is not None:
-        data = reject_peaks(data, remove=reject)
-    if delete is not None:
-        data = delete_peaks(data, remove=delete)
+    if len(edits.rejected) > 0:
+        data = reject_peaks(data, remove=sorted(edits.rejected))
+    if len(edits.deleted) > 0:
+        data = delete_peaks(data, remove=sorted(edits.deleted))
+    if len(edits.included) > 0:
+        data = add_peaks(data, add=sorted(edits.included))
 
     return data
 
diff --git a/peakdet/physio.py b/peakdet/physio.py
@@ -21,6 +21,8 @@ class Physio():
         Functions performed on `data`. Default: None
     metadata : dict, optional
         Metadata associated with `data`. Default: None
+    suppdata : array_like, optional
+        Support data array. Default: None
 
     Attributes
     ----------
@@ -35,9 +37,11 @@ class Physio():
         Indices of peaks in `data`
     troughs : :obj:`numpy.ndarray`
         Indices of troughs in `data`
+    suppdata : :obj:`numpy.ndarray`
+        Secondary physiological waveform
     """
 
-    def __init__(self, data, fs=None, history=None, metadata=None):
+    def __init__(self, data, fs=None, history=None, metadata=None, suppdata=None):
         self._data = np.asarray(data).squeeze()
         if self.data.ndim > 1:
             raise ValueError('Provided data dimensionality {} > 1.'
@@ -68,6 +72,7 @@ def __init__(self, data, fs=None, history=None, metadata=None):
             self._metadata = dict(peaks=np.empty(0, dtype=int),
                                   troughs=np.empty(0, dtype=int),
                                   reject=np.empty(0, dtype=int))
+        self._suppdata = None if suppdata is None else np.asarray(suppdata).squeeze()
 
     def __array__(self):
         return self.data
diff --git a/peakdet/utils.py b/peakdet/utils.py
@@ -141,12 +141,13 @@ def check_physio(data, ensure_fs=True, copy=False):
     if copy is True:
         return new_physio_like(data, data.data,
                                copy_history=True,
-                               copy_metadata=True)
+                               copy_metadata=True,
+                               copy_suppdata=True)
     return data
 
 
-def new_physio_like(ref_physio, data, *, fs=None, dtype=None,
-                    copy_history=True, copy_metadata=True):
+def new_physio_like(ref_physio, data, *, fs=None, suppdata=None, dtype=None,
+                    copy_history=True, copy_metadata=True, copy_suppdata=True):
     """
     Makes `data` into physio object like `ref_data`
 
@@ -159,10 +160,16 @@ def new_physio_like(ref_physio, data, *, fs=None, dtype=None,
     fs : float, optional
         Sampling rate of `data`. If not supplied, assumed to be the same as
         in `ref_physio`
+    suppdata : array_like, optional
+        New supplementary data. If not supplied, assumed to be the same.
     dtype : data_type, optional
         Data type to convert `data` to, if conversion needed. Default: None
     copy_history : bool, optional
         Copy history from `ref_physio` to new physio object. Default: True
+    copy_metadata : bool, optional
+        Copy metadata from `ref_physio` to new physio object. Default: True
+    copy_suppdata : bool, optional
+        Copy suppdata from `ref_physio` to new physio object. Default: True
 
     Returns
     -------
@@ -177,9 +184,13 @@ def new_physio_like(ref_physio, data, *, fs=None, dtype=None,
     history = list(ref_physio.history) if copy_history else []
     metadata = dict(**ref_physio._metadata) if copy_metadata else None
 
+    if suppdata is None:
+        suppdata = ref_physio._suppdata if copy_suppdata else None
+
     # make new class
     out = ref_physio.__class__(np.array(data, dtype=dtype),
-                               fs=fs, history=history, metadata=metadata)
+                               fs=fs, history=history, metadata=metadata,
+                               suppdata=suppdata)
     return out
 
 
