Test analyzer in lazy mode

spikeinterface_gui/controller.py

@@ -6,19 +6,15 @@
 
 from copy import deepcopy
 
-from spikeinterface import compute_sparsity
-from spikeinterface.core import get_template_extremum_channel, BaseEvent
-from spikeinterface.core.sorting_tools import spike_vector_to_indices
+from spikeinterface import compute_sparsity, get_template_extremum_channel
+from spikeinterface.core.base import minimum_spike_dtype
 from spikeinterface.curation import validate_curation_dict
 from spikeinterface.curation.curation_model import Curation
 from spikeinterface.widgets.utils import make_units_table_from_analyzer
 
 from .curation_tools import add_merge, default_label_definitions, empty_curation_data
 from .event_tools import parse_events
 
-spike_dtype =[('sample_index', 'int64'), ('unit_index', 'int64'), 
-    ('channel_index', 'int64'), ('segment_index', 'int64'),
-    ('visible', 'bool'), ('selected', 'bool'), ('rand_selected', 'bool')]
 
 
 _default_main_settings = dict(
@@ -275,8 +271,12 @@ def __init__(
 
         t0 = time.perf_counter()
 
-        self._extremum_channel = get_template_extremum_channel(self.analyzer,
-                                    mode="extremum", peak_sign='both', outputs='index')
+        self._extremum_channel = get_template_extremum_channel(
+            self.analyzer,
+            mode="extremum",
+            peak_sign='both',
+            outputs='index'
+        )
 
         # spikeinterface handle colors in matplotlib style tuple values in range (0,1)
         self.refresh_colors()
@@ -295,43 +295,56 @@ def __init__(
         unit_ids = self.analyzer.unit_ids
         num_seg = self.analyzer.get_num_segments()
         self.num_spikes = self.analyzer.sorting.count_num_spikes_per_unit(outputs="dict")
-        # print("self.num_spikes", self.num_spikes)
 
-        spike_vector = self.analyzer.sorting.to_spike_vector(concatenated=True, extremum_channel_inds=self._extremum_channel)
-        # spike_vector = self.analyzer.sorting.to_spike_vector(concatenated=True)
-
+        if self.analyzer._lazy:
+            # If the analyzer is lazy, avoid materializing the full spike vector, which can be very large.
+            self.spikes = self.analyzer.sorting.to_spike_vector()
+        else:
+            # In this case we make a copy with align=True, which is required for np.searchsorted
+            # (and therefore trace views) to be fast.
+            spike_vector = self.analyzer.sorting.to_spike_vector()
+            self.spikes = np.zeros(spike_vector.size, dtype=np.dtype(minimum_spike_dtype, align=True))
+            self.spikes['sample_index'] = spike_vector['sample_index']
+            self.spikes['unit_index'] = spike_vector['unit_index']
+            self.spikes['segment_index'] = spike_vector['segment_index']
+
         self.random_spikes_indices = self.analyzer.get_extension("random_spikes").get_data()
+        self._random_spikes_set = set(int(i) for i in self.random_spikes_indices)
 
-        # align=True is required for np.searchsorted (and therefore trace
-        # views) to be fast.
-        self.spikes = np.zeros(spike_vector.size, dtype=np.dtype(spike_dtype, align=True))
-        self.spikes['sample_index'] = spike_vector['sample_index']
-        self.spikes['unit_index'] = spike_vector['unit_index']
-        self.spikes['segment_index'] = spike_vector['segment_index']
-        self.spikes['channel_index'] = spike_vector['channel_index']
-        self.spikes['rand_selected'][:] = False
-        self.spikes['rand_selected'][self.random_spikes_indices] = True
-
-        # self.num_spikes = self.analyzer.sorting.count_num_spikes_per_unit(outputs="dict")
-        seg_limits = np.searchsorted(self.spikes["segment_index"], np.arange(num_seg + 1))
-        self.segment_slices = {segment_index: slice(seg_limits[segment_index], seg_limits[segment_index + 1]) for segment_index in range(num_seg)}
-
-        spike_vector2 = self.analyzer.sorting.to_spike_vector(concatenated=False)
-        self.final_spike_samples = [segment_spike_vector[-1][0] for segment_spike_vector in spike_vector2]
-        # this is dict of list because per segment spike_indices[segment_index][unit_id]
-        spike_indices_abs = spike_vector_to_indices(spike_vector2, unit_ids, absolute_index=True)
-        spike_indices = spike_vector_to_indices(spike_vector2, unit_ids)
-        # this is flatten
-        spike_per_seg = [s.size for s in spike_vector2]
-        # dict[unit_id] -> all indices for this unit across segments
-        self._spike_index_by_units = {}
-        # dict[segment_index][unit_id] -> all indices for this unit for one segment
-        self._spike_index_by_segment_and_units = spike_indices_abs
-        for unit_id in unit_ids:
-            inds = []
-            for seg_ind in range(num_seg):
-                inds.append(spike_indices[seg_ind][unit_id] + int(np.sum(spike_per_seg[:seg_ind])))
-            self._spike_index_by_units[unit_id] = np.concatenate(inds)
+        self._ext_channel_inds = np.array([self._extremum_channel[unit_id] for unit_id in self.unit_ids])
+
+        cached = self.analyzer.sorting._cached_spike_vector_segment_slices
+        if cached is not None:
+            # shape (num_seg, 2): columns are [start, stop]
+            self.segment_slices = {seg: slice(int(cached[seg, 0]), int(cached[seg, 1])) for seg in range(num_seg)}
+        else:
+            bounds = np.searchsorted(np.asarray(self.spikes["segment_index"]), np.arange(num_seg + 1))
+            self.segment_slices = {seg: slice(int(bounds[seg]), int(bounds[seg + 1])) for seg in range(num_seg)}
+
+        # Load unit_index once to build per-unit lookup structures, avoiding a
+        # second to_spike_vector() call that would materialise full structured
+        # arrays from zarr for every segment.
+        unit_index_all = np.asarray(self.spikes["unit_index"])
+
+        # last sample per segment: one cheap element read instead of full materialisation
+        sample_index_arr = self.spikes["sample_index"]
+        self.final_spike_samples = [int(sample_index_arr[self.segment_slices[seg].stop - 1]) for seg in range(num_seg)]
+
+        # dict[segment_index][unit_id] -> absolute spike indices for that unit in that segment
+        num_units = len(unit_ids)
+        self._spike_index_by_segment_and_units = {}
+        for seg_ind in range(num_seg):
+            sl = self.segment_slices[seg_ind]
+            seg_unit_idx = unit_index_all[sl]
+            abs_offset = sl.start
+            order = np.argsort(seg_unit_idx, stable=True)
+            sorted_unit = seg_unit_idx[order]
+            sorted_abs = (order + abs_offset).astype(np.int64)
+            boundaries = np.searchsorted(sorted_unit, np.arange(num_units + 1, dtype=np.int64))
+            self._spike_index_by_segment_and_units[seg_ind] = {
+                uid: sorted_abs[boundaries[u]:boundaries[u + 1]].copy()
+                for u, uid in enumerate(unit_ids)
+            }
 
         t1 = time.perf_counter()
         if verbose:
@@ -664,9 +677,10 @@ def get_dict_unit_visible(self):
 
     def update_visible_spikes(self):
         inds = []
-        for unit_index, unit_id in self.iter_visible_units():
-            inds.append(self._spike_index_by_units[unit_id])
-
+        for _, unit_id in self.iter_visible_units():
+            for seg_ind in self._spike_index_by_segment_and_units:
+                inds.append(self._spike_index_by_segment_and_units[seg_ind][unit_id])
+
         if len(inds) > 0:
             inds = np.concatenate(inds)
             inds = np.sort(inds)
@@ -693,10 +707,11 @@ def set_indices_spike_selected(self, inds):
 
     def get_spike_indices(self, unit_id, segment_index=None):
         if segment_index is None:
-            # dict[unit_id] -> all indices for this unit across segments
-            return self._spike_index_by_units[unit_id]
+            return np.concatenate([
+                self._spike_index_by_segment_and_units[seg][unit_id]
+                for seg in sorted(self._spike_index_by_segment_and_units)
+            ])
         else:
-            # dict[segment_index][unit_id] -> all indices for this unit for one segment
             return self._spike_index_by_segment_and_units[segment_index][unit_id]
 
     def get_num_samples(self, segment_index):

spikeinterface_gui/spikelistview.py

@@ -53,6 +53,8 @@ def data(self, index, role):
 
             abs_ind = self.visible_ind[row]
             spike = self.controller.spikes[abs_ind]
+            channel_index = self.controller._ext_channel_inds[spike['unit_index']]
+            rand_selected = int(abs_ind) in self.controller._random_spikes_set
             unit_id = self.controller.unit_ids[spike['unit_index']]
 
             if role ==QT.Qt.DisplayRole :
@@ -65,9 +67,9 @@ def data(self, index, role):
                 elif col == 3:
                     return '{}'.format(spike['sample_index'])
                 elif col == 4:
-                    return '{}'.format(spike['channel_index'])
+                    return '{}'.format(channel_index)
                 elif col == 5:
-                    return '{}'.format(spike['rand_selected'])
+                    return '{}'.format(rand_selected)
                 else:
                     return None
             elif role == QT.Qt.DecorationRole :
@@ -316,14 +318,17 @@ def _panel_refresh_table(self):
             color = mcolors.to_hex(self.controller.get_unit_color(unit_id))
             spike_unit_ids.append({"id": unit_id, "color": color})
 
+        channel_inds = self.controller._ext_channel_inds[spikes['unit_index']]
+        rand_selected = np.isin(visible_inds, self.controller.random_spikes_indices)
+
         # Prepare data for tabulator
         data = {
             '#': visible_inds,
             'unit_id': spike_unit_ids,
             'segment_index': spikes['segment_index'],
             'sample_index': spikes['sample_index'],
-            'channel_index': spikes['channel_index'],
-            'rand_selected': spikes['rand_selected']
+            'channel_index': channel_inds,
+            'rand_selected': rand_selected
         }
 
         # Update table data without replacing entire dataframe

spikeinterface_gui/spikerateview.py

@@ -77,7 +77,7 @@ def _qt_refresh(self):
         for r, unit_id in enumerate(visible_unit_ids):
 
             spike_inds = self.controller.get_spike_indices(unit_id, segment_index=segment_index)
-            spikes = self.controller.spikes[spike_inds]['sample_index']
+            spikes = self.controller.spikes['sample_index'][spike_inds]
 
             count, _ = np.histogram(spikes, bins=bin_edges)
 
@@ -161,7 +161,7 @@ def _panel_refresh(self):
         for unit_id in visible_unit_ids:
 
             spike_inds = self.controller.get_spike_indices(unit_id, segment_index=segment_index)
-            spikes = self.controller.spikes[spike_inds]['sample_index']
+            spikes = self.controller.spikes['sample_index'][spike_inds]
 
             count, _ = np.histogram(spikes, bins=bin_edges)
 

spikeinterface_gui/tests/test_mainwindow_panel.py

@@ -22,8 +22,7 @@ def teardown_module():
 
 def test_mainwindow(start_app=False, verbose=True, curation=False, only_some_extensions=False, events=False, port=0):
 
-
-    analyzer = load_sorting_analyzer(test_folder / "sorting_analyzer")
+    analyzer = load_sorting_analyzer(test_folder / "sorting_analyzer", load_extensions=False)
     # analyzer = load_analyzer(test_folder / "sorting_analyzer.zarr")
 
     print(analyzer)

spikeinterface_gui/tests/test_mainwindow_qt.py

@@ -35,15 +35,12 @@ def teardown_module():
     clean_all(test_folder)
 
 
-def test_mainwindow(start_app=False, verbose=True, curation=False, only_some_extensions=False, events=False):
+def test_mainwindow(start_app=False, verbose=True, curation=False, only_some_extensions=False, events=False, lazy_load=False):
 
 
-    analyzer = load_sorting_analyzer(test_folder / "sorting_analyzer")
+    analyzer = load_sorting_analyzer(test_folder / "sorting_analyzer", load_extensions=False, lazy=lazy_load)
     # analyzer = load_analyzer(test_folder / "sorting_analyzer.zarr")
 
-    tm = analyzer.get_extension("template_metrics").get_data().iloc[0, :]
-    # print(tm)
-    # return
 
     print(analyzer)
 
@@ -109,7 +106,8 @@ def test_mainwindow(start_app=False, verbose=True, curation=False, only_some_ext
         displayed_unit_properties=None,
         extra_unit_properties=extra_unit_properties,
         layout_preset='default',
-        events=events_dict
+        events=events_dict,
+        lazy_load=lazy_load
         # user_settings={"mainsettings": {"color_mode": "color_by_visibility", "max_visible_units": 5}}
     )
 
@@ -144,6 +142,8 @@ def test_launcher(verbose=True):
 parser = ArgumentParser()
 parser.add_argument('--dataset', default="small", help='Path to the dataset folder')
 parser.add_argument('--events', action="store_true", help='Simulate and add events')
+parser.add_argument('--lazy', action="store_true", help='Lazy load')
+
 
 if __name__ == '__main__':
     args = parser.parse_args()
@@ -155,7 +155,7 @@ def test_launcher(verbose=True):
     if not test_folder.is_dir():
         setup_module()
 
-    win = test_mainwindow(start_app=True, verbose=True, curation=True, events=args.events)
+    win = test_mainwindow(start_app=True, verbose=True, curation=True, events=args.events, lazy_load=args.lazy)
     # win = test_mainwindow(start_app=True, verbose=True, curation=False)
 
     # test_launcher(verbose=True)

spikeinterface_gui/traceview.py

@@ -38,6 +38,7 @@ def get_data_in_chunk(self, t1, t2, 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_channel_chunk = self.controller._ext_channel_inds[spikes_chunk["unit_index"]]
             spikes_chunk["sample_index"] -= ind1
 
             # for trace map view, this returns the channels ordered by depth
@@ -73,7 +74,7 @@ def get_data_in_chunk(self, t1, t2, segment_index):
 
                 # Get spikes for this unit
                 unit_spikes = spikes_chunk[inds]
-                channel_inds = unit_spikes["channel_index"]
+                channel_inds = spikes_channel_chunk[inds]
                 sample_inds = unit_spikes["sample_index"]
 
                 chan_mask = np.isin(channel_inds, visible_channel_inds)
@@ -933,14 +934,15 @@ def find_nearest_spike(controller, x, segment_index, max_distance_samples=None):
         max_distance_samples = controller.sampling_frequency // 30
 
     ind_click = controller.time_to_sample_index(x)
-    (in_seg,) = np.nonzero(controller.spikes["segment_index"] == segment_index)
+    sl = controller.segment_slices[segment_index]
 
-    if len(in_seg) == 0:
+    if sl.start == sl.stop:
         return None
 
-    nearest = np.argmin(np.abs(controller.spikes[in_seg]["sample_index"] - ind_click))
-    ind_spike_nearest = in_seg[nearest]
-    sample_index = controller.spikes[ind_spike_nearest]["sample_index"]
+    sample_index_seg = np.asarray(controller.spikes["sample_index"][sl])
+    nearest_in_seg = np.argmin(np.abs(sample_index_seg - ind_click))
+    ind_spike_nearest = sl.start + nearest_in_seg
+    sample_index = sample_index_seg[nearest_in_seg]
 
     if np.abs(ind_click - sample_index) > max_distance_samples:
         return None