Add context manager for busy indicator and use for traces, merges, compute

Author: alejoe91

spikeinterface_gui/mergeview.py

@@ -198,7 +198,7 @@ def _qt_make_layout(self):
         row_layout = QT.QHBoxLayout()
 
         but = QT.QPushButton('Calculate merges')
-        but.clicked.connect(self._qt_calculate_potential_automerge)
+        but.clicked.connect(self._compute_merges)
         row_layout.addWidget(but)
 
         if self.controller.curation:
@@ -269,8 +269,11 @@ def _qt_refresh(self):
             self.table.resizeColumnToContents(i)
         self.table.setSortingEnabled(True)
 
-    def _qt_calculate_potential_automerge(self):
-        self.get_potential_merges()
+    def _compute_merges(self):
+        with self.busy_cursor():
+            self.get_potential_merges()
+        if len(self.proposed_merge_unit_groups) == 0:
+            self.warning(f"No potential merges found with method {self.method}")
         self.refresh()
 
     def _qt_on_spike_selection_changed(self):
@@ -318,7 +321,7 @@ def _panel_make_layout(self):
         self.table_area = pn.pane.Placeholder("No merges computed yet.", height=400)
 
         self.caluculate_merges_button = pn.widgets.Button(name="Calculate merges", button_type="primary", sizing_mode="stretch_width")
-        self.caluculate_merges_button.on_click(self._panel_calculate_merges)
+        self.caluculate_merges_button.on_click(self._panel_compute_merges)
 
         calculate_list = [self.caluculate_merges_button]
 
@@ -378,11 +381,8 @@ def _panel_refresh(self):
         self.table.on_click(self._panel_on_click)
         self.table_area.update(self.table)
 
-    def _panel_calculate_merges(self, event):
-        import panel as pn
-        self.table_area.update(pn.indicators.LoadingSpinner(size=50, value=True))
-        self.get_potential_merges()
-        self.refresh()
+    def _panel_compute_merges(self, event):
+        self._compute_merges()
 
     def _panel_on_method_change(self, event):
         self.method = event.new

spikeinterface_gui/tracemapview.py

@@ -1,6 +1,5 @@
 import numpy as np
-
-import matplotlib.cm
+from contextlib import nullcontext
 import matplotlib.colors
 
 from .view_base import ViewBase
@@ -29,9 +28,12 @@ def __init__(self, controller=None, parent=None, backend="qt"):
         self.channel_order_reverse = np.argsort(self.channel_order, kind="stable")
         self.color_limit = None
         self.last_data_curves = None
+        self.factor = None
 
         self.xsize = 0.5
         self._block_auto_refresh_and_notify = False
+        self._retrieve_traces_time_checked = None
+        self.trace_context = nullcontext
 
         ViewBase.__init__(self, controller=controller, parent=parent, backend=backend)
         MixinViewTrace.__init__(self)
@@ -62,60 +64,8 @@ def make_color_lut(self):
         if self.settings['reverse_colormap']:
             self.lut = self.lut[::-1]
 
-
-    def get_data_in_chunk(self, t1, t2, segment_index):
-        t_start = 0.0
-        sr = self.controller.sampling_frequency
-
-        ind1 = max(0, int((t1 - t_start) * sr))
-        ind2 = min(self.controller.get_num_samples(segment_index), int((t2 - t_start) * sr))
-
-        traces_chunk = self.controller.get_traces(segment_index=segment_index, start_frame=ind1, end_frame=ind2)
-
-        sl = self.controller.segment_slices[segment_index]
-        spikes_seg = self.controller.spikes[sl]
-        i1, i2 = np.searchsorted(spikes_seg["sample_index"], [ind1, ind2])
-        spikes_chunk = spikes_seg[i1:i2].copy()
-        spikes_chunk["sample_index"] -= ind1
-
-        data_curves = traces_chunk[:, self.channel_order]
-
-        if data_curves.dtype != "float32":
-            data_curves = data_curves.astype("float32")
-
-        times_chunk = np.arange(traces_chunk.shape[0], dtype='float64')/self.controller.sampling_frequency+max(t1, 0)
-
-        scatter_x = []
-        scatter_y = []
-        scatter_colors = []
-        scatter_unit_ids = []
-
-        for unit_index, unit_id in self.controller.iter_visible_units():
-
-            inds = np.flatnonzero(spikes_chunk["unit_index"] == unit_index)
-            if inds.size == 0:
-                continue
-
-            # Get spikes for this unit
-            unit_spikes = spikes_chunk[inds]
-            channel_inds = unit_spikes["channel_index"]
-            sample_inds = unit_spikes["sample_index"]
-
-            x = times_chunk[sample_inds]
-            y = self.channel_order_reverse[channel_inds] + 0.5
-
-            # This should both for qt (QTColor) and panel (html color)
-            color = self.get_unit_color(unit_id)
-
-            scatter_x.extend(x)
-            scatter_y.extend(y)
-            scatter_colors.extend([color] * len(x))
-            scatter_unit_ids.extend([str(unit_id)] * len(x))
-
-        # used for auto scaled
-        self.last_data_curves = data_curves
-
-        return times_chunk, data_curves, scatter_x, scatter_y, scatter_colors, scatter_unit_ids
+    def get_visible_channel_inds(self):
+        return np.arange(self.controller.analyzer.get_num_channels())
 
 
     ## Qt ##
@@ -199,6 +149,7 @@ def _qt_seek(self, t):
         seg_index = self.controller.get_time()[1]
         times_chunk, data_curves, scatter_x, scatter_y, scatter_colors, scatter_unit_ids = \
             self.get_data_in_chunk(t1, t2, seg_index)
+        self.last_data_curves = data_curves
 
         if self.color_limit is None:
             self.color_limit = np.max(np.abs(data_curves))

spikeinterface_gui/traceview.py

@@ -1,5 +1,6 @@
 import numpy as np
 import time
+from contextlib import nullcontext
 
 from .view_base import ViewBase
 
@@ -13,6 +14,94 @@
 #   * 
 
 class MixinViewTrace:
+
+    MAX_RETRIEVE_TIME_FOR_BUSY_CURSOR = 0.5  # seconds
+
+    def get_data_in_chunk(self, t1, t2, segment_index):
+        with self.trace_context():
+            if not self._retrieve_traces_time_checked:
+                t_traces_start = time.perf_counter()
+            t_start = 0.0
+            sr = self.controller.sampling_frequency
+
+            ind1 = max(0, int((t1 - t_start) * sr))
+            ind2 = min(self.controller.get_num_samples(segment_index), int((t2 - t_start) * sr))
+
+            traces_chunk = self.controller.get_traces(segment_index=segment_index, start_frame=ind1, end_frame=ind2)
+
+            sl = self.controller.segment_slices[segment_index]
+            spikes_seg = self.controller.spikes[sl]
+            i1, i2 = np.searchsorted(spikes_seg["sample_index"], [ind1, ind2])
+            spikes_chunk = spikes_seg[i1:i2].copy()
+            spikes_chunk["sample_index"] -= ind1
+
+            visible_channel_inds = self.get_visible_channel_inds()
+
+            data_curves = traces_chunk[:, visible_channel_inds].T.copy()
+
+            if data_curves.dtype != "float32":
+                data_curves = data_curves.astype("float32")
+
+            if self.factor is not None:
+                n = visible_channel_inds.size
+                gains = np.ones(n, dtype=float) * 1.0 / (self.factor * max(self.mad[visible_channel_inds]))
+                offsets = np.arange(n)[::-1] - self.med[visible_channel_inds] * gains
+
+                data_curves *= gains[:, None]
+                data_curves += offsets[:, None]
+
+            times_chunk = np.arange(traces_chunk.shape[0], dtype='float64') / self.controller.sampling_frequency+max(t1, 0)
+
+            scatter_x = []
+            scatter_y = []
+            scatter_colors = []
+            scatter_unit_ids = []
+
+            global_to_local_chan_inds = np.zeros(self.controller.channel_ids.size, dtype='int64')
+            global_to_local_chan_inds[visible_channel_inds] = np.arange(visible_channel_inds.size, dtype='int64')
+
+            for unit_index, unit_id in self.controller.iter_visible_units():
+
+                inds = np.flatnonzero(spikes_chunk["unit_index"] == unit_index)
+                if inds.size == 0:
+                    continue
+
+                # Get spikes for this unit
+                unit_spikes = spikes_chunk[inds]
+                channel_inds = unit_spikes["channel_index"]
+                sample_inds = unit_spikes["sample_index"]
+
+                chan_mask = np.isin(channel_inds, visible_channel_inds)
+                if not np.any(chan_mask):
+                    continue
+
+                sample_inds = sample_inds[chan_mask]
+                channel_inds = channel_inds[chan_mask]
+                # Map channel indices to their positions in visible_channel_inds
+                local_channel_inds = global_to_local_chan_inds[channel_inds]
+
+
+                # Calculate y values using signal values
+                x = times_chunk[sample_inds]
+                y = data_curves[local_channel_inds, sample_inds]
+
+                # This should both for qt (QTColor) and panel (html color)
+                color = self.get_unit_color(unit_id)
+
+                scatter_x.extend(x)
+                scatter_y.extend(y)
+                scatter_colors.extend([color] * len(x))
+                scatter_unit_ids.extend([str(unit_id)] * len(x))
+            if not self._retrieve_traces_time_checked:
+                t_traces_end = time.perf_counter()
+                elapsed = t_traces_end - t_traces_start
+                if elapsed > self.MAX_RETRIEVE_TIME_FOR_BUSY_CURSOR:
+                    print(f"Trace retrieval took {elapsed:.3f} seconds. Enabling busy cursor.")
+                    self.trace_context = self.busy_cursor
+                self._retrieve_traces_time_checked = True
+
+        return times_chunk, data_curves, scatter_x, scatter_y, scatter_colors, scatter_unit_ids
+
     ## Qt ##
     def _qt_create_toolbar(self):
         from .myqt import QT
@@ -284,6 +373,8 @@ def __init__(self, controller=None, parent=None, backend="qt"):
         self.factor = 15.0
         self.xsize = 0.5
         self._block_auto_refresh_and_notify = False
+        self._retrieve_traces_time_checked = False
+        self.trace_context = nullcontext
 
         ViewBase.__init__(self, controller=controller, parent=parent,  backend=backend)
         MixinViewTrace.__init__(self)
@@ -308,82 +399,6 @@ def apply_gain_zoom(self, factor_ratio):
         self.factor *= factor_ratio
         self.refresh()
 
-    def get_data_in_chunk(self, t1, t2, segment_index):
-        t_start = 0.0
-        sr = self.controller.sampling_frequency
-
-        ind1 = max(0, int((t1 - t_start) * sr))
-        ind2 = min(self.controller.get_num_samples(segment_index), int((t2 - t_start) * sr))
-
-        traces_chunk = self.controller.get_traces(segment_index=segment_index, start_frame=ind1, end_frame=ind2)
-
-        sl = self.controller.segment_slices[segment_index]
-        spikes_seg = self.controller.spikes[sl]
-        i1, i2 = np.searchsorted(spikes_seg["sample_index"], [ind1, ind2])
-        spikes_chunk = spikes_seg[i1:i2].copy()
-        spikes_chunk["sample_index"] -= ind1
-
-        visible_channel_inds = self.get_visible_channel_inds()
-        
-        data_curves = traces_chunk[:, visible_channel_inds].T.copy()
-
-        if data_curves.dtype != "float32":
-            data_curves = data_curves.astype("float32")
-        
-        n = visible_channel_inds.size
-        gains = np.ones(n, dtype=float) * 1.0 / (self.factor * max(self.mad[visible_channel_inds]))
-        offsets = np.arange(n)[::-1] - self.med[visible_channel_inds] * gains
-        
-        data_curves *= gains[:, None]
-        data_curves += offsets[:, None]
-
-        times_chunk = np.arange(traces_chunk.shape[0], dtype='float64')/self.controller.sampling_frequency+max(t1, 0)
-
-        scatter_x = []
-        scatter_y = []
-        scatter_colors = []
-        scatter_unit_ids = []
-
-        global_to_local_chan_inds = np.zeros(self.controller.channel_ids.size, dtype='int64')
-        global_to_local_chan_inds[visible_channel_inds] = np.arange(visible_channel_inds.size, dtype='int64')
-        
-
-        for unit_index, unit_id in self.controller.iter_visible_units():
-
-            inds = np.flatnonzero(spikes_chunk["unit_index"] == unit_index)
-            if inds.size == 0:
-                continue
-
-            # Get spikes for this unit
-            unit_spikes = spikes_chunk[inds]
-            channel_inds = unit_spikes["channel_index"]
-            sample_inds = unit_spikes["sample_index"]
-
-            
-            chan_mask = np.isin(channel_inds, visible_channel_inds)
-            if not np.any(chan_mask):
-                continue
-
-            sample_inds = sample_inds[chan_mask]
-            channel_inds = channel_inds[chan_mask]
-            # Map channel indices to their positions in visible_channel_inds
-            local_channel_inds = global_to_local_chan_inds[channel_inds]
-
-
-            # Calculate y values using signal values
-            x = times_chunk[sample_inds]
-            y = data_curves[local_channel_inds, sample_inds]
-
-            # This should both for qt (QTColor) and panel (html color)
-            color = self.get_unit_color(unit_id)
-
-            scatter_x.extend(x)
-            scatter_y.extend(y)
-            scatter_colors.extend([color] * len(x))
-            scatter_unit_ids.extend([str(unit_id)] * len(x))
-
-        return times_chunk, data_curves, scatter_x, scatter_y, scatter_colors, scatter_unit_ids
-
     ## qt ##
     def _qt_make_layout(self):
         from .myqt import QT