Improve merging and iterative merging

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=["precision"], 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/tests/test_benchmark_merging.py

@@ -0,0 +1,74 @@
+import pytest
+from pathlib import Path
+import numpy as np
+
+import shutil
+
+from spikeinterface.benchmark.benchmark_merging import MergingStudy
+from spikeinterface.benchmark.tests.common_benchmark_testing import make_dataset
+
+from spikeinterface.generation.splitting_tools import split_sorting_by_amplitudes, split_sorting_by_times
+
+
+@pytest.mark.skip()
+def test_benchmark_merging(create_cache_folder):
+    cache_folder = create_cache_folder
+    job_kwargs = dict(n_jobs=0.8, chunk_duration="1s")
+
+    recording, gt_sorting, gt_analyzer = make_dataset()
+
+    # create study
+    study_folder = cache_folder / "study_clustering"
+    # datasets = {"toy": (recording, gt_sorting)}
+    datasets = {"toy": gt_analyzer}
+
+    gt_analyzer.compute(["random_spikes", "templates", "spike_amplitudes"])
+
+    splitted_sorting = {}
+    splitted_sorting["times"] = split_sorting_by_times(gt_analyzer)
+    splitted_sorting["amplitudes"] = split_sorting_by_amplitudes(gt_analyzer)
+
+    cases = {}
+    for splits in ["times", "amplitudes"]:
+        for method in ["circus", "lussac"]:
+            cases[(method, splits)] = {
+                "label": f"{method}",
+                "dataset": "toy",
+                "init_kwargs": {"gt_sorting": gt_sorting, "splitted_cells": splitted_sorting[splits][1]},
+                "params": {"method": method, "splitted_sorting": splitted_sorting[splits][0], "method_kwargs": {}},
+            }
+
+    if study_folder.exists():
+        shutil.rmtree(study_folder)
+    study = MergingStudy.create(study_folder, datasets=datasets, cases=cases)
+    print(study)
+
+    # this study needs analyzer
+    # study.create_sorting_analyzer_gt(**job_kwargs)
+    study.compute_metrics()
+
+    study = MergingStudy(study_folder)
+
+    # run and result
+    study.run(**job_kwargs)
+    study.compute_results()
+
+    # load study to check persistency
+    study = MergingStudy(study_folder)
+    print(study)
+
+    # plots
+    # study.plot_performances_vs_snr()
+    study.plot_agreements()
+    study.plot_unit_counts()
+    # study.plot_error_metrics()
+    # study.plot_metrics_vs_snr()
+    # study.plot_run_times()
+    # study.plot_metrics_vs_snr("cosine")
+    # study.homogeneity_score(ignore_noise=False)
+    # import matplotlib.pyplot as plt
+    # plt.show()
+
+
+if __name__ == "__main__":
+    test_benchmark_merging()

src/spikeinterface/curation/__init__.py

@@ -3,7 +3,11 @@
 from .remove_redundant import remove_redundant_units, find_redundant_units
 from .remove_duplicated_spikes import remove_duplicated_spikes
 from .remove_excess_spikes import remove_excess_spikes
-from .auto_merge import compute_merge_unit_groups, auto_merge_units, get_potential_auto_merge
+from .auto_merge import (
+    compute_merge_unit_groups,
+    auto_merge_units,
+    get_potential_auto_merge,
+)
 
 
 # manual sorting,