Improve and benchmark auto-merging functions

src/spikeinterface/curation/curation_tools.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 from typing import Optional
 import numpy as np
+from spikeinterface import SortingAnalyzer
 
 
 try:
@@ -133,3 +134,44 @@ def find_duplicated_spikes(
         return _find_duplicated_spikes_keep_last_iterative(spike_train.astype(np.int64), censored_period)
     else:
         raise ValueError(f"Method '{method}' isn't a valid method for find_duplicated_spikes. Use one of {_methods}")
+
+
+def remove_empty_units(sorting_or_sorting_analyzer, minimum_spikes=10):
+    if isinstance(sorting_or_sorting_analyzer, SortingAnalyzer):
+        sorting = sorting_or_sorting_analyzer.sorting
+        counts = sorting.get_total_num_spikes()
+        ids_to_select = []
+        for id, num_spikes in counts.items():
+            if num_spikes >= minimum_spikes:
+                ids_to_select += [id]
+        return sorting_or_sorting_analyzer.select_units(ids_to_select)
+    else:
+        counts = sorting_or_sorting_analyzer.get_total_num_spikes()
+        ids_to_select = []
+        for id, num_spikes in counts.items():
+            if num_spikes >= minimum_spikes:
+                ids_to_select += [id]
+        return sorting_or_sorting_analyzer.select_units(ids_to_select)
+
+
+def resolve_merging_graph(sorting, potential_merges):
+    """
+    Function to provide, given a list of potential_merges, a resolved merging
+    graph based on the connected components.
+    """
+    from scipy.sparse.csgraph import connected_components
+    from scipy.sparse import lil_matrix
+
+    n = len(sorting.unit_ids)
+    graph = lil_matrix((n, n))
+    for i, j in potential_merges:
+        graph[sorting.id_to_index(i), sorting.id_to_index(j)] = 1
+
+    n_components, labels = connected_components(graph, directed=True, connection="weak", return_labels=True)
+    final_merges = []
+    for i in range(n_components):
+        merges = labels == i
+        if merges.sum() > 1:
+            final_merges += [list(sorting.unit_ids[np.flatnonzero(merges)])]
+
+    return final_merges

src/spikeinterface/curation/merge_temporal_splits.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+import numpy as np
+
+
+def presence_distance(sorting, unit1, unit2, bin_duration_s=2, bins=None, num_samples=None):
+    """
+    Compute the presence distance between two units.
+
+    The presence distance is defined as the Wasserstein distance between the two histograms of
+    the firing activity over time.
+
+    Parameters
+    ----------
+    sorting : Sorting
+        The sorting object.
+    unit1 : int or str
+        The id of the first unit.
+    unit2 : int or str
+        The id of the second unit.
+    bin_duration_s : float
+        The duration of the bin in seconds.
+    bins : array-like
+        The bins used to compute the firing rate.
+    num_samples : list | int | None, default: None
+        The number of samples for each segment. Required if the sorting doesn't have a recording
+        attached.
+
+    Returns
+    -------
+    d : float
+        The presence distance between the two units.
+    """
+    import scipy
+
+    distances = []
+    if num_samples is not None:
+        if isinstance(num_samples, int):
+            num_samples = [num_samples]
+
+    if not sorting.has_recording():
+        if num_samples is None:
+            raise ValueError("num_samples must be provided if sorting has no recording")
+        if len(num_samples) != sorting.get_num_segments():
+            raise ValueError("num_samples must have the same length as the number of segments")
+
+    for segment_index in range(sorting.get_num_segments()):
+        if bins is None:
+            bin_size = bin_duration_s * sorting.sampling_frequency
+            if sorting.has_recording():
+                ns = sorting.get_num_samples(segment_index)
+            else:
+                ns = num_samples[segment_index]
+            bins = np.arange(0, ns, bin_size)
+
+        st1 = sorting.get_unit_spike_train(unit_id=unit1)
+        st2 = sorting.get_unit_spike_train(unit_id=unit2)
+
+        h1, _ = np.histogram(st1, bins)
+        h1 = h1.astype(float)
+
+        h2, _ = np.histogram(st2, bins)
+        h2 = h2.astype(float)
+
+        xaxis = bins[1:] / sorting.sampling_frequency
+        d = scipy.stats.wasserstein_distance(xaxis, xaxis, h1, h2)
+        distances.append(d)
+
+    return np.mean(d)
+
+
+def compute_presence_distance(sorting, pair_mask, num_samples=None, **presence_distance_kwargs):
+    """
+    Get the potential drift-related merges based on similarity and presence completeness.
+
+    Parameters
+    ----------
+    sorting : Sorting
+        The sorting object
+    pair_mask : None or boolean array
+        A bool matrix of size (num_units, num_units) to select
+        which pair to compute.
+    num_samples : list | int | None, default: None
+        The number of samples for each segment. Required if the sorting doesn't have a recording
+        attached.
+    presence_distance_threshold : float
+        The presence distance threshold used to consider two units as similar
+    presence_distance_kwargs : A dictionary of kwargs to be passed to compute_presence_distance().
+
+    Returns
+    -------
+    potential_merges : list
+        The list of potential merges
+
+    """
+
+    unit_ids = sorting.unit_ids
+    n = len(unit_ids)
+
+    if pair_mask is None:
+        pair_mask = np.ones((n, n), dtype="bool")
+
+    presence_distances = np.ones((sorting.get_num_units(), sorting.get_num_units()))
+
+    for unit_ind1 in range(n):
+        for unit_ind2 in range(unit_ind1 + 1, n):
+            if not pair_mask[unit_ind1, unit_ind2]:
+                continue
+            unit1 = unit_ids[unit_ind1]
+            unit2 = unit_ids[unit_ind2]
+            d = presence_distance(sorting, unit1, unit2, num_samples=num_samples, **presence_distance_kwargs)
+            presence_distances[unit_ind1, unit_ind2] = d
+
+    return presence_distances

src/spikeinterface/curation/tests/test_auto_merge.py

@@ -12,8 +12,10 @@
 from spikeinterface.curation.tests.common import make_sorting_analyzer, sorting_analyzer_for_curation
 
 
-def test_get_auto_merge_list(sorting_analyzer_for_curation):
+@pytest.mark.parametrize("preset", ["lussac", "knn", "temporal_splits", None])
+def test_get_auto_merge_list(sorting_analyzer_for_curation, preset):
 
+    print(sorting_analyzer_for_curation)
     sorting = sorting_analyzer_for_curation.sorting
     recording = sorting_analyzer_for_curation.recording
     num_unit_splited = 1
@@ -34,32 +36,43 @@ def test_get_auto_merge_list(sorting_analyzer_for_curation):
     sorting_analyzer.compute("random_spikes")
     sorting_analyzer.compute("waveforms", **job_kwargs)
     sorting_analyzer.compute("templates")
-
-    potential_merges, outs = get_potential_auto_merge(
-        sorting_analyzer,
-        minimum_spikes=1000,
-        maximum_distance_um=150.0,
-        peak_sign="neg",
-        bin_ms=0.25,
-        window_ms=100.0,
-        corr_diff_thresh=0.16,
-        template_diff_thresh=0.25,
-        censored_period_ms=0.0,
-        refractory_period_ms=4.0,
-        sigma_smooth_ms=0.6,
-        contamination_threshold=0.2,
-        adaptative_window_threshold=0.5,
-        num_channels=5,
-        num_shift=5,
-        firing_contamination_balance=1.5,
-        extra_outputs=True,
-    )
-
-    assert len(potential_merges) == num_unit_splited
-    for true_pair in other_ids.values():
-        true_pair = tuple(true_pair)
-        assert true_pair in potential_merges
-
+    sorting_analyzer.compute(["spike_amplitudes", "spike_locations"])
+
+    if preset is not None:
+        potential_merges, outs = get_potential_auto_merge(
+            sorting_analyzer,
+            preset=preset,
+            minimum_spikes=1000,
+            maximum_distance_um=150.0,
+            peak_sign="neg",
+            bin_ms=0.25,
+            window_ms=100.0,
+            corr_diff_thresh=0.16,
+            template_diff_thresh=0.25,
+            censored_period_ms=0.0,
+            refractory_period_ms=4.0,
+            sigma_smooth_ms=0.6,
+            contamination_threshold=0.2,
+            adaptative_window_threshold=0.5,
+            num_channels=5,
+            num_shift=5,
+            firing_contamination_balance=1.5,
+            extra_outputs=True,
+        )
+        if preset == "lussac":
+            assert len(potential_merges) == num_unit_splited
+            for true_pair in other_ids.values():
+                true_pair = tuple(true_pair)
+                assert true_pair in potential_merges
+    else:
+        # when preset is None you have to specify the steps
+        with pytest.raises(ValueError):
+            potential_merges = get_potential_auto_merge(sorting_analyzer, preset=preset)
+        potential_merges = get_potential_auto_merge(
+            sorting_analyzer, preset=preset, steps=["min_spikes", "min_snr", "remove_contaminated", "unit_positions"]
+        )
+
+    # DEBUG
     # import matplotlib.pyplot as plt
     # from spikeinterface.curation.auto_merge import normalize_correlogram
     # templates_diff = outs['templates_diff']

src/spikeinterface/generation/drift_tools.py

@@ -577,3 +577,104 @@ def get_traces(
 
     def get_num_samples(self) -> int:
         return self.num_samples
+
+
+def split_sorting_by_times(
+    sorting_analyzer, splitting_probability=0.5, partial_split_prob=0.95, unit_ids=None, min_snr=None, seed=None
+):
+    sa = sorting_analyzer
+    sorting = sa.sorting
+    rng = np.random.RandomState(seed)
+
+    sorting_split = sorting.select_units(sorting.unit_ids)
+    split_units = []
+    original_units = []
+    nb_splits = int(splitting_probability * len(sorting.unit_ids))
+    if unit_ids is None:
+        select_from = sorting.unit_ids
+        if min_snr is not None:
+            if sa.get_extension("noise_levels") is None:
+                sa.compute("noise_levels")
+            if sa.get_extension("quality_metrics") is None:
+                sa.compute("quality_metrics", metric_names=["snr"])
+
+            snr = sa.get_extension("quality_metrics").get_data()["snr"].values
+            select_from = select_from[snr > min_snr]
+
+        to_split_ids = rng.choice(select_from, nb_splits, replace=False)
+    else:
+        to_split_ids = unit_ids
+
+    import spikeinterface.curation as scur
+
+    for unit in to_split_ids:
+        num_spikes = len(sorting_split.get_unit_spike_train(unit))
+        indices = np.zeros(num_spikes, dtype=int)
+        indices[: num_spikes // 2] = (rng.rand(num_spikes // 2) < partial_split_prob).astype(int)
+        indices[num_spikes // 2 :] = (rng.rand(num_spikes - num_spikes // 2) < 1 - partial_split_prob).astype(int)
+        sorting_split = scur.split_unit_sorting(
+            sorting_split, split_unit_id=unit, indices_list=indices, properties_policy="remove"
+        )
+        split_units.append(sorting_split.unit_ids[-2:])
+        original_units.append(unit)
+    return sorting_split, split_units
+
+
+def split_sorting_by_amplitudes(sorting_analyzer, splitting_probability=0.5, unit_ids=None, min_snr=None, seed=None):
+    """
+    Fonction used to split a sorting based on the amplitudes of the units. This
+    might be used for benchmarking meta merging step (see components)
+    """
+
+    sa = sorting_analyzer
+    if sa.get_extension("spike_amplitudes") is None:
+        sa.compute("spike_amplitudes")
+
+    rng = np.random.RandomState(seed)
+
+    from spikeinterface.core.numpyextractors import NumpySorting
+    from spikeinterface.core.template_tools import get_template_extremum_channel
+
+    extremum_channel_inds = get_template_extremum_channel(sa, outputs="index")
+    spikes = sa.sorting.to_spike_vector(extremum_channel_inds=extremum_channel_inds)
+    new_spikes = spikes.copy()
+    amplitudes = sa.get_extension("spike_amplitudes").get_data()
+    nb_splits = int(splitting_probability * len(sa.sorting.unit_ids))
+
+    if unit_ids is None:
+        select_from = sa.sorting.unit_ids
+        if min_snr is not None:
+            if sa.get_extension("noise_levels") is None:
+                sa.compute("noise_levels")
+            if sa.get_extension("quality_metrics") is None:
+                sa.compute("quality_metrics", metric_names=["snr"])
+
+            snr = sa.get_extension("quality_metrics").get_data()["snr"].values
+            select_from = select_from[snr > min_snr]
+        to_split_ids = rng.choice(select_from, nb_splits, replace=False)
+    else:
+        to_split_ids = unit_ids
+
+    max_index = np.max(spikes["unit_index"])
+    new_unit_ids = list(sa.sorting.unit_ids.copy())
+    splitted_pairs = []
+    for unit_id in to_split_ids:
+        ind_mask = spikes["unit_index"] == sa.sorting.id_to_index(unit_id)
+
+        m = amplitudes[ind_mask].mean()
+        s = amplitudes[ind_mask].std()
+        thresh = m + 0.2 * s
+
+        amplitude_mask = amplitudes > thresh
+        mask = ind_mask & amplitude_mask
+        new_spikes["unit_index"][mask] = max_index + 1
+
+        amplitude_mask = (amplitudes > m) * (amplitudes < thresh)
+        mask = ind_mask & amplitude_mask
+        new_spikes["unit_index"][mask] = (max_index + 1) * rng.rand(np.sum(mask)) > 0.5
+        max_index += 1
+        new_unit_ids += [max(new_unit_ids) + 1]
+        splitted_pairs += [(unit_id, new_unit_ids[-1])]
+
+    new_sorting = NumpySorting(new_spikes, sampling_frequency=sa.sampling_frequency, unit_ids=new_unit_ids)
+    return new_sorting, splitted_pairs