diff --git a/spikeinterface_gui/controller.py b/spikeinterface_gui/controller.py index 1c5f179..c8c8083 100644 --- a/spikeinterface_gui/controller.py +++ b/spikeinterface_gui/controller.py @@ -6,9 +6,8 @@ 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 @@ -16,9 +15,6 @@ 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): diff --git a/spikeinterface_gui/spikelistview.py b/spikeinterface_gui/spikelistview.py index c8cd73c..1149d8a 100644 --- a/spikeinterface_gui/spikelistview.py +++ b/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 diff --git a/spikeinterface_gui/spikerateview.py b/spikeinterface_gui/spikerateview.py index d76ddf7..dbf9201 100644 --- a/spikeinterface_gui/spikerateview.py +++ b/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) diff --git a/spikeinterface_gui/tests/test_mainwindow_panel.py b/spikeinterface_gui/tests/test_mainwindow_panel.py index 05d5e19..38f0b70 100644 --- a/spikeinterface_gui/tests/test_mainwindow_panel.py +++ b/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) diff --git a/spikeinterface_gui/tests/test_mainwindow_qt.py b/spikeinterface_gui/tests/test_mainwindow_qt.py index 0fbc3ef..6656dfd 100644 --- a/spikeinterface_gui/tests/test_mainwindow_qt.py +++ b/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) diff --git a/spikeinterface_gui/traceview.py b/spikeinterface_gui/traceview.py index e79f29d..a742e4a 100644 --- a/spikeinterface_gui/traceview.py +++ b/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