Merge pull request #1908 from JoeZiminski/change_kilosort_tmp_file_cl…

…eanup_signature

Change the signature on kilosort's delete intermediate files parameters.

src/spikeinterface/sorters/external/kilosort.py

@@ -42,7 +42,7 @@ class KilosortSorter(KilosortBase, BaseSorter):
         "Nfilt": None,
         "NT": None,
         "wave_length": 61,
-        "delete_tmp_files": True,
+        "delete_tmp_files": ("matlab_files",),
         "delete_recording_dat": False,
     }
 
@@ -56,7 +56,9 @@ class KilosortSorter(KilosortBase, BaseSorter):
         "Nfilt": "Number of clusters to use (if None it is automatically computed)",
         "NT": "Batch size (if None it is automatically computed)",
         "wave_length": "size of the waveform extracted around each detected peak, (Default 61, maximum 81)",
-        "delete_tmp_files": "Whether to delete all temporary files after a successful run",
+        "delete_tmp_files": "Delete temporary files created during sorting (matlab files and the `temp_wh.dat` file that "
+        "contains kilosort-preprocessed data). Accepts `False` (deletes no files), `True` (deletes all files) "
+        "or a Tuple containing the files to delete. Options are: ('temp_wh.dat', 'matlab_files')",
         "delete_recording_dat": "Whether to delete the 'recording.dat' file after a successful run",
     }
 

src/spikeinterface/sorters/external/kilosort2.py

@@ -50,7 +50,7 @@ class Kilosort2Sorter(KilosortBase, BaseSorter):
         "skip_kilosort_preprocessing": False,
         "scaleproc": None,
         "save_rez_to_mat": False,
-        "delete_tmp_files": True,
+        "delete_tmp_files": ("matlab_files",),
         "delete_recording_dat": False,
     }
 
@@ -73,7 +73,9 @@ class Kilosort2Sorter(KilosortBase, BaseSorter):
         "skip_kilosort_preprocessing": "Can optionaly skip the internal kilosort preprocessing",
         "scaleproc": "int16 scaling of whitened data, if None set to 200.",
         "save_rez_to_mat": "Save the full rez internal struc to mat file",
-        "delete_tmp_files": "Whether to delete all temporary files after a successful run",
+        "delete_tmp_files": "Delete temporary files created during sorting (matlab files and the `temp_wh.dat` file that "
+        "contains kilosort-preprocessed data). Accepts `False` (deletes no files), `True` (deletes all files) "
+        "or a Tuple containing the files to delete. Options are: ('temp_wh.dat', 'matlab_files')",
         "delete_recording_dat": "Whether to delete the 'recording.dat' file after a successful run",
     }
 

src/spikeinterface/sorters/external/kilosort2_5.py

@@ -57,7 +57,7 @@ class Kilosort2_5Sorter(KilosortBase, BaseSorter):
         "skip_kilosort_preprocessing": False,
         "scaleproc": None,
         "save_rez_to_mat": False,
-        "delete_tmp_files": True,
+        "delete_tmp_files": ("matlab_files",),
         "delete_recording_dat": False,
     }
 
@@ -83,7 +83,9 @@ class Kilosort2_5Sorter(KilosortBase, BaseSorter):
         "skip_kilosort_preprocessing": "Can optionaly skip the internal kilosort preprocessing",
         "scaleproc": "int16 scaling of whitened data, if None set to 200.",
         "save_rez_to_mat": "Save the full rez internal struc to mat file",
-        "delete_tmp_files": "Whether to delete all temporary files after a successful run",
+        "delete_tmp_files": "Delete temporary files created during sorting (matlab files and the `temp_wh.dat` file that "
+        "contains kilosort-preprocessed data). Accepts `False` (deletes no files), `True` (deletes all files) "
+        "or a Tuple containing the files to delete. Options are: ('temp_wh.dat', 'matlab_files') ",
         "delete_recording_dat": "Whether to delete the 'recording.dat' file after a successful run",
     }
 

src/spikeinterface/sorters/external/kilosort3.py

@@ -54,7 +54,7 @@ class Kilosort3Sorter(KilosortBase, BaseSorter):
         "skip_kilosort_preprocessing": False,
         "scaleproc": None,
         "save_rez_to_mat": False,
-        "delete_tmp_files": True,
+        "delete_tmp_files": ("matlab_files",),
         "delete_recording_dat": False,
     }
 
@@ -80,7 +80,9 @@ class Kilosort3Sorter(KilosortBase, BaseSorter):
         "skip_kilosort_preprocessing": "Can optionaly skip the internal kilosort preprocessing",
         "scaleproc": "int16 scaling of whitened data, if None set to 200.",
         "save_rez_to_mat": "Save the full rez internal struc to mat file",
-        "delete_tmp_files": "Whether to delete all temporary files after a successful run",
+        "delete_tmp_files": "Delete temporary files created during sorting (matlab files and the `temp_wh.dat` file that "
+        "contains kilosort-preprocessed data). Accepts `False` (deletes no files), `True` (deletes all files) "
+        "or a Tuple containing the files to delete. Options are: ('temp_wh.dat', 'matlab_files')",
         "delete_recording_dat": "Whether to delete the 'recording.dat' file after a successful run",
     }
 

src/spikeinterface/sorters/external/kilosortbase.py

@@ -215,16 +215,35 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
             raise Exception(f"{cls.sorter_name} returned a non-zero exit code")
 
         # Clean-up temporary files
-        if params["delete_tmp_files"]:
-            for temp_file in sorter_output_folder.glob("*.m"):
-                temp_file.unlink()
-            for temp_file in sorter_output_folder.glob("*.mat"):
-                temp_file.unlink()
-            if (sorter_output_folder / "temp_wh.dat").exists():
-                (sorter_output_folder / "temp_wh.dat").unlink()
         if params["delete_recording_dat"] and (recording_file := sorter_output_folder / "recording.dat").exists():
             recording_file.unlink()
 
+        all_temp_files = ("matlab_files", "temp_wh.dat")
+
+        if isinstance(params["delete_tmp_files"], bool):
+            if params["delete_tmp_files"]:
+                tmp_files_to_remove = all_tmp_files
+            else:
+                tmp_files_to_remove = ()
+        else:
+            assert isinstance(
+                params["delete_tmp_files"], (tuple, list)
+            ), "`delete_tmp_files` must be a `Bool`, `Tuple` or `List`."
+
+            for name in params["delete_tmp_files"]:
+                assert name in all_tmp_files, f"{name} is not a valid option, must be one of: {all_tmp_files}"
+
+            tmp_files_to_remove = params["delete_tmp_files"]
+
+        if "temp_wh.dat" in tmp_files_to_remove:
+            if (temp_wh_file := sorter_output_folder / "temp_wh.dat").exists():
+                temp_wh_file.unlink()
+
+        if "matlab_files" in tmp_files_to_remove:
+            for ext in ["*.m", "*.mat"]:
+                for temp_file in sorter_output_folder.glob(ext):
+                    temp_file.unlink()
+
     @classmethod
     def _get_result_from_folder(cls, sorter_output_folder):
         sorter_output_folder = Path(sorter_output_folder)