From ba947d869808f1bd684850c00961bd2174ef147b Mon Sep 17 00:00:00 2001 From: KohakuKirisame Date: Tue, 6 Oct 2026 13:35:35 +0200 Subject: [PATCH 1/9] Add a VanillaSort external sorter backed by its public Python API Register VanillaSortSorter and map model selection, local checkpoint paths, CPU/CUDA devices, seed, mixture count, profile and batch sizes to vanillasort.sort(). Keep all detection and clustering code in the external package. Use BaseSorter's recording serialization and metadata flow, preserve multi-segment results and main_channel_id with native sorting-folder storage, and capture package progress in the standard runtime log. Report package installation/version and optional GPU support without loading checkpoints during discovery. Reuse SorterCommonTestSuite with a short synthetic recording and add one direct-versus-wrapper integration test. Check exact spike-vector/unit/property equivalence, result reload, recording metadata, version reporting, output cleanup and standard failure logging. Validation: the new external tests and existing sorter registry/dependency checks pass (9 passed, 1 platform-specific skip). Black, targeted Ruff, compileall and git diff --check pass. The optional VanillaSort package is tested from its local 0.1.0 inference-package branch. --- .../external/tests/test_vanillasort.py | 74 ++++++++++++ .../sorters/external/vanillasort.py | 113 ++++++++++++++++++ src/spikeinterface/sorters/sorterlist.py | 2 + 3 files changed, 189 insertions(+) create mode 100644 src/spikeinterface/sorters/external/tests/test_vanillasort.py create mode 100644 src/spikeinterface/sorters/external/vanillasort.py 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..ef055724f3 --- /dev/null +++ b/src/spikeinterface/sorters/external/tests/test_vanillasort.py @@ -0,0 +1,74 @@ +import json +import unittest + +import numpy as np +import pytest + +from spikeinterface.core import BaseSorting, NumpyRecording +from spikeinterface.sorters import VanillaSortSorter, available_sorters, run_sorter +from spikeinterface.sorters.tests.common_tests import SorterCommonTestSuite + + +def make_recording(multi_segment=False): + """Short deterministic four-channel fixture with enough events for default K.""" + rng = np.random.default_rng(42) + traces = rng.normal(size=(7500, 4)).astype("float32") + t = np.arange(-30, 30) + waveform = -np.exp(-((t / 3.0) ** 2)) + 0.35 * np.exp(-(((t - 9) / 6.0) ** 2)) + for index, center in enumerate(range(100, len(traces) - 100, 150)): + amplitudes = np.roll([8.0, 5.0, 2.0, 1.0], index % 4) + traces[center - 30 : center + 30] += (waveform[:, None] * amplitudes).astype("float32") + segments = [traces, traces[:3750].copy()] if multi_segment else [traces] + recording = NumpyRecording(segments, 30000, channel_ids=["a", "b", "c", "d"]) + recording.set_dummy_probe_from_locations(np.array([[0, 0], [20, 0], [0, 20], [20, 20]])) + return recording + + +@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort >= 0.1.0 not installed") +class VanillaSortCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): + SorterClass = VanillaSortSorter + + def setUp(self): + # Keep the common tests, using 0.25 s instead of the suite's 60 s fixture. + self.recording = make_recording().save(folder=self.cache_folder / "rec", verbose=False, format="binary") + + +@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort >= 0.1.0 not installed") +def test_vanillasort_direct_and_wrapper(tmp_path): + import vanillasort + + assert "vanillasort" in available_sorters() + assert VanillaSortSorter.get_sorter_version() == vanillasort.__version__ + assert set(VanillaSortSorter.default_params()) == set(VanillaSortSorter.params_description()) + assert not VanillaSortSorter.use_gpu({"device": "cpu"}) + recording = make_recording(multi_segment=True).save(folder=tmp_path / "rec", verbose=False) + params = dict(components=4, device="cpu", seed=0) + direct = vanillasort.sort(recording, verbose=False, **params) + folder = tmp_path / "wrapped" + wrapped = run_sorter("vanillasort", recording, folder=folder, verbose=True, **params) + assert isinstance(wrapped, BaseSorting) + assert wrapped.get_num_segments() == 2 + assert wrapped.to_spike_vector().size > 0 + np.testing.assert_array_equal(wrapped.unit_ids, direct.unit_ids) + np.testing.assert_array_equal(wrapped.to_spike_vector(), direct.to_spike_vector()) + np.testing.assert_array_equal(wrapped.get_property("main_channel_id"), direct.get_property("main_channel_id")) + restored = VanillaSortSorter.get_result_from_folder(folder) + np.testing.assert_array_equal(restored.to_spike_vector(), direct.to_spike_vector()) + np.testing.assert_array_equal(restored.get_property("main_channel_id"), direct.get_property("main_channel_id")) + assert restored.has_recording() + assert restored.sorting_info["params"]["sorter_params"]["seed"] == 0 + log = json.loads((folder / "spikeinterface_log.json").read_text()) + assert log["sorter_version"] == vanillasort.__version__ + assert log["error"] is False and log["run_time"] > 0 + assert any("Completed:" in line for line in log["runtime_trace"]) + + # BaseSorter retains normal error logs for incompatible model input. + bad = NumpyRecording([np.zeros((200, 4), dtype="float32")], 20000) + bad.set_dummy_probe_from_locations(recording.get_channel_locations()) + bad = bad.save(folder=tmp_path / "bad_rec", verbose=False) + failed_folder = tmp_path / "failed" + failed = run_sorter("vanillasort", bad, folder=failed_folder, raise_error=False, with_output=False, **params) + assert failed is None + failed_log = json.loads((failed_folder / "spikeinterface_log.json").read_text()) + assert failed_log["error"] is True + assert "30 kHz" in "\n".join(failed_log["error_trace"]) diff --git a/src/spikeinterface/sorters/external/vanillasort.py b/src/spikeinterface/sorters/external/vanillasort.py new file mode 100644 index 0000000000..35ac69c708 --- /dev/null +++ b/src/spikeinterface/sorters/external/vanillasort.py @@ -0,0 +1,113 @@ +from contextlib import redirect_stdout +from pathlib import Path + +from ..basesorter import BaseSorter +from ...core import load + + +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 = "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", + "detector_batch": 8, + "embedding_batch": 128, + } + + _params_description = { + "model": "Model name; 'default' uses the bundled pretrained checkpoints without downloading.", + "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; choose K for the recording (no automatic estimation).", + "profile": "Published preprocessing/threshold profile: 'd1' or 'canonical'.", + "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.", + } + + sorter_description = """VanillaSort combines VanillaDet detection, HuiduRep waveform embeddings and + relative-amplitude features with Gaussian-mixture clustering and template-residual refinement. + The pretrained models require 30 kHz recordings and 2D channel locations, with at least four channels + per group/shank. Larger probes use experimental local neighborhoods. Segments have independent unit IDs. + See https://github.com/IgarashiAkatuki/VanillaSort and https://doi.org/10.64898/2026.09.18.752552. + """ + + installation_mesg = """ + Install VanillaSort >= 0.1.0 from its repository: + + git clone https://github.com/IgarashiAkatuki/VanillaSort.git + cd VanillaSort + pip install -e . + + Follow the repository README for CPU/CUDA PyTorch installation instructions. + """ + + @classmethod + def is_installed(cls): + try: + import vanillasort + from packaging.version import parse + + return callable(getattr(vanillasort, "sort", None)) and parse(vanillasort.__version__) >= parse("0.1.0") + except (ImportError, AttributeError): + return False + + @staticmethod + def get_sorter_version(): + import vanillasort + + return vanillasort.__version__ + + @classmethod + def use_gpu(cls, params): + device = params.get("device", "auto") + if device == "auto": + import torch + + return torch.cuda.is_available() + return str(device).startswith("cuda") + + @classmethod + def _check_params(cls, recording, output_folder, params): + 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") + # BaseSorter includes this trace in spikeinterface_log.json, including on failure. + with (sorter_output_folder / "vanillasort.log").open("w", encoding="utf-8") as log, redirect_stdout(log): + 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, From b72b998529f5c21d85a25c9c54f432587a54b40f Mon Sep 17 00:00:00 2001 From: KohakuKirisame Date: Tue, 6 Oct 2026 13:35:43 +0200 Subject: [PATCH 2/9] Document VanillaSort installation, invocation and model constraints Add VanillaSort to the external sorter list and document source installation, the public package API, repository and paper citation, and a basic run_sorter example. Describe 30 kHz and 2D geometry requirements, four-channel local neighborhoods, optional CPU/CUDA execution, configurable GMM K and local models. Explicitly identify experimental larger-probe behavior, independent segment units and full-segment memory requirements. Validation: documentation changes are limited to the existing sorter list and installation guide; whitespace and diff checks pass. Submission and upstream discussion remain for human contributors under AGENTS.md. --- doc/get_started/install_sorters.rst | 42 +++++++++++++++++++++++++++++ doc/modules/sorters.rst | 1 + 2 files changed, 43 insertions(+) diff --git a/doc/get_started/install_sorters.rst b/doc/get_started/install_sorters.rst index fa3d883da9..0150aabd22 100644 --- a/doc/get_started/install_sorters.rst +++ b/doc/get_started/install_sorters.rst @@ -203,6 +203,48 @@ Mountainsort5 pip install mountainsort5 +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 (requires the VanillaSort package, version 0.1.0 or newer):: + + git clone https://github.com/IgarashiAkatuki/VanillaSort.git + cd VanillaSort + pip install -e . + +VanillaSort combines VanillaDet spike detection, HuiduRep waveform embeddings, +relative-amplitude features, Gaussian-mixture clustering, and template-residual +refinement. The wrapper calls the package's public ``vanillasort.sort()`` API; +the algorithm and pretrained checkpoints are maintained in VanillaSort. +Checkpoints are bundled with the package and are not downloaded at import time. +See the repository README for platform-specific PyTorch installation. + +The pretrained model requires raw 30 kHz recordings with 2D channel locations in +micrometres and at least four channels in each group/shank. Larger probes use +experimental nearest-four-channel neighborhoods. Segments are clustered +independently with distinct unit IDs. Full-segment filtering and normalization +require sufficient host RAM even though detector and embedding inference are +batched. The package documents these limitations and its validation scope. + +.. code-block:: python + + import spikeinterface.sorters as ss + + sorting = ss.run_sorter( + "vanillasort", recording, folder="vanillasort_output", + device="auto", components=22, seed=0, + ) + +Choose ``components`` (GMM K per neighborhood/segment) for your data; the default +22 is a historical preset, not an automatic unit-count estimate. ``model_path`` +can select a local checkpoint bundle or directory. Results retain segment sample +indices and the ``main_channel_id`` property. + + SpyKING CIRCUS ^^^^^^^^^^^^^^ diff --git a/doc/modules/sorters.rst b/doc/modules/sorters.rst index 5c67296794..c07970d107 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')` * **Wave clus** :code:`run_sorter(sorter_name='waveclus')` * **Combinato** :code:`run_sorter(sorter_name='combinato')` * **HDSort** :code:`run_sorter(sorter_name='hdsort')` From 0cd8318f4f468a162d4ffe679da0c672c756f807 Mon Sep 17 00:00:00 2001 From: KohakuKirisame Date: Tue, 6 Oct 2026 13:48:07 +0200 Subject: [PATCH 3/9] Keep VanillaSort integration tests offline with local tiny checkpoints Document that default pretrained models are downloaded lazily from a pinned Hugging Face release and cached separately from the Python package. Update the model parameter description and link the published model repository. Generate small local VanillaDet/HuiduRep checkpoints for the existing common-suite and equivalence tests. Exercise the real inference path and local-model parameter without downloading pretrained weights or adding test cases. Validation: HF_HUB_OFFLINE=1 run of the external sorter, registry and runner dependency tests reports 9 passed and 1 platform-specific skip. A separate smoke test with the actual Hub model and installed lightweight wheel confirms exact direct/wrapper equivalence. Black, targeted Ruff and git diff --check pass. --- doc/get_started/install_sorters.rst | 4 +- .../external/tests/test_vanillasort.py | 47 ++++++++++++++++++- .../sorters/external/vanillasort.py | 2 +- 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/doc/get_started/install_sorters.rst b/doc/get_started/install_sorters.rst index 0150aabd22..15344d84a4 100644 --- a/doc/get_started/install_sorters.rst +++ b/doc/get_started/install_sorters.rst @@ -220,7 +220,9 @@ VanillaSort combines VanillaDet spike detection, HuiduRep waveform embeddings, relative-amplitude features, Gaussian-mixture clustering, and template-residual refinement. The wrapper calls the package's public ``vanillasort.sort()`` API; the algorithm and pretrained checkpoints are maintained in VanillaSort. -Checkpoints are bundled with the package and are not downloaded at import time. +The default checkpoints are hosted at https://huggingface.co/Kohaku2580/VanillaSort. +They download from a pinned revision on first inference and are then cached locally; +installation and ordinary import do not download weights. See the repository README for platform-specific PyTorch installation. The pretrained model requires raw 30 kHz recordings with 2D channel locations in diff --git a/src/spikeinterface/sorters/external/tests/test_vanillasort.py b/src/spikeinterface/sorters/external/tests/test_vanillasort.py index ef055724f3..0e9be0784b 100644 --- a/src/spikeinterface/sorters/external/tests/test_vanillasort.py +++ b/src/spikeinterface/sorters/external/tests/test_vanillasort.py @@ -1,5 +1,6 @@ import json import unittest +from unittest.mock import patch import numpy as np import pytest @@ -9,6 +10,34 @@ from spikeinterface.sorters.tests.common_tests import SorterCommonTestSuite +@pytest.fixture +def tiny_model(tmp_path): + """Use real small model architectures; external CI never downloads weights.""" + import torch + from vanillasort.checkpoints import load_config + from vanillasort.models import VanillaDet, HuiduRep + + config = load_config() + config["detector_architecture"].update( + d_model=8, nhead=2, num_layers=1, dim_ff=16, frontend_channels=8, max_len=5000 + ) + config["huidurep_architecture"].update(embedding_dim=16, n_heads=2, ff_dim=16, num_layers=3) + config["chunk_len"] = 128 + config["profiles"]["d1"].update(base_threshold=0.0001, direct_keep=0.0001) + model_path = tmp_path / "tiny.pt" + with torch.random.fork_rng(): + torch.manual_seed(0) + torch.save( + { + "config": config, + "detector_state_dict": VanillaDet(**config["detector_architecture"]).state_dict(), + "huidurep_state_dict": HuiduRep(**config["huidurep_architecture"]).state_dict(), + }, + model_path, + ) + return model_path + + def make_recording(multi_segment=False): """Short deterministic four-channel fixture with enough events for default K.""" rng = np.random.default_rng(42) @@ -28,13 +57,27 @@ def make_recording(multi_segment=False): class VanillaSortCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): SorterClass = VanillaSortSorter + @pytest.fixture(autouse=True) + def configure_model(self, tiny_model): + self.model_path = tiny_model + def setUp(self): + defaults = patch.dict( + self.SorterClass._default_params, + { + "model_path": str(self.model_path), + "components": 2, + "device": "cpu", + }, + ) + defaults.start() + self.addCleanup(defaults.stop) # Keep the common tests, using 0.25 s instead of the suite's 60 s fixture. self.recording = make_recording().save(folder=self.cache_folder / "rec", verbose=False, format="binary") @pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort >= 0.1.0 not installed") -def test_vanillasort_direct_and_wrapper(tmp_path): +def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model): import vanillasort assert "vanillasort" in available_sorters() @@ -42,7 +85,7 @@ def test_vanillasort_direct_and_wrapper(tmp_path): assert set(VanillaSortSorter.default_params()) == set(VanillaSortSorter.params_description()) assert not VanillaSortSorter.use_gpu({"device": "cpu"}) recording = make_recording(multi_segment=True).save(folder=tmp_path / "rec", verbose=False) - params = dict(components=4, device="cpu", seed=0) + params = dict(components=4, device="cpu", seed=0, model_path=str(tiny_model)) direct = vanillasort.sort(recording, verbose=False, **params) folder = tmp_path / "wrapped" wrapped = run_sorter("vanillasort", recording, folder=folder, verbose=True, **params) diff --git a/src/spikeinterface/sorters/external/vanillasort.py b/src/spikeinterface/sorters/external/vanillasort.py index 35ac69c708..4e32ccd595 100644 --- a/src/spikeinterface/sorters/external/vanillasort.py +++ b/src/spikeinterface/sorters/external/vanillasort.py @@ -26,7 +26,7 @@ class VanillaSortSorter(BaseSorter): } _params_description = { - "model": "Model name; 'default' uses the bundled pretrained checkpoints without downloading.", + "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.", From 263b570667b91368347d2f9d6c0d323e4c38328a Mon Sep 17 00:00:00 2001 From: KohakuKirisame Date: Wed, 7 Oct 2026 09:41:06 +0200 Subject: [PATCH 4/9] Use PyPI installation instructions for VanillaSort Replace the repository clone and editable-install commands with pip install vanillasort in the external sorter's installation message and the sorter installation guide. Retain the minimum supported version and the link to platform-specific PyTorch guidance. Validation: Black and git diff --check pass. A targeted source/documentation search confirms both VanillaSort installation instructions now use the PyPI command. This changes instructional text only. --- doc/get_started/install_sorters.rst | 4 +--- src/spikeinterface/sorters/external/vanillasort.py | 6 ++---- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/doc/get_started/install_sorters.rst b/doc/get_started/install_sorters.rst index 15344d84a4..dc79caeb3d 100644 --- a/doc/get_started/install_sorters.rst +++ b/doc/get_started/install_sorters.rst @@ -212,9 +212,7 @@ VanillaSort * Reference: `Spike Sorting with VanillaSort `_ (bioRxiv, 2026) * Installation (requires the VanillaSort package, version 0.1.0 or newer):: - git clone https://github.com/IgarashiAkatuki/VanillaSort.git - cd VanillaSort - pip install -e . + pip install vanillasort VanillaSort combines VanillaDet spike detection, HuiduRep waveform embeddings, relative-amplitude features, Gaussian-mixture clustering, and template-residual diff --git a/src/spikeinterface/sorters/external/vanillasort.py b/src/spikeinterface/sorters/external/vanillasort.py index 4e32ccd595..dd077e6b09 100644 --- a/src/spikeinterface/sorters/external/vanillasort.py +++ b/src/spikeinterface/sorters/external/vanillasort.py @@ -44,11 +44,9 @@ class VanillaSortSorter(BaseSorter): """ installation_mesg = """ - Install VanillaSort >= 0.1.0 from its repository: + Install VanillaSort >= 0.1.0 from PyPI: - git clone https://github.com/IgarashiAkatuki/VanillaSort.git - cd VanillaSort - pip install -e . + pip install vanillasort Follow the repository README for CPU/CUDA PyTorch installation instructions. """ From 475f425730123b7827dfcbffb43e019caf7b650d Mon Sep 17 00:00:00 2001 From: Kohaku Kirisame Date: Wed, 14 Oct 2026 11:47:17 +0200 Subject: [PATCH 5/9] Avoid heavy imports during VanillaSort sorter discovery Address the maintainer review by checking module availability with importlib.util.find_spec and reading the installed distribution version with importlib.metadata.version. Preserve the minimum supported version without importing the optional VanillaSort package. Make use_gpu report the requested device intent without importing Torch on the host. Use the recognized nvidia-optional capability so SpikeInterface's container runner handles hardware availability, while VanillaSort resolves auto in the inference environment. Explain how redirect_stdout captures package progress in the standard runtime trace and restores stdout after success or failure. Add one focused import-guard regression test and extend the existing integration test to verify stdout restoration. Validation: the VanillaSort external suite and existing registry/dependency checks pass with HF_HUB_OFFLINE=1 (10 passed, 1 platform-specific skip) against the published vanillasort 0.1.0 package. A fresh-process check confirms discovery, version reporting and GPU intent leave both vanillasort and torch unloaded. Black, targeted Ruff, compileall and git diff --check pass. --- .../external/tests/test_vanillasort.py | 35 +++++++++++++++++++ .../sorters/external/vanillasort.py | 32 ++++++++--------- 2 files changed, 51 insertions(+), 16 deletions(-) diff --git a/src/spikeinterface/sorters/external/tests/test_vanillasort.py b/src/spikeinterface/sorters/external/tests/test_vanillasort.py index 0e9be0784b..177ff33b01 100644 --- a/src/spikeinterface/sorters/external/tests/test_vanillasort.py +++ b/src/spikeinterface/sorters/external/tests/test_vanillasort.py @@ -1,4 +1,6 @@ +import builtins import json +import sys import unittest from unittest.mock import patch @@ -10,6 +12,36 @@ from spikeinterface.sorters.tests.common_tests import SorterCommonTestSuite +def test_vanillasort_discovery_without_heavy_imports(monkeypatch): + from spikeinterface.sorters.external import vanillasort + + original_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + assert name.split(".")[0] not in {"vanillasort", "torch"}, f"Discovery imported {name}" + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guarded_import) + monkeypatch.setattr(vanillasort.importlib.util, "find_spec", lambda name: None) + assert not VanillaSortSorter.is_installed() + + monkeypatch.setattr(vanillasort.importlib.util, "find_spec", lambda name: object()) + monkeypatch.setattr(vanillasort, "version", lambda name: "0.0.9") + assert not VanillaSortSorter.is_installed() + monkeypatch.setattr(vanillasort, "version", lambda name: "0.1.0") + assert VanillaSortSorter.is_installed() + assert VanillaSortSorter.get_sorter_version() == "0.1.0" + assert VanillaSortSorter.gpu_capability == "nvidia-optional" + for params, expected in [ + ({}, True), + ({"device": "auto"}, True), + ({"device": "cpu"}, False), + ({"device": "cuda"}, True), + ({"device": "cuda:1"}, True), + ]: + assert VanillaSortSorter.use_gpu(params) is expected + + @pytest.fixture def tiny_model(tmp_path): """Use real small model architectures; external CI never downloads weights.""" @@ -88,7 +120,9 @@ def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model): params = dict(components=4, device="cpu", seed=0, model_path=str(tiny_model)) direct = vanillasort.sort(recording, verbose=False, **params) folder = tmp_path / "wrapped" + stdout = sys.stdout wrapped = run_sorter("vanillasort", recording, folder=folder, verbose=True, **params) + assert sys.stdout is stdout assert isinstance(wrapped, BaseSorting) assert wrapped.get_num_segments() == 2 assert wrapped.to_spike_vector().size > 0 @@ -111,6 +145,7 @@ def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model): bad = bad.save(folder=tmp_path / "bad_rec", verbose=False) failed_folder = tmp_path / "failed" failed = run_sorter("vanillasort", bad, folder=failed_folder, raise_error=False, with_output=False, **params) + assert sys.stdout is stdout assert failed is None failed_log = json.loads((failed_folder / "spikeinterface_log.json").read_text()) assert failed_log["error"] is True diff --git a/src/spikeinterface/sorters/external/vanillasort.py b/src/spikeinterface/sorters/external/vanillasort.py index dd077e6b09..dc3ace0401 100644 --- a/src/spikeinterface/sorters/external/vanillasort.py +++ b/src/spikeinterface/sorters/external/vanillasort.py @@ -1,6 +1,10 @@ +import importlib.util 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 @@ -11,7 +15,7 @@ class VanillaSortSorter(BaseSorter): sorter_name = "vanillasort" requires_locations = True handle_multi_segment = True - gpu_capability = "optional" + gpu_capability = "nvidia-optional" compatible_with_parallel = {"loky": False, "multiprocessing": False, "threading": False} _default_params = { @@ -53,28 +57,21 @@ class VanillaSortSorter(BaseSorter): @classmethod def is_installed(cls): - try: - import vanillasort - from packaging.version import parse - - return callable(getattr(vanillasort, "sort", None)) and parse(vanillasort.__version__) >= parse("0.1.0") - except (ImportError, AttributeError): + if importlib.util.find_spec("vanillasort") is None: return False + return parse(cls.get_sorter_version()) >= parse("0.1.0") @staticmethod def get_sorter_version(): - import vanillasort - - return vanillasort.__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") - if device == "auto": - import torch - - return torch.cuda.is_available() - return str(device).startswith("cuda") + return device == "auto" or str(device).startswith("cuda") @classmethod def _check_params(cls, recording, output_folder, params): @@ -98,7 +95,10 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose): 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") - # BaseSorter includes this trace in spikeinterface_log.json, including on failure. + # Capture the package's print-based progress in vanillasort.log, which + # BaseSorter copies into spikeinterface_log.json's runtime_trace. + # redirect_stdout restores sys.stdout on exit, including exceptions. + # Its process-wide effect is why this wrapper disables threaded runs. with (sorter_output_folder / "vanillasort.log").open("w", encoding="utf-8") as log, redirect_stdout(log): print(f"VanillaSort {vanillasort.__version__}", flush=True) sorting = vanillasort.sort( From 72a80d3199f4672de65c13cd097513af6b1e5c10 Mon Sep 17 00:00:00 2001 From: IgarashiAkatuki Date: Sun, 18 Oct 2026 15:35:15 +0200 Subject: [PATCH 6/9] Show VanillaSort progress during verbose sorting Echo package progress to stdout while retaining vanillasort.log and the standard runtime trace. Preserve quiet execution and restore stdout when sorting fails. Use pip install --upgrade vanillasort without a hard-coded version requirement. Extend the existing integration test to check visible and quiet progress, log retention, and exception handling without adding tests. Validation: targeted sorter, discovery, and dependency tests passed (10 passed, 1 Linux-only skip); direct/wrapper and discovery compatibility checks also passed against the older package (2 passed). Targeted Ruff, formatting, and git diff --check passed. --- doc/get_started/install_sorters.rst | 4 +- .../external/tests/test_vanillasort.py | 30 +++++++++--- .../sorters/external/vanillasort.py | 49 ++++++++++++------- 3 files changed, 56 insertions(+), 27 deletions(-) diff --git a/doc/get_started/install_sorters.rst b/doc/get_started/install_sorters.rst index dc79caeb3d..ca4575678a 100644 --- a/doc/get_started/install_sorters.rst +++ b/doc/get_started/install_sorters.rst @@ -210,9 +210,9 @@ VanillaSort * Url: https://github.com/IgarashiAkatuki/VanillaSort * Authors: Zishuo Feng and Feng Cao * Reference: `Spike Sorting with VanillaSort `_ (bioRxiv, 2026) -* Installation (requires the VanillaSort package, version 0.1.0 or newer):: +* Installation:: - pip install vanillasort + pip install --upgrade vanillasort VanillaSort combines VanillaDet spike detection, HuiduRep waveform embeddings, relative-amplitude features, Gaussian-mixture clustering, and template-residual diff --git a/src/spikeinterface/sorters/external/tests/test_vanillasort.py b/src/spikeinterface/sorters/external/tests/test_vanillasort.py index 177ff33b01..2f1b06d653 100644 --- a/src/spikeinterface/sorters/external/tests/test_vanillasort.py +++ b/src/spikeinterface/sorters/external/tests/test_vanillasort.py @@ -26,11 +26,10 @@ def guarded_import(name, *args, **kwargs): assert not VanillaSortSorter.is_installed() monkeypatch.setattr(vanillasort.importlib.util, "find_spec", lambda name: object()) - monkeypatch.setattr(vanillasort, "version", lambda name: "0.0.9") - assert not VanillaSortSorter.is_installed() - monkeypatch.setattr(vanillasort, "version", lambda name: "0.1.0") + installed_version = "test-version" + monkeypatch.setattr(vanillasort, "version", lambda name: installed_version) assert VanillaSortSorter.is_installed() - assert VanillaSortSorter.get_sorter_version() == "0.1.0" + assert VanillaSortSorter.get_sorter_version() == installed_version assert VanillaSortSorter.gpu_capability == "nvidia-optional" for params, expected in [ ({}, True), @@ -85,7 +84,7 @@ def make_recording(multi_segment=False): return recording -@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort >= 0.1.0 not installed") +@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort not installed") class VanillaSortCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): SorterClass = VanillaSortSorter @@ -108,8 +107,8 @@ def setUp(self): self.recording = make_recording().save(folder=self.cache_folder / "rec", verbose=False, format="binary") -@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort >= 0.1.0 not installed") -def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model): +@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort not installed") +def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model, capsys): import vanillasort assert "vanillasort" in available_sorters() @@ -119,10 +118,14 @@ def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model): recording = make_recording(multi_segment=True).save(folder=tmp_path / "rec", verbose=False) params = dict(components=4, device="cpu", seed=0, model_path=str(tiny_model)) direct = vanillasort.sort(recording, verbose=False, **params) + assert capsys.readouterr().out == "" folder = tmp_path / "wrapped" stdout = sys.stdout wrapped = run_sorter("vanillasort", recording, folder=folder, verbose=True, **params) assert sys.stdout is stdout + progress = capsys.readouterr().out + for message in ("VanillaSort ", "Preprocessing ", "Detector:", "Completed:"): + assert message in progress assert isinstance(wrapped, BaseSorting) assert wrapped.get_num_segments() == 2 assert wrapped.to_spike_vector().size > 0 @@ -138,13 +141,24 @@ def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model): assert log["sorter_version"] == vanillasort.__version__ assert log["error"] is False and log["run_time"] > 0 assert any("Completed:" in line for line in log["runtime_trace"]) + assert all(line in progress for line in log["runtime_trace"]) + + quiet_folder = tmp_path / "quiet" + quiet = run_sorter("vanillasort", recording, folder=quiet_folder, verbose=False, **params) + assert sys.stdout is stdout + assert capsys.readouterr().out == "" + np.testing.assert_array_equal(quiet.to_spike_vector(), direct.to_spike_vector()) + quiet_log = json.loads((quiet_folder / "spikeinterface_log.json").read_text()) + assert any("Completed:" in line for line in quiet_log["runtime_trace"]) # BaseSorter retains normal error logs for incompatible model input. bad = NumpyRecording([np.zeros((200, 4), dtype="float32")], 20000) bad.set_dummy_probe_from_locations(recording.get_channel_locations()) bad = bad.save(folder=tmp_path / "bad_rec", verbose=False) failed_folder = tmp_path / "failed" - failed = run_sorter("vanillasort", bad, folder=failed_folder, raise_error=False, with_output=False, **params) + failed = run_sorter( + "vanillasort", bad, folder=failed_folder, verbose=True, raise_error=False, with_output=False, **params + ) assert sys.stdout is stdout assert failed is None failed_log = json.loads((failed_folder / "spikeinterface_log.json").read_text()) diff --git a/src/spikeinterface/sorters/external/vanillasort.py b/src/spikeinterface/sorters/external/vanillasort.py index dc3ace0401..e8b35834a9 100644 --- a/src/spikeinterface/sorters/external/vanillasort.py +++ b/src/spikeinterface/sorters/external/vanillasort.py @@ -1,14 +1,30 @@ 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.""" @@ -48,18 +64,16 @@ class VanillaSortSorter(BaseSorter): """ installation_mesg = """ - Install VanillaSort >= 0.1.0 from PyPI: + Install or upgrade VanillaSort from PyPI: - pip install vanillasort + pip install --upgrade vanillasort Follow the repository README for CPU/CUDA PyTorch installation instructions. """ @classmethod def is_installed(cls): - if importlib.util.find_spec("vanillasort") is None: - return False - return parse(cls.get_sorter_version()) >= parse("0.1.0") + return importlib.util.find_spec("vanillasort") is not None @staticmethod def get_sorter_version(): @@ -95,16 +109,17 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose): 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") - # Capture the package's print-based progress in vanillasort.log, which - # BaseSorter copies into spikeinterface_log.json's runtime_trace. - # redirect_stdout restores sys.stdout on exit, including exceptions. - # Its process-wide effect is why this wrapper disables threaded runs. - with (sorter_output_folder / "vanillasort.log").open("w", encoding="utf-8") as log, redirect_stdout(log): - 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) + # 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): From 48e8f9e70cc25446d4f26118c177f9794b95a2b7 Mon Sep 17 00:00:00 2001 From: Kohaku Kirisame Date: Mon, 19 Oct 2026 22:04:49 +0200 Subject: [PATCH 7/9] Document VanillaSort parameters Add comprehensive VanillaSort 0.2.1 documentation covering wrapper parameters, direct API options, gating, custom callbacks, fixed-K postprocessing, outputs, and usage examples. Link the new guide from sorter listings and installation instructions. --- doc/get_started/install_sorters.rst | 8 + doc/modules/sorters.rst | 7 +- doc/modules/vanillasort.rst | 346 ++++++++++++++++++++++++++++ 3 files changed, 360 insertions(+), 1 deletion(-) create mode 100644 doc/modules/vanillasort.rst diff --git a/doc/get_started/install_sorters.rst b/doc/get_started/install_sorters.rst index ca4575678a..b96390272e 100644 --- a/doc/get_started/install_sorters.rst +++ b/doc/get_started/install_sorters.rst @@ -203,6 +203,8 @@ Mountainsort5 pip install mountainsort5 +.. _install-vanillasort: + VanillaSort ^^^^^^^^^^^ @@ -244,6 +246,12 @@ Choose ``components`` (GMM K per neighborhood/segment) for your data; the defaul can select a local checkpoint bundle or directory. Results retain segment sample indices and the ``main_channel_id`` property. +See :ref:`vanillasort-parameters` for the complete wrapper parameter table and the +additional VanillaSort 0.2.1 Python API options, including detection thresholds, +SNR gates, custom callbacks, neighborhood selection, and fixed-K postprocessing. +The guide distinguishes parameters registered by ``run_sorter()`` from options +that currently require calling ``vanillasort.sort()`` directly. + SpyKING CIRCUS ^^^^^^^^^^^^^^ diff --git a/doc/modules/sorters.rst b/doc/modules/sorters.rst index c07970d107..239068b6ba 100644 --- a/doc/modules/sorters.rst +++ b/doc/modules/sorters.rst @@ -484,11 +484,16 @@ 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')` +* **VanillaSort** :code:`run_sorter(sorter_name='vanillasort')` (:ref:`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')` +.. toctree:: + :hidden: + + vanillasort + Here is a list of internal sorter based on :py:mod:`spikeinterface.sortingcomponents`: * **Lupin** :code:`run_sorter(sorter_name='lupin')` diff --git a/doc/modules/vanillasort.rst b/doc/modules/vanillasort.rst new file mode 100644 index 0000000000..80c95e0dae --- /dev/null +++ b/doc/modules/vanillasort.rst @@ -0,0 +1,346 @@ +.. _vanillasort-parameters: + +VanillaSort parameters and examples +=================================== + +VanillaSort accepts a SpikeInterface recording and returns a +:py:class:`~spikeinterface.core.BaseSorting`. This page describes the VanillaSort +0.2.1 API and the parameters exposed by the current SpikeInterface wrapper. +See :ref:`install-vanillasort` for installation and model downloads. + +Choose the calling interface +---------------------------- + +``ss.run_sorter("vanillasort", ...)`` accepts the eight sorter parameters in the +first table below. SpikeInterface validates parameter names against the wrapper's +registered defaults before calling VanillaSort. Extra package parameters such as +``threshold``, ``gate_fn``, and ``target_units`` currently raise ``ValueError`` +when passed to this wrapper. Upgrading the VanillaSort package alone does not +extend that registered parameter list. + +Use ``vanillasort.sort(recording, ...)`` directly for the additional 0.2.1 options. +Its result is also a SpikeInterface sorting object and can be used with the same +analysis, comparison, and export tools. Custom gate callbacks belong to this +direct Python interface; SpikeInterface saves wrapper parameters as JSON and +does not serialize arbitrary Python functions. + +Inspect the wrapper available in your installed SpikeInterface version: + +.. code-block:: python + + import spikeinterface.sorters as ss + + defaults = ss.get_default_sorter_params("vanillasort") + descriptions = ss.get_sorter_params_description("vanillasort") + for name, default in defaults.items(): + print(name, default, descriptions[name]) + +The examples below assume ``recording`` contains raw traces sampled at 30,000 Hz +(tolerance 1 Hz), at least 100 samples per segment, and finite 2D channel locations +in micrometres. Each group/shank needs at least four channels. Channel gains and +offsets are applied when available; otherwise traces need a common voltage scale. +Filtering and normalization run inside VanillaSort. + +Parameters accepted by run_sorter +--------------------------------- + +.. list-table:: Registered VanillaSort wrapper parameters + :header-rows: 1 + :widths: 20 15 65 + + * - Parameter + - Default + - Meaning + * - ``model`` + - ``"default"`` + - Published model bundle. Pinned Hugging Face checkpoints are downloaded on + first inference, verified by SHA-256, and reused from the cache. + * - ``model_path`` + - ``None`` + - Local combined ``.pt`` checkpoint or directory containing ``config.json`` + and both model checkpoints. Overrides the default model download. + * - ``device`` + - ``"auto"`` + - ``"auto"`` selects CUDA when available and CPU otherwise. Explicit values + are ``"cpu"``, ``"cuda"``, and ``"cuda:N"``. MPS is not supported. + * - ``seed`` + - ``0`` + - Seed for model initialization and GMM clustering. The standalone + VanillaSort CLI has a separate default of ``30``. + * - ``components`` + - ``22`` + - Integer GMM component count, at least 2, for each neighborhood and segment. + Choose it for the recording; 22 is a historical preset. A nonempty + neighborhood with fewer sortable events than this count raises an error. + * - ``profile`` + - ``"d1"`` + - ``"d1"`` or ``"canonical"`` preprocessing and threshold profile from the + loaded model configuration. + * - ``detector_batch`` + - ``8`` + - Positive number of detector chunks per inference batch. The published + model uses 2,500-sample chunks. CUDA memory retries reduce this batch size. + * - ``embedding_batch`` + - ``128`` + - Positive number of waveforms per HuiduRep batch. CUDA memory retries reduce + this batch size. + +.. code-block:: python + + sorting = ss.run_sorter( + "vanillasort", + recording, + folder="vanillasort_output", + device="auto", + components=22, + profile="d1", + seed=0, + detector_batch=8, + embedding_batch=128, + verbose=True, + ) + +``folder`` and ``verbose`` are SpikeInterface runner options. The wrapper manages +VanillaSort's ``output_folder`` and saves package outputs under +``vanillasort_output/sorter_output/inference/``. Do not pass ``output_folder`` as +a sorter parameter. With ``verbose=True``, progress is shown in the terminal and +retained in ``sorter_output/vanillasort.log``; ``spikeinterface_log.json`` is +finalized when sorting finishes. + +Additional parameters in vanillasort.sort +----------------------------------------- + +The direct API accepts all eight parameters above plus the following options. +The options in this table require VanillaSort 0.2.1 or later and are not currently +registered by the SpikeInterface wrapper. + +.. list-table:: Additional direct Python API options + :header-rows: 1 + :widths: 23 15 62 + + * - Parameter + - Default + - Meaning and accepted values + * - ``enable_gate`` + - ``True`` + - Boolean enabling the built-in SNR gate or a supplied custom gate. + ``False`` keeps score thresholding, peak selection, refractory NMS, and + main-channel ownership checks while disabling the SNR gate. + * - ``threshold`` + - ``None`` + - Finite detector sigmoid score in [0, 1], inclusive. Candidates must satisfy + ``score >= threshold``. ``None`` uses the loaded profile's value: + published ``d1``: 0.039; ``canonical``: 0.038. This is a score threshold, + independent of voltage amplitude. + * - ``direct_keep`` + - ``None`` + - Finite score in [0, 1]. In the built-in gate, eligible candidates with + ``score >= direct_keep`` bypass the SNR checks. Published defaults are + 0.045 for ``d1`` and 0.041 for ``canonical``. A custom gate replaces this + bypass together with the SNR rule. + * - ``primary_snr`` + - ``None`` + - Finite, nonnegative minimum SNR on the strongest channel for the built-in + gate. ``None`` uses the model setting, published default 3. + * - ``secondary_snr`` + - ``None`` + - Finite, nonnegative minimum SNR on additional channels for the built-in + gate. ``None`` uses the model setting, published default 2. + * - ``min_secondary_channels`` + - ``None`` + - Integer from 1 to 3: how many additional channels must meet + ``secondary_snr``. ``None`` uses the model setting, published default 1. + * - ``gate_fn`` + - ``None`` + - Callable ``gate_fn(context)`` returning one boolean per candidate. + Replaces the built-in SNR/direct-acceptance rule and requires + ``enable_gate=True``. See the callback example below. + * - ``neighborhood_mode`` + - ``"dense"`` + - ``"dense"`` retains the original nearest-four neighborhoods. + ``"cover"`` selects a covering subset with fewer full-record detector + passes. Every channel remains an input and owns exactly one selected + patch; group/shank boundaries are respected. Sparser contexts can reduce + accuracy, so cover mode requires explicit selection. + * - ``target_units`` + - ``None`` + - Optional positive integer enabling fixed-K merging after cross-patch + event deduplication. The target counts nonempty output units per segment, + independently of the local GMM ``components``. K=1 is supported here. + * - ``postprocessing`` + - ``None`` + - Dictionary of fixed-K policies and thresholds listed below. Requires + ``target_units``. Set ``target_units`` as its own argument, outside this + dictionary. Unknown dictionary keys are rejected. + +The direct API also accepts ``output_folder=None`` and ``verbose=True``. +``output_folder=None`` keeps the sorting in memory; a supplied path must be new +or empty. ``verbose=False`` suppresses progress. The wrapper sets these two +options through the runner's ``folder`` and ``verbose`` arguments instead. + +Thresholds and SNR gating +------------------------- + +After peak selection and refractory NMS, the score threshold is applied first. +With the built-in gate enabled, candidates above ``direct_keep`` pass directly. +Remaining candidates need the strongest-channel SNR to meet ``primary_snr`` and +the required number of other channels to meet ``secondary_snr``. These are +per-event negative-peak SNR estimates on filtered, MAD-normalized traces. + +Disable the SNR gate while retaining a detector-score threshold: + +.. code-block:: python + + import vanillasort + + sorting = vanillasort.sort( + recording, components=22, threshold=0.139, enable_gate=False, + output_folder="vanillasort_no_gate", device="auto", + ) + +``threshold`` and ``direct_keep`` are independent. If ``threshold >= direct_keep``, +every threshold-passing candidate bypasses the built-in SNR checks. To retain SNR +selection for scores between 0.139 and 0.2, for example, set +``threshold=0.139, direct_keep=0.2, enable_gate=True``. These example values do not +change the defaults. Omitting overrides preserves the loaded model settings. + +Custom Python gate +------------------ + +.. code-block:: python + + import numpy as np + import vanillasort + + def keep_events(context): + ranked = np.sort(context.snrs, axis=1) + return (context.scores >= 0.2) | ( + (ranked[:, -1] >= 3) & (ranked[:, -2] >= 2) + ) + + sorting = vanillasort.sort( + recording, components=22, threshold=0.139, + enable_gate=True, gate_fn=keep_events, + ) + +``context`` is a ``vanillasort.gating.GateContext`` for one segment/neighborhood. +It contains read-only CPU array copies: ``samples`` and ``scores`` with shape +``(N,)``, ``snrs`` with shape ``(N, 4)``, ``channel_ids`` with shape ``(4,)``, and +``channel_locations`` with shape ``(4, 2)``. Sample indices are relative to the +segment start; SNR columns follow the supplied channel order. Other fields are +``fs_hz``, ``segment_index``, ``neighborhood_index``, ``profile``, and +``event_snr_radius`` in samples. + +Return a boolean array or list of shape ``(N,)``. The callback runs once for each +nonempty set of threshold-passing candidates; empty input skips it. Invalid +return types/shapes and callback exceptions stop the run. Callbacks can reject +candidates but cannot restore events removed by thresholding, peak selection, +or NMS. Ownership checks and event deduplication still follow the custom gate. +Only the callable's name is logged; retain its implementation and captured +parameters to reproduce the run. + +Fixed-K postprocessing +---------------------- + +``components`` controls the GMM fitted independently in each neighborhood. +``target_units`` controls subsequent merging into nonempty output units for each +segment, using waveforms on common channels. The target is independent of the +local GMM size. Different segments retain separate unit identities; a target of +10 does not establish 10 shared identities across an entire multisegment record. + +.. code-block:: python + + sorting = vanillasort.sort( + recording, + components=22, + target_units=10, + neighborhood_mode="dense", + postprocessing={"on_shortfall": "return_available"}, + output_folder="vanillasort_fixed_k", + device="auto", + ) + +The default ``on_shortfall="raise"`` raises +``vanillasort.postprocessing.InsufficientAnchorsError`` when fewer than the target +number of reliable anchors are available. The exception's ``audit`` attribute +describes the shortfall. ``"return_available"`` allows fewer units, including +zero, and records ``exact_k=False``. ``anchor_policy="supplement"`` permits +observed weak or duplicate clusters and marks quality relaxation; it cannot +create missing units. Noise labels and degenerate templates are excluded. + +.. list-table:: Keys accepted inside the postprocessing dictionary + :header-rows: 1 + :widths: 25 20 55 + + * - Key + - Default + - Meaning + * - ``on_shortfall`` + - ``"raise"`` + - ``"raise"`` or ``"return_available"`` controls insufficient-anchor handling. + * - ``anchor_policy`` + - ``"reliable_only"`` + - ``"reliable_only"`` or ``"supplement"`` controls whether weak/duplicate + observed clusters may supplement reliable anchors. + * - ``core_filter`` + - ``True`` + - Boolean enabling reliability criteria: at least 20 events, normalized + template peak at least 2.5, and absolute peak offset at most 8 samples. + ``False`` relaxes these criteria and records reduced quality requirements. + * - ``anchor_only`` + - ``False`` + - Boolean restricting output to selected anchor clusters when true. + Otherwise compatible fragments may merge into anchors. + * - ``rank_quality`` + - ``"coherence"`` + - ``"coherence"`` ranks anchors by event count times waveform coherence. + ``"none"`` uses event count, treating usable clusters as fully coherent. + * - ``anchor_duplicate_cosine`` + - ``0.995`` + - Finite value at least -1. Anchors meeting this cosine threshold and having + an amplitude ratio within 1.15 are considered duplicates. Values above 1 + disable this cosine-based duplicate selection criterion. + * - ``amp_weight`` + - ``0.25`` + - Finite, nonnegative weight of absolute log amplitude ratio in the distance + used to assign fragments to anchors. + * - ``min_fragment_coherence`` + - ``0.5`` + - Finite fraction in [0, 1], the minimum waveform coherence for merging a + non-anchor fragment. + * - ``min_cosine`` + - ``0.9`` + - Finite value in [-1, 1], the minimum template cosine similarity for merging + a fragment into its selected anchor. + * - ``max_amplitude_ratio`` + - ``1.5`` + - Finite value at least 1, bounding the fragment/anchor amplitude ratio in + either direction. + * - ``min_margin`` + - ``0.03`` + - Finite, nonnegative minimum difference between the second-best and best + anchor distances. This criterion does not apply when there is one anchor. + +These defaults describe version 0.2.1. Changing them affects anchor selection, +fragment acceptance, and quality; choose settings using the recording's needs. +The independent array utility ``vanillasort.postprocessing.postprocess_fixed_k`` +also exists. The recording API prepares its common-channel waveforms automatically. + +Outputs and practical constraints +--------------------------------- + +Spike trains use zero-based segment sample indices, and the ``main_channel_id`` +property retains original recording channel IDs. Use +``sorting.get_annotation("detection_gate")`` for resolved gate settings and +``sorting.get_annotation("fixed_k_postprocessing")`` after enabling fixed-K +postprocessing. Saved runs contain ``run.json``, ``events.npz``, and a reloadable +``sorting/`` directory; fixed-K runs also save ``postprocessing_segment*.npz/json``. +On a strict fixed-K shortfall, the JSON diagnostic is written before the error. + +Full-segment filtering and normalization require host RAM even with small +inference batches. Runtime grows with duration and neighborhood count because +each neighborhood runs a full detector pass. ``"cover"`` reduces that count while +changing detector contexts. Drift correction and cross-segment identity matching +are not implemented. See the +`VanillaSort 0.2.1 documentation `_ +for model configuration, saved diagnostics, and validation scope. From 9f44444ff4fa3f44c9688e63d3d376695bb1f109 Mon Sep 17 00:00:00 2001 From: KohakuKirisame Date: Mon, 19 Oct 2026 22:15:51 +0200 Subject: [PATCH 8/9] Expose VanillaSort 0.2.1 parameters in the sorter wrapper Register the nine JSON-serializable detection, neighborhood, and fixed-K postprocessing options with the same defaults as vanillasort.sort(). Keep gate_fn on the direct Python API, while preserving the existing inference implementation and stdout logging behavior. Check the installed distribution version before inference and provide an upgrade message for VanillaSort versions older than 0.2.1. Version checks, sorter discovery, and GPU intent checks continue to avoid heavy imports. Update the installation and parameter guides for all 17 wrapper options, including run_sorter examples for threshold control and per-segment fixed-K merging, shortfall policies, and the direct-only callback boundary. Validation: - Focused offline wrapper suite: 5 passed. - Eight-channel, two-segment direct/wrapper results match exactly. - Parameter JSON, empty results, and strict-shortfall diagnostics verified. - Fresh-process discovery and version checks import neither torch nor vanillasort. - Black, targeted Ruff, diff checks, and focused Sphinx build pass. - Documentation defaults and Python examples match the package API. --- doc/get_started/install_sorters.rst | 19 ++-- doc/modules/vanillasort.rst | 87 +++++++++++-------- .../external/tests/test_vanillasort.py | 87 +++++++++++++++++++ .../sorters/external/vanillasort.py | 36 +++++++- 4 files changed, 184 insertions(+), 45 deletions(-) diff --git a/doc/get_started/install_sorters.rst b/doc/get_started/install_sorters.rst index b96390272e..2fd8ba17ce 100644 --- a/doc/get_started/install_sorters.rst +++ b/doc/get_started/install_sorters.rst @@ -218,7 +218,8 @@ VanillaSort VanillaSort combines VanillaDet spike detection, HuiduRep waveform embeddings, relative-amplitude features, Gaussian-mixture clustering, and template-residual -refinement. The wrapper calls the package's public ``vanillasort.sort()`` API; +refinement. The wrapper requires VanillaSort >= 0.2.1 and calls the package's +public ``vanillasort.sort()`` API; the algorithm and pretrained checkpoints are maintained in VanillaSort. The default checkpoints are hosted at https://huggingface.co/Kohaku2580/VanillaSort. They download from a pinned revision on first inference and are then cached locally; @@ -246,11 +247,17 @@ Choose ``components`` (GMM K per neighborhood/segment) for your data; the defaul can select a local checkpoint bundle or directory. Results retain segment sample indices and the ``main_channel_id`` property. -See :ref:`vanillasort-parameters` for the complete wrapper parameter table and the -additional VanillaSort 0.2.1 Python API options, including detection thresholds, -SNR gates, custom callbacks, neighborhood selection, and fixed-K postprocessing. -The guide distinguishes parameters registered by ``run_sorter()`` from options -that currently require calling ``vanillasort.sort()`` directly. +The wrapper also accepts ``enable_gate``, ``threshold``, ``direct_keep``, +``primary_snr``, ``secondary_snr``, ``min_secondary_channels``, +``neighborhood_mode``, ``target_units``, and ``postprocessing``. For example, +``target_units=10`` enables merging into a requested count of nonempty units per +segment; it is independent of the per-neighborhood ``components``. The default +shortfall policy raises when too few reliable anchors are available. + +See :ref:`vanillasort-parameters` for all 17 wrapper parameters, fixed-K policies, +and examples. Only the Python callback ``gate_fn`` requires calling +``vanillasort.sort()`` directly. Update both VanillaSort and the SpikeInterface +wrapper to use the newly exposed parameters. SpyKING CIRCUS diff --git a/doc/modules/vanillasort.rst b/doc/modules/vanillasort.rst index 80c95e0dae..c446fa0b1e 100644 --- a/doc/modules/vanillasort.rst +++ b/doc/modules/vanillasort.rst @@ -6,23 +6,27 @@ VanillaSort parameters and examples VanillaSort accepts a SpikeInterface recording and returns a :py:class:`~spikeinterface.core.BaseSorting`. This page describes the VanillaSort 0.2.1 API and the parameters exposed by the current SpikeInterface wrapper. +The wrapper requires VanillaSort >= 0.2.1; upgrade with +``pip install --upgrade vanillasort`` and use a SpikeInterface version containing +the updated wrapper. See :ref:`install-vanillasort` for installation and model downloads. Choose the calling interface ---------------------------- -``ss.run_sorter("vanillasort", ...)`` accepts the eight sorter parameters in the -first table below. SpikeInterface validates parameter names against the wrapper's -registered defaults before calling VanillaSort. Extra package parameters such as -``threshold``, ``gate_fn``, and ``target_units`` currently raise ``ValueError`` -when passed to this wrapper. Upgrading the VanillaSort package alone does not -extend that registered parameter list. +``ss.run_sorter("vanillasort", ...)`` accepts all 17 parameters in the model and +detection/merging tables below, including ``threshold``, ``enable_gate``, +``neighborhood_mode``, ``target_units``, and the ``postprocessing`` dictionary. +SpikeInterface validates parameter names against the wrapper's registered +defaults, saves their values as JSON, and forwards them to ``vanillasort.sort()``. +The default settings preserve the existing gate, dense neighborhoods, and disabled +fixed-K postprocessing. Older wrappers expose only the eight model/inference +parameters, so upgrading VanillaSort alone does not update the wrapper. -Use ``vanillasort.sort(recording, ...)`` directly for the additional 0.2.1 options. -Its result is also a SpikeInterface sorting object and can be used with the same -analysis, comparison, and export tools. Custom gate callbacks belong to this -direct Python interface; SpikeInterface saves wrapper parameters as JSON and -does not serialize arbitrary Python functions. +Use ``vanillasort.sort(recording, gate_fn=...)`` directly for custom Python gate +callbacks. ``gate_fn`` remains outside the wrapper's JSON parameter interface; +passing it to ``run_sorter()`` raises ``ValueError``. Both interfaces return a +SpikeInterface sorting object for the same analysis, comparison, and export tools. Inspect the wrapper available in your installed SpikeInterface version: @@ -41,10 +45,10 @@ in micrometres. Each group/shank needs at least four channels. Channel gains and offsets are applied when available; otherwise traces need a common voltage scale. Filtering and normalization run inside VanillaSort. -Parameters accepted by run_sorter ---------------------------------- +Model and inference parameters +------------------------------ -.. list-table:: Registered VanillaSort wrapper parameters +.. list-table:: Model and inference parameters accepted by run_sorter :header-rows: 1 :widths: 20 15 65 @@ -107,14 +111,14 @@ a sorter parameter. With ``verbose=True``, progress is shown in the terminal and retained in ``sorter_output/vanillasort.log``; ``spikeinterface_log.json`` is finalized when sorting finishes. -Additional parameters in vanillasort.sort ------------------------------------------ +Detection and merging parameters +-------------------------------- -The direct API accepts all eight parameters above plus the following options. -The options in this table require VanillaSort 0.2.1 or later and are not currently -registered by the SpikeInterface wrapper. +The following nine options are accepted by both ``ss.run_sorter("vanillasort", +...)`` and the direct ``vanillasort.sort()`` API. Their defaults match the package; +``None`` threshold/SNR overrides preserve the loaded model configuration. -.. list-table:: Additional direct Python API options +.. list-table:: Detection and merging parameters accepted by run_sorter :header-rows: 1 :widths: 23 15 62 @@ -123,7 +127,7 @@ registered by the SpikeInterface wrapper. - Meaning and accepted values * - ``enable_gate`` - ``True`` - - Boolean enabling the built-in SNR gate or a supplied custom gate. + - Boolean enabling the built-in SNR gate. ``False`` keeps score thresholding, peak selection, refractory NMS, and main-channel ownership checks while disabling the SNR gate. * - ``threshold`` @@ -136,8 +140,7 @@ registered by the SpikeInterface wrapper. - ``None`` - Finite score in [0, 1]. In the built-in gate, eligible candidates with ``score >= direct_keep`` bypass the SNR checks. Published defaults are - 0.045 for ``d1`` and 0.041 for ``canonical``. A custom gate replaces this - bypass together with the SNR rule. + 0.045 for ``d1`` and 0.041 for ``canonical``. * - ``primary_snr`` - ``None`` - Finite, nonnegative minimum SNR on the strongest channel for the built-in @@ -150,11 +153,6 @@ registered by the SpikeInterface wrapper. - ``None`` - Integer from 1 to 3: how many additional channels must meet ``secondary_snr``. ``None`` uses the model setting, published default 1. - * - ``gate_fn`` - - ``None`` - - Callable ``gate_fn(context)`` returning one boolean per candidate. - Replaces the built-in SNR/direct-acceptance rule and requires - ``enable_gate=True``. See the callback example below. * - ``neighborhood_mode`` - ``"dense"`` - ``"dense"`` retains the original nearest-four neighborhoods. @@ -191,11 +189,12 @@ Disable the SNR gate while retaining a detector-score threshold: .. code-block:: python - import vanillasort + import spikeinterface.sorters as ss - sorting = vanillasort.sort( - recording, components=22, threshold=0.139, enable_gate=False, - output_folder="vanillasort_no_gate", device="auto", + sorting = ss.run_sorter( + "vanillasort", recording, components=22, + threshold=0.139, enable_gate=False, + folder="vanillasort_no_gate", device="auto", verbose=True, ) ``threshold`` and ``direct_keep`` are independent. If ``threshold >= direct_keep``, @@ -207,6 +206,10 @@ change the defaults. Omitting overrides preserves the loaded model settings. Custom Python gate ------------------ +``gate_fn=None`` is available only through the direct Python API. A supplied +callable ``gate_fn(context)`` replaces the built-in SNR rule and its +``direct_keep`` bypass, and requires ``enable_gate=True``. + .. code-block:: python import numpy as np @@ -250,20 +253,28 @@ local GMM size. Different segments retain separate unit identities; a target of .. code-block:: python - sorting = vanillasort.sort( + import spikeinterface.sorters as ss + + sorting = ss.run_sorter( + "vanillasort", recording, components=22, target_units=10, neighborhood_mode="dense", postprocessing={"on_shortfall": "return_available"}, - output_folder="vanillasort_fixed_k", + folder="vanillasort_fixed_k", device="auto", + verbose=True, ) -The default ``on_shortfall="raise"`` raises +With the default ``on_shortfall="raise"``, the package raises ``vanillasort.postprocessing.InsufficientAnchorsError`` when fewer than the target number of reliable anchors are available. The exception's ``audit`` attribute -describes the shortfall. ``"return_available"`` allows fewer units, including +describes the shortfall. Through ``run_sorter()``, SpikeInterface records the +underlying exception in ``spikeinterface_log.json`` and raises +``SpikeSortingError`` unless ``raise_error=False``. The per-segment diagnostic +remains in ``sorter_output/inference/postprocessing_segment*.json``. +``"return_available"`` allows fewer units, including zero, and records ``exact_k=False``. ``anchor_policy="supplement"`` permits observed weak or duplicate clusters and marks quality relaxation; it cannot create missing units. Noise labels and degenerate templates are excluded. @@ -335,7 +346,9 @@ property retains original recording channel IDs. Use ``sorting.get_annotation("fixed_k_postprocessing")`` after enabling fixed-K postprocessing. Saved runs contain ``run.json``, ``events.npz``, and a reloadable ``sorting/`` directory; fixed-K runs also save ``postprocessing_segment*.npz/json``. -On a strict fixed-K shortfall, the JSON diagnostic is written before the error. +These files are under ``folder/sorter_output/inference/`` for ``run_sorter()`` or +the direct API's ``output_folder``. On a strict fixed-K shortfall, the JSON +diagnostic is written before the error. Full-segment filtering and normalization require host RAM even with small inference batches. Runtime grows with duration and neighborhood count because diff --git a/src/spikeinterface/sorters/external/tests/test_vanillasort.py b/src/spikeinterface/sorters/external/tests/test_vanillasort.py index 2f1b06d653..ab13d9ce44 100644 --- a/src/spikeinterface/sorters/external/tests/test_vanillasort.py +++ b/src/spikeinterface/sorters/external/tests/test_vanillasort.py @@ -40,6 +40,14 @@ def guarded_import(name, *args, **kwargs): ]: assert VanillaSortSorter.use_gpu(params) is expected + # Version checks must also stay lightweight and fail before inference. + for installed_version in ("0.1.0", "0.2.0"): + with pytest.raises(RuntimeError, match="VanillaSort >= 0.2.1.*pip install --upgrade vanillasort"): + VanillaSortSorter._check_params(None, None, VanillaSortSorter.default_params()) + for installed_version in ("0.2.1", "0.10.0"): + params = VanillaSortSorter.default_params() + assert VanillaSortSorter._check_params(None, None, params) == params + @pytest.fixture def tiny_model(tmp_path): @@ -164,3 +172,82 @@ def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model, capsys): failed_log = json.loads((failed_folder / "spikeinterface_log.json").read_text()) assert failed_log["error"] is True assert "30 kHz" in "\n".join(failed_log["error_trace"]) + + +@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort not installed") +def test_vanillasort_extended_params_roundtrip(tmp_path, tiny_model): + import vanillasort + + base = make_recording(multi_segment=True) + segments = [] + for segment in range(base.get_num_segments()): + traces = base.get_traces(segment_index=segment) + segments.append(np.column_stack([traces, 0.8 * np.roll(traces, 13, axis=0)])) + recording = NumpyRecording(segments, 30000, channel_ids=list("abcdefgh")) + recording.set_dummy_probe_from_locations(np.column_stack([np.tile([0, 20], 4), np.repeat([0, 20, 40, 60], 2)])) + recording = recording.save(folder=tmp_path / "rec", verbose=False) + options = dict( + enable_gate=False, + threshold=0.0002, + direct_keep=0.75, + primary_snr=4, + secondary_snr=2.5, + min_secondary_channels=2, + neighborhood_mode="cover", + target_units=1, + postprocessing={"core_filter": False, "on_shortfall": "return_available", "min_cosine": 0.8}, + ) + params = dict(components=4, device="cpu", seed=0, model_path=str(tiny_model), **options) + direct = vanillasort.sort(recording, verbose=False, **params) + folder = tmp_path / "wrapped" + wrapped = run_sorter("vanillasort", recording, folder=folder, verbose=False, **params) + np.testing.assert_array_equal(wrapped.unit_ids, direct.unit_ids) + np.testing.assert_array_equal(wrapped.to_spike_vector(), direct.to_spike_vector()) + np.testing.assert_array_equal(wrapped.get_property("main_channel_id"), direct.get_property("main_channel_id")) + assert wrapped.get_num_segments() == 2 and wrapped.get_num_units() == 2 + for segment in range(2): + active_units = [u for u in wrapped.unit_ids if len(wrapped.get_unit_spike_train(u, segment_index=segment))] + assert len(active_units) == 1 + + saved = json.loads((folder / "spikeinterface_params.json").read_text())["sorter_params"] + assert {name: saved[name] for name in options} == options + inference = folder / "sorter_output" / "inference" + audit = json.loads((inference / "run.json").read_text()) + assert audit["neighborhood_mode"] == "cover" and audit["target_units"] == 1 + assert len(audit["runs"]) == 4 # two covering patches in each segment + gate = wrapped.get_annotation("detection_gate") + for name in ("enable_gate", "threshold", "direct_keep", "primary_snr", "secondary_snr", "min_secondary_channels"): + assert gate[name] == options[name] + assert gate["mode"] == "off" + for segment, diagnostic in enumerate(wrapped.get_annotation("fixed_k_postprocessing")): + assert diagnostic["requested_units"] == diagnostic["output_units"] == 1 + for name, value in options["postprocessing"].items(): + assert diagnostic["config"][name] == value + assert (inference / f"postprocessing_segment{segment}.json").is_file() + + # A high score threshold produces empty results with an explicit shortfall policy. + empty_params = dict(params, threshold=1.0) + empty = run_sorter("vanillasort", recording, folder=tmp_path / "empty", verbose=False, **empty_params) + assert empty.get_num_units() == 0 and empty.get_num_segments() == 2 + assert all(d["status"] == "partial" for d in empty.get_annotation("fixed_k_postprocessing")) + + failed_folder = tmp_path / "failed" + failed = run_sorter( + "vanillasort", + recording, + folder=failed_folder, + verbose=False, + raise_error=False, + with_output=False, + **dict(empty_params, postprocessing={"core_filter": False}), + ) + assert failed is None + failure = json.loads((failed_folder / "spikeinterface_log.json").read_text()) + assert failure["error"] and "InsufficientAnchorsError" in "\n".join(failure["error_trace"]) + diagnostic = json.loads( + (failed_folder / "sorter_output" / "inference" / "postprocessing_segment0.json").read_text() + ) + assert diagnostic["failure_reason"] == "empty_input" + + with pytest.raises(ValueError, match="Invalid parameters.*gate_fn"): + VanillaSortSorter.set_params_to_folder(recording, tmp_path, {"gate_fn": lambda context: []}, verbose=False) diff --git a/src/spikeinterface/sorters/external/vanillasort.py b/src/spikeinterface/sorters/external/vanillasort.py index e8b35834a9..b90e9aa7a3 100644 --- a/src/spikeinterface/sorters/external/vanillasort.py +++ b/src/spikeinterface/sorters/external/vanillasort.py @@ -4,6 +4,8 @@ from importlib.metadata import version from pathlib import Path +from packaging.version import parse + from ..basesorter import BaseSorter from ...core import load @@ -41,8 +43,17 @@ class VanillaSortSorter(BaseSorter): "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 = { @@ -50,10 +61,25 @@ class VanillaSortSorter(BaseSorter): "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; choose K for the recording (no automatic estimation).", + "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 @@ -64,7 +90,7 @@ class VanillaSortSorter(BaseSorter): """ installation_mesg = """ - Install or upgrade VanillaSort from PyPI: + This wrapper requires VanillaSort >= 0.2.1. Install or upgrade from PyPI: pip install --upgrade vanillasort @@ -89,6 +115,12 @@ def use_gpu(cls, params): @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 From 45181a74cdc95d9d66aae473a41a5479f4dbbdcd Mon Sep 17 00:00:00 2001 From: KohakuKirisame Date: Tue, 20 Oct 2026 12:16:08 +0200 Subject: [PATCH 9/9] Keep VanillaSort docs and specialized tests in its own repository Follow the review on PR #4832: remove the dedicated parameter page and its hidden toctree, link installation and sorter settings to the VanillaSort-owned guide, and reduce the test module to the shared SorterCommonTestSuite subclass. The migrated tiny-model regressions remain in VanillaSort. Keep the 17-parameter wrapper and common test infrastructure unchanged; align the sorter description with the current multichannel package API. Validation: both common tests pass on the default 60-second, 25 kHz fixture with published checkpoints and current VanillaSort main; all 21 package/integration tests pass in VanillaSort; focused Sphinx, parameter-default consistency, Black/Ruff, and diff checks pass. --- doc/get_started/install_sorters.rst | 45 +-- doc/modules/sorters.rst | 7 +- doc/modules/vanillasort.rst | 359 ------------------ .../external/tests/test_vanillasort.py | 244 +----------- .../sorters/external/vanillasort.py | 5 +- 5 files changed, 8 insertions(+), 652 deletions(-) delete mode 100644 doc/modules/vanillasort.rst diff --git a/doc/get_started/install_sorters.rst b/doc/get_started/install_sorters.rst index 2fd8ba17ce..90412f488d 100644 --- a/doc/get_started/install_sorters.rst +++ b/doc/get_started/install_sorters.rst @@ -216,48 +216,9 @@ VanillaSort pip install --upgrade vanillasort -VanillaSort combines VanillaDet spike detection, HuiduRep waveform embeddings, -relative-amplitude features, Gaussian-mixture clustering, and template-residual -refinement. The wrapper requires VanillaSort >= 0.2.1 and calls the package's -public ``vanillasort.sort()`` API; -the algorithm and pretrained checkpoints are maintained in VanillaSort. -The default checkpoints are hosted at https://huggingface.co/Kohaku2580/VanillaSort. -They download from a pinned revision on first inference and are then cached locally; -installation and ordinary import do not download weights. -See the repository README for platform-specific PyTorch installation. - -The pretrained model requires raw 30 kHz recordings with 2D channel locations in -micrometres and at least four channels in each group/shank. Larger probes use -experimental nearest-four-channel neighborhoods. Segments are clustered -independently with distinct unit IDs. Full-segment filtering and normalization -require sufficient host RAM even though detector and embedding inference are -batched. The package documents these limitations and its validation scope. - -.. code-block:: python - - import spikeinterface.sorters as ss - - sorting = ss.run_sorter( - "vanillasort", recording, folder="vanillasort_output", - device="auto", components=22, seed=0, - ) - -Choose ``components`` (GMM K per neighborhood/segment) for your data; the default -22 is a historical preset, not an automatic unit-count estimate. ``model_path`` -can select a local checkpoint bundle or directory. Results retain segment sample -indices and the ``main_channel_id`` property. - -The wrapper also accepts ``enable_gate``, ``threshold``, ``direct_keep``, -``primary_snr``, ``secondary_snr``, ``min_secondary_channels``, -``neighborhood_mode``, ``target_units``, and ``postprocessing``. For example, -``target_units=10`` enables merging into a requested count of nonempty units per -segment; it is independent of the per-neighborhood ``components``. The default -shortfall policy raises when too few reliable anchors are available. - -See :ref:`vanillasort-parameters` for all 17 wrapper parameters, fixed-K policies, -and examples. Only the Python callback ``gate_fn`` requires calling -``vanillasort.sort()`` directly. Update both VanillaSort and the SpikeInterface -wrapper to use the newly exposed parameters. +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 239068b6ba..b3d0bc2b6d 100644 --- a/doc/modules/sorters.rst +++ b/doc/modules/sorters.rst @@ -484,16 +484,11 @@ 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')` (:ref:`parameters and examples `) +* **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')` -.. toctree:: - :hidden: - - vanillasort - Here is a list of internal sorter based on :py:mod:`spikeinterface.sortingcomponents`: * **Lupin** :code:`run_sorter(sorter_name='lupin')` diff --git a/doc/modules/vanillasort.rst b/doc/modules/vanillasort.rst deleted file mode 100644 index c446fa0b1e..0000000000 --- a/doc/modules/vanillasort.rst +++ /dev/null @@ -1,359 +0,0 @@ -.. _vanillasort-parameters: - -VanillaSort parameters and examples -=================================== - -VanillaSort accepts a SpikeInterface recording and returns a -:py:class:`~spikeinterface.core.BaseSorting`. This page describes the VanillaSort -0.2.1 API and the parameters exposed by the current SpikeInterface wrapper. -The wrapper requires VanillaSort >= 0.2.1; upgrade with -``pip install --upgrade vanillasort`` and use a SpikeInterface version containing -the updated wrapper. -See :ref:`install-vanillasort` for installation and model downloads. - -Choose the calling interface ----------------------------- - -``ss.run_sorter("vanillasort", ...)`` accepts all 17 parameters in the model and -detection/merging tables below, including ``threshold``, ``enable_gate``, -``neighborhood_mode``, ``target_units``, and the ``postprocessing`` dictionary. -SpikeInterface validates parameter names against the wrapper's registered -defaults, saves their values as JSON, and forwards them to ``vanillasort.sort()``. -The default settings preserve the existing gate, dense neighborhoods, and disabled -fixed-K postprocessing. Older wrappers expose only the eight model/inference -parameters, so upgrading VanillaSort alone does not update the wrapper. - -Use ``vanillasort.sort(recording, gate_fn=...)`` directly for custom Python gate -callbacks. ``gate_fn`` remains outside the wrapper's JSON parameter interface; -passing it to ``run_sorter()`` raises ``ValueError``. Both interfaces return a -SpikeInterface sorting object for the same analysis, comparison, and export tools. - -Inspect the wrapper available in your installed SpikeInterface version: - -.. code-block:: python - - import spikeinterface.sorters as ss - - defaults = ss.get_default_sorter_params("vanillasort") - descriptions = ss.get_sorter_params_description("vanillasort") - for name, default in defaults.items(): - print(name, default, descriptions[name]) - -The examples below assume ``recording`` contains raw traces sampled at 30,000 Hz -(tolerance 1 Hz), at least 100 samples per segment, and finite 2D channel locations -in micrometres. Each group/shank needs at least four channels. Channel gains and -offsets are applied when available; otherwise traces need a common voltage scale. -Filtering and normalization run inside VanillaSort. - -Model and inference parameters ------------------------------- - -.. list-table:: Model and inference parameters accepted by run_sorter - :header-rows: 1 - :widths: 20 15 65 - - * - Parameter - - Default - - Meaning - * - ``model`` - - ``"default"`` - - Published model bundle. Pinned Hugging Face checkpoints are downloaded on - first inference, verified by SHA-256, and reused from the cache. - * - ``model_path`` - - ``None`` - - Local combined ``.pt`` checkpoint or directory containing ``config.json`` - and both model checkpoints. Overrides the default model download. - * - ``device`` - - ``"auto"`` - - ``"auto"`` selects CUDA when available and CPU otherwise. Explicit values - are ``"cpu"``, ``"cuda"``, and ``"cuda:N"``. MPS is not supported. - * - ``seed`` - - ``0`` - - Seed for model initialization and GMM clustering. The standalone - VanillaSort CLI has a separate default of ``30``. - * - ``components`` - - ``22`` - - Integer GMM component count, at least 2, for each neighborhood and segment. - Choose it for the recording; 22 is a historical preset. A nonempty - neighborhood with fewer sortable events than this count raises an error. - * - ``profile`` - - ``"d1"`` - - ``"d1"`` or ``"canonical"`` preprocessing and threshold profile from the - loaded model configuration. - * - ``detector_batch`` - - ``8`` - - Positive number of detector chunks per inference batch. The published - model uses 2,500-sample chunks. CUDA memory retries reduce this batch size. - * - ``embedding_batch`` - - ``128`` - - Positive number of waveforms per HuiduRep batch. CUDA memory retries reduce - this batch size. - -.. code-block:: python - - sorting = ss.run_sorter( - "vanillasort", - recording, - folder="vanillasort_output", - device="auto", - components=22, - profile="d1", - seed=0, - detector_batch=8, - embedding_batch=128, - verbose=True, - ) - -``folder`` and ``verbose`` are SpikeInterface runner options. The wrapper manages -VanillaSort's ``output_folder`` and saves package outputs under -``vanillasort_output/sorter_output/inference/``. Do not pass ``output_folder`` as -a sorter parameter. With ``verbose=True``, progress is shown in the terminal and -retained in ``sorter_output/vanillasort.log``; ``spikeinterface_log.json`` is -finalized when sorting finishes. - -Detection and merging parameters --------------------------------- - -The following nine options are accepted by both ``ss.run_sorter("vanillasort", -...)`` and the direct ``vanillasort.sort()`` API. Their defaults match the package; -``None`` threshold/SNR overrides preserve the loaded model configuration. - -.. list-table:: Detection and merging parameters accepted by run_sorter - :header-rows: 1 - :widths: 23 15 62 - - * - Parameter - - Default - - Meaning and accepted values - * - ``enable_gate`` - - ``True`` - - Boolean enabling the built-in SNR gate. - ``False`` keeps score thresholding, peak selection, refractory NMS, and - main-channel ownership checks while disabling the SNR gate. - * - ``threshold`` - - ``None`` - - Finite detector sigmoid score in [0, 1], inclusive. Candidates must satisfy - ``score >= threshold``. ``None`` uses the loaded profile's value: - published ``d1``: 0.039; ``canonical``: 0.038. This is a score threshold, - independent of voltage amplitude. - * - ``direct_keep`` - - ``None`` - - Finite score in [0, 1]. In the built-in gate, eligible candidates with - ``score >= direct_keep`` bypass the SNR checks. Published defaults are - 0.045 for ``d1`` and 0.041 for ``canonical``. - * - ``primary_snr`` - - ``None`` - - Finite, nonnegative minimum SNR on the strongest channel for the built-in - gate. ``None`` uses the model setting, published default 3. - * - ``secondary_snr`` - - ``None`` - - Finite, nonnegative minimum SNR on additional channels for the built-in - gate. ``None`` uses the model setting, published default 2. - * - ``min_secondary_channels`` - - ``None`` - - Integer from 1 to 3: how many additional channels must meet - ``secondary_snr``. ``None`` uses the model setting, published default 1. - * - ``neighborhood_mode`` - - ``"dense"`` - - ``"dense"`` retains the original nearest-four neighborhoods. - ``"cover"`` selects a covering subset with fewer full-record detector - passes. Every channel remains an input and owns exactly one selected - patch; group/shank boundaries are respected. Sparser contexts can reduce - accuracy, so cover mode requires explicit selection. - * - ``target_units`` - - ``None`` - - Optional positive integer enabling fixed-K merging after cross-patch - event deduplication. The target counts nonempty output units per segment, - independently of the local GMM ``components``. K=1 is supported here. - * - ``postprocessing`` - - ``None`` - - Dictionary of fixed-K policies and thresholds listed below. Requires - ``target_units``. Set ``target_units`` as its own argument, outside this - dictionary. Unknown dictionary keys are rejected. - -The direct API also accepts ``output_folder=None`` and ``verbose=True``. -``output_folder=None`` keeps the sorting in memory; a supplied path must be new -or empty. ``verbose=False`` suppresses progress. The wrapper sets these two -options through the runner's ``folder`` and ``verbose`` arguments instead. - -Thresholds and SNR gating -------------------------- - -After peak selection and refractory NMS, the score threshold is applied first. -With the built-in gate enabled, candidates above ``direct_keep`` pass directly. -Remaining candidates need the strongest-channel SNR to meet ``primary_snr`` and -the required number of other channels to meet ``secondary_snr``. These are -per-event negative-peak SNR estimates on filtered, MAD-normalized traces. - -Disable the SNR gate while retaining a detector-score threshold: - -.. code-block:: python - - import spikeinterface.sorters as ss - - sorting = ss.run_sorter( - "vanillasort", recording, components=22, - threshold=0.139, enable_gate=False, - folder="vanillasort_no_gate", device="auto", verbose=True, - ) - -``threshold`` and ``direct_keep`` are independent. If ``threshold >= direct_keep``, -every threshold-passing candidate bypasses the built-in SNR checks. To retain SNR -selection for scores between 0.139 and 0.2, for example, set -``threshold=0.139, direct_keep=0.2, enable_gate=True``. These example values do not -change the defaults. Omitting overrides preserves the loaded model settings. - -Custom Python gate ------------------- - -``gate_fn=None`` is available only through the direct Python API. A supplied -callable ``gate_fn(context)`` replaces the built-in SNR rule and its -``direct_keep`` bypass, and requires ``enable_gate=True``. - -.. code-block:: python - - import numpy as np - import vanillasort - - def keep_events(context): - ranked = np.sort(context.snrs, axis=1) - return (context.scores >= 0.2) | ( - (ranked[:, -1] >= 3) & (ranked[:, -2] >= 2) - ) - - sorting = vanillasort.sort( - recording, components=22, threshold=0.139, - enable_gate=True, gate_fn=keep_events, - ) - -``context`` is a ``vanillasort.gating.GateContext`` for one segment/neighborhood. -It contains read-only CPU array copies: ``samples`` and ``scores`` with shape -``(N,)``, ``snrs`` with shape ``(N, 4)``, ``channel_ids`` with shape ``(4,)``, and -``channel_locations`` with shape ``(4, 2)``. Sample indices are relative to the -segment start; SNR columns follow the supplied channel order. Other fields are -``fs_hz``, ``segment_index``, ``neighborhood_index``, ``profile``, and -``event_snr_radius`` in samples. - -Return a boolean array or list of shape ``(N,)``. The callback runs once for each -nonempty set of threshold-passing candidates; empty input skips it. Invalid -return types/shapes and callback exceptions stop the run. Callbacks can reject -candidates but cannot restore events removed by thresholding, peak selection, -or NMS. Ownership checks and event deduplication still follow the custom gate. -Only the callable's name is logged; retain its implementation and captured -parameters to reproduce the run. - -Fixed-K postprocessing ----------------------- - -``components`` controls the GMM fitted independently in each neighborhood. -``target_units`` controls subsequent merging into nonempty output units for each -segment, using waveforms on common channels. The target is independent of the -local GMM size. Different segments retain separate unit identities; a target of -10 does not establish 10 shared identities across an entire multisegment record. - -.. code-block:: python - - import spikeinterface.sorters as ss - - sorting = ss.run_sorter( - "vanillasort", - recording, - components=22, - target_units=10, - neighborhood_mode="dense", - postprocessing={"on_shortfall": "return_available"}, - folder="vanillasort_fixed_k", - device="auto", - verbose=True, - ) - -With the default ``on_shortfall="raise"``, the package raises -``vanillasort.postprocessing.InsufficientAnchorsError`` when fewer than the target -number of reliable anchors are available. The exception's ``audit`` attribute -describes the shortfall. Through ``run_sorter()``, SpikeInterface records the -underlying exception in ``spikeinterface_log.json`` and raises -``SpikeSortingError`` unless ``raise_error=False``. The per-segment diagnostic -remains in ``sorter_output/inference/postprocessing_segment*.json``. -``"return_available"`` allows fewer units, including -zero, and records ``exact_k=False``. ``anchor_policy="supplement"`` permits -observed weak or duplicate clusters and marks quality relaxation; it cannot -create missing units. Noise labels and degenerate templates are excluded. - -.. list-table:: Keys accepted inside the postprocessing dictionary - :header-rows: 1 - :widths: 25 20 55 - - * - Key - - Default - - Meaning - * - ``on_shortfall`` - - ``"raise"`` - - ``"raise"`` or ``"return_available"`` controls insufficient-anchor handling. - * - ``anchor_policy`` - - ``"reliable_only"`` - - ``"reliable_only"`` or ``"supplement"`` controls whether weak/duplicate - observed clusters may supplement reliable anchors. - * - ``core_filter`` - - ``True`` - - Boolean enabling reliability criteria: at least 20 events, normalized - template peak at least 2.5, and absolute peak offset at most 8 samples. - ``False`` relaxes these criteria and records reduced quality requirements. - * - ``anchor_only`` - - ``False`` - - Boolean restricting output to selected anchor clusters when true. - Otherwise compatible fragments may merge into anchors. - * - ``rank_quality`` - - ``"coherence"`` - - ``"coherence"`` ranks anchors by event count times waveform coherence. - ``"none"`` uses event count, treating usable clusters as fully coherent. - * - ``anchor_duplicate_cosine`` - - ``0.995`` - - Finite value at least -1. Anchors meeting this cosine threshold and having - an amplitude ratio within 1.15 are considered duplicates. Values above 1 - disable this cosine-based duplicate selection criterion. - * - ``amp_weight`` - - ``0.25`` - - Finite, nonnegative weight of absolute log amplitude ratio in the distance - used to assign fragments to anchors. - * - ``min_fragment_coherence`` - - ``0.5`` - - Finite fraction in [0, 1], the minimum waveform coherence for merging a - non-anchor fragment. - * - ``min_cosine`` - - ``0.9`` - - Finite value in [-1, 1], the minimum template cosine similarity for merging - a fragment into its selected anchor. - * - ``max_amplitude_ratio`` - - ``1.5`` - - Finite value at least 1, bounding the fragment/anchor amplitude ratio in - either direction. - * - ``min_margin`` - - ``0.03`` - - Finite, nonnegative minimum difference between the second-best and best - anchor distances. This criterion does not apply when there is one anchor. - -These defaults describe version 0.2.1. Changing them affects anchor selection, -fragment acceptance, and quality; choose settings using the recording's needs. -The independent array utility ``vanillasort.postprocessing.postprocess_fixed_k`` -also exists. The recording API prepares its common-channel waveforms automatically. - -Outputs and practical constraints ---------------------------------- - -Spike trains use zero-based segment sample indices, and the ``main_channel_id`` -property retains original recording channel IDs. Use -``sorting.get_annotation("detection_gate")`` for resolved gate settings and -``sorting.get_annotation("fixed_k_postprocessing")`` after enabling fixed-K -postprocessing. Saved runs contain ``run.json``, ``events.npz``, and a reloadable -``sorting/`` directory; fixed-K runs also save ``postprocessing_segment*.npz/json``. -These files are under ``folder/sorter_output/inference/`` for ``run_sorter()`` or -the direct API's ``output_folder``. On a strict fixed-K shortfall, the JSON -diagnostic is written before the error. - -Full-segment filtering and normalization require host RAM even with small -inference batches. Runtime grows with duration and neighborhood count because -each neighborhood runs a full detector pass. ``"cover"`` reduces that count while -changing detector contexts. Drift correction and cross-segment identity matching -are not implemented. See the -`VanillaSort 0.2.1 documentation `_ -for model configuration, saved diagnostics, and validation scope. diff --git a/src/spikeinterface/sorters/external/tests/test_vanillasort.py b/src/spikeinterface/sorters/external/tests/test_vanillasort.py index ab13d9ce44..efbb3f2bf8 100644 --- a/src/spikeinterface/sorters/external/tests/test_vanillasort.py +++ b/src/spikeinterface/sorters/external/tests/test_vanillasort.py @@ -1,253 +1,11 @@ -import builtins -import json -import sys import unittest -from unittest.mock import patch -import numpy as np import pytest -from spikeinterface.core import BaseSorting, NumpyRecording -from spikeinterface.sorters import VanillaSortSorter, available_sorters, run_sorter +from spikeinterface.sorters import VanillaSortSorter from spikeinterface.sorters.tests.common_tests import SorterCommonTestSuite -def test_vanillasort_discovery_without_heavy_imports(monkeypatch): - from spikeinterface.sorters.external import vanillasort - - original_import = builtins.__import__ - - def guarded_import(name, *args, **kwargs): - assert name.split(".")[0] not in {"vanillasort", "torch"}, f"Discovery imported {name}" - return original_import(name, *args, **kwargs) - - monkeypatch.setattr(builtins, "__import__", guarded_import) - monkeypatch.setattr(vanillasort.importlib.util, "find_spec", lambda name: None) - assert not VanillaSortSorter.is_installed() - - monkeypatch.setattr(vanillasort.importlib.util, "find_spec", lambda name: object()) - installed_version = "test-version" - monkeypatch.setattr(vanillasort, "version", lambda name: installed_version) - assert VanillaSortSorter.is_installed() - assert VanillaSortSorter.get_sorter_version() == installed_version - assert VanillaSortSorter.gpu_capability == "nvidia-optional" - for params, expected in [ - ({}, True), - ({"device": "auto"}, True), - ({"device": "cpu"}, False), - ({"device": "cuda"}, True), - ({"device": "cuda:1"}, True), - ]: - assert VanillaSortSorter.use_gpu(params) is expected - - # Version checks must also stay lightweight and fail before inference. - for installed_version in ("0.1.0", "0.2.0"): - with pytest.raises(RuntimeError, match="VanillaSort >= 0.2.1.*pip install --upgrade vanillasort"): - VanillaSortSorter._check_params(None, None, VanillaSortSorter.default_params()) - for installed_version in ("0.2.1", "0.10.0"): - params = VanillaSortSorter.default_params() - assert VanillaSortSorter._check_params(None, None, params) == params - - -@pytest.fixture -def tiny_model(tmp_path): - """Use real small model architectures; external CI never downloads weights.""" - import torch - from vanillasort.checkpoints import load_config - from vanillasort.models import VanillaDet, HuiduRep - - config = load_config() - config["detector_architecture"].update( - d_model=8, nhead=2, num_layers=1, dim_ff=16, frontend_channels=8, max_len=5000 - ) - config["huidurep_architecture"].update(embedding_dim=16, n_heads=2, ff_dim=16, num_layers=3) - config["chunk_len"] = 128 - config["profiles"]["d1"].update(base_threshold=0.0001, direct_keep=0.0001) - model_path = tmp_path / "tiny.pt" - with torch.random.fork_rng(): - torch.manual_seed(0) - torch.save( - { - "config": config, - "detector_state_dict": VanillaDet(**config["detector_architecture"]).state_dict(), - "huidurep_state_dict": HuiduRep(**config["huidurep_architecture"]).state_dict(), - }, - model_path, - ) - return model_path - - -def make_recording(multi_segment=False): - """Short deterministic four-channel fixture with enough events for default K.""" - rng = np.random.default_rng(42) - traces = rng.normal(size=(7500, 4)).astype("float32") - t = np.arange(-30, 30) - waveform = -np.exp(-((t / 3.0) ** 2)) + 0.35 * np.exp(-(((t - 9) / 6.0) ** 2)) - for index, center in enumerate(range(100, len(traces) - 100, 150)): - amplitudes = np.roll([8.0, 5.0, 2.0, 1.0], index % 4) - traces[center - 30 : center + 30] += (waveform[:, None] * amplitudes).astype("float32") - segments = [traces, traces[:3750].copy()] if multi_segment else [traces] - recording = NumpyRecording(segments, 30000, channel_ids=["a", "b", "c", "d"]) - recording.set_dummy_probe_from_locations(np.array([[0, 0], [20, 0], [0, 20], [20, 20]])) - return recording - - @pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort not installed") class VanillaSortCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): SorterClass = VanillaSortSorter - - @pytest.fixture(autouse=True) - def configure_model(self, tiny_model): - self.model_path = tiny_model - - def setUp(self): - defaults = patch.dict( - self.SorterClass._default_params, - { - "model_path": str(self.model_path), - "components": 2, - "device": "cpu", - }, - ) - defaults.start() - self.addCleanup(defaults.stop) - # Keep the common tests, using 0.25 s instead of the suite's 60 s fixture. - self.recording = make_recording().save(folder=self.cache_folder / "rec", verbose=False, format="binary") - - -@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort not installed") -def test_vanillasort_direct_and_wrapper(tmp_path, tiny_model, capsys): - import vanillasort - - assert "vanillasort" in available_sorters() - assert VanillaSortSorter.get_sorter_version() == vanillasort.__version__ - assert set(VanillaSortSorter.default_params()) == set(VanillaSortSorter.params_description()) - assert not VanillaSortSorter.use_gpu({"device": "cpu"}) - recording = make_recording(multi_segment=True).save(folder=tmp_path / "rec", verbose=False) - params = dict(components=4, device="cpu", seed=0, model_path=str(tiny_model)) - direct = vanillasort.sort(recording, verbose=False, **params) - assert capsys.readouterr().out == "" - folder = tmp_path / "wrapped" - stdout = sys.stdout - wrapped = run_sorter("vanillasort", recording, folder=folder, verbose=True, **params) - assert sys.stdout is stdout - progress = capsys.readouterr().out - for message in ("VanillaSort ", "Preprocessing ", "Detector:", "Completed:"): - assert message in progress - assert isinstance(wrapped, BaseSorting) - assert wrapped.get_num_segments() == 2 - assert wrapped.to_spike_vector().size > 0 - np.testing.assert_array_equal(wrapped.unit_ids, direct.unit_ids) - np.testing.assert_array_equal(wrapped.to_spike_vector(), direct.to_spike_vector()) - np.testing.assert_array_equal(wrapped.get_property("main_channel_id"), direct.get_property("main_channel_id")) - restored = VanillaSortSorter.get_result_from_folder(folder) - np.testing.assert_array_equal(restored.to_spike_vector(), direct.to_spike_vector()) - np.testing.assert_array_equal(restored.get_property("main_channel_id"), direct.get_property("main_channel_id")) - assert restored.has_recording() - assert restored.sorting_info["params"]["sorter_params"]["seed"] == 0 - log = json.loads((folder / "spikeinterface_log.json").read_text()) - assert log["sorter_version"] == vanillasort.__version__ - assert log["error"] is False and log["run_time"] > 0 - assert any("Completed:" in line for line in log["runtime_trace"]) - assert all(line in progress for line in log["runtime_trace"]) - - quiet_folder = tmp_path / "quiet" - quiet = run_sorter("vanillasort", recording, folder=quiet_folder, verbose=False, **params) - assert sys.stdout is stdout - assert capsys.readouterr().out == "" - np.testing.assert_array_equal(quiet.to_spike_vector(), direct.to_spike_vector()) - quiet_log = json.loads((quiet_folder / "spikeinterface_log.json").read_text()) - assert any("Completed:" in line for line in quiet_log["runtime_trace"]) - - # BaseSorter retains normal error logs for incompatible model input. - bad = NumpyRecording([np.zeros((200, 4), dtype="float32")], 20000) - bad.set_dummy_probe_from_locations(recording.get_channel_locations()) - bad = bad.save(folder=tmp_path / "bad_rec", verbose=False) - failed_folder = tmp_path / "failed" - failed = run_sorter( - "vanillasort", bad, folder=failed_folder, verbose=True, raise_error=False, with_output=False, **params - ) - assert sys.stdout is stdout - assert failed is None - failed_log = json.loads((failed_folder / "spikeinterface_log.json").read_text()) - assert failed_log["error"] is True - assert "30 kHz" in "\n".join(failed_log["error_trace"]) - - -@pytest.mark.skipif(not VanillaSortSorter.is_installed(), reason="vanillasort not installed") -def test_vanillasort_extended_params_roundtrip(tmp_path, tiny_model): - import vanillasort - - base = make_recording(multi_segment=True) - segments = [] - for segment in range(base.get_num_segments()): - traces = base.get_traces(segment_index=segment) - segments.append(np.column_stack([traces, 0.8 * np.roll(traces, 13, axis=0)])) - recording = NumpyRecording(segments, 30000, channel_ids=list("abcdefgh")) - recording.set_dummy_probe_from_locations(np.column_stack([np.tile([0, 20], 4), np.repeat([0, 20, 40, 60], 2)])) - recording = recording.save(folder=tmp_path / "rec", verbose=False) - options = dict( - enable_gate=False, - threshold=0.0002, - direct_keep=0.75, - primary_snr=4, - secondary_snr=2.5, - min_secondary_channels=2, - neighborhood_mode="cover", - target_units=1, - postprocessing={"core_filter": False, "on_shortfall": "return_available", "min_cosine": 0.8}, - ) - params = dict(components=4, device="cpu", seed=0, model_path=str(tiny_model), **options) - direct = vanillasort.sort(recording, verbose=False, **params) - folder = tmp_path / "wrapped" - wrapped = run_sorter("vanillasort", recording, folder=folder, verbose=False, **params) - np.testing.assert_array_equal(wrapped.unit_ids, direct.unit_ids) - np.testing.assert_array_equal(wrapped.to_spike_vector(), direct.to_spike_vector()) - np.testing.assert_array_equal(wrapped.get_property("main_channel_id"), direct.get_property("main_channel_id")) - assert wrapped.get_num_segments() == 2 and wrapped.get_num_units() == 2 - for segment in range(2): - active_units = [u for u in wrapped.unit_ids if len(wrapped.get_unit_spike_train(u, segment_index=segment))] - assert len(active_units) == 1 - - saved = json.loads((folder / "spikeinterface_params.json").read_text())["sorter_params"] - assert {name: saved[name] for name in options} == options - inference = folder / "sorter_output" / "inference" - audit = json.loads((inference / "run.json").read_text()) - assert audit["neighborhood_mode"] == "cover" and audit["target_units"] == 1 - assert len(audit["runs"]) == 4 # two covering patches in each segment - gate = wrapped.get_annotation("detection_gate") - for name in ("enable_gate", "threshold", "direct_keep", "primary_snr", "secondary_snr", "min_secondary_channels"): - assert gate[name] == options[name] - assert gate["mode"] == "off" - for segment, diagnostic in enumerate(wrapped.get_annotation("fixed_k_postprocessing")): - assert diagnostic["requested_units"] == diagnostic["output_units"] == 1 - for name, value in options["postprocessing"].items(): - assert diagnostic["config"][name] == value - assert (inference / f"postprocessing_segment{segment}.json").is_file() - - # A high score threshold produces empty results with an explicit shortfall policy. - empty_params = dict(params, threshold=1.0) - empty = run_sorter("vanillasort", recording, folder=tmp_path / "empty", verbose=False, **empty_params) - assert empty.get_num_units() == 0 and empty.get_num_segments() == 2 - assert all(d["status"] == "partial" for d in empty.get_annotation("fixed_k_postprocessing")) - - failed_folder = tmp_path / "failed" - failed = run_sorter( - "vanillasort", - recording, - folder=failed_folder, - verbose=False, - raise_error=False, - with_output=False, - **dict(empty_params, postprocessing={"core_filter": False}), - ) - assert failed is None - failure = json.loads((failed_folder / "spikeinterface_log.json").read_text()) - assert failure["error"] and "InsufficientAnchorsError" in "\n".join(failure["error_trace"]) - diagnostic = json.loads( - (failed_folder / "sorter_output" / "inference" / "postprocessing_segment0.json").read_text() - ) - assert diagnostic["failure_reason"] == "empty_input" - - with pytest.raises(ValueError, match="Invalid parameters.*gate_fn"): - VanillaSortSorter.set_params_to_folder(recording, tmp_path, {"gate_fn": lambda context: []}, verbose=False) diff --git a/src/spikeinterface/sorters/external/vanillasort.py b/src/spikeinterface/sorters/external/vanillasort.py index b90e9aa7a3..b3299c69c2 100644 --- a/src/spikeinterface/sorters/external/vanillasort.py +++ b/src/spikeinterface/sorters/external/vanillasort.py @@ -84,9 +84,10 @@ class VanillaSortSorter(BaseSorter): sorter_description = """VanillaSort combines VanillaDet detection, HuiduRep waveform embeddings and relative-amplitude features with Gaussian-mixture clustering and template-residual refinement. - The pretrained models require 30 kHz recordings and 2D channel locations, with at least four channels + 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. - See https://github.com/IgarashiAkatuki/VanillaSort and https://doi.org/10.64898/2026.09.18.752552. + 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 = """