Add support for segment annotation

Author: me-pic

peakdet/editor.py

@@ -32,7 +32,7 @@ def __init__(self, data):
 
         # 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, self.included = set(), set(), set()
+        self.deleted, self.rejected, self.included, self.annotated = set(), set(), set(), set()
 
         # make main plot objects depending on supplementary data
         if self.suppdata is None:
@@ -57,8 +57,8 @@ def __init__(self, data):
         # three selectors for:
         #    1. rejection (central mouse),
         #    2. addition (right mouse), and
-        #    3. deletion (left mouse)
-        delete = functools.partial(self.on_edit, method="delete")
+        #    3. deletion or annotation (left mouse)
+        self.mode = "delete"
         reject = functools.partial(self.on_edit, method="reject")
         insert = functools.partial(self.on_edit, method="insert")
 
@@ -71,7 +71,7 @@ def __init__(self, data):
 
         self.span2 = SpanSelector(
             self.ax,
-            delete,
+            self._span_callback,
             "horizontal",
             button=1,
             useblit=True,
@@ -118,6 +118,21 @@ def plot_signals(self, plot=True):
             ".g",
         )
 
+        if self.data.segments.shape[0]>0:
+            for segment in self.data.segments:
+                ymin, ymax = self.ax.get_ylim()
+                height = ymax - ymin
+                width = self.time[int(segment[1])] - self.time[int(segment[0])]
+
+                rect = plt.Rectangle(
+                    [self.time[int(segment[0])], ymin], 
+                    width, 
+                    height,
+                    alpha=0.3,
+                    facecolor='y'
+                )
+                self.ax.add_patch(rect)
+
         if self.suppdata is not None:
             self._ax[1].plot(self.time, self.suppdata, "k", linewidth=0.7)
             self._ax[1].set_ylim(-0.5, 0.5)
@@ -142,6 +157,13 @@ def on_key(self, event):
             self.undo()
         elif event.key in ["ctrl+q", "super+d"]:
             self.quit()
+        elif event.key in ["ctrl+a"]:
+            self.annotate_mode()
+        elif event.key in ["ctrl+e"]:
+            self.edit_mode()
+
+    def _span_callback(self, xmin, xmax):
+        self.on_edit(xmin, xmax, method=self.mode)
 
     def on_edit(self, xmin, xmax, *, method):
         """
@@ -150,10 +172,10 @@ def on_edit(self, xmin, xmax, *, method):
         Removes specified peaks by either rejection / deletion, OR
         Include one peak by finding the max in the selection.
 
-        method accepts 'insert', 'reject', 'delete'
+        method accepts 'insert', 'reject', 'delete', 'annotate'
         """
         logger.debug("Edited peaks with action: {}", method)
-        if method not in ["insert", "reject", "delete"]:
+        if method not in ["insert", "reject", "delete", "annotate"]:
             raise ValueError(f'Action "{method}" not supported.')
 
         tmin, tmax = np.searchsorted(self.time, (xmin, xmax))
@@ -165,6 +187,8 @@ def on_edit(self, xmin, xmax, *, method):
             if newpeak == tmin:
                 self.plot_signals()
                 return
+        elif method == "annotate":
+            segment = (tmin, tmax)
         else:
             bad = np.arange(pmin, pmax, dtype=int)
             if len(bad) == 0:
@@ -180,6 +204,11 @@ def on_edit(self, xmin, xmax, *, method):
         if method == "insert":
             self.included.add(newpeak)
             self.data = operations.add_peaks(self.data, newpeak)
+        elif method == "annotate":
+            self.data = operations.annotate_segments(self.data, segment)
+            self.annotated = set()
+            for s in self.data.segments:
+                self.annotated.add((s[0], s[1]))
         else:
             rej.update(self.data.peaks[bad].tolist())
             self.data = fcn(self.data, self.data.peaks[bad])
@@ -189,7 +218,7 @@ def on_edit(self, xmin, xmax, *, method):
     def undo(self):
         """Reset last span select peak removal."""
         # check if last history entry was a manual reject / delete
-        relevant = ["reject_peaks", "delete_peaks", "add_peaks"]
+        relevant = ["reject_peaks", "delete_peaks", "add_peaks", "annotate_segments"]
         if self.data._history[-1][0] not in relevant:
             return
 
@@ -215,7 +244,20 @@ def undo(self):
                 np.searchsorted(self.data._metadata["peaks"], peaks["add"]),
             )
             self.included.remove(peaks["add"])
+        elif func == "annotate_segments":
+            self.data._metadata["segments"] = np.delete(
+                self.data._metadata["segments"],
+                -1,
+                axis=0
+            )
+            self.annotated.pop()
         self.data._metadata["troughs"] = utils.check_troughs(
             self.data, self.data.peaks, self.data.troughs
         )
         self.plot_signals()
+
+    def annotate_mode(self):
+        self.mode = "annotate"
+
+    def edit_mode(self):
+        self.mode = "delete"

peakdet/operations.py

@@ -249,6 +249,19 @@ def add_peaks(data, add):
     return data
 
 
+@utils.make_operation()
+def annotate_segments(data, segment):
+    if isinstance(segment, tuple):
+        segment = [segment]
+    for s in segment:
+        data = utils.check_physio(data, ensure_fs=False, copy=True)
+        segments = utils.check_segment(data, np.array([s], dtype=data._metadata["segments"].dtype))
+    
+        data._metadata["segments"] = np.unique(segments)
+
+    return data
+
+
 def edit_physio(data):
     """
     Open interactive plot with `data` to permit manual editing of time series.
@@ -281,6 +294,8 @@ def edit_physio(data):
         data = delete_peaks(data, remove=sorted(edits.deleted))
     if len(edits.included) > 0:
         data = add_peaks(data, add=sorted(edits.included))
+    if len(edits.annotated) > 0:
+        data = annotate_segments(data, segment=sorted(edits.annotated))
 
     return data
 

peakdet/physio.py

@@ -74,6 +74,7 @@ def __init__(self, data, fs=None, history=None, metadata=None, suppdata=None):
                 peaks=np.empty(0, dtype=int),
                 troughs=np.empty(0, dtype=int),
                 reject=np.empty(0, dtype=int),
+                segments=np.empty(0, dtype=[('onset', 'f8'), ('offset', 'f8')])
             )
         self._suppdata = None if suppdata is None else np.asarray(suppdata).squeeze()
 
@@ -115,6 +116,11 @@ def peaks(self):
     def troughs(self):
         """Indices of detected troughs in `data`."""
         return self._metadata["troughs"]
+    
+    @property
+    def segments(self):
+        """Indices marking the onset and offset of annotated segments"""
+        return self._metadata["segments"]
 
     @property
     def _masked(self):

peakdet/utils.py

@@ -60,6 +60,21 @@ def wrapper(data, *args, **kwargs):
             # append everything to data instance history
             data._history += [(name, provided)]
 
+            # clean duplicates from history
+            tmp_history, duplicate = [], set()
+            for name, value in data._history:
+                if name == "annotate_segments":
+                    segments = value.get("segment", [])
+                    new_segments = [s for s in segments if s not in duplicate]
+                    duplicate.update(new_segments)
+
+                    tmp_history.append((name, {"segment": new_segments}))
+                else:
+                    tmp_history.append((name, value))
+
+            if len(tmp_history) > 0:
+                data._history = tmp_history
+
             return data
 
         return wrapper
@@ -244,6 +259,36 @@ def check_troughs(data, peaks, troughs=None):
     return all_troughs
 
 
+def check_segment(data, segment):
+    """
+    Confirms that segments are not overlapping. The last segment specified 
+    as precedence over the overlapping ones.
+
+    Parameters
+    ----------
+    data : array-like
+        Input data for which `troughs` and `peaks` were detected
+    segments : array-like
+        Array containing the onset and offset of segments to annotate
+    
+    Returns
+    -------
+    segments : np.ndarray
+        Array containing validated (non-overlapping) segments
+    """
+    tmp_segment = data.segments
+    if tmp_segment.shape[0] == 0:
+        return segment
+    else:
+        new_segment = segment[0]
+        for idx, s in enumerate(tmp_segment):
+            if (new_segment[1] > s[0] and new_segment[1] <= s[1]) or (new_segment[0] >= s[0] and new_segment[0] < s[1]):
+                tmp_segment[idx] = new_segment
+            else:
+                tmp_segment = np.append(tmp_segment, segment)
+        return tmp_segment
+
+
 def enable_logger(loglevel="INFO", diagnose=True, backtrace=True):
     """
     Toggle the use of the module's logger and configures it.