Support released phylib spike selector

Author: rossant

docs/changelog.md

@@ -40,6 +40,8 @@ behavior they verify rather than listed separately.
 
 ### Fixed
 
+- Start the GUI with released phylib versions that do not yet expose the
+  disjoint-spike selection optimization hint.
 - Keep dataset-local view settings isolated from global GUI state. In
   particular, a Firing Rate time range saved or leaked from another recording
   no longer clips spikes in a fresh dataset.

phy/apps/base.py

@@ -113,6 +113,20 @@ def _sample_spikes_evenly(spike_ids, n_spikes):
     return np.asarray(spike_ids[indices], dtype=np.int64)
 
 
+def _select_spikes_evenly(selector, n_spikes, cluster_ids, **kwargs):
+    """Use phylib's even selector when available, with a 2.7-compatible fallback."""
+    if 'sample_evenly' in inspect.signature(selector).parameters:
+        return selector(n_spikes, cluster_ids, sample_evenly=True, **kwargs)
+
+    selected = [
+        _sample_spikes_evenly(selector(None, [cluster_id], **kwargs), n_spikes)
+        for cluster_id in cluster_ids
+    ]
+    if not selected:
+        return np.array([], dtype=np.int64)
+    return np.concatenate(selected)
+
+
 def _spike_budget_fields(per_cluster, total, max_n_clusters, background=None):
     """Return dialog fields for a per-cluster budget and optional shared cap."""
     per_cluster_default = per_cluster if per_cluster is not None else total or 100000
@@ -318,8 +332,8 @@ def _get_waveforms_with_n_spikes(self, cluster_id, n_spikes_waveforms, current_f
         # Or keep spikes from a subset of the chunks for performance reasons (decompression will
         # happen on the fly here).
         else:
-            spike_ids = self.selector(
-                n_spikes_waveforms, [cluster_id], subset_chunks=True, sample_evenly=True
+            spike_ids = _select_spikes_evenly(
+                self.selector, n_spikes_waveforms, [cluster_id], subset_chunks=True
             )
 
         # Get the best channels.
@@ -1269,13 +1283,17 @@ def spikes_per_cluster(cluster_id):
         except AttributeError:
             chunk_bounds = [0.0, self.model.spike_samples[-1] + 1]
 
-        self.selector = SpikeSelector(
-            get_spikes_per_cluster=spikes_per_cluster,
-            spike_times=self.model.spike_samples,  # NOTE: chunk_bounds is in samples, not seconds
-            chunk_bounds=chunk_bounds,
-            n_chunks_kept=self.n_chunks_kept,
-            spikes_are_disjoint=True,
-        )
+        selector_kwargs = {
+            'get_spikes_per_cluster': spikes_per_cluster,
+            # NOTE: chunk_bounds is in samples, not seconds.
+            'spike_times': self.model.spike_samples,
+            'chunk_bounds': chunk_bounds,
+            'n_chunks_kept': self.n_chunks_kept,
+        }
+        # phylib 2.7 does not yet expose this optimization hint.
+        if 'spikes_are_disjoint' in inspect.signature(SpikeSelector).parameters:
+            selector_kwargs['spikes_are_disjoint'] = True
+        self.selector = SpikeSelector(**selector_kwargs)
 
     def _cache_methods(self):
         """Cache methods as specified in `self._memcached` and `self._cached`."""

phy/apps/tests/test_base.py

@@ -46,6 +46,7 @@
     TraceMixin,
     WaveformMixin,
     _allocate_spike_counts,
+    _select_spikes_evenly,
     _spike_budget_fields,
     _spike_budget_values,
 )
@@ -215,6 +216,64 @@ def test_controller_close(tempdir):
     assert all(handler.stream is None for handler in handlers)
 
 
+def test_set_selector_supports_released_and_newer_phylib():
+    controller = object.__new__(BaseController)
+    controller.model = Bunch(spike_samples=np.arange(4))
+    controller.supervisor = Bunch(clustering=Clustering(np.array([0, 0, 1, 1])))
+    controller.n_chunks_kept = 2
+
+    class ReleasedSpikeSelector:
+        def __init__(
+            self,
+            get_spikes_per_cluster=None,
+            spike_times=None,
+            chunk_bounds=None,
+            n_chunks_kept=None,
+        ):
+            self.spikes_are_disjoint = None
+
+    with patch('phy.apps.base.SpikeSelector', ReleasedSpikeSelector):
+        controller._set_selector()
+    assert controller.selector.spikes_are_disjoint is None
+
+    class NewerSpikeSelector(ReleasedSpikeSelector):
+        def __init__(self, *args, spikes_are_disjoint=False, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.spikes_are_disjoint = spikes_are_disjoint
+
+    with patch('phy.apps.base.SpikeSelector', NewerSpikeSelector):
+        controller._set_selector()
+    assert controller.selector.spikes_are_disjoint is True
+
+
+def test_select_spikes_evenly_supports_released_and_newer_phylib():
+    spikes = {0: np.arange(5), 1: np.arange(5, 10)}
+
+    def released_selector(n_spikes, cluster_ids, subset_chunks=False):
+        assert subset_chunks
+        assert n_spikes is None
+        return np.concatenate([spikes[cluster_id] for cluster_id in cluster_ids])
+
+    np.testing.assert_array_equal(
+        _select_spikes_evenly(released_selector, 2, [0, 1], subset_chunks=True),
+        [0, 4, 5, 9],
+    )
+
+    def newer_selector(n_spikes, cluster_ids, subset_chunks=False, sample_evenly=False):
+        assert (n_spikes, cluster_ids, subset_chunks, sample_evenly) == (
+            2,
+            [0, 1],
+            True,
+            True,
+        )
+        return np.array([1, 8])
+
+    np.testing.assert_array_equal(
+        _select_spikes_evenly(newer_selector, 2, [0, 1], subset_chunks=True),
+        [1, 8],
+    )
+
+
 def test_get_firing_rate_fast_path():
     spike_times = np.array([0.1, 0.2, 0.4, 0.8, 1.6, 3.2])
     clustering = Clustering(np.array([0, 1, 1, 2, 2, 3]))