Merge pull request SpikeInterface#3556 from yger/bencharmk_remove

Fix use of keep during run_sorter benchmark
samuelgarcia · Jan 8, 2025 · a64aed9 · a64aed9
2 parents 1c43ef1 + 692ed25
commit a64aed9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/src/spikeinterface/benchmark/benchmark_base.py b/src/spikeinterface/benchmark/benchmark_base.py
@@ -208,10 +208,11 @@ def run(self, case_keys=None, keep=True, verbose=False, **job_kwargs):
         for key in case_keys:
 
             result_folder = self.folder / "results" / self.key_to_str(key)
+            sorter_folder = self.folder / "sorters" / self.key_to_str(key)
 
             if keep and result_folder.exists():
                 continue
-            elif not keep and result_folder.exists():
+            elif not keep and (result_folder.exists() or sorter_folder.exists()):
                 self.remove_benchmark(key)
             job_keys.append(key)
 

diff --git a/src/spikeinterface/benchmark/benchmark_sorter.py b/src/spikeinterface/benchmark/benchmark_sorter.py
@@ -56,6 +56,15 @@ def create_benchmark(self, key):
         benchmark = SorterBenchmark(recording, gt_sorting, params, sorter_folder)
         return benchmark
 
+    def remove_benchmark(self, key):
+        BenchmarkStudy.remove_benchmark(self, key)
+
+        sorter_folder = self.folder / "sorters" / self.key_to_str(key)
+        import shutil
+
+        if sorter_folder.exists():
+            shutil.rmtree(sorter_folder)
+
     def get_performance_by_unit(self, case_keys=None):
         import pandas as pd