Improve merging and iterative merging (#3487)

Author: yger

doc/modules/curation.rst

@@ -120,31 +120,31 @@ the unit (among the redundant ones), with a better template alignment.
 Auto-merging units
 ^^^^^^^^^^^^^^^^^^
 
-The :py:func:`~spikeinterface.curation.get_potential_auto_merge` function returns a list of potential merges.
+The :py:func:`~spikeinterface.curation.compute_merge_unit_groups` function returns a list of potential merges.
 The list of potential merges can be then applied to the sorting output.
-:py:func:`~spikeinterface.curation.get_potential_auto_merge` has many internal tricks and steps to identify potential
+:py:func:`~spikeinterface.curation.compute_merge_unit_groups` has many internal tricks and steps to identify potential
 merges. It offers multiple "presets" and the flexibility to apply individual steps, with different parameters.
 **Read the function documentation carefully and do not apply it blindly!**
 
 
 .. code-block:: python
 
     from spikeinterface import create_sorting_analyzer
-    from spikeinterface.curation import get_potential_auto_merge
+    from spikeinterface.curation import compute_merge_unit_groups
 
     analyzer = create_sorting_analyzer(sorting=sorting, recording=recording)
 
     # some extensions are required
     analyzer.compute(["random_spikes", "templates", "template_similarity", "correlograms"])
 
     # merges is a list of unit pairs, with unit_ids to be merged.
-    merge_unit_pairs = get_potential_auto_merge(
+    merge_unit_pairs = compute_merge_unit_groups(
         analyzer=analyzer,
         preset="similarity_correlograms",
     )
     # with resolve_graph=True, merges_resolved is a list of merge groups,
     # which can contain more than two units
-    merge_unit_groups = get_potential_auto_merge(
+    merge_unit_groups = compute_merge_unit_groups(
         analyzer=analyzer,
         preset="similarity_correlograms",
         resolve_graph=True
@@ -153,6 +153,44 @@ merges. It offers multiple "presets" and the flexibility to apply individual ste
     # here we apply the merges
     analyzer_merged = analyzer.merge_units(merge_unit_groups=merge_unit_groups)
 
+There is also the convenient :py:func:`~spikeinterface.curation.auto_merge_units` function that combines the
+:py:func:`~spikeinterface.curation.compute_merge_unit_groups` and :py:func:`~spikeinterface.core.SortingAnalyzer.merge_units` functions.
+This is a high level function that allows you to apply either one or several presets/lists of steps in one go. For example, let's
+assume you want to apply the "x_contamination" preset, but iteratively and with slightly different parameters: first,
+you want to focus on the templates that are very similar, according to their template similarities, before
+considering those that might be more distant. Such a greedy and iterative scheme has been proved to be less
+prone to wrong merges. To do so, you'll need to do the following:
+
+.. code-block:: python
+
+    from spikeinterface import create_sorting_analyzer
+    from spikeinterface.curation import auto_merge_units
+
+    analyzer = create_sorting_analyzer(sorting=sorting, recording=recording)
+
+    # some extensions are required
+    analyzer.compute(["random_spikes", "templates", "template_similarity", "correlograms"])
+    analyzer.compute("unit_locations", method="monopolar_triangulation")
+
+    template_diff_thresh = [0.05, 0.15, 0.25]
+    presets = ["x_contaminations"] * len(template_diff_thresh)
+    steps_params = [
+        {"template_similarity": {"template_diff_thresh": i}}
+        for i in template_diff_thresh
+    ]
+
+    analyzer_merged = auto_merge_units(
+        analyzer,
+        presets=presets,
+        steps_params=steps_params,
+        recursive=True,
+        **job_kwargs,
+    )
+
+The extra keyword ``recursive`` specifies that for each presets/sequences of steps, merges are performed
+until no further merges are possible. The ``job_kwargs`` are the parameters for the parallelization.
+**Be careful that the merges can not be reverted, so be sure to not erase your analyzer and create a new variable**
+
 
 Manual curation
 ---------------

src/spikeinterface/benchmark/benchmark_base.py

@@ -10,7 +10,7 @@
 
 
 from spikeinterface.core import SortingAnalyzer
-
+from spikeinterface.core.job_tools import fix_job_kwargs, split_job_kwargs
 from spikeinterface import load, create_sorting_analyzer, load_sorting_analyzer
 from spikeinterface.widgets import get_some_colors
 
@@ -219,7 +219,7 @@ def run(self, case_keys=None, keep=True, verbose=False, **job_kwargs):
         for key in job_keys:
             benchmark = self.create_benchmark(key)
             t0 = time.perf_counter()
-            benchmark.run()
+            benchmark.run(**job_kwargs)
             t1 = time.perf_counter()
             self.benchmarks[key] = benchmark
             bench_folder = self.folder / "results" / self.key_to_str(key)
@@ -264,6 +264,7 @@ def plot_run_times(self, case_keys=None):
         return plot_run_times(self, case_keys=case_keys)
 
     def compute_results(self, case_keys=None, verbose=False, **result_params):
+
         if case_keys is None:
             case_keys = list(self.cases.keys())
 
@@ -309,7 +310,7 @@ def get_templates(self, key, operator="average"):
         templates = sorting_analyzer.get_extenson("templates").get_data(operator=operator)
         return templates
 
-    def compute_metrics(self, case_keys=None, metric_names=["snr", "firing_rate"], force=False):
+    def compute_metrics(self, case_keys=None, metric_names=["snr", "firing_rate"], force=False, **job_kwargs):
         if case_keys is None:
             case_keys = self.cases.keys()
 
@@ -329,7 +330,7 @@ def compute_metrics(self, case_keys=None, metric_names=["snr", "firing_rate"], f
             sorting_analyzer = self.get_sorting_analyzer(key)
             qm_ext = sorting_analyzer.get_extension("quality_metrics")
             if qm_ext is None or force:
-                qm_ext = sorting_analyzer.compute("quality_metrics", metric_names=metric_names)
+                qm_ext = sorting_analyzer.compute("quality_metrics", metric_names=metric_names, **job_kwargs)
 
             # TODO remove this metics CSV file!!!!
             metrics = qm_ext.get_data()

src/spikeinterface/benchmark/benchmark_clustering.py

@@ -10,7 +10,7 @@
 
 
 import numpy as np
-
+from spikeinterface.core.job_tools import fix_job_kwargs, split_job_kwargs
 from .benchmark_base import Benchmark, BenchmarkStudy
 from spikeinterface.core.sortinganalyzer import create_sorting_analyzer
 from spikeinterface.core.template_tools import get_template_extremum_channel
@@ -36,6 +36,8 @@ def run(self, **job_kwargs):
         self.result["peak_labels"] = peak_labels
 
     def compute_result(self, **result_params):
+        result_params, job_kwargs = split_job_kwargs(result_params)
+        job_kwargs = fix_job_kwargs(job_kwargs)
         self.noise = self.result["peak_labels"] < 0
         spikes = self.gt_sorting.to_spike_vector()
         self.result["sliced_gt_sorting"] = NumpySorting(
@@ -47,8 +49,11 @@ def compute_result(self, **result_params):
         gt_unit_locations = self.gt_sorting.get_property("gt_unit_locations")
         if gt_unit_locations is None:
             print("'gt_unit_locations' is not a property of the sorting so compute it")
-            gt_analyzer = create_sorting_analyzer(self.gt_sorting, self.recording, format="memory", sparse=True)
-            gt_analyzer.compute(["random_spikes", "templates"])
+            gt_analyzer = create_sorting_analyzer(
+                self.gt_sorting, self.recording, format="memory", sparse=True, **job_kwargs
+            )
+            gt_analyzer.compute("random_spikes")
+            gt_analyzer.compute("templates", **job_kwargs)
             ext = gt_analyzer.compute("unit_locations", method="monopolar_triangulation")
             gt_unit_locations = ext.get_data()
             self.gt_sorting.set_property("gt_unit_locations", gt_unit_locations)
@@ -64,17 +69,17 @@ def compute_result(self, **result_params):
         )
 
         sorting_analyzer = create_sorting_analyzer(
-            self.result["sliced_gt_sorting"], self.recording, format="memory", sparse=False
+            self.result["sliced_gt_sorting"], self.recording, format="memory", sparse=False, **job_kwargs
         )
         sorting_analyzer.compute("random_spikes")
-        ext = sorting_analyzer.compute("templates")
+        ext = sorting_analyzer.compute("templates", **job_kwargs)
         self.result["sliced_gt_templates"] = ext.get_data(outputs="Templates")
 
         sorting_analyzer = create_sorting_analyzer(
-            self.result["clustering"], self.recording, format="memory", sparse=False
+            self.result["clustering"], self.recording, format="memory", sparse=False, **job_kwargs
         )
         sorting_analyzer.compute("random_spikes")
-        ext = sorting_analyzer.compute("templates")
+        ext = sorting_analyzer.compute("templates", **job_kwargs)
         self.result["clustering_templates"] = ext.get_data(outputs="Templates")
 
     _run_key_saved = [("peak_labels", "npy")]

src/spikeinterface/benchmark/benchmark_matching.py

@@ -133,10 +133,10 @@ def get_count_units(self, case_keys=None, well_detected_score=None, redundant_sc
 
         return count_units
 
-    def plot_unit_counts(self, case_keys=None, figsize=None):
-        from spikeinterface.widgets.widget_list import plot_study_unit_counts
+    def plot_unit_counts(self, case_keys=None, **kwargs):
+        from .benchmark_plot_tools import plot_unit_counts
 
-        plot_study_unit_counts(self, case_keys, figsize=figsize)
+        return plot_unit_counts(self, case_keys, **kwargs)
 
     def plot_unit_losses(self, before, after, metric=["accuracy"], figsize=None):
         import matplotlib.pyplot as plt

src/spikeinterface/benchmark/benchmark_merging.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+from spikeinterface.curation.auto_merge import auto_merge_units
+from spikeinterface.comparison import compare_sorter_to_ground_truth
+from spikeinterface.core.sortinganalyzer import create_sorting_analyzer
+from spikeinterface.widgets import (
+    plot_unit_templates,
+    plot_amplitudes,
+    plot_crosscorrelograms,
+)
+
+import numpy as np
+from .benchmark_base import Benchmark, BenchmarkStudy
+
+
+class MergingBenchmark(Benchmark):
+
+    def __init__(self, recording, splitted_sorting, params, gt_sorting, splitted_cells=None):
+        self.recording = recording
+        self.splitted_sorting = splitted_sorting
+        self.gt_sorting = gt_sorting
+        self.splitted_cells = splitted_cells
+        self.method_kwargs = params["method_kwargs"]
+        self.result = {}
+
+    def run(self, **job_kwargs):
+        sorting_analyzer = create_sorting_analyzer(
+            self.splitted_sorting, self.recording, format="memory", sparse=True, **job_kwargs
+        )
+        # sorting_analyzer.compute(['random_spikes', 'templates'])
+        # sorting_analyzer.compute('template_similarity', max_lag_ms=0.1, method="l2", **job_kwargs)
+        merged_analyzer, self.result["merged_pairs"], self.result["merges"], self.result["outs"] = auto_merge_units(
+            sorting_analyzer, extra_outputs=True, **self.method_kwargs, **job_kwargs
+        )
+
+        self.result["sorting"] = merged_analyzer.sorting
+
+    def compute_result(self, **result_params):
+        sorting = self.result["sorting"]
+        comp = compare_sorter_to_ground_truth(self.gt_sorting, sorting, exhaustive_gt=True)
+        self.result["gt_comparison"] = comp
+
+    _run_key_saved = [("sorting", "sorting"), ("merges", "pickle"), ("merged_pairs", "pickle"), ("outs", "pickle")]
+    _result_key_saved = [("gt_comparison", "pickle")]
+
+
+class MergingStudy(BenchmarkStudy):
+
+    benchmark_class = MergingBenchmark
+
+    def create_benchmark(self, key):
+        dataset_key = self.cases[key]["dataset"]
+        recording, gt_sorting = self.datasets[dataset_key]
+        params = self.cases[key]["params"]
+        init_kwargs = self.cases[key]["init_kwargs"]
+        benchmark = MergingBenchmark(recording, gt_sorting, params, **init_kwargs)
+        return benchmark
+
+    def get_count_units(self, case_keys=None, well_detected_score=None, redundant_score=None, overmerged_score=None):
+        import pandas as pd
+
+        if case_keys is None:
+            case_keys = list(self.cases.keys())
+
+        if isinstance(case_keys[0], str):
+            index = pd.Index(case_keys, name=self.levels)
+        else:
+            index = pd.MultiIndex.from_tuples(case_keys, names=self.levels)
+
+        columns = ["num_gt", "num_sorter", "num_well_detected"]
+        comp = self.get_result(case_keys[0])["gt_comparison"]
+        if comp.exhaustive_gt:
+            columns.extend(["num_false_positive", "num_redundant", "num_overmerged", "num_bad"])
+        count_units = pd.DataFrame(index=index, columns=columns, dtype=int)
+
+        for key in case_keys:
+            comp = self.get_result(key)["gt_comparison"]
+            assert comp is not None, "You need to do study.run_comparisons() first"
+
+            gt_sorting = comp.sorting1
+            sorting = comp.sorting2
+
+            count_units.loc[key, "num_gt"] = len(gt_sorting.get_unit_ids())
+            count_units.loc[key, "num_sorter"] = len(sorting.get_unit_ids())
+            count_units.loc[key, "num_well_detected"] = comp.count_well_detected_units(well_detected_score)
+
+            if comp.exhaustive_gt:
+                count_units.loc[key, "num_redundant"] = comp.count_redundant_units(redundant_score)
+                count_units.loc[key, "num_overmerged"] = comp.count_overmerged_units(overmerged_score)
+                count_units.loc[key, "num_false_positive"] = comp.count_false_positive_units(redundant_score)
+                count_units.loc[key, "num_bad"] = comp.count_bad_units()
+
+        return count_units
+
+    def plot_agreement_matrix(self, **kwargs):
+        from .benchmark_plot_tools import plot_agreement_matrix
+
+        return plot_agreement_matrix(self, **kwargs)
+
+    def plot_unit_counts(self, case_keys=None, **kwargs):
+        from .benchmark_plot_tools import plot_unit_counts
+
+        return plot_unit_counts(self, case_keys, **kwargs)
+
+    def get_splitted_pairs(self, case_key):
+        return self.benchmarks[case_key].splitted_cells
+
+    def get_splitted_pairs_index(self, case_key, pair):
+        for count, i in enumerate(self.benchmarks[case_key].splitted_cells):
+            if i == pair:
+                return count
+
+    def plot_splitted_amplitudes(self, case_key, pair_index=0, backend="ipywidgets"):
+        analyzer = self.get_sorting_analyzer(case_key)
+        if analyzer.get_extension("spike_amplitudes") is None:
+            analyzer.compute(["spike_amplitudes"])
+        plot_amplitudes(analyzer, unit_ids=self.get_splitted_pairs(case_key)[pair_index], backend=backend)
+
+    def plot_splitted_correlograms(self, case_key, pair_index=0, backend="ipywidgets"):
+        analyzer = self.get_sorting_analyzer(case_key)
+        if analyzer.get_extension("correlograms") is None:
+            analyzer.compute(["correlograms"])
+        if analyzer.get_extension("template_similarity") is None:
+            analyzer.compute(["template_similarity"])
+        plot_crosscorrelograms(analyzer, unit_ids=self.get_splitted_pairs(case_key)[pair_index])
+
+    def plot_splitted_templates(self, case_key, pair_index=0, backend="ipywidgets"):
+        analyzer = self.get_sorting_analyzer(case_key)
+        if analyzer.get_extension("spike_amplitudes") is None:
+            analyzer.compute(["spike_amplitudes"])
+        plot_unit_templates(analyzer, unit_ids=self.get_splitted_pairs(case_key)[pair_index], backend=backend)
+
+    def plot_potential_merges(self, case_key, min_snr=None, backend="ipywidgets"):
+        analyzer = self.get_sorting_analyzer(case_key)
+        mylist = self.get_splitted_pairs(case_key)
+
+        if analyzer.get_extension("spike_amplitudes") is None:
+            analyzer.compute(["spike_amplitudes"])
+        if analyzer.get_extension("correlograms") is None:
+            analyzer.compute(["correlograms"])
+
+        if min_snr is not None:
+            select_from = analyzer.sorting.unit_ids
+            if analyzer.get_extension("noise_levels") is None:
+                analyzer.compute("noise_levels")
+            if analyzer.get_extension("quality_metrics") is None:
+                analyzer.compute("quality_metrics", metric_names=["snr"])
+
+            snr = analyzer.get_extension("quality_metrics").get_data()["snr"].values
+            select_from = select_from[snr > min_snr]
+            mylist_selection = []
+            for i in mylist:
+                if (i[0] in select_from) or (i[1] in select_from):
+                    mylist_selection += [i]
+            mylist = mylist_selection
+
+        from spikeinterface.widgets import plot_potential_merges
+
+        plot_potential_merges(analyzer, mylist, backend=backend)
+
+    def plot_performed_merges(self, case_key, backend="ipywidgets"):
+        analyzer = self.get_sorting_analyzer(case_key)
+
+        if analyzer.get_extension("spike_amplitudes") is None:
+            analyzer.compute(["spike_amplitudes"])
+        if analyzer.get_extension("correlograms") is None:
+            analyzer.compute(["correlograms"])
+
+        all_merges = list(self.benchmarks[case_key].result["merged_pairs"].values())
+
+        from spikeinterface.widgets import plot_potential_merges
+
+        plot_potential_merges(analyzer, all_merges, backend=backend)

src/spikeinterface/benchmark/benchmark_motion_estimation.py

@@ -86,7 +86,7 @@ def __init__(
     def run(self, **job_kwargs):
         p = self.params
 
-        noise_levels = get_noise_levels(self.recording, return_scaled=False)
+        noise_levels = get_noise_levels(self.recording, return_scaled=False, **job_kwargs)
 
         t0 = time.perf_counter()
         peaks = detect_peaks(self.recording, noise_levels=noise_levels, **p["detect_kwargs"], **job_kwargs)