Merge branch 'select_sorting_periods' of github.com:alejoe91/spikeinterface into select_sorting_periods

Author: alejoe91

src/spikeinterface/core/analyzer_extension_core.py

@@ -901,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.

src/spikeinterface/metrics/quality/misc_metrics.py

@@ -1129,13 +1129,20 @@ def compute_drift_metrics(
         unit_ids = sorting.unit_ids
 
     spike_locations_ext = sorting_analyzer.get_extension("spike_locations")
-    spike_locations_array = spike_locations_ext.get_data(periods=periods)
+    spike_locations_by_unit_and_segments = spike_locations_ext.get_data(
+        outputs="by_unit", concatenated=False, periods=periods
+    )
     spike_locations_by_unit = spike_locations_ext.get_data(outputs="by_unit", concatenated=True, periods=periods)
 
     segment_samples = [sorting_analyzer.get_num_samples(i) for i in range(sorting_analyzer.get_num_segments())]
-    assert direction in spike_locations_array.dtype.names, (
-        f"Direction {direction} is invalid. Available directions: " f"{spike_locations_array.dtype.names}"
+    data = spike_locations_by_unit[unit_ids[0]]
+    assert direction in data.dtype.names, (
+        f"Direction {direction} is invalid. Available directions: " f"{data.dtype.names}"
+    )
+    bin_edges_for_units = compute_bin_edges_per_unit(
+        sorting, segment_samples=segment_samples, periods=periods, bin_duration_s=interval_s, concatenated=False
     )
+    failed_units = []
 
     # we need
     drift_ptps = {}
@@ -1144,62 +1151,43 @@ def compute_drift_metrics(
 
     # reference positions are the medians across segments
     reference_positions = {}
+    median_position_segments = {unit_id: np.array([]) for unit_id in unit_ids}
+
     for unit_id in unit_ids:
         reference_positions[unit_id] = np.median(spike_locations_by_unit[unit_id][direction])
 
-    # now compute median positions and concatenate them over segments
-    spike_vector = sorting.to_spike_vector()
-    spike_sample_indices = spike_vector["sample_index"].copy()
-    # we need to add the cumulative sum of segment samples to have global sample indices
-    cumulative_segment_samples = np.cumsum([0] + segment_samples[:-1])
     for segment_index in range(sorting_analyzer.get_num_segments()):
-        segment_slice = sorting._get_spike_vector_segment_slices()[segment_index]
-        spike_sample_indices[segment_slice[0] : segment_slice[1]] += cumulative_segment_samples[segment_index]
-
-    bin_edges_for_units = compute_bin_edges_per_unit(
-        sorting,
-        segment_samples=segment_samples,
-        periods=periods,
-        bin_duration_s=interval_s,
-    )
-
-    failed_units = []
-    median_positions_per_unit = {}
+        for unit_id in unit_ids:
+            bins = bin_edges_for_units[unit_id][segment_index]
+            num_bin_edges = len(bins)
+            if (num_bin_edges - 1) < min_num_bins:
+                failed_units.append(unit_id)
+                continue
+            median_positions = np.nan * np.zeros((num_bin_edges - 1))
+            spikes_in_segment_of_unit = sorting.get_unit_spike_train(unit_id, segment_index)
+            bounds = np.searchsorted(spikes_in_segment_of_unit, bins, side="left")
+            for bin_index, (i0, i1) in enumerate(zip(bounds[:-1], bounds[1:])):
+                spike_locations_in_bin = spike_locations_by_unit_and_segments[segment_index][unit_id][i0:i1][direction]
+                if (i1 - i0) >= min_spikes_per_interval:
+                    median_positions[bin_index] = np.median(spike_locations_in_bin)
+            median_position_segments[unit_id] = np.concatenate((median_position_segments[unit_id], median_positions))
+
+    # finally, compute deviations and drifts
     for unit_id in unit_ids:
-        bins = bin_edges_for_units[unit_id]
-        num_bins = len(bins) - 1
-        if num_bins < min_num_bins:
+        # Skip units that already failed because not enough bins in at least one segment
+        if unit_id in failed_units:
             drift_ptps[unit_id] = np.nan
             drift_stds[unit_id] = np.nan
             drift_mads[unit_id] = np.nan
-            failed_units.append(unit_id)
             continue
-
-        # bin_edges are global across segments, so we have to use spike_sample_indices,
-        # since we offseted them to be global
-        bin_spike_indices = np.searchsorted(spike_sample_indices, bins)
-        median_positions = np.nan * np.zeros(num_bins)
-        for bin_index, (i0, i1) in enumerate(zip(bin_spike_indices[:-1], bin_spike_indices[1:])):
-            spikes_in_bin = spike_vector[i0:i1]
-            spike_locations_in_bin = spike_locations_array[i0:i1][direction]
-
-            unit_index = sorting_analyzer.sorting.id_to_index(unit_id)
-            mask = spikes_in_bin["unit_index"] == unit_index
-            if np.sum(mask) >= min_spikes_per_interval:
-                median_positions[bin_index] = np.median(spike_locations_in_bin[mask])
-            else:
-                median_positions[bin_index] = np.nan
-        median_positions_per_unit[unit_id] = median_positions
-
-        # now compute deviations and drifts for this unit
-        position_diff = median_positions - reference_positions[unit_id]
+        position_diff = median_position_segments[unit_id] - reference_positions[unit_id]
         if np.any(np.isnan(position_diff)):
             # deal with nans: if more than 50% nans --> set to nan
             if np.sum(np.isnan(position_diff)) > min_fraction_valid_intervals * len(position_diff):
-                failed_units.append(unit_id)
                 ptp_drift = np.nan
                 std_drift = np.nan
                 mad_drift = np.nan
+                failed_units.append(unit_id)
             else:
                 ptp_drift = np.nanmax(position_diff) - np.nanmin(position_diff)
                 std_drift = np.nanstd(np.abs(position_diff))
@@ -1219,7 +1207,7 @@ def compute_drift_metrics(
         )
 
     if return_positions:
-        outs = res(drift_ptps, drift_stds, drift_mads), median_positions_per_unit
+        outs = res(drift_ptps, drift_stds, drift_mads), median_positions
     else:
         outs = res(drift_ptps, drift_stds, drift_mads)
     return outs

src/spikeinterface/metrics/quality/tests/test_metrics_functions.py

@@ -13,10 +13,7 @@
 
 from spikeinterface.metrics.utils import create_ground_truth_pc_distributions, create_regular_periods
 
-from spikeinterface.metrics.quality import (
-    get_quality_metric_list,
-    compute_quality_metrics,
-)
+from spikeinterface.metrics.quality import get_quality_metric_list, compute_quality_metrics, ComputeQualityMetrics
 from spikeinterface.metrics.quality.misc_metrics import (
     misc_metrics_list,
     compute_amplitude_cutoffs,
@@ -657,37 +654,9 @@ def test_save_quality_metrics(small_sorting_analyzer, create_cache_folder):
 
     # can't use _misc_metric_name_to_func as some functions compute several qms
     # e.g. isi_violation and synchrony
-    quality_metrics = [
-        "num_spikes",
-        "firing_rate",
-        "presence_ratio",
-        "snr",
-        "isi_violations_ratio",
-        "isi_violations_count",
-        "rp_contamination",
-        "rp_violations",
-        "sliding_rp_violation",
-        "amplitude_cutoff",
-        "amplitude_median",
-        "amplitude_cv_median",
-        "amplitude_cv_range",
-        "sync_spike_2",
-        "sync_spike_4",
-        "sync_spike_8",
-        "firing_range",
-        "drift_ptp",
-        "drift_std",
-        "drift_mad",
-        "sd_ratio",
-        "isolation_distance",
-        "l_ratio",
-        "d_prime",
-        "silhouette",
-        "nn_hit_rate",
-        "nn_miss_rate",
-    ]
-
-    small_sorting_analyzer.compute("quality_metrics")
+    quality_metric_columns = ComputeQualityMetrics.get_metric_columns()
+    all_metrics = ComputeQualityMetrics.get_available_metric_names()
+    small_sorting_analyzer.compute("quality_metrics", metric_names=all_metrics)
 
     cache_folder = create_cache_folder
     output_folder = cache_folder / "sorting_analyzer"
@@ -699,7 +668,7 @@ def test_save_quality_metrics(small_sorting_analyzer, create_cache_folder):
         saved_metrics = csv.reader(metrics_file)
         metric_names = next(saved_metrics)
 
-    for metric_name in quality_metrics:
+    for metric_name in quality_metric_columns:
         assert metric_name in metric_names
 
     folder_analyzer.compute("quality_metrics", metric_names=["snr"], delete_existing_metrics=False)
@@ -708,7 +677,7 @@ def test_save_quality_metrics(small_sorting_analyzer, create_cache_folder):
         saved_metrics = csv.reader(metrics_file)
         metric_names = next(saved_metrics)
 
-    for metric_name in quality_metrics:
+    for metric_name in quality_metric_columns:
         assert metric_name in metric_names
 
     folder_analyzer.compute("quality_metrics", metric_names=["snr"], delete_existing_metrics=True)
@@ -717,7 +686,7 @@ def test_save_quality_metrics(small_sorting_analyzer, create_cache_folder):
         saved_metrics = csv.reader(metrics_file)
         metric_names = next(saved_metrics)
 
-    for metric_name in quality_metrics:
+    for metric_name in quality_metric_columns:
         if metric_name == "snr":
             assert metric_name in metric_names
         else:

src/spikeinterface/metrics/utils.py

@@ -4,7 +4,7 @@
 from spikeinterface.core.base import unit_period_dtype
 
 
-def compute_bin_edges_per_unit(sorting, segment_samples, bin_duration_s=1.0, periods=None):
+def compute_bin_edges_per_unit(sorting, segment_samples, bin_duration_s=1.0, periods=None, concatenated=True):
     """
     Compute bin edges for units, optionally taking into account periods.
 
@@ -18,6 +18,16 @@ def compute_bin_edges_per_unit(sorting, segment_samples, bin_duration_s=1.0, per
         Duration of each bin in seconds
     periods : array of unit_period_dtype, default: None
         Periods to consider for each unit
+    concatenated : bool, default: True
+        Wheter the bins are concatenated across segments or not.
+        If False, the bin edges are computed per segment and the first index of each segment is 0.
+        If True, the bin edges are computed on the concatenated segments, with the correct offsets.
+
+    Returns
+    -------
+    dict
+        Bin edges for each unit. If concatenated is True, the bin edges are a 1D array.
+        If False, the bin edges are a list of arrays, one per segment.
     """
     bin_edges_for_units = {}
     num_segments = len(segment_samples)
@@ -31,27 +41,38 @@ def compute_bin_edges_per_unit(sorting, segment_samples, bin_duration_s=1.0, per
             for seg_index in range(num_segments):
                 seg_periods = periods_unit[periods_unit["segment_index"] == seg_index]
                 if len(seg_periods) == 0:
+                    if not concatenated:
+                        bin_edges.append(np.array([]))
                     continue
-                seg_start = np.sum(segment_samples[:seg_index])
+                seg_start = np.sum(segment_samples[:seg_index]) if concatenated else 0
+                bin_edges_segment = []
                 for period in seg_periods:
                     start_sample = seg_start + period["start_sample_index"]
                     end_sample = seg_start + period["end_sample_index"]
                     end_sample = end_sample // bin_duration_samples * bin_duration_samples + 1  # align to bin
-                    bin_edges.extend(np.arange(start_sample, end_sample, bin_duration_samples))
-            bin_edges_for_units[unit_id] = np.unique(np.array(bin_edges))
+                    bin_edges_segment.extend(np.arange(start_sample, end_sample, bin_duration_samples))
+                bin_edges_segment = np.unique(np.array(bin_edges_segment))
+                if concatenated:
+                    bin_edges.extend(bin_edges_segment)
+                else:
+                    bin_edges.append(bin_edges_segment)
+            bin_edges_for_units[unit_id] = bin_edges
     else:
         for unit_id in sorting.unit_ids:
             bin_edges = []
             for seg_index in range(num_segments):
-                seg_start = np.sum(segment_samples[:seg_index])
+                seg_start = np.sum(segment_samples[:seg_index]) if concatenated else 0
                 seg_end = seg_start + segment_samples[seg_index]
                 # for segments which are not the last, we don't need to correct the end
                 # since the first index of the next segment will be the end of the current segment
                 if seg_index == num_segments - 1:
                     seg_end = seg_end // bin_duration_samples * bin_duration_samples + 1  # align to bin
-                bins = np.arange(seg_start, seg_end, bin_duration_samples)
-                bin_edges.extend(bins)
-            bin_edges_for_units[unit_id] = np.array(bin_edges)
+                bin_edges_segment = np.arange(seg_start, seg_end, bin_duration_samples)
+                if concatenated:
+                    bin_edges.extend(bin_edges_segment)
+                else:
+                    bin_edges.append(bin_edges_segment)
+            bin_edges_for_units[unit_id] = bin_edges
     return bin_edges_for_units
 
 
@@ -82,7 +103,8 @@ def compute_total_samples_per_unit(sorting_analyzer, periods=None):
                 num_samples_in_period += period["end_sample_index"] - period["start_sample_index"]
             total_samples[unit_id] = num_samples_in_period
     else:
-        total_samples = {unit_id: sorting_analyzer.get_total_samples() for unit_id in sorting_analyzer.unit_ids}
+        total = sorting_analyzer.get_total_samples()
+        total_samples = {unit_id: total for unit_id in sorting_analyzer.unit_ids}
     return total_samples