From e1817f23bf3ff327581ae492a60ba347530e76b3 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Tue, 24 Oct 2023 09:30:41 +0200 Subject: [PATCH 1/9] Initial dartsort wrapper implementation. --- .../sorters/external/dartsort.py | 126 ++++++++++++++++++ .../sorters/external/tests/test_dartsort.py | 16 +++ src/spikeinterface/sorters/sorterlist.py | 2 + 3 files changed, 144 insertions(+) create mode 100644 src/spikeinterface/sorters/external/dartsort.py create mode 100644 src/spikeinterface/sorters/external/tests/test_dartsort.py diff --git a/src/spikeinterface/sorters/external/dartsort.py b/src/spikeinterface/sorters/external/dartsort.py new file mode 100644 index 0000000000..be31b733fd --- /dev/null +++ b/src/spikeinterface/sorters/external/dartsort.py @@ -0,0 +1,126 @@ +from pathlib import Path +from packaging.version import parse + +from ..basesorter import BaseSorter +from ...core import NumpyFolderSorting + +class DartsortSorter(BaseSorter): + """Dasrtsort wrapper""" + + sorter_name = "dartsort" + requires_locations = False + compatible_with_parallel = {"loky": False, "multiprocessing": False, "threading": False} + + # @charlie @julien @cole: tell which parameters you want to propagate here + _default_params = { + "n_jobs": -1, + "device": None, + "waveform": { + "ms_before": 1.4, + "ms_after": 2.6, + }, + "featurization":{ + "do_nn_denoise": True, + "do_tpca_denoise": True, + "do_enforce_decrease": True, + "denoise_only":False, + # ... more params are available + }, + "subtraction":{ + "spike_length_samples": 121, + "detection_thresholds": [12, 10, 8, 6, 5, 4], + "chunk_length_samples": 30_000, + "peak_sign": "neg", + "spatial_dedup_radius": 150.0, + "extract_radius": 200.0, + "n_chunks_fit": 40, + "fit_subsampling_random_state": 0, + "residnorm_decrease_threshold": 3.162, + # ... more params are available + + }, + "template": { + "spikes_per_unit": 500, + # ... more params are available + }, + "matching": { + "threshold": 50., + # ... more params are available + } + + } + + _params_description = { + "n_jobs": "number of worker", + "device": "Torch device used. None is auto." + } + + sorter_description = "Dartsort is the Columbia university sorter made with love by Charlie Windolf, Julien Boussard, Cole Hurwitz, Chris Langfield and Hyun Dong Lee from Liam Paninski team." + + installation_mesg = """\nTo use dartsort run:\n + >>> pip install dartsort + + More information on mountainsort5 at: + * https://github.com/cwindolf/dartsort + """ + + @classmethod + def is_installed(cls): + try: + import dartsort + + HAVE_DARTSORT = True + except ImportError: + HAVE_DARTSORT = False + + return HAVE_DARTSORT + + @staticmethod + def get_sorter_version(): + import dartsort + + if hasattr(dartsort, "__version__"): + return dartsort.__version__ + return "unknown" + + @classmethod + def _setup_recording(cls, recording, sorter_output_folder, params, verbose): + pass + + @classmethod + def _run_from_folder(cls, sorter_output_folder, params, verbose): + from dartsort.main import dartsort as dartsort_main + from dartsort.config import WaveformConfig, FeaturizationConfig, SubtractionConfig, TemplateConfig, MatchingConfig + + recording = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) + + # dartsort config are set using dataclass we need to map this + + waveform_config = WaveformConfig(**p["featurization"]) + trough_offset_samples = waveform_config.trough_offset_samples + featurization_config = FeaturizationConfig(**p["featurization"]) + subtraction_config = SubtractionConfig(trough_offset_samples=trough_offset_samples, **p["subtraction"]) + template_config = TemplateConfig(trough_offset_samples=trough_offset_samples, **p["template"]) + matching_config = MatchingConfig(trough_offset_samples=trough_offset_samples, **p["matching"]) + + sorting= dartsort_main( + recording, + sorter_output_folder, + + featurization_config=featurization_config, + subtraction_config=subtraction_config, + + + n_jobs=p["n_jobs"], + overwrite=False, + show_progress=verbose, + device=p["device"], + ) + + NumpyFolderSorting.write_sorting(sorting, sorter_output_folder / "final_sorting") + + @classmethod + def _get_result_from_folder(cls, sorter_output_folder): + sorter_output_folder = Path(sorter_output_folder) + sorting = NumpyFolderSorting(sorter_output_folder / "final_sorting") + return sorting diff --git a/src/spikeinterface/sorters/external/tests/test_dartsort.py b/src/spikeinterface/sorters/external/tests/test_dartsort.py new file mode 100644 index 0000000000..bb0601b896 --- /dev/null +++ b/src/spikeinterface/sorters/external/tests/test_dartsort.py @@ -0,0 +1,16 @@ +import unittest +import pytest + +from spikeinterface.sorters import DartsortSorter +from spikeinterface.sorters.tests.common_tests import SorterCommonTestSuite + + +@pytest.mark.skipif(not DartsortSorter.is_installed(), reason="dartsort not installed") +class DartsortCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): + SorterClass = Mountainsort4Sorter + + +if __name__ == "__main__": + test = DartsortCommonTestSuite() + test.setUp() + test.test_with_run() diff --git a/src/spikeinterface/sorters/sorterlist.py b/src/spikeinterface/sorters/sorterlist.py index 761bb6d716..bf065bff83 100644 --- a/src/spikeinterface/sorters/sorterlist.py +++ b/src/spikeinterface/sorters/sorterlist.py @@ -1,4 +1,5 @@ from .external.combinato import CombinatoSorter +from .external.dartsort import DartsortSorter from .external.hdsort import HDSortSorter from .external.herdingspikes import HerdingspikesSorter from .external.ironclust import IronClustSorter @@ -23,6 +24,7 @@ sorter_full_list = [ # external CombinatoSorter, + DartsortSorter, HDSortSorter, HerdingspikesSorter, IronClustSorter, From 86800eb08193e5144d0c37a0a415bf6b0b8826d9 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Wed, 11 Feb 2026 18:16:01 +0100 Subject: [PATCH 2/9] wip --- .../sorters/external/dartsort.py | 55 ++++++------------- 1 file changed, 16 insertions(+), 39 deletions(-) diff --git a/src/spikeinterface/sorters/external/dartsort.py b/src/spikeinterface/sorters/external/dartsort.py index be31b733fd..87701e2fa1 100644 --- a/src/spikeinterface/sorters/external/dartsort.py +++ b/src/spikeinterface/sorters/external/dartsort.py @@ -11,51 +11,28 @@ class DartsortSorter(BaseSorter): requires_locations = False compatible_with_parallel = {"loky": False, "multiprocessing": False, "threading": False} - # @charlie @julien @cole: tell which parameters you want to propagate here - _default_params = { - "n_jobs": -1, - "device": None, - "waveform": { - "ms_before": 1.4, - "ms_after": 2.6, - }, - "featurization":{ - "do_nn_denoise": True, - "do_tpca_denoise": True, - "do_enforce_decrease": True, - "denoise_only":False, - # ... more params are available - }, - "subtraction":{ - "spike_length_samples": 121, - "detection_thresholds": [12, 10, 8, 6, 5, 4], - "chunk_length_samples": 30_000, - "peak_sign": "neg", - "spatial_dedup_radius": 150.0, - "extract_radius": 200.0, - "n_chunks_fit": 40, - "fit_subsampling_random_state": 0, - "residnorm_decrease_threshold": 3.162, - # ... more params are available - - }, - "template": { - "spikes_per_unit": 500, - # ... more params are available - }, - "matching": { - "threshold": 50., - # ... more params are available - } + sorter_description = "Dartsort is the Columbia university sorter made with love by Charlie Windolf and Liam Paninski's team." + _default_params = { } _params_description = { - "n_jobs": "number of worker", - "device": "Torch device used. None is auto." } - sorter_description = "Dartsort is the Columbia university sorter made with love by Charlie Windolf, Julien Boussard, Cole Hurwitz, Chris Langfield and Hyun Dong Lee from Liam Paninski team." + @classmethod + def _dynamic_params(cls): + from darsort import DARTsortUserConfig + from pydantic import RootModel + # the trick is to transform the DARTsortUserConfig (a pydantic.dataclass) into a pydantic model + model = RootModel[DARTsortUserConfig](DARTsortUserConfig()) + default_params = model.model_dump(mode='python') + # default_params_descriptions = + + return default_params, default_params_descriptions + + + + installation_mesg = """\nTo use dartsort run:\n >>> pip install dartsort From c60b6a6d142e72e23fde25f6198c52968fe5ca43 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Thu, 12 Feb 2026 12:01:24 +0100 Subject: [PATCH 3/9] wip dartsort wrapper --- .../sorters/external/dartsort.py | 66 ++++++++----------- 1 file changed, 26 insertions(+), 40 deletions(-) diff --git a/src/spikeinterface/sorters/external/dartsort.py b/src/spikeinterface/sorters/external/dartsort.py index 87701e2fa1..a077529c48 100644 --- a/src/spikeinterface/sorters/external/dartsort.py +++ b/src/spikeinterface/sorters/external/dartsort.py @@ -5,13 +5,18 @@ from ...core import NumpyFolderSorting class DartsortSorter(BaseSorter): - """Dasrtsort wrapper""" + """Dartsort wrapper""" sorter_name = "dartsort" requires_locations = False compatible_with_parallel = {"loky": False, "multiprocessing": False, "threading": False} - sorter_description = "Dartsort is the Columbia university sorter made with love by Charlie Windolf and Liam Paninski's team." + installation_mesg = """\nTo use dartsort run:\n + >>> pip install dartsort + + More information on mountainsort5 at: + * https://github.com/cwindolf/dartsort + """ _default_params = { } @@ -21,31 +26,25 @@ class DartsortSorter(BaseSorter): @classmethod def _dynamic_params(cls): - from darsort import DARTsortUserConfig + from dartsort import DARTsortUserConfig from pydantic import RootModel # the trick is to transform the DARTsortUserConfig (a pydantic.dataclass) into a pydantic model - model = RootModel[DARTsortUserConfig](DARTsortUserConfig()) - default_params = model.model_dump(mode='python') - # default_params_descriptions = + Model = RootModel[DARTsortUserConfig] + # so we can dump to dict + cfg = Model(DARTsortUserConfig()) + default_params = cfg.model_dump(mode='python') + # and retrieve properties + schema = Model.model_json_schema() + default_params_descriptions = {} + for k, props in schema['$defs']['DARTsortUserConfig']['properties'].items(): + default_params_descriptions[k] = props['title'] return default_params, default_params_descriptions - - - - - installation_mesg = """\nTo use dartsort run:\n - >>> pip install dartsort - - More information on mountainsort5 at: - * https://github.com/cwindolf/dartsort - """ - @classmethod def is_installed(cls): try: import dartsort - HAVE_DARTSORT = True except ImportError: HAVE_DARTSORT = False @@ -55,7 +54,6 @@ def is_installed(cls): @staticmethod def get_sorter_version(): import dartsort - if hasattr(dartsort, "__version__"): return dartsort.__version__ return "unknown" @@ -66,38 +64,26 @@ def _setup_recording(cls, recording, sorter_output_folder, params, verbose): @classmethod def _run_from_folder(cls, sorter_output_folder, params, verbose): - from dartsort.main import dartsort as dartsort_main - from dartsort.config import WaveformConfig, FeaturizationConfig, SubtractionConfig, TemplateConfig, MatchingConfig + from dartsort import dartsort as dartsort_main + from dartsort import DARTsortUserConfig recording = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) # dartsort config are set using dataclass we need to map this + print(params) + cfg = DARTsortUserConfig(**params) + print(cfg) - waveform_config = WaveformConfig(**p["featurization"]) - trough_offset_samples = waveform_config.trough_offset_samples - featurization_config = FeaturizationConfig(**p["featurization"]) - subtraction_config = SubtractionConfig(trough_offset_samples=trough_offset_samples, **p["subtraction"]) - template_config = TemplateConfig(trough_offset_samples=trough_offset_samples, **p["template"]) - matching_config = MatchingConfig(trough_offset_samples=trough_offset_samples, **p["matching"]) - - sorting= dartsort_main( + sorting = dartsort_main( recording, sorter_output_folder, - - featurization_config=featurization_config, - subtraction_config=subtraction_config, - - - n_jobs=p["n_jobs"], - overwrite=False, - show_progress=verbose, - device=p["device"], + cfg, ) - NumpyFolderSorting.write_sorting(sorting, sorter_output_folder / "final_sorting") + NumpyFolderSorting.write_sorting(sorting, sorter_output_folder / "final_darsort_sorting") @classmethod def _get_result_from_folder(cls, sorter_output_folder): sorter_output_folder = Path(sorter_output_folder) - sorting = NumpyFolderSorting(sorter_output_folder / "final_sorting") + sorting = NumpyFolderSorting(sorter_output_folder / "final_darsort_sorting") return sorting From 2b76ead2b2c3b001eaac590181696e5d092d9373 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Thu, 12 Feb 2026 12:27:50 +0100 Subject: [PATCH 4/9] output of darsort --- src/spikeinterface/sorters/external/dartsort.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/sorters/external/dartsort.py b/src/spikeinterface/sorters/external/dartsort.py index a077529c48..f3366c2e16 100644 --- a/src/spikeinterface/sorters/external/dartsort.py +++ b/src/spikeinterface/sorters/external/dartsort.py @@ -74,11 +74,12 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose): cfg = DARTsortUserConfig(**params) print(cfg) - sorting = dartsort_main( + ret = dartsort_main( recording, sorter_output_folder, cfg, ) + sorting = ret['sorting'] NumpyFolderSorting.write_sorting(sorting, sorter_output_folder / "final_darsort_sorting") From d4459174b818b7b968486c82eff63a495adb33ff Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Fri, 13 Feb 2026 15:01:10 +0100 Subject: [PATCH 5/9] debug darsort output structure --- src/spikeinterface/sorters/external/dartsort.py | 15 +++++++++++---- .../sorters/external/tests/test_dartsort.py | 4 +++- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/src/spikeinterface/sorters/external/dartsort.py b/src/spikeinterface/sorters/external/dartsort.py index f3366c2e16..353791f4b1 100644 --- a/src/spikeinterface/sorters/external/dartsort.py +++ b/src/spikeinterface/sorters/external/dartsort.py @@ -2,7 +2,7 @@ from packaging.version import parse from ..basesorter import BaseSorter -from ...core import NumpyFolderSorting +from ...core import NumpyFolderSorting, NumpySorting class DartsortSorter(BaseSorter): """Dartsort wrapper""" @@ -70,16 +70,23 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose): recording = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) # dartsort config are set using dataclass we need to map this - print(params) cfg = DARTsortUserConfig(**params) - print(cfg) ret = dartsort_main( recording, sorter_output_folder, cfg, ) - sorting = ret['sorting'] + # the dartsort_sorting is not the spikeinterface sorting!!! + dartsort_sorting = ret['sorting'] + + times_samples = dartsort_sorting.times_samples + labels = dartsort_sorting.labels + mask = labels >= 0 + + sorting = NumpySorting.from_samples_and_labels( + [times_samples[mask]], [labels[mask]], dartsort_sorting.sampling_frequency + ) NumpyFolderSorting.write_sorting(sorting, sorter_output_folder / "final_darsort_sorting") diff --git a/src/spikeinterface/sorters/external/tests/test_dartsort.py b/src/spikeinterface/sorters/external/tests/test_dartsort.py index bb0601b896..4f87f63261 100644 --- a/src/spikeinterface/sorters/external/tests/test_dartsort.py +++ b/src/spikeinterface/sorters/external/tests/test_dartsort.py @@ -7,10 +7,12 @@ @pytest.mark.skipif(not DartsortSorter.is_installed(), reason="dartsort not installed") class DartsortCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): - SorterClass = Mountainsort4Sorter + SorterClass = DartsortSorter if __name__ == "__main__": + from pathlib import Path test = DartsortCommonTestSuite() + test.cache_folder = Path(__file__).resolve().parents[4] / "cache_folder" / "sorters" test.setUp() test.test_with_run() From a25835c395a61ea4ebfd4ea5666dd09a51299b94 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Thu, 23 Apr 2026 11:52:50 +0200 Subject: [PATCH 6/9] Improve backward compatibility for dartsort after the chunkable PR (4472) --- src/spikeinterface/core/recording_tools.py | 38 +++++++++++++++++++- src/spikeinterface/core/time_series_tools.py | 2 +- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py index 14fd921992..355093d93d 100644 --- a/src/spikeinterface/core/recording_tools.py +++ b/src/spikeinterface/core/recording_tools.py @@ -18,7 +18,7 @@ split_job_kwargs, ) -from .time_series_tools import get_random_sample_slices, get_chunks, get_chunk_with_margin +from .time_series_tools import get_random_sample_slices, get_chunks, get_time_series_chunk_with_margin from .time_series_tools import write_binary as _write_binary from .time_series_tools import write_memory as _write_memory from .time_series_tools import _write_time_series_to_zarr @@ -804,3 +804,39 @@ def do_recording_attributes_match( exception_str = "" return attributes_match, exception_str + + +def get_chunk_with_margin( + rec_segment, + start_frame, + end_frame, + channel_indices, + margin, + add_zeros=False, + add_reflect_padding=False, + window_on_margin=False, + dtype=None, +): + """ + Helper to get chunk with margin + + The margin is extracted from the recording when possible. If + at the edge of the recording, no margin is used unless one + of `add_zeros` or `add_reflect_padding` is True. In the first + case zero padding is used, in the second case np.pad is called + with mod="reflect". + """ + + # wrapper to keep backward compatibility with the previous naming and signature + # this help dartsort for instance + return get_time_series_chunk_with_margin( + rec_segment, + start_frame, + end_frame, + channel_indices, + margin, + add_zeros=add_zeros, + add_reflect_padding=add_reflect_padding, + window_on_margin=window_on_margin, + dtype=dtype, + ) \ No newline at end of file diff --git a/src/spikeinterface/core/time_series_tools.py b/src/spikeinterface/core/time_series_tools.py index 1c15daed21..fad697f94e 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -567,7 +567,7 @@ def get_chunks(time_series: TimeSeries, concatenated=True, get_data_kwargs=None, return chunk_list -def get_chunk_with_margin( +def get_time_series_chunk_with_margin( chunkable_segment: TimeSeriesSegment, start_frame, end_frame, From 2e679d0573919429ca00650c33e9e1b876abd0c1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 23 Apr 2026 09:55:51 +0000 Subject: [PATCH 7/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/spikeinterface/core/recording_tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py index 355093d93d..74b3ccb56e 100644 --- a/src/spikeinterface/core/recording_tools.py +++ b/src/spikeinterface/core/recording_tools.py @@ -815,7 +815,7 @@ def get_chunk_with_margin( add_zeros=False, add_reflect_padding=False, window_on_margin=False, - dtype=None, + dtype=None, ): """ Helper to get chunk with margin @@ -839,4 +839,4 @@ def get_chunk_with_margin( add_reflect_padding=add_reflect_padding, window_on_margin=window_on_margin, dtype=dtype, - ) \ No newline at end of file + ) From ed6c7e52e0d4a53e908cd282e96dd0fb626ea995 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Thu, 23 Apr 2026 18:07:04 +0200 Subject: [PATCH 8/9] oups --- src/spikeinterface/preprocessing/highpass_spatial_filter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/preprocessing/highpass_spatial_filter.py b/src/spikeinterface/preprocessing/highpass_spatial_filter.py index 15eb924fdc..d2d03dd517 100644 --- a/src/spikeinterface/preprocessing/highpass_spatial_filter.py +++ b/src/spikeinterface/preprocessing/highpass_spatial_filter.py @@ -221,7 +221,7 @@ def get_traces(self, start_frame, end_frame, channel_indices): self.parent_recording_segment, start_frame=start_frame, end_frame=end_frame, - last_dimension_indices=slice(None), + channel_indices=slice(None) margin=margin, ) # apply sorting by depth From 8de0656fa7a8318af7ef64e951e2a401087eeda9 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Thu, 23 Apr 2026 18:08:15 +0200 Subject: [PATCH 9/9] oups --- src/spikeinterface/preprocessing/highpass_spatial_filter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/preprocessing/highpass_spatial_filter.py b/src/spikeinterface/preprocessing/highpass_spatial_filter.py index d2d03dd517..497bbdd482 100644 --- a/src/spikeinterface/preprocessing/highpass_spatial_filter.py +++ b/src/spikeinterface/preprocessing/highpass_spatial_filter.py @@ -221,7 +221,7 @@ def get_traces(self, start_frame, end_frame, channel_indices): self.parent_recording_segment, start_frame=start_frame, end_frame=end_frame, - channel_indices=slice(None) + channel_indices=slice(None), margin=margin, ) # apply sorting by depth