Faster unit selection sorting.

src/spikeinterface/core/basesorting.py

@@ -166,6 +166,8 @@ def get_unit_spike_train(
             per unit and per segment compact in memory.
             Using the cache makes the first call quite slow but then future calls are very fast.
 
+        Note: if use_cache=False, but the lexsorted cache is already computed then it will be used anyway.
+
         Returns
         -------
         spike_train : np.ndarray
@@ -188,9 +190,14 @@ def get_unit_spike_train(
             )
 
         segment_index = self._check_segment_index(segment_index)
+
+        lexsort_key = ("sample_index", "segment_index", "unit_index")
+        if lexsort_key in self._cached_lexsorted_spike_vector.keys():
+            use_cache = True
+
         if use_cache:
             ordered_spike_vector, slices = self.to_reordered_spike_vector(
-                lexsort=("sample_index", "segment_index", "unit_index"),
+                lexsort=lexsort_key,
                 return_order=False,
                 return_slices=True,
             )

src/spikeinterface/core/unitsselectionsorting.py

@@ -59,11 +59,23 @@ def _compute_and_cache_spike_vector(self) -> None:
             all_old_unit_ids=self._parent_sorting.unit_ids,
             all_new_unit_ids=self._unit_ids,
         )
-        # lexsort by segment_index, sample_index, unit_index
-        sort_indices = np.lexsort(
-            (spike_vector["unit_index"], spike_vector["sample_index"], spike_vector["segment_index"])
-        )
-        self._cached_spike_vector = spike_vector[sort_indices]
+
+        # check if order is preserved
+        pos = np.searchsorted(self._parent_sorting.unit_ids, self.unit_ids)
+        order_is_preserved = np.all(np.diff(pos) > 0)
+
+        if not order_is_preserved:
+            # Note: this can be a very high cost and make big dataset very slow
+            # the only goal of this is to ensure the unit_index order when the sample is the same
+            # TODO: maybe we can remove it, if we don't guarantee the order of unit_index
+            # when sample_index is the same, but it can be a problem for some downstream analysis
+
+            # lexsort by segment_index, sample_index, unit_index
+            sort_indices = np.lexsort(
+                (spike_vector["unit_index"], spike_vector["sample_index"], spike_vector["segment_index"])
+            )
+            spike_vector = spike_vector[sort_indices]
+        self._cached_spike_vector = spike_vector
 
 
 class UnitsSelectionSortingSegment(BaseSortingSegment):

src/spikeinterface/extractors/phykilosortextractors.py

@@ -149,6 +149,7 @@ def __init__(
             del cluster_info["id"]
 
         if remove_empty_units:
+
             unique_unit_ids_list = [int(clust) for clust in unique_unit_ids]
             cluster_info = cluster_info.query(f"cluster_id in {unique_unit_ids_list}")