Valid unit periods extension

src/spikeinterface/core/analyzer_extension_core.py

@@ -13,6 +13,7 @@
 import numpy as np
 from collections import namedtuple
 
+from .numpyextractors import NumpySorting
 from .sortinganalyzer import SortingAnalyzer, AnalyzerExtension, register_result_extension
 from .waveform_tools import extract_waveforms_to_single_buffer, estimate_templates_with_accumulator
 from .recording_tools import get_noise_levels
@@ -823,10 +824,9 @@ class BaseMetric:
     metric_columns = {}  # column names and their dtypes of the dataframe
     metric_descriptions = {}  # descriptions of each metric column
     needs_recording = False  # whether the metric needs recording
-    needs_tmp_data = (
-        False  # whether the metric needs temporary data comoputed with _prepare_data at the MetricExtension level
-    )
-    needs_job_kwargs = False
+    needs_tmp_data = False  # whether the metric needs temporary data computed with MetricExtension._prepare_data
+    needs_job_kwargs = False  # whether the metric needs job_kwargs
+    supports_periods = False  # whether the metric function supports periods
     depend_on = []  # extensions the metric depends on
 
     # the metric function must have the signature:
@@ -839,7 +839,7 @@ class BaseMetric:
     metric_function = None  # to be defined in subclass
 
     @classmethod
-    def compute(cls, sorting_analyzer, unit_ids, metric_params, tmp_data, job_kwargs):
+    def compute(cls, sorting_analyzer, unit_ids, metric_params, tmp_data, job_kwargs, periods=None):
         """Compute the metric.
 
         Parameters
@@ -854,6 +854,8 @@ def compute(cls, sorting_analyzer, unit_ids, metric_params, tmp_data, job_kwargs
             Temporary data to pass to the metric function
         job_kwargs : dict
             Job keyword arguments to control parallelization
+        periods : np.ndarray | None
+            Numpy array of unit periods of unit_period_dtype if supports_periods is True
 
         Returns
         -------
@@ -865,6 +867,8 @@ def compute(cls, sorting_analyzer, unit_ids, metric_params, tmp_data, job_kwargs
             args += (tmp_data,)
         if cls.needs_job_kwargs:
             args += (job_kwargs,)
+        if cls.supports_periods:
+            args += (periods,)
 
         results = cls.metric_function(*args, **metric_params)
 
@@ -897,6 +901,17 @@ class BaseMetricExtension(AnalyzerExtension):
     need_backward_compatibility_on_load = False
     metric_list: list[BaseMetric] = None  # list of BaseMetric
 
+    @classmethod
+    def get_available_metric_names(cls):
+        """Get the available metric names.
+
+        Returns
+        -------
+        available_metric_names : list[str]
+            List of available metric names.
+        """
+        return [m.metric_name for m in cls.metric_list]
+
     @classmethod
     def get_default_metric_params(cls):
         """Get the default metric parameters.
@@ -988,6 +1003,7 @@ def _set_params(
         metric_params: dict | None = None,
         delete_existing_metrics: bool = False,
         metrics_to_compute: list[str] | None = None,
+        periods: np.ndarray | None = None,
         **other_params,
     ):
         """
@@ -1004,6 +1020,8 @@ def _set_params(
             If True, existing metrics in the extension will be deleted before computing new ones.
         metrics_to_compute : list[str] | None
             List of metric names to compute. If None, all metrics in `metric_names` are computed.
+        periods : np.ndarray | None
+            Numpy array of unit_period_dtype defining periods to compute metrics over.
         other_params : dict
             Additional parameters for metric computation.
 
@@ -1079,6 +1097,7 @@ def _set_params(
             metrics_to_compute=metrics_to_compute,
             delete_existing_metrics=delete_existing_metrics,
             metric_params=metric_params,
+            periods=periods,
             **other_params,
         )
         return params
@@ -1129,6 +1148,8 @@ def _compute_metrics(
         if metric_names is None:
             metric_names = self.params["metric_names"]
 
+        periods = self.params.get("periods", None)
+
         column_names_dtypes = {}
         for metric_name in metric_names:
             metric = [m for m in self.metric_list if m.metric_name == metric_name][0]
@@ -1153,6 +1174,7 @@ def _compute_metrics(
                     metric_params=metric_params,
                     tmp_data=tmp_data,
                     job_kwargs=job_kwargs,
+                    periods=periods,
                 )
             except Exception as e:
                 warnings.warn(f"Error computing metric {metric_name}: {e}")
@@ -1179,6 +1201,7 @@ def _run(self, **job_kwargs):
 
         metrics_to_compute = self.params["metrics_to_compute"]
         delete_existing_metrics = self.params["delete_existing_metrics"]
+        periods = self.params.get("periods", None)
 
         _, job_kwargs = split_job_kwargs(job_kwargs)
         job_kwargs = fix_job_kwargs(job_kwargs)
@@ -1445,13 +1468,22 @@ def _get_data(self, outputs="numpy", concatenated=False, return_data_name=None,
                 ), f"return_data_name {return_data_name} not in nodepipeline_variables {self.nodepipeline_variables}"
 
         all_data = self.data[return_data_name]
-        keep_mask = None
         if periods is not None:
             keep_mask = select_sorting_periods_mask(
                 self.sorting_analyzer.sorting,
                 periods,
             )
             all_data = all_data[keep_mask]
+            # since we have the mask already, we can use it directly to avoid double computation
+            spike_vector = self.sorting_analyzer.sorting.to_spike_vector(concatenated=True)
+            sliced_spike_vector = spike_vector[keep_mask]
+            sorting = NumpySorting(
+                sliced_spike_vector,
+                sampling_frequency=self.sorting_analyzer.sampling_frequency,
+                unit_ids=self.sorting_analyzer.unit_ids,
+            )
+        else:
+            sorting = self.sorting_analyzer.sorting
 
         if outputs == "numpy":
             if copy:
@@ -1460,16 +1492,16 @@ def _get_data(self, outputs="numpy", concatenated=False, return_data_name=None,
                 return all_data
         elif outputs == "by_unit":
             unit_ids = self.sorting_analyzer.unit_ids
-            spike_vector = self.sorting_analyzer.sorting.to_spike_vector(concatenated=False)
+
             if keep_mask is not None:
                 # since we are filtering spikes, we need to recompute the spike indices
-                spike_vector = spike_vector[keep_mask]
+                spike_vector = sorting.to_spike_vector(concatenated=False)
                 spike_indices = spike_vector_to_indices(spike_vector, unit_ids, absolute_index=True)
             else:
                 # use the cache of indices
                 spike_indices = self.sorting_analyzer.sorting.get_spike_vector_to_indices()
             data_by_units = {}
-            for segment_index in range(self.sorting_analyzer.sorting.get_num_segments()):
+            for segment_index in range(self.sorting_analyzer.get_num_segments()):
                 data_by_units[segment_index] = {}
                 for unit_id in unit_ids:
                     inds = spike_indices[segment_index][unit_id]

src/spikeinterface/core/sortinganalyzer.py

@@ -2702,6 +2702,7 @@ def _save_data(self):
                     extension_group.create_dataset(
                         name=ext_data_name, data=np.array([ext_data], dtype=object), object_codec=numcodecs.JSON()
                     )
+                    extension_group[ext_data_name].attrs["dict"] = True
                 elif isinstance(ext_data, np.ndarray):
                     extension_group.create_dataset(name=ext_data_name, data=ext_data, **saving_options)
                 elif HAS_PANDAS and isinstance(ext_data, pd.DataFrame):
@@ -2884,6 +2885,7 @@ def set_data(self, ext_data_name, ext_data):
     "spike_locations": "spikeinterface.postprocessing",
     "template_similarity": "spikeinterface.postprocessing",
     "unit_locations": "spikeinterface.postprocessing",
+    "valid_unit_periods": "spikeinterface.postprocessing",
     # from metrics
     "quality_metrics": "spikeinterface.metrics",
     "template_metrics": "spikeinterface.metrics",

src/spikeinterface/metrics/conftest.py

@@ -1,8 +1,85 @@
 import pytest
 
-from spikeinterface.postprocessing.tests.conftest import _small_sorting_analyzer
+from spikeinterface.core import (
+    generate_ground_truth_recording,
+    create_sorting_analyzer,
+)
+
+job_kwargs = dict(n_jobs=2, progress_bar=True, chunk_duration="1s")
+
+
+def make_small_analyzer():
+    recording, sorting = generate_ground_truth_recording(
+        durations=[10.0],
+        num_units=10,
+        seed=1205,
+    )
+
+    channel_ids_as_integers = [id for id in range(recording.get_num_channels())]
+    unit_ids_as_integers = [id for id in range(sorting.get_num_units())]
+    recording = recording.rename_channels(new_channel_ids=channel_ids_as_integers)
+    sorting = sorting.rename_units(new_unit_ids=unit_ids_as_integers)
+
+    sorting = sorting.select_units([2, 7, 0], ["#3", "#9", "#4"])
+
+    sorting_analyzer = create_sorting_analyzer(recording=recording, sorting=sorting, format="memory")
+
+    extensions_to_compute = {
+        "random_spikes": {"seed": 1205},
+        "noise_levels": {"seed": 1205},
+        "waveforms": {},
+        "templates": {"operators": ["average", "median"]},
+        "spike_amplitudes": {},
+        "spike_locations": {},
+        "principal_components": {},
+    }
+
+    sorting_analyzer.compute(extensions_to_compute)
+
+    return sorting_analyzer
 
 
 @pytest.fixture(scope="module")
 def small_sorting_analyzer():
-    return _small_sorting_analyzer()
+    return make_small_analyzer()
+
+
+@pytest.fixture(scope="module")
+def sorting_analyzer_simple():
+    # we need high firing rate for amplitude_cutoff
+    recording, sorting = generate_ground_truth_recording(
+        durations=[
+            120.0,
+        ],
+        sampling_frequency=30_000.0,
+        num_channels=6,
+        num_units=10,
+        generate_sorting_kwargs=dict(firing_rates=10.0, refractory_period_ms=4.0),
+        generate_unit_locations_kwargs=dict(
+            margin_um=5.0,
+            minimum_z=5.0,
+            maximum_z=20.0,
+        ),
+        generate_templates_kwargs=dict(
+            unit_params=dict(
+                alpha=(200.0, 500.0),
+            )
+        ),
+        noise_kwargs=dict(noise_levels=5.0, strategy="tile_pregenerated"),
+        seed=1205,
+    )
+
+    channel_ids_as_integers = [id for id in range(recording.get_num_channels())]
+    unit_ids_as_integers = [id for id in range(sorting.get_num_units())]
+    recording = recording.rename_channels(new_channel_ids=channel_ids_as_integers)
+    sorting = sorting.rename_units(new_unit_ids=unit_ids_as_integers)
+
+    sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=True)
+
+    sorting_analyzer.compute("random_spikes", max_spikes_per_unit=300, seed=1205)
+    sorting_analyzer.compute("noise_levels")
+    sorting_analyzer.compute("waveforms", **job_kwargs)
+    sorting_analyzer.compute("templates")
+    sorting_analyzer.compute(["spike_amplitudes", "spike_locations"], **job_kwargs)
+
+    return sorting_analyzer

src/spikeinterface/metrics/quality/__init__.py

@@ -20,4 +20,5 @@
     compute_sliding_rp_violations,
     compute_sd_ratio,
     compute_synchrony_metrics,
+    compute_refrac_period_violations,
 )