diff --git a/doc/get_started/install_sorters.rst b/doc/get_started/install_sorters.rst index fa3d883da9..90412f488d 100644 --- a/doc/get_started/install_sorters.rst +++ b/doc/get_started/install_sorters.rst @@ -203,6 +203,24 @@ Mountainsort5 pip install mountainsort5 +.. _install-vanillasort: + +VanillaSort +^^^^^^^^^^^ + +* Python / PyTorch, CPU or optional CUDA +* Url: https://github.com/IgarashiAkatuki/VanillaSort +* Authors: Zishuo Feng and Feng Cao +* Reference: `Spike Sorting with VanillaSort `_ (bioRxiv, 2026) +* Installation:: + + pip install --upgrade vanillasort + +The wrapper requires VanillaSort >= 0.2.1. See the +`VanillaSort documentation `_ +for input requirements, model downloads, parameters, and examples. + + SpyKING CIRCUS ^^^^^^^^^^^^^^ diff --git a/doc/modules/sorters.rst b/doc/modules/sorters.rst index 5c67296794..b3d0bc2b6d 100644 --- a/doc/modules/sorters.rst +++ b/doc/modules/sorters.rst @@ -484,6 +484,7 @@ Here is the list of external sorters accessible using the run_sorter wrapper: * **RTSort** :code:`run_sorter(sorter_name='rtsort')` * **SpyKING Circus** :code:`run_sorter(sorter_name='spykingcircus')` * **Tridesclous** :code:`run_sorter(sorter_name='tridesclous')` +* **VanillaSort** :code:`run_sorter(sorter_name='vanillasort')` (`parameters and examples `_) * **Wave clus** :code:`run_sorter(sorter_name='waveclus')` * **Combinato** :code:`run_sorter(sorter_name='combinato')` * **HDSort** :code:`run_sorter(sorter_name='hdsort')` diff --git a/src/spikeinterface/sorters/external/tests/test_vanillasort.py b/src/spikeinterface/sorters/external/tests/test_vanillasort.py new file mode 100644 index 0000000000..efbb3f2bf8 --- /dev/null +++ b/src/spikeinterface/sorters/external/tests/test_vanillasort.py @@ -0,0 +1,11 @@ +import unittest + +import pytest + +from spikeinterface.sorters import VanillaSortSorter +from spikeinterface.sorters.tests.common_tests import SorterCommonTestSuite + + +@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort not installed") +class VanillaSortCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): + SorterClass = VanillaSortSorter diff --git a/src/spikeinterface/sorters/external/vanillasort.py b/src/spikeinterface/sorters/external/vanillasort.py new file mode 100644 index 0000000000..b3299c69c2 --- /dev/null +++ b/src/spikeinterface/sorters/external/vanillasort.py @@ -0,0 +1,159 @@ +import importlib.util +import sys +from contextlib import redirect_stdout +from importlib.metadata import version +from pathlib import Path + +from packaging.version import parse + +from ..basesorter import BaseSorter +from ...core import load + + +class _LogAndStdout: + """Keep the runtime log while echoing verbose package progress.""" + + def __init__(self, log, stdout): + self.log = log + self.stdout = stdout + + def write(self, text): + written = self.log.write(text) + self.stdout.write(text) + return written + + def flush(self): + self.log.flush() + self.stdout.flush() + + +class VanillaSortSorter(BaseSorter): + """Run the VanillaSort package's public recording-to-sorting API.""" + + sorter_name = "vanillasort" + requires_locations = True + handle_multi_segment = True + gpu_capability = "nvidia-optional" + compatible_with_parallel = {"loky": False, "multiprocessing": False, "threading": False} + + _default_params = { + "model": "default", + "model_path": None, + "device": "auto", + "seed": 0, + "components": 22, + "profile": "d1", + "enable_gate": True, + "threshold": None, + "direct_keep": None, + "primary_snr": None, + "secondary_snr": None, + "min_secondary_channels": None, + "detector_batch": 8, + "embedding_batch": 128, + "neighborhood_mode": "dense", + "target_units": None, + "postprocessing": None, + } + + _params_description = { + "model": "Model name; 'default' downloads pinned Hugging Face checkpoints once and reuses the local cache.", + "model_path": "Optional local combined .pt checkpoint or directory containing config.json and both models.", + "device": "'auto' selects CUDA when available, otherwise CPU; also accepts 'cpu', 'cuda', or 'cuda:N'.", + "seed": "Random seed for model initialization and GMM clustering.", + "components": ( + "GMM component count per neighborhood/segment, at least 2; " + "independent of the output target_units. No automatic K estimation." + ), + "profile": "Published preprocessing/threshold profile: 'd1' or 'canonical'.", + "enable_gate": "Enable the built-in SNR gate; False retains score thresholding, peak/NMS and ownership checks.", + "threshold": "Minimum detector sigmoid score in [0, 1]; None uses the loaded profile's default.", + "direct_keep": "Eligible scores at or above this value bypass SNR checks; None uses the profile's default.", + "primary_snr": "Nonnegative strongest-channel event SNR minimum; None uses the model's default.", + "secondary_snr": "Nonnegative additional-channel event SNR minimum; None uses the model's default.", + "min_secondary_channels": "Number of additional channels meeting secondary_snr, from 1 to 3; None uses the model.", + "detector_batch": "Number of 2500-sample detector chunks per batch; reduced on CUDA out-of-memory errors.", + "embedding_batch": "Number of waveforms per HuiduRep batch; reduced on CUDA out-of-memory errors.", + "neighborhood_mode": "'dense' preserves original neighborhoods; 'cover' uses fewer passes and may change accuracy.", + "target_units": "Optional positive output unit count per segment after merging; None disables fixed-K processing.", + "postprocessing": ( + "Optional fixed-K options dictionary; requires target_units. " + "on_shortfall='return_available' permits fewer units when reliable anchors are insufficient." + ), + } + + sorter_description = """VanillaSort combines VanillaDet detection, HuiduRep waveform embeddings and + relative-amplitude features with Gaussian-mixture clustering and template-residual refinement. + The pretrained models use multichannel recordings with 2D channel locations and at least four channels + per group/shank. Larger probes use experimental local neighborhoods. Segments have independent unit IDs. + Parameters and examples: https://github.com/IgarashiAkatuki/VanillaSort/blob/main/docs/spikeinterface.md + Reference: https://doi.org/10.64898/2026.09.18.752552. + """ + + installation_mesg = """ + This wrapper requires VanillaSort >= 0.2.1. Install or upgrade from PyPI: + + pip install --upgrade vanillasort + + Follow the repository README for CPU/CUDA PyTorch installation instructions. + """ + + @classmethod + def is_installed(cls): + return importlib.util.find_spec("vanillasort") is not None + + @staticmethod + def get_sorter_version(): + return version("vanillasort") + + @classmethod + def use_gpu(cls, params): + # Express GPU intent without importing torch on the host. The container + # runner checks NVIDIA availability for nvidia-optional sorters, and + # VanillaSort resolves "auto" in the environment that runs inference. + device = params.get("device", "auto") + return device == "auto" or str(device).startswith("cuda") + + @classmethod + def _check_params(cls, recording, output_folder, params): + installed_version = cls.get_sorter_version() + if parse(installed_version) < parse("0.2.1"): + raise RuntimeError( + f"VanillaSort >= 0.2.1 is required; found {installed_version}. " + "Upgrade with 'pip install --upgrade vanillasort'." + ) + if params["model_path"] is not None: + params["model_path"] = str(Path(params["model_path"]).expanduser().resolve()) + return params + + @classmethod + def _check_apply_filter_in_params(cls, params): + return True + + @classmethod + def _setup_recording(cls, recording, sorter_output_folder, params, verbose): + # BaseSorter serializes the recording; the package reads its traces and geometry. + pass + + @classmethod + def _run_from_folder(cls, sorter_output_folder, params, verbose): + import vanillasort + + recording = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) + if recording is None: + raise RuntimeError("Unable to load the recording for VanillaSort") + # Keep progress in the runtime trace and echo it for verbose runs. + # redirect_stdout restores sys.stdout even on failure. Its process-wide + # effect is why this wrapper disables threaded runs. + with (sorter_output_folder / "vanillasort.log").open("w", encoding="utf-8", buffering=1) as log: + stream = _LogAndStdout(log, sys.stdout) if verbose else log + with redirect_stdout(stream): + print(f"VanillaSort {vanillasort.__version__}", flush=True) + sorting = vanillasort.sort( + recording, output_folder=sorter_output_folder / "inference", verbose=verbose, **params + ) + print(f"Completed: {sorting.get_num_units()} units, {sorting.get_num_segments()} segments", flush=True) + + @classmethod + def _get_result_from_folder(cls, sorter_output_folder): + return load(Path(sorter_output_folder) / "inference" / "sorting") diff --git a/src/spikeinterface/sorters/sorterlist.py b/src/spikeinterface/sorters/sorterlist.py index 81279e8ebf..d0fe912109 100644 --- a/src/spikeinterface/sorters/sorterlist.py +++ b/src/spikeinterface/sorters/sorterlist.py @@ -15,6 +15,7 @@ from .external.rt_sort import RTSortSorter from .external.spyking_circus import SpykingcircusSorter from .external.tridesclous import TridesclousSorter +from .external.vanillasort import VanillaSortSorter from .external.waveclus import WaveClusSorter from .external.waveclus_snippets import WaveClusSnippetsSorter from .external.yass import YassSorter @@ -43,6 +44,7 @@ RTSortSorter, SpykingcircusSorter, TridesclousSorter, + VanillaSortSorter, WaveClusSorter, WaveClusSnippetsSorter, YassSorter,