diff --git a/src/spikeinterface/sorters/external/kilosort.py b/src/spikeinterface/sorters/external/kilosort.py index 62b1e8b9e2..f1d656644b 100644 --- a/src/spikeinterface/sorters/external/kilosort.py +++ b/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", } diff --git a/src/spikeinterface/sorters/external/kilosort2.py b/src/spikeinterface/sorters/external/kilosort2.py index 267ff38e36..00ab3fbde5 100644 --- a/src/spikeinterface/sorters/external/kilosort2.py +++ b/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", } diff --git a/src/spikeinterface/sorters/external/kilosort2_5.py b/src/spikeinterface/sorters/external/kilosort2_5.py index 0c9e36177e..dd9130b9ae 100644 --- a/src/spikeinterface/sorters/external/kilosort2_5.py +++ b/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", } diff --git a/src/spikeinterface/sorters/external/kilosort3.py b/src/spikeinterface/sorters/external/kilosort3.py index 77e83e35b9..77267620fa 100644 --- a/src/spikeinterface/sorters/external/kilosort3.py +++ b/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", } diff --git a/src/spikeinterface/sorters/external/kilosortbase.py b/src/spikeinterface/sorters/external/kilosortbase.py index 9918d73edc..2b8eb621b9 100644 --- a/src/spikeinterface/sorters/external/kilosortbase.py +++ b/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)