diff --git a/.github/workflows/all-tests.yml b/.github/workflows/all-tests.yml index 25b220a8a6..e4a2888161 100644 --- a/.github/workflows/all-tests.yml +++ b/.github/workflows/all-tests.yml @@ -25,7 +25,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.13"] # Lower and higher versions we support + python-version: ["3.10", "3.14"] # Lower and higher versions we support os: [macos-latest, windows-latest, ubuntu-latest] steps: - uses: actions/checkout@v6 @@ -94,16 +94,20 @@ jobs: - name: Pip list run: uv pip list --system - - name: Test core - run: pytest -m "core" - shell: bash - - name: Install Dependencies for Timing Display run: | uv pip install --system tabulate uv pip install --system pandas shell: bash + - name: Set execute permissions on run_tests.sh + shell: bash + run: chmod +x .github/run_tests.sh + + - name: Test core + shell: bash + run: ./.github/run_tests.sh core --no-virtual-env + - name: Get current hash (SHA) of the ephy_testing_data repo shell: bash id: repo_hash @@ -125,7 +129,8 @@ jobs: run: | uv pip install --system datalad-installer if [ ${{ runner.os }} = 'Linux' ]; then - datalad-installer --sudo ok git-annex --method datalad/packages + # TODO: remove the pin once datalad uploads the git-annex 10.20261006 package (its latest-version points to a missing file) + datalad-installer --sudo ok git-annex=10.20260901 --method datalad/packages elif [ ${{ runner.os }} = 'macOS' ]; then datalad-installer --sudo ok git-annex --method brew elif [ ${{ runner.os }} = 'Windows' ]; then @@ -137,10 +142,6 @@ jobs: if: env.INSTALL_PLEXON_DEPENDENCIES == 'true' uses: ./.github/actions/install-wine - - name: Set execute permissions on run_tests.sh - shell: bash - run: chmod +x .github/run_tests.sh - - name: Test extractors shell: bash env: diff --git a/.github/workflows/core-test.yml b/.github/workflows/core-test.yml deleted file mode 100644 index e64c49851e..0000000000 --- a/.github/workflows/core-test.yml +++ /dev/null @@ -1,48 +0,0 @@ -name: Testing core - -on: - pull_request: - types: [synchronize, opened, reopened] - branches: - - main - -concurrency: # Cancel previous workflows on the same pull request - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: true - -jobs: - build-and-test: - name: Test on ${{ matrix.os }} OS - runs-on: ${{ matrix.os }} - strategy: - fail-fast: false - matrix: - os: ["ubuntu-latest", "macos-latest", "windows-latest"] - steps: - - uses: actions/checkout@v6 - - uses: actions/setup-python@v6 - with: - python-version: "3.10" - - uses: astral-sh/setup-uv@v7 - with: - python-version: "3.10" - enable-cache: true - - name: Install dependencies - run: | - git config --global user.email "CI@example.com" - git config --global user.name "CI Almighty" - uv pip install --system -e . --group test-core - - name: Test core with pytest - run: | - pytest -m "core" -vv -ra --durations=0 --durations-min=0.001 | tee report.txt; test $? -eq 0 || exit 1 - shell: bash # Necessary for pipeline to work on windows - - name: Build test summary - run: | - uv pip install --system pandas - uv pip install --system tabulate - echo "# Timing profile of core tests in ${{matrix.os}}" >> $GITHUB_STEP_SUMMARY - # Outputs markdown summary to standard output - python ./.github/scripts/build_job_summary.py report.txt >> $GITHUB_STEP_SUMMARY - cat $GITHUB_STEP_SUMMARY - rm report.txt - shell: bash # Necessary for pipeline to work on windows diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9fd06a2be6..03c2a8c178 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,7 +6,7 @@ repos: - id: end-of-file-fixer - id: trailing-whitespace - repo: https://github.com/psf/black-pre-commit-mirror - rev: 26.5.1 + rev: 26.10.0 hooks: - id: black files: ^src/ diff --git a/conftest.py b/conftest.py index e326c60d71..e4a08c170a 100644 --- a/conftest.py +++ b/conftest.py @@ -1,3 +1,4 @@ +import platform import pytest from pathlib import Path @@ -21,6 +22,24 @@ def pytest_addoption(parser): default=False, help="Enable debug plots during tests", ) + # Users on Linux get fork by default but the tests run with forkserver (the default since Python 3.14) + parser.addoption( + "--mp-context", + default="forkserver" if platform.system() == "Linux" else None, + help="Multiprocessing context used by the tests instead of the default one (forkserver on Linux)", + ) + + +def pytest_configure(config): + mp_context = config.getoption("--mp-context") + if mp_context is not None: + import multiprocessing + import spikeinterface.core.globals as si_globals + + # The default is patched as well so it survives reset_global_job_kwargs() + multiprocessing.set_start_method(mp_context, force=True) + si_globals._default_job_kwargs["mp_context"] = mp_context + si_globals.global_job_kwargs["mp_context"] = mp_context def pytest_collection_modifyitems(config, items): @@ -29,9 +48,18 @@ def pytest_collection_modifyitems(config, items): Marking them in turn allows the tests to be run by using the pytest -m marker_name option. """ + from spikeinterface.core.core_tools import _is_zarr_write_supported + rootdir = Path(config.rootdir) modules_location = rootdir / "src" / "spikeinterface" for item in items: + # TODO: remove once writing to zarr is supported with zarr>=3 + if item.get_closest_marker("requires_zarr_write") and not _is_zarr_write_supported(): + item.add_marker(pytest.mark.skip(reason="Writing to zarr is not supported yet with zarr>=3")) + + if config.getoption("--mp-context") is not None and item.name == "test_global_job_kwargs": + item.add_marker(pytest.mark.skip(reason="--mp-context changes the default job kwargs")) + try: rel_path = Path(item.fspath).relative_to(modules_location) except: diff --git a/doc/api.rst b/doc/api.rst index 335bc186aa..c0ace59d3b 100755 --- a/doc/api.rst +++ b/doc/api.rst @@ -111,6 +111,7 @@ NEO-based .. autofunction:: read_neuralynx .. autofunction:: read_neuralynx_sorting .. autofunction:: read_neuroexplorer + .. autofunction:: read_neuroexplorer_sorting .. autofunction:: read_neuroscope .. autofunction:: read_nix .. autofunction:: read_openephys diff --git a/doc/how_to/process_by_channel_group.rst b/doc/how_to/process_by_channel_group.rst index f1dbb80a42..e92be4eff1 100644 --- a/doc/how_to/process_by_channel_group.rst +++ b/doc/how_to/process_by_channel_group.rst @@ -188,9 +188,9 @@ to different groups. for group, sub_recording in split_preprocessed_recording.items(): sorting = run_sorter( sorter_name='kilosort2', - recording=split_preprocessed_recording, + recording=sub_recording, folder=f"folder_KS2_group{group}" - ) + ) sortings[group] = sorting diff --git a/doc/modules/extractors.rst b/doc/modules/extractors.rst index 89ea3dbc06..3179da863f 100644 --- a/doc/modules/extractors.rst +++ b/doc/modules/extractors.rst @@ -175,6 +175,7 @@ For sorted data formats, we currently support: * **Neurodata Without Borders** :py:func:`~spikeinterface.extractors.read_nwb_sorting()` * **Neuroscope** :py:func:`~spikeinterface.extractors.read_neuroscope_sorting()` * **Neuralynx spikes** :py:func:`~spikeinterface.extractors.read_neuralynx_sorting()` +* **NeuroExplorer spikes** :py:func:`~spikeinterface.extractors.read_neuroexplorer_sorting()` * **NPZ (created by SpikeInterface)** :py:func:`~spikeinterface.core.read_npz_sorting()` * **Plexon spikes** :py:func:`~spikeinterface.extractors.read_plexon_sorting()` * **Plexon 2 spikes** :py:func:`~spikeinterface.extractors.read_plexon2_sorting()` diff --git a/doc/releases/0.105.1.rst b/doc/releases/0.105.1.rst new file mode 100644 index 0000000000..fa72a6f9b2 --- /dev/null +++ b/doc/releases/0.105.1.rst @@ -0,0 +1,85 @@ +.. _release0.105.1: + +SpikeInterface 0.105.1 release notes +------------------------------------ + +October 9th 2026 + +Minor release with bug fixes and performance improvements + + +core: + +* Override search_cached_spikes_sorted for identity selections (#4845) +* Fix group (and other property) propagation in channel slice (#4838) +* Fix: remove ``set_times`` warning (#4834) +* Enable case omission on ``get_unit_spike_train_in_seconds`` (#4811) +* Optimize ``SpikeRetriever`` and RAM usage for extensions (#4809) +* Fix: provenance handling in ``SortingAnalyzer`` (#4805) +* Fix: drop new_unit_ids of non-mergeable units (#4804) +* Fix: zarr pickling in ``SortingAnalyzer`` with ``try``/``except`` (#4801) +* Fix: make seed JSON serializable (#4798) +* Fix sorting times segment propagation (#4796) +* UnitSelectionSorting improvements: remove lexsort, fix bug, save memory by sharing parent cache when possible (#4783) +* Extract waveforms to zarr dataset (#4708) + +extractors: + +* Add ``NeuroExplorerSortingExtractor`` (#4825) +* Don't mount remote paths when running a sorter in a container (#4821) +* Build one probe per electrode group in NwbRecordingExtractor (#4820) +* Read the last spike from the unit end offsets in ``NwbSortingExtractor`` (#4819) +* Don't follow symlinks when storing extractor paths (#4812) + +preprocessing: + +* Fix: artifact detection crossing chunks (#4803) +* Speed up FFT phase shifts for Neuropixels recordings (#4802) +* Fix: frame slicing after resample/decimation (#4799) + +sorters: + +* Don't mount remote paths when running a sorter in a container (#4821) + +postprocessing: + +* Test empty postprocessing and metric extension outputs (#4807) +* Precompute unit-pair spatial overlap once in find_collisions (#4797) +* Avoid dense waveform allocation when exporting PC features (#4789) +* Speed up template similarity with pair-level parallelism (#4785) + +curation: + +* Fix: cast infinites to NaN is overflow error (#4830) + +sortingcomponents: + +* Speed up Wobble peak search with a linear maximum filter (#4782) + +documentation: + +* Fix manual split example in process_by_channel_group docs (#4846) +* Clarify PSD threshold units and recording scaling (#4808) + +continuous integration: + +* Drop the separate core test workflow (#4822) +* Run the tests with ``forkserver`` on Linux (#4818) + +packaging: + +* Support python 3.14 (#4173) + +testing: + +* Test empty postprocessing and metric extension outputs (#4807) + +Contributors: + +* @alejoe91 +* @AtomicGlance +* @grahamfindlay +* @guptaishaan +* @h-mayorquin +* @JESUSROYETH +* @JoeZiminski diff --git a/doc/whatisnew.rst b/doc/whatisnew.rst index 9d79a612c4..a6524f7a43 100644 --- a/doc/whatisnew.rst +++ b/doc/whatisnew.rst @@ -8,6 +8,7 @@ Release notes .. toctree:: :maxdepth: 1 + releases/0.105.1.rst releases/0.105.0.rst releases/0.104.9.rst releases/0.104.8.rst @@ -63,6 +64,11 @@ Release notes releases/0.9.1.rst +Versions 0.105.1 +================ + +* Minor releases with bug fixes and performance improvements + Version 0.105.0 =============== diff --git a/examples/tutorials/core/plot_4_sorting_analyzer.py b/examples/tutorials/core/plot_4_sorting_analyzer.py index 7c89088c09..c52426251c 100644 --- a/examples/tutorials/core/plot_4_sorting_analyzer.py +++ b/examples/tutorials/core/plot_4_sorting_analyzer.py @@ -33,6 +33,8 @@ # First let's generate a simulated recording and sorting recording, sorting = generate_ground_truth_recording() +# Make the recording JSON serializable, so it can be reloaded later +recording = recording.save(folder="recording_folder") print(recording) print(sorting) diff --git a/pyproject.toml b/pyproject.toml index fc44c7030d..5fb13c9edd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "spikeinterface" -version = "0.105.0" +version = "0.105.1" authors = [ { name="Alessio Buccino", email="alessiop.buccino@gmail.com" }, { name="Samuel Garcia", email="sam.garcia.die@gmail.com" }, @@ -20,10 +20,12 @@ classifiers = [ dependencies = [ - "numpy>=2.0.0", + "numpy>=2.0.0;python_version<'3.14'", + "numpy>=2.3.2;python_version>='3.14'", "threadpoolctl>=3.0.0", "tqdm", - "zarr>=2.18,<3", + "zarr>=2.18,<3;python_version<'3.14'", + "zarr>=3;python_version>='3.14'", # Zarr writing is not yet supported with zarr 3 # for release we need pypi, so this needs to be commented # "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", @@ -31,7 +33,8 @@ dependencies = [ "probeinterface>=0.4.0", "packaging", "pydantic", - "numcodecs<0.16.0", # For supporting zarr < 3 + "numcodecs<0.16.0;python_version<'3.14'", # For supporting zarr < 3 + "numcodecs>=0.16.5;python_version>='3.14'", ] [build-system] @@ -340,6 +343,12 @@ test-all = [ {include-group = "test-generation"}, ] +# Tools for the pre-commit hooks. uv installs this group by default. +dev = [ + "black", + "pre-commit", +] + [tool.pytest.ini_options] testpaths = ["src"] markers = [ @@ -358,6 +367,7 @@ markers = [ "widgets", "sortingcomponents", "streaming_extractors: extractors that require streaming such as ross and fsspec", + "requires_zarr_write: tests that write to zarr, skipped with zarr>=3 where writing is not supported yet", ] filterwarnings =[ 'ignore:.*distutils Version classes are deprecated.*:DeprecationWarning', diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index c2261fb390..070176244f 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/src/spikeinterface/core/analyzer_extension_core.py @@ -182,7 +182,6 @@ def _run(self, verbose=False, **job_kwargs): recording = self.sorting_analyzer.recording sorting = self.sorting_analyzer.sorting - unit_ids = sorting.unit_ids # retrieve spike vector and the sampling some_spikes = self.sorting_analyzer.get_extension("random_spikes").get_random_spikes() @@ -192,6 +191,10 @@ def _run(self, verbose=False, **job_kwargs): file_path = self._get_binary_extension_folder() / "waveforms.npy" mode = "memmap" copy = False + elif self.format == "zarr": + file_path = self._get_binary_extension_folder() / "waveforms" + mode = "zarr" + copy = False else: file_path = None mode = "shared_memory" @@ -205,7 +208,6 @@ def _run(self, verbose=False, **job_kwargs): all_waveforms = extract_waveforms_to_single_buffer( recording, some_spikes, - unit_ids, self.nbefore, self.nafter, mode=mode, @@ -218,6 +220,8 @@ def _run(self, verbose=False, **job_kwargs): verbose=verbose, **job_kwargs, ) + if not self.sorting_analyzer._lazy and not isinstance(all_waveforms, np.ndarray): + all_waveforms = materialize_array(all_waveforms) self.data["waveforms"] = all_waveforms diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index 1e136f2dff..fdd15e08e6 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -149,6 +149,44 @@ def get_total_duration(self) -> float: ), "This methods requires an associated recording. Call self.register_recording() first." return self._recording.get_total_duration() + def search_cached_spikes_sorted( + self, + indices: list[int], + segment_index: int | None = None, + ): + """ + Search sample indices (frames) in the cached spike vector of one segment. + + Equivalent to `np.searchsorted(segment_sample_index, indices, side="left")`. + Sortings with a lazy spike vector override this to avoid materialising it. + + Parameters + ---------- + indices : list[int] + The sample indices (frames) to search. + segment_index : int | None, default: None + The segment to search. Can be None for mono-segment sortings. + + Returns + ------- + positions : np.ndarray + The insertion positions, relative to the start of the segment in the spike vector. + """ + if self._cached_spike_vector is None: + self._compute_and_cache_spike_vector() + spikes = self._cached_spike_vector + if not isinstance(spikes, np.ndarray): # np.memmap is an ndarray + # np.searchsorted would materialise a lazy vector on every call + raise TypeError( + f"{type(self).__name__} holds a lazy spike vector ({type(spikes).__name__}) " + "and must override search_cached_spikes_sorted()" + ) + if segment_index is None: + assert self.get_num_segments() == 1, "segment_index is required for multi-segment sortings" + segment_index = 0 + start, stop = self._get_spike_vector_segment_slices()[segment_index] + return np.searchsorted(spikes["sample_index"][start:stop], indices) + def get_unit_spike_train( self, unit_id: str | int, @@ -200,6 +238,7 @@ def get_unit_spike_train( segment_index=segment_index, start_time=start_time, end_time=end_time, + use_cache=use_cache, ) segment_index = self._check_segment_index(segment_index) @@ -237,6 +276,7 @@ def get_unit_spike_train_in_seconds( segment_index: int | None = None, start_time: float | None = None, end_time: float | None = None, + use_cache: bool = True, ) -> np.ndarray: """ Get spike train for a unit in seconds. @@ -261,6 +301,9 @@ def get_unit_spike_train_in_seconds( The start time in seconds for spike train extraction end_time : float or None, default: None The end time in seconds for spike train extraction + use_cache : bool, default: True + Passed to `get_unit_spike_train` when the times are computed from frames. + Ignored when the segment returns native times. Returns ------- @@ -283,7 +326,7 @@ def get_unit_spike_train_in_seconds( start_frame=start_frame, end_frame=end_frame, return_times=False, - use_cache=True, + use_cache=use_cache, ) spike_times = self.sample_index_to_time(spike_frames, segment_index=segment_index) @@ -322,7 +365,7 @@ def get_unit_spike_train_in_seconds( start_frame=start_frame, end_frame=end_frame, return_times=False, - use_cache=True, + use_cache=use_cache, ) t_start = segment._t_start if segment._t_start is not None else 0 @@ -535,6 +578,8 @@ def save(self, format="numpy_folder", **save_kwargs): * "numpy_folder" format: - folder : str or Path The folder where the files will be saved. + - mmap_mode : str or None, default: None + The memory-mapping mode to use when saving numpy files. If None, no memory mapping is used. - overwrite : bool, default: False If True, existing files in the folder will be overwritten. * "zarr" format: @@ -1076,11 +1121,16 @@ def _get_spike_vector_segment_slices(self): if self._cached_spike_vector_segment_slices is None: # compute the, this is needed when spikevector is loaded from format and not computed num_seg = self.get_num_segments() - slices = np.searchsorted(self._cached_spike_vector["segment_index"], np.arange(num_seg + 1)) - self._cached_spike_vector_segment_slices = np.zeros((num_seg, 2), dtype="int64") - for seg_index in range(num_seg): - self._cached_spike_vector_segment_slices[seg_index, 0] = slices[seg_index] - self._cached_spike_vector_segment_slices[seg_index, 1] = slices[seg_index + 1] + if num_seg == 1: + self._cached_spike_vector_segment_slices = np.array( + [[0, self._cached_spike_vector.size]], dtype="int64" + ) + else: + slices = np.searchsorted(self._cached_spike_vector["segment_index"], np.arange(num_seg + 1)) + self._cached_spike_vector_segment_slices = np.zeros((num_seg, 2), dtype="int64") + for seg_index in range(num_seg): + self._cached_spike_vector_segment_slices[seg_index, 0] = slices[seg_index] + self._cached_spike_vector_segment_slices[seg_index, 1] = slices[seg_index + 1] return self._cached_spike_vector_segment_slices def to_reordered_spike_vector( diff --git a/src/spikeinterface/core/binaryfolder.py b/src/spikeinterface/core/binaryfolder.py index c0090719ca..c3c4e07cb8 100644 --- a/src/spikeinterface/core/binaryfolder.py +++ b/src/spikeinterface/core/binaryfolder.py @@ -15,7 +15,7 @@ make_paths_absolute, load_properties_from_folder, save_properties_to_folder, - save_extractor_provenance, + save_provenance_to_folder, save_annotations_to_folder, load_annotations_from_folder, ) @@ -160,7 +160,7 @@ def write_recording( ) save_properties_to_folder(folder_path / "properties", recording) - save_extractor_provenance(folder_path, recording) + save_provenance_to_folder(folder_path, recording, relative_to=folder_path) # new in version 0.105.0, before that annotations were handle by "si_folder.json" file save_annotations_to_folder(folder_path, recording) diff --git a/src/spikeinterface/core/channelslice.py b/src/spikeinterface/core/channelslice.py index 4630e0be24..b4478360f7 100644 --- a/src/spikeinterface/core/channelslice.py +++ b/src/spikeinterface/core/channelslice.py @@ -58,8 +58,8 @@ def __init__(self, recording, channel_ids=None, renamed_channel_ids=None): self.add_recording_segment(sub_segment) # copy annotation and properties - recording.copy_metadata(self, only_main=False, ids=self._channel_ids) self._parent = recording + recording.copy_metadata(self, only_main=False, ids=self._channel_ids) # change the wiring of the probe if self._parent.has_probe(): @@ -67,6 +67,8 @@ def __init__(self, recording, channel_ids=None, renamed_channel_ids=None): sliced_probegroup = parent_probegroup.get_slice(self._parent_channel_indices) sliced_probegroup.set_global_device_channel_indices(np.arange(len(self._channel_ids))) self.set_probegroup(sliced_probegroup) + # Reset channel groups to original ones to avoid remapping + self.set_channel_groups(recording.get_channel_groups()[self._parent_channel_indices]) # update dump dict self._kwargs = { diff --git a/src/spikeinterface/core/core_tools.py b/src/spikeinterface/core/core_tools.py index b5ee518208..58d53dca91 100644 --- a/src/spikeinterface/core/core_tools.py +++ b/src/spikeinterface/core/core_tools.py @@ -763,6 +763,27 @@ def measure_memory_allocation(measure_in_process: bool = True) -> float: return memory +def _is_zarr_write_supported() -> bool: + """ + Whether writing to zarr is supported with the installed zarr, which is not yet the case for zarr>=3. + """ + return int(zarr.__version__.split(".")[0]) < 3 + + +def _check_zarr_write_is_supported() -> None: + """ + Raise an informative error when writing to zarr with zarr>=3, which is not supported yet. + + zarr>=3 is what gets installed on Python 3.14, where reading existing zarr folders works but writing does not. + """ + if not _is_zarr_write_supported(): + raise NotImplementedError( + f"Writing to zarr is not supported yet with zarr {zarr.__version__}, which is the version installed " + "on Python 3.14. Use Python 3.13 or lower to save in zarr format, or save in another format such as " + "'binary_folder'." + ) + + def is_path_remote(path: str | Path) -> bool: """ Returns True if the path is a remote path (e.g., s3:// or gcs://). @@ -777,7 +798,8 @@ def is_path_remote(path: str | Path) -> bool: is_remote: bool Whether the path is a remote path. """ - return "s3://" in str(path) or "gcs://" in str(path) + remote_schemes = ("s3://", "gcs://", "http://", "https://") + return any(scheme in str(path) for scheme in remote_schemes) def ms_to_samples(ms: float, sampling_frequency: float) -> int: @@ -934,11 +956,11 @@ def load_annotations_from_folder(folder: str | Path, extractor: "BaseExtractor") extractor._annotations.update(annotations) -def save_extractor_provenance(folder: str | Path, extractor: "BaseExtractor"): +def save_provenance_to_folder(folder: str | Path, extractor: "BaseExtractor", relative_to: str | Path | None = None): folder = Path(folder) if extractor.check_serializability("json"): provenance_file_path = folder / f"provenance.json" - extractor.dump_to_json(file_path=provenance_file_path, relative_to=folder) + extractor.dump_to_json(file_path=provenance_file_path, relative_to=relative_to) elif extractor.check_serializability("pickle"): provenance_file = folder / f"provenance.pkl" extractor.dump_to_pickle(provenance_file, relative_to=folder) diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py index 6ecf66918d..c66d68a2e4 100644 --- a/src/spikeinterface/core/node_pipeline.py +++ b/src/spikeinterface/core/node_pipeline.py @@ -12,7 +12,7 @@ from spikeinterface.core import BaseRecording, get_chunk_with_margin from spikeinterface.core.job_tools import TimeSeriesChunkExecutor, fix_job_kwargs, _shared_job_kwargs_doc from spikeinterface.core import get_channel_distances -from spikeinterface.core.core_tools import ms_to_samples, samples_to_ms +from spikeinterface.core.core_tools import ms_to_samples, samples_to_ms, _check_zarr_write_is_supported class PipelineNode: @@ -51,7 +51,11 @@ def __init__( parents = [parents] self.parents = parents - self._kwargs = dict() + self._kwargs = dict( + time_series=time_series, + return_output=return_output, + parents=parents, + ) def get_margin(self): # can optionally be overwritten @@ -107,13 +111,18 @@ class PeakRetriever(PeakSource): def __init__(self, recording, peaks): PipelineNode.__init__(self, recording, return_output=False) - self.peaks = peaks + self._peaks = peaks # precompute segment slice - self.segment_slices = [] - for segment_index in range(recording.get_num_segments()): - i0, i1 = np.searchsorted(peaks["segment_index"], [segment_index, segment_index + 1]) - self.segment_slices.append(slice(i0, i1)) + if recording.get_num_segments() > 1: + self.segment_slices = [] + for segment_index in range(recording.get_num_segments()): + i0, i1 = np.searchsorted(peaks["segment_index"], [segment_index, segment_index + 1]) + self.segment_slices.append(slice(i0, i1)) + else: + self.segment_slices = None + + self._kwargs.update(dict(peaks=peaks)) def get_margin(self): return 0 @@ -122,15 +131,22 @@ def get_dtype(self): return base_peak_dtype def get_peak_slice(self, segment_index, start_frame, end_frame, max_margin): - sl = self.segment_slices[segment_index] - peaks_in_segment = self.peaks[sl] + if self.segment_slices is not None: + sl = self.segment_slices[segment_index] + peaks_in_segment = self._peaks[sl] + else: + peaks_in_segment = self._peaks i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame]) return i0, i1 def compute(self, traces, start_frame, end_frame, segment_index, max_margin, peak_slice): # get local peaks - sl = self.segment_slices[segment_index] - peaks_in_segment = self.peaks[sl] + if self.segment_slices is not None: + sl = self.segment_slices[segment_index] + peaks_in_segment = self._peaks[sl] + else: + peaks_in_segment = self._peaks + # i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame]) i0, i1 = peak_slice local_peaks = peaks_in_segment[i0:i1] @@ -195,7 +211,6 @@ def __init__( category=FutureWarning, stacklevel=2, ) - self._dtype = spike_peak_dtype self.include_spikes_in_margin = include_spikes_in_margin @@ -204,19 +219,36 @@ def __init__( main_channel_ids = sorting.get_property("main_channel_id") assert main_channel_ids is not None, "SpikeRetriever needs the sorting to have `main_channel_id`s." - main_channel_indices = recording.ids_to_indices(main_channel_ids) - self.peaks = sorting_to_peaks(sorting, main_channel_indices, self._dtype) + self.main_channel_indices = recording.ids_to_indices(main_channel_ids) + self.spike_vector, segment_slices = sorting.to_spike_vector(return_slices=True) + self.sorting = sorting + self._peaks = None + + # get_peak_slice() returns positions relative to the segment start + if sorting.get_num_segments() == 1: + self.segment_slices = None + else: + self.segment_slices = [slice(int(s0), int(s1)) for s0, s1 in segment_slices] + + # build any lazy search index now, in the parent, so forked workers inherit it + # instead of each building its own + sorting.search_cached_spikes_sorted([0], segment_index=0) if not channel_from_template: channel_distance = get_channel_distances(recording) self.neighbours_mask = channel_distance <= radius_um self.peak_sign = peak_sign - # precompute segment slice - self.segment_slices = [] - for segment_index in range(recording.get_num_segments()): - i0, i1 = np.searchsorted(self.peaks["segment_index"], [segment_index, segment_index + 1]) - self.segment_slices.append(slice(i0, i1)) + self._kwargs.update( + dict( + sorting=sorting, + channel_from_template=channel_from_template, + extremum_channel_inds=extremum_channel_inds, + radius_um=radius_um, + peak_sign=peak_sign, + include_spikes_in_margin=include_spikes_in_margin, + ) + ) def get_margin(self): return 0 @@ -225,26 +257,32 @@ def get_dtype(self): return self._dtype def get_peak_slice(self, segment_index, start_frame, end_frame, max_margin): - sl = self.segment_slices[segment_index] - peaks_in_segment = self.peaks[sl] if self.include_spikes_in_margin: - i0, i1 = np.searchsorted( - peaks_in_segment["sample_index"], [start_frame - max_margin, end_frame + max_margin] - ) + indices = [start_frame - max_margin, end_frame + max_margin] else: - i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame]) - return i0, i1 + indices = [start_frame, end_frame] + return self.sorting.search_cached_spikes_sorted( + indices=indices, + segment_index=segment_index, + ) def compute(self, traces, start_frame, end_frame, segment_index, max_margin, peak_slice): # get local peaks - sl = self.segment_slices[segment_index] - peaks_in_segment = self.peaks[sl] i0, i1 = peak_slice + if self.segment_slices is not None: + s0 = self.segment_slices[segment_index].start + spikes = self.spike_vector[s0 + i0 : s0 + i1] + else: + spikes = self.spike_vector[i0:i1] - local_peaks = peaks_in_segment[i0:i1] + local_peaks = np.zeros(spikes.size, dtype=self._dtype) + local_peaks["sample_index"] = spikes["sample_index"] + local_peaks["channel_index"] = self.main_channel_indices[spikes["unit_index"]] + local_peaks["amplitude"] = 0.0 + local_peaks["segment_index"] = spikes["segment_index"] + local_peaks["unit_index"] = spikes["unit_index"] # make sample index local to traces - local_peaks = local_peaks.copy() local_peaks["sample_index"] -= start_frame - max_margin # handle flag for margin @@ -347,6 +385,15 @@ def __init__( self.nafter = ms_to_samples(ms_after, sampling_frequency) self.neighbours_mask = None + self._kwargs.update( + dict( + ms_before=ms_before, + ms_after=ms_after, + nbefore=nbefore, + nafter=nafter, + ) + ) + class ExtractDenseWaveforms(WaveformsNode): def __init__( @@ -470,6 +517,13 @@ def __init__( self.neighbours_mask = self.channel_distance <= radius_um self.max_num_chans = np.max(np.sum(self.neighbours_mask, axis=1)) + self._kwargs.update( + dict( + radius_um=radius_um, + sparsity_mask=sparsity_mask, + ) + ) + def get_margin(self): return max(self.nbefore, self.nafter) @@ -1044,6 +1098,7 @@ def __init__( from spikeinterface.core.zarrextractors import get_default_zarr_compressor + _check_zarr_write_is_supported() if compressor == "default": compressor = get_default_zarr_compressor() self.compressor = compressor diff --git a/src/spikeinterface/core/numpyextractors.py b/src/spikeinterface/core/numpyextractors.py index b58a98eddb..7f5f5321b9 100644 --- a/src/spikeinterface/core/numpyextractors.py +++ b/src/spikeinterface/core/numpyextractors.py @@ -276,8 +276,7 @@ def __init__(self, spikes, sampling_frequency, unit_ids): self._serializability["memory"] = True self._serializability["json"] = False - # theorically this should be False but for simplicity make generators simples we still need this. - self._serializability["pickle"] = True + self._serializability["pickle"] = False if spikes.size == 0: nseg = 1 @@ -304,6 +303,9 @@ def from_sorting(source_sorting: BaseSorting, with_metadata=False, copy_spike_ve sorting = NumpySorting(spike_vector, source_sorting.get_sampling_frequency(), source_sorting.unit_ids.copy()) if source_sorting.has_recording(): sorting._recording = source_sorting._recording + for source_segment, target_segment in zip(source_sorting._sorting_segments, sorting._sorting_segments): + target_segment._t_start = source_segment._t_start + target_segment._native_t_start = source_segment._native_t_start if with_metadata: source_sorting.copy_metadata(sorting) return sorting diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 3e80b9fab4..b702237445 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -27,6 +27,7 @@ retrieve_importing_provenance, is_path_remote, clean_zarr_folder_name, + _check_zarr_write_is_supported, ) from .sorting_tools import ( generate_unit_ids_for_merge_group, @@ -714,7 +715,6 @@ def create_binary_folder( info_file = folder / "spikeinterface_info.json" settings_file = folder / "settings.json" sorting_folder = folder / "sorting" - sorting_provenance_file = folder / "sorting_provenance.json" sparsity_file = folder / "sparsity_mask.npy" recording_provenance_file = folder / "recording" # .json (default) or .pickle (if not json-serializable) rec_attributes_file = recording_info_folder / "recording_attributes.json" @@ -731,17 +731,8 @@ def create_binary_folder( json.dump(check_json(settings), f, indent=4) # Save the sorting output - sorting = sorting.save(folder=sorting_folder) - - # Dump sorting provenance - if sorting.check_serializability("json"): - sorting.dump(sorting_provenance_file, relative_to=folder) - elif sorting.check_serializability("pickle"): - sorting.dump(sorting_provenance_file.with_suffix(".pickle"), relative_to=folder) - else: - warnings.warn( - "The sorting provenance is not serializable! The sorting provenance link will be lost for future load" - ) + mmap_mode = "r" if lazy else None + sorting_cached = sorting.save(folder=sorting_folder, relative_to=folder, mmap_mode=mmap_mode) # Save sparsity if sparsity is not None: @@ -776,7 +767,7 @@ def create_binary_folder( # Create SortingAnalyzer sorting_analyzer = SortingAnalyzer( - sorting=sorting, + sorting=sorting_cached, recording=recording, rec_attributes={**rec_attributes_to_save, "probegroup": probegroup}, format="binary_folder", @@ -1023,6 +1014,8 @@ def load_from_binary_folder( def _get_zarr_root(self, mode="r+"): assert mode in ("r+", "a", "r"), "mode must be 'r+', 'a' or 'r'" + if mode != "r": + _check_zarr_write_is_supported() storage_options = self._backend_options.get("storage_options", {}) zarr_root = super_zarr_open(self.folder, mode=mode, storage_options=storage_options) @@ -1042,10 +1035,9 @@ def create_zarr( backend_options: dict | None, lazy: bool = False, ) -> "SortingAnalyzer": - # used by create and save_as - - from .zarrextractors import add_sorting_to_zarr_group + from .zarrextractors import add_sorting_to_zarr_group, ZarrSortingExtractor + _check_zarr_write_is_supported() is_remote = is_path_remote(folder) if not is_remote: folder = clean_zarr_folder_name(folder) @@ -1071,27 +1063,10 @@ def create_zarr( settings = {"return_in_uV": return_in_uV, "peak_sign": peak_sign, "peak_mode": peak_mode} zarr_root.attrs["settings"] = check_json(settings) - # Save the sorting output - add_sorting_to_zarr_group(sorting, sorting_group, **saving_options) - - # Dump sorting provenance relative_to = None if is_remote else folder - sort_dict = sorting.to_dict(relative_to=relative_to, recursive=True) - if sorting.check_serializability("json"): - _write_object_array(zarr_root, "sorting_provenance", check_json(sort_dict), codec="json") - elif sorting.check_serializability("pickle"): - try: - _write_object_array(zarr_root, "sorting_provenance", sort_dict, codec="pickle") - except: - warnings.warn( - "Failed to serialize sorting provenance with Pickle Codec! " - "The sorting provenance link will be lost for future load" - ) - else: - warnings.warn( - "The sorting provenance is not serializable! " - "The sorting provenance link will be lost for future load" - ) + # Save the sorting output + add_sorting_to_zarr_group(sorting, sorting_group, relative_to=folder, **saving_options) + sorting_cached = ZarrSortingExtractor(folder / "sorting", lazy_spike_vector=lazy) # Save sparsity if sparsity is not None: @@ -1103,7 +1078,13 @@ def create_zarr( if recording.check_serializability("json"): _write_object_array(zarr_root, "recording", check_json(rec_dict), codec="json") elif recording.check_serializability("pickle"): - _write_object_array(zarr_root, "recording", rec_dict, codec="pickle") + try: + _write_object_array(zarr_root, "recording", rec_dict, codec="pickle") + except: + warnings.warn( + "Failed to serialize recording with Pickle Codec! " + "The recording link will be lost for future load" + ) else: warnings.warn("The Recording is not serializable! The recording link will be lost for future load") else: @@ -1129,7 +1110,7 @@ def create_zarr( # Create SortingAnalyzer sorting_analyzer = SortingAnalyzer( - sorting=NumpySorting.from_sorting(sorting, with_metadata=True, copy_spike_vector=True), + sorting=sorting_cached, recording=recording, rec_attributes={**rec_attributes_to_save, "probegroup": probegroup}, format="zarr", @@ -1194,22 +1175,13 @@ def load_from_zarr( settings = cls._handle_backward_compatibility_settings_pre_init(settings) # Load sorting (in memory or lazy) + zarr_sorting = ZarrSortingExtractor( + folder, zarr_group="sorting", storage_options=storage_options, lazy_spike_vector=lazy + ) if lazy: - copy_spike_vector = False - lazy_spike_vector = True + sorting = zarr_sorting else: - copy_spike_vector = True - lazy_spike_vector = False - sorting = NumpySorting.from_sorting( - ZarrSortingExtractor( - folder, - zarr_group="sorting", - storage_options=storage_options, - lazy_spike_vector=lazy_spike_vector, - ), - with_metadata=True, - copy_spike_vector=copy_spike_vector, - ) + sorting = NumpySorting.from_sorting(zarr_sorting, with_metadata=True, copy_spike_vector=True) # Load recording (if available) if recording is None: @@ -2000,9 +1972,12 @@ def merge_units( mergeable_unit_groups = [] unmergeable_unit_groups = [] - for merge_unit_group, mergeable in zip(merge_unit_groups, mergeable.values()): + new_unit_ids_mergeable = [] if new_unit_ids is not None else None + for i, (merge_unit_group, mergeable) in enumerate(zip(merge_unit_groups, mergeable.values())): if mergeable: mergeable_unit_groups.append(merge_unit_group) + if new_unit_ids_mergeable is not None: + new_unit_ids_mergeable.append(new_unit_ids[i]) else: unmergeable_unit_groups.append(merge_unit_group) @@ -2015,9 +1990,14 @@ def merge_units( warnings.warn(warning_message) else: mergeable_unit_groups = merge_unit_groups + new_unit_ids_mergeable = new_unit_ids + + if len(mergeable_unit_groups) == 0: + warnings.warn("No mergeable unit groups found.") + return self if not return_new_unit_ids else (self, []) new_unit_ids = generate_unit_ids_for_merge_group( - self.unit_ids, mergeable_unit_groups, new_unit_ids, new_id_strategy + self.unit_ids, mergeable_unit_groups, new_unit_ids_mergeable, new_id_strategy ) all_unit_ids = _get_ids_after_merging(self.unit_ids, mergeable_unit_groups, new_unit_ids=new_unit_ids) @@ -2169,34 +2149,36 @@ def get_sorting_provenance(self): """ from .loading import load + sorting_provenance = None if self.format == "memory": # the original sorting provenance is not kept in that case - sorting_provenance = None - + pass elif self.format == "binary_folder": for type in ("json", "pickle"): - filename = self.folder / f"sorting_provenance.{type}" + filename = self.folder / "sorting" / f"provenance.{type}" sorting_provenance = None + if not filename.is_file(): + # Legacy location + filename = self.folder / "sorting_provenance.json" + # try-except here is because it's not required to be able + # to load the sorting provenance, as the user might have deleted + # the original sorting folder if filename.is_file(): - # try-except here is because it's not required to be able - # to load the sorting provenance, as the user might have deleted - # the original sorting folder try: sorting_provenance = load(filename, base_folder=self.folder) break except: pass - # sorting_provenance = None - elif self.format == "zarr": zarr_root = self._get_zarr_root(mode="r") sorting_provenance = None - if "sorting_provenance" in zarr_root.keys(): - # try-except here is because it's not required to be able - # to load the sorting provenance, as the user might have deleted - # the original sorting folder - try: + sort_dict = zarr_root["sorting"].attrs.get("provenance", None) + if sort_dict is None: + # Legacy + if "sorting_provenance" in zarr_root.keys(): sort_dict = zarr_root["sorting_provenance"][0] + if sort_dict is not None: + try: sorting_provenance = load(sort_dict, base_folder=self.folder) except: pass @@ -2415,7 +2397,6 @@ def compute_one_extension(self, extension_name, save=True, verbose=False, **kwar else: extension_instance.run(save=save, verbose=verbose) if not self._lazy: - for variable_name in extension_instance.data.keys(): # Materialize the data if not in lazy mode if isinstance(extension_instance.data[variable_name], (np.memmap, zarr.Array)): @@ -3402,6 +3383,8 @@ def load_data(self, lazy=False): ext_data.loc[:, col] = ext_data_[col][:] ext_data = ext_data.convert_dtypes() elif "object" in ext_data_.attrs: + if lazy: + continue ext_data = ext_data_[0] else: ext_data = ext_data_ if lazy else np.array(ext_data_[:]) diff --git a/src/spikeinterface/core/sortingfolder.py b/src/spikeinterface/core/sortingfolder.py index 3a068e0384..fefe5c1229 100644 --- a/src/spikeinterface/core/sortingfolder.py +++ b/src/spikeinterface/core/sortingfolder.py @@ -15,7 +15,7 @@ save_properties_to_folder, save_annotations_to_folder, load_annotations_from_folder, - save_extractor_provenance, + save_provenance_to_folder, ) @@ -61,7 +61,13 @@ def __init__(self, folder_path, mmap_mode: str | None = None): self._kwargs = dict(folder_path=str(folder_path.absolute()), mmap_mode=mmap_mode) @staticmethod - def write_sorting(sorting, folder_path, overwrite: bool = False): + def write_sorting( + sorting, + folder_path, + overwrite: bool = False, + mmap_mode: str | None = None, + relative_to: str | Path | None = None, + ): # the folder can already exists but not contaning numpysorting_info.json folder_path = Path(folder_path) if folder_path.is_dir(): @@ -81,14 +87,15 @@ def write_sorting(sorting, folder_path, overwrite: bool = False): np.save(folder_path / "spikes.npy", sorting.to_spike_vector()) save_properties_to_folder(folder_path / "properties", sorting) - save_extractor_provenance(folder_path, sorting) + relative_to = folder_path if relative_to is None else relative_to + save_provenance_to_folder(folder_path, sorting, relative_to=relative_to) # new in version 0.105.0, before that annotations were handle by "si_folder.json" file save_annotations_to_folder(folder_path, sorting) # Create the si_folder file to make the load() easier until version 0.105.0 # All properties, annotations, and probe information are already saved in the folder, # so we don't need to include them in the si_folder.json - cached = NumpyFolderSorting(folder_path=folder_path) + cached = NumpyFolderSorting(folder_path=folder_path, mmap_mode=mmap_mode) si_folder_path = folder_path / f"si_folder.json" cached.dump_to_json( file_path=si_folder_path, diff --git a/src/spikeinterface/core/template.py b/src/spikeinterface/core/template.py index b661b2d83f..0f0922e10b 100644 --- a/src/spikeinterface/core/template.py +++ b/src/spikeinterface/core/template.py @@ -321,6 +321,10 @@ def add_templates_to_zarr_group(self, zarr_group: "zarr.Group") -> None: to optimize read/write operations for individual units. """ + from .core_tools import _check_zarr_write_is_supported + + _check_zarr_write_is_supported() + # Saves one chunk per unit arrays_chunk = (1, None, None) zarr_group.create_dataset("templates_array", data=self.templates_array, chunks=arrays_chunk) @@ -353,6 +357,9 @@ def to_zarr(self, folder_path: str | Path) -> None: """ import zarr + from .core_tools import _check_zarr_write_is_supported + + _check_zarr_write_is_supported() zarr_group = zarr.open_group(folder_path, mode="w") self.add_templates_to_zarr_group(zarr_group) diff --git a/src/spikeinterface/core/tests/test_analyzer_extension_core.py b/src/spikeinterface/core/tests/test_analyzer_extension_core.py index f606ca34ac..8436d7c7a3 100644 --- a/src/spikeinterface/core/tests/test_analyzer_extension_core.py +++ b/src/spikeinterface/core/tests/test_analyzer_extension_core.py @@ -11,6 +11,9 @@ from spikeinterface.core.sortinganalyzer import _extension_children, _get_children_dependencies import numpy as np +from spikeinterface.core.core_tools import _is_zarr_write_supported + +analyzer_formats = ("memory", "binary_folder", "zarr") if _is_zarr_write_supported() else ("memory", "binary_folder") def get_sorting_analyzer(cache_folder, format="memory", sparse=True): @@ -51,7 +54,7 @@ def get_sorting_analyzer(cache_folder, format="memory", sparse=True): def _check_result_extension(sorting_analyzer, extension_name, cache_folder): # select unit_ids to several format - for format in ("memory", "binary_folder", "zarr"): + for format in analyzer_formats: # for format in ("memory", ): if format != "memory": if format == "zarr": @@ -72,7 +75,9 @@ def _check_result_extension(sorting_analyzer, extension_name, cache_folder): # print(k, arr.shape) -@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +@pytest.mark.parametrize( + "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) @pytest.mark.parametrize( "sparse", [ @@ -102,7 +107,9 @@ def test_ComputeRandomSpikes(format, sparse, create_cache_folder): _check_result_extension(sorting_analyzer, "random_spikes", cache_folder) -@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +@pytest.mark.parametrize( + "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) @pytest.mark.parametrize("sparse", [True, False]) def test_ComputeWaveforms(format, sparse, create_cache_folder): cache_folder = create_cache_folder @@ -115,7 +122,65 @@ def test_ComputeWaveforms(format, sparse, create_cache_folder): _check_result_extension(sorting_analyzer, "waveforms", cache_folder) -@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +@pytest.mark.requires_zarr_write +@pytest.mark.parametrize("sparse", [False, True]) +def test_ComputeWaveforms_consistent_across_formats(create_cache_folder, sparse): + # Computing waveforms (and templates) on memory / binary_folder / zarr analyzers with the same + # data and fixed seeds must produce identical results, and the saved on-disk files (the npy for + # binary_folder and the zarr dataset) must match the in-memory computation. + import zarr + + cache_folder = create_cache_folder + recording, sorting = generate_ground_truth_recording( + durations=[15.0], + sampling_frequency=16000.0, + num_channels=8, + num_units=5, + seed=2406, + ) + job_kwargs = dict(n_jobs=2, chunk_duration="1s", progress_bar=False) + + analyzers = {} + for fmt in ["memory", "binary_folder", "zarr"]: + if fmt == "memory": + folder = None + elif fmt == "binary_folder": + folder = cache_folder / f"consistency_across_formats_{sparse}_binary" + else: + folder = cache_folder / f"consistency_across_formats_{sparse}.zarr" + if folder is not None and folder.exists(): + shutil.rmtree(folder) + + sorting_analyzer = create_sorting_analyzer( + sorting, recording, format=fmt, folder=folder, sparse=sparse, sparsity=None + ) + sorting_analyzer.compute("random_spikes", max_spikes_per_unit=50, seed=2205) + sorting_analyzer.compute("waveforms", **job_kwargs) + sorting_analyzer.compute("templates", operators=["average", "std"]) + analyzers[fmt] = sorting_analyzer + + # waveforms and templates must be identical across formats + wfs_mem = np.asarray(analyzers["memory"].get_extension("waveforms").get_data()) + for fmt in ["binary_folder", "zarr"]: + wfs = np.asarray(analyzers[fmt].get_extension("waveforms").get_data()) + assert np.array_equal(wfs, wfs_mem), f"waveforms differ for format {fmt}" + for operator in ["average", "std"]: + template_mem = analyzers["memory"].get_extension("templates").get_templates(operator=operator) + template_fmt = analyzers[fmt].get_extension("templates").get_templates(operator=operator) + assert np.allclose(template_fmt, template_mem), f"templates '{operator}' differ for format {fmt}" + + # the files saved on disk must match the in-memory waveforms + binary_waveforms = np.load(analyzers["binary_folder"].folder / "extensions" / "waveforms" / "waveforms.npy") + assert np.array_equal(binary_waveforms, wfs_mem) + + zarr_root = zarr.open(str(analyzers["zarr"].folder), mode="r") + zarr_waveforms = zarr_root["extensions"]["waveforms"]["waveforms"][:] + assert np.array_equal(zarr_waveforms, wfs_mem) + + +@pytest.mark.parametrize( + "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) @pytest.mark.parametrize("sparse", [True, False]) def test_ComputeTemplates(format, sparse, create_cache_folder): cache_folder = create_cache_folder @@ -204,7 +269,9 @@ def test_ComputeTemplates(format, sparse, create_cache_folder): _check_result_extension(sorting_analyzer, "templates", cache_folder) -@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +@pytest.mark.parametrize( + "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) @pytest.mark.parametrize("sparse", [True, False]) def test_ComputeNoiseLevels(format, sparse, create_cache_folder): cache_folder = create_cache_folder diff --git a/src/spikeinterface/core/tests/test_baserecording.py b/src/spikeinterface/core/tests/test_baserecording.py index 63fe18f899..8fd38ce4bf 100644 --- a/src/spikeinterface/core/tests/test_baserecording.py +++ b/src/spikeinterface/core/tests/test_baserecording.py @@ -24,6 +24,7 @@ from spikeinterface.core.testing import check_recordings_equal from spikeinterface.core import generate_recording +from spikeinterface.core.core_tools import _is_zarr_write_supported def test_BaseRecording(create_cache_folder): @@ -364,32 +365,34 @@ def test_BaseRecording(create_cache_folder): rec_2d = rec_3d.planarize(axes="zy") assert np.allclose(rec_2d.get_channel_locations(), locations_3d[:, [2, 1]]) - # test save to zarr - compressor = get_default_zarr_compressor() - rec_zarr = rec2.save(format="zarr", folder=cache_folder / "recording", compressor=compressor) - rec_zarr_loaded = load(cache_folder / "recording.zarr") - # annotations is False because Zarr adds compression ratios - check_recordings_equal(rec2, rec_zarr, return_in_uV=False, check_annotations=False, check_properties=True) - check_recordings_equal( - rec_zarr, rec_zarr_loaded, return_in_uV=False, check_annotations=False, check_properties=True - ) - for annotation_name in rec2.get_annotation_keys(): - assert rec2.get_annotation(annotation_name) == rec_zarr.get_annotation(annotation_name) - assert rec2.get_annotation(annotation_name) == rec_zarr_loaded.get_annotation(annotation_name) - - rec_zarr2 = rec2.save( - format="zarr", folder=cache_folder / "recording_channel_chunk", compressor=compressor, channel_chunk_size=2 - ) - rec_zarr2_loaded = load(cache_folder / "recording_channel_chunk.zarr") - - # annotations is False because Zarr adds compression ratios - check_recordings_equal(rec2, rec_zarr2, return_in_uV=False, check_annotations=False, check_properties=True) - check_recordings_equal( - rec_zarr2, rec_zarr2_loaded, return_in_uV=False, check_annotations=False, check_properties=True - ) - for annotation_name in rec2.get_annotation_keys(): - assert rec2.get_annotation(annotation_name) == rec_zarr2.get_annotation(annotation_name) - assert rec2.get_annotation(annotation_name) == rec_zarr2_loaded.get_annotation(annotation_name) + # TODO: remove once writing to zarr is supported with zarr>=3 + if _is_zarr_write_supported(): + # test save to zarr + compressor = get_default_zarr_compressor() + rec_zarr = rec2.save(format="zarr", folder=cache_folder / "recording", compressor=compressor) + rec_zarr_loaded = load(cache_folder / "recording.zarr") + # annotations is False because Zarr adds compression ratios + check_recordings_equal(rec2, rec_zarr, return_in_uV=False, check_annotations=False, check_properties=True) + check_recordings_equal( + rec_zarr, rec_zarr_loaded, return_in_uV=False, check_annotations=False, check_properties=True + ) + for annotation_name in rec2.get_annotation_keys(): + assert rec2.get_annotation(annotation_name) == rec_zarr.get_annotation(annotation_name) + assert rec2.get_annotation(annotation_name) == rec_zarr_loaded.get_annotation(annotation_name) + + rec_zarr2 = rec2.save( + format="zarr", folder=cache_folder / "recording_channel_chunk", compressor=compressor, channel_chunk_size=2 + ) + rec_zarr2_loaded = load(cache_folder / "recording_channel_chunk.zarr") + + # annotations is False because Zarr adds compression ratios + check_recordings_equal(rec2, rec_zarr2, return_in_uV=False, check_annotations=False, check_properties=True) + check_recordings_equal( + rec_zarr2, rec_zarr2_loaded, return_in_uV=False, check_annotations=False, check_properties=True + ) + for annotation_name in rec2.get_annotation_keys(): + assert rec2.get_annotation(annotation_name) == rec_zarr2.get_annotation(annotation_name) + assert rec2.get_annotation(annotation_name) == rec_zarr2_loaded.get_annotation(annotation_name) def test_json_pickle_equivalence(create_cache_folder): @@ -477,7 +480,7 @@ def test_time_slice_with_time_vector(): sampling_frequency = 10_000.0 recording = generate_recording(durations=[1.0], num_channels=3, sampling_frequency=sampling_frequency) times = 1 + np.arange(0, 10_000) / sampling_frequency - recording.set_times(times=times, segment_index=0, with_warning=False) + recording.set_times(times=times, segment_index=0) sliced_recording_times = recording.time_slice(start_time=1.1, end_time=1.8) sliced_recording_frames = recording.frame_slice(start_frame=1000, end_frame=8000) @@ -504,7 +507,7 @@ def test_save_load_binary_with_time_vector(create_cache_folder, mp_context): times = rec.get_times(segment_index=0) + 100.0 # Set time vector - rec.set_times(times=times, segment_index=0, with_warning=False) + rec.set_times(times=times, segment_index=0) # Save rec_saved = rec.save(folder=cache_folder / f"recording_with_time_vector_{mp_context}", format="binary") assert np.allclose(rec.get_times(segment_index=0), rec_saved.get_times(segment_index=0)) diff --git a/src/spikeinterface/core/tests/test_basesorting.py b/src/spikeinterface/core/tests/test_basesorting.py index 25ba1a3c2d..ec7c2cefb6 100644 --- a/src/spikeinterface/core/tests/test_basesorting.py +++ b/src/spikeinterface/core/tests/test_basesorting.py @@ -25,6 +25,7 @@ from spikeinterface.core.base import BaseExtractor, minimum_spike_dtype, unit_period_dtype from spikeinterface.core.basesorting import LEXSORT_UNIT_COMPACT from spikeinterface.core.testing import check_sorted_arrays_equal, check_sortings_equal +from spikeinterface.core.core_tools import _is_zarr_write_supported def test_BaseSorting(create_cache_folder): @@ -143,16 +144,18 @@ def test_BaseSorting(create_cache_folder): del sorting6 del sorting5 - # test save to zarr - # compressor = get_default_zarr_compressor() - sorting_zarr = sorting.save(format="zarr", folder=cache_folder / "sorting.zarr") - sorting_zarr_loaded = load(cache_folder / "sorting.zarr") - # annotations is False because Zarr adds compression ratios - check_sortings_equal(sorting, sorting_zarr, check_annotations=False, check_properties=True) - check_sortings_equal(sorting_zarr, sorting_zarr_loaded, check_annotations=False, check_properties=True) - for annotation_name in sorting.get_annotation_keys(): - assert sorting.get_annotation(annotation_name) == sorting_zarr.get_annotation(annotation_name) - assert sorting.get_annotation(annotation_name) == sorting_zarr_loaded.get_annotation(annotation_name) + # TODO: remove once writing to zarr is supported with zarr>=3 + if _is_zarr_write_supported(): + # test save to zarr + # compressor = get_default_zarr_compressor() + sorting_zarr = sorting.save(format="zarr", folder=cache_folder / "sorting.zarr") + sorting_zarr_loaded = load(cache_folder / "sorting.zarr") + # annotations is False because Zarr adds compression ratios + check_sortings_equal(sorting, sorting_zarr, check_annotations=False, check_properties=True) + check_sortings_equal(sorting_zarr, sorting_zarr_loaded, check_annotations=False, check_properties=True) + for annotation_name in sorting.get_annotation_keys(): + assert sorting.get_annotation(annotation_name) == sorting_zarr.get_annotation(annotation_name) + assert sorting.get_annotation(annotation_name) == sorting_zarr_loaded.get_annotation(annotation_name) def _make_sorting_with_shuffled_ties(num_units, num_segments, seed=42): diff --git a/src/spikeinterface/core/tests/test_channelslicerecording.py b/src/spikeinterface/core/tests/test_channelslicerecording.py index c563888e07..0f8247e85e 100644 --- a/src/spikeinterface/core/tests/test_channelslicerecording.py +++ b/src/spikeinterface/core/tests/test_channelslicerecording.py @@ -1,6 +1,3 @@ -import shutil -from pathlib import Path - import pytest import numpy as np @@ -129,5 +126,15 @@ def test_select_channels_preserves_probe_metadata(): assert probes[0].annotations.get("manufacturer") == "vendor_Y" +def test_channel_slice_retain_original_groups(): + rec = generate_recording(num_channels=8, durations=[1.0]) # has a probe attached + rec.set_property("group", [0, 0, 1, 1, 2, 2, 3, 3]) + for g, sub in rec.split_by("group").items(): + assert np.all(sub.get_property("group") == g) + + rec_sliced = rec.remove_channels(remove_channel_ids=["1", "3", "5", "7"]) + assert np.all(rec_sliced.get_property("group") == [0, 1, 2, 3]) + + if __name__ == "__main__": test_ChannelSliceRecording() diff --git a/src/spikeinterface/core/tests/test_loading.py b/src/spikeinterface/core/tests/test_loading.py index 257895fcc2..364c5bd910 100644 --- a/src/spikeinterface/core/tests/test_loading.py +++ b/src/spikeinterface/core/tests/test_loading.py @@ -1,4 +1,5 @@ import pytest +import zarr import numpy as np from spikeinterface import ( @@ -89,7 +90,7 @@ def generate_motion_object(): return motion -@pytest.mark.parametrize("output_format", ["binary", "zarr"]) +@pytest.mark.parametrize("output_format", ["binary", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) def test_load_binary_recording(generate_recording_sorting, tmp_path, output_format): rec, _ = generate_recording_sorting _ = rec.save(folder=tmp_path / "test_recording", format=output_format, overwrite=True) @@ -102,7 +103,7 @@ def test_load_binary_recording(generate_recording_sorting, tmp_path, output_form check_recordings_equal(rec, rec_loaded) -@pytest.mark.parametrize("output_format", ["numpy_folder", "zarr"]) +@pytest.mark.parametrize("output_format", ["numpy_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) def test_load_binary_sorting(generate_recording_sorting, tmp_path, output_format): _, sort = generate_recording_sorting _ = sort.save(folder=tmp_path / "test_sorting", format=output_format, overwrite=True) @@ -140,7 +141,9 @@ def test_load_ext_extractors(generate_recording_sorting, tmp_path, extension): check_sortings_equal(sort, sort_loaded, check_properties=False) -@pytest.mark.parametrize("output_format", ["binary_folder", "zarr"]) +@pytest.mark.parametrize( + "output_format", ["binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) def test_load_sorting_analyzer(generate_sorting_analyzer, tmp_path, output_format): analyzer = generate_sorting_analyzer _ = analyzer.save_as(folder=tmp_path / "analyzer", format=output_format) @@ -157,6 +160,7 @@ def test_load_sorting_analyzer(generate_sorting_analyzer, tmp_path, output_forma assert ext in analyzer_loaded.extensions +@pytest.mark.requires_zarr_write def test_load_templates(tmp_path, generate_templates_object): templates = generate_templates_object templates_dict = templates.to_dict() @@ -212,6 +216,8 @@ def test_remote_recording(): @pytest.mark.streaming_extractors @pytest.mark.skipif(not HAVE_S3, reason="s3fs not installed") +# TODO: remove once streaming zarr is supported with zarr>=3 +@pytest.mark.skipif(int(zarr.__version__.split(".")[0]) >= 3, reason="Streaming zarr is not supported yet with zarr>=3") def test_remote_analyzer(): s3_path = "s3://spikeinterface-sorting-analyzer-test/analyzer_remote_test.zarr/" analyzer = load(s3_path) diff --git a/src/spikeinterface/core/tests/test_node_pipeline.py b/src/spikeinterface/core/tests/test_node_pipeline.py index a721fc3367..533966d703 100644 --- a/src/spikeinterface/core/tests/test_node_pipeline.py +++ b/src/spikeinterface/core/tests/test_node_pipeline.py @@ -16,6 +16,7 @@ ExtractDenseWaveforms, sorting_to_peaks, ) +from spikeinterface.core.core_tools import _is_zarr_write_supported class AmplitudeExtractionNode(PipelineNode): @@ -180,34 +181,36 @@ def test_run_node_pipeline(cache_folder_creation): assert np.array_equal(denoised_waveforms_rms, denoised_waveforms_rms2) assert np.array_equal(denoised_waveforms_rms2, denoised_waveforms_rms3) - # gather zarr mode - import zarr - - zarr_folder = cache_folder / f"pipeline_folder_{loop}.zarr" - if zarr_folder.is_dir(): - shutil.rmtree(zarr_folder) - output = run_node_pipeline( - recording, - nodes, - job_kwargs, - gather_mode="zarr", - dest=zarr_folder, - names=["amplitudes", "waveforms_rms", "denoised_waveforms_rms"], - ) - amplitudes_z, waveforms_rms_z, denoised_waveforms_rms_z = output - - # values must match the memory gather - assert np.array_equal(amplitudes, amplitudes_z[:]) - assert np.array_equal(waveforms_rms, waveforms_rms_z[:]) - assert np.array_equal(denoised_waveforms_rms, denoised_waveforms_rms_z[:]) - - # arrays must be persisted on disk and re-openable - zarr_root = zarr.open(str(zarr_folder), mode="r") - for name in ("amplitudes", "waveforms_rms", "denoised_waveforms_rms"): - assert name in zarr_root - assert np.array_equal(amplitudes, zarr_root["amplitudes"][:]) - assert np.array_equal(waveforms_rms, zarr_root["waveforms_rms"][:]) - assert np.array_equal(denoised_waveforms_rms, zarr_root["denoised_waveforms_rms"][:]) + # TODO: remove once writing to zarr is supported with zarr>=3 + if _is_zarr_write_supported(): + # gather zarr mode + import zarr + + zarr_folder = cache_folder / f"pipeline_folder_{loop}.zarr" + if zarr_folder.is_dir(): + shutil.rmtree(zarr_folder) + output = run_node_pipeline( + recording, + nodes, + job_kwargs, + gather_mode="zarr", + dest=zarr_folder, + names=["amplitudes", "waveforms_rms", "denoised_waveforms_rms"], + ) + amplitudes_z, waveforms_rms_z, denoised_waveforms_rms_z = output + + # values must match the memory gather + assert np.array_equal(amplitudes, amplitudes_z[:]) + assert np.array_equal(waveforms_rms, waveforms_rms_z[:]) + assert np.array_equal(denoised_waveforms_rms, denoised_waveforms_rms_z[:]) + + # arrays must be persisted on disk and re-openable + zarr_root = zarr.open(str(zarr_folder), mode="r") + for name in ("amplitudes", "waveforms_rms", "denoised_waveforms_rms"): + assert name in zarr_root + assert np.array_equal(amplitudes, zarr_root["amplitudes"][:]) + assert np.array_equal(waveforms_rms, zarr_root["waveforms_rms"][:]) + assert np.array_equal(denoised_waveforms_rms, zarr_root["denoised_waveforms_rms"][:]) # gather npy mode with an explicit list of file paths (final location) npy_files_folder = cache_folder / f"pipeline_npy_files_{loop}" @@ -232,36 +235,38 @@ def test_run_node_pipeline(cache_folder_creation): assert np.array_equal(waveforms_rms, waveforms_rms_f) assert np.array_equal(denoised_waveforms_rms, denoised_waveforms_rms_f) - # gather zarr mode with an explicit list of dataset paths, created on the fly - # inside an existing store (final location, e.g. an analyzer extension group) - datasets_store = cache_folder / f"pipeline_zarr_datasets_{loop}.zarr" - if datasets_store.is_dir(): - shutil.rmtree(datasets_store) - # pre-existing store that must not be wiped - root = zarr.open(str(datasets_store), mode="w") - root.attrs["preexisting"] = True - dataset_paths = [ - datasets_store / "extensions" / "amplitudes", - datasets_store / "extensions" / "waveforms_rms", - datasets_store / "extensions" / "denoised_waveforms_rms", - ] - output = run_node_pipeline( - recording, - nodes, - job_kwargs, - gather_mode="zarr", - dest=dataset_paths, - ) - amplitudes_d, waveforms_rms_d, denoised_waveforms_rms_d = output - assert np.array_equal(amplitudes, amplitudes_d[:]) - assert np.array_equal(waveforms_rms, waveforms_rms_d[:]) - assert np.array_equal(denoised_waveforms_rms, denoised_waveforms_rms_d[:]) - # data must be persisted at the passed final location and the store not wiped - root_reopen = zarr.open(str(datasets_store), mode="r") - assert root_reopen.attrs.get("preexisting", False) - assert np.array_equal(amplitudes, root_reopen["extensions"]["amplitudes"][:]) - assert np.array_equal(waveforms_rms, root_reopen["extensions"]["waveforms_rms"][:]) - assert np.array_equal(denoised_waveforms_rms, root_reopen["extensions"]["denoised_waveforms_rms"][:]) + # TODO: remove once writing to zarr is supported with zarr>=3 + if _is_zarr_write_supported(): + # gather zarr mode with an explicit list of dataset paths, created on the fly + # inside an existing store (final location, e.g. an analyzer extension group) + datasets_store = cache_folder / f"pipeline_zarr_datasets_{loop}.zarr" + if datasets_store.is_dir(): + shutil.rmtree(datasets_store) + # pre-existing store that must not be wiped + root = zarr.open(str(datasets_store), mode="w") + root.attrs["preexisting"] = True + dataset_paths = [ + datasets_store / "extensions" / "amplitudes", + datasets_store / "extensions" / "waveforms_rms", + datasets_store / "extensions" / "denoised_waveforms_rms", + ] + output = run_node_pipeline( + recording, + nodes, + job_kwargs, + gather_mode="zarr", + dest=dataset_paths, + ) + amplitudes_d, waveforms_rms_d, denoised_waveforms_rms_d = output + assert np.array_equal(amplitudes, amplitudes_d[:]) + assert np.array_equal(waveforms_rms, waveforms_rms_d[:]) + assert np.array_equal(denoised_waveforms_rms, denoised_waveforms_rms_d[:]) + # data must be persisted at the passed final location and the store not wiped + root_reopen = zarr.open(str(datasets_store), mode="r") + assert root_reopen.attrs.get("preexisting", False) + assert np.array_equal(amplitudes, root_reopen["extensions"]["amplitudes"][:]) + assert np.array_equal(waveforms_rms, root_reopen["extensions"]["waveforms_rms"][:]) + assert np.array_equal(denoised_waveforms_rms, root_reopen["extensions"]["denoised_waveforms_rms"][:]) # Test pickle mechanism for node in nodes: @@ -271,6 +276,7 @@ def test_run_node_pipeline(cache_folder_creation): unpickled_node = pickle.loads(pickled_node) +@pytest.mark.requires_zarr_write def test_gather_to_zarr_chunking(tmp_path): # the zarr chunk size along the first axis must be picked from a byte target (not from the # size of the first gathered buffer), so it stays sensible for billions of spikes and never @@ -319,6 +325,7 @@ def test_gather_to_zarr_chunking(tmp_path): assert np.array_equal(waveforms[:], waveforms2[:]) +@pytest.mark.requires_zarr_write def test_gather_to_zarr_chunk_bytes_per_name(tmp_path): # `zarr_target_chunk_bytes` can also be a dict to use a different byte target per array recording, sorting = generate_ground_truth_recording(num_channels=8, num_units=5, durations=[20.0], seed=7) diff --git a/src/spikeinterface/core/tests/test_numpy_extractors.py b/src/spikeinterface/core/tests/test_numpy_extractors.py index daaf99d72a..d0a8c5be70 100644 --- a/src/spikeinterface/core/tests/test_numpy_extractors.py +++ b/src/spikeinterface/core/tests/test_numpy_extractors.py @@ -90,6 +90,22 @@ def test_NumpySorting(setup_NumpyRecording): sorting = NumpySorting.from_sorting(other_sorting) # print(sorting) + # Verify recording segment offsets and shifted sorting times survive conversion to NumpySorting. + recording = generate_recording(num_channels=2, durations=[1.0, 1.0]) + recording.shift_times(shift=5.0, segment_index=0) + recording.shift_times(shift=120.0, segment_index=1) + other_sorting.register_recording(recording) + other_sorting.shift_times(2.0) + sorting_with_times = NumpySorting.from_sorting(other_sorting) + + for segment_index in range(other_sorting.get_num_segments()): + assert sorting_with_times.get_start_time(segment_index) == other_sorting.get_start_time(segment_index) + for unit_id in other_sorting.unit_ids: + assert np.array_equal( + sorting_with_times.get_unit_spike_train(unit_id, segment_index=segment_index, return_times=True), + other_sorting.get_unit_spike_train(unit_id, segment_index=segment_index, return_times=True), + ) + # construct back from kwargs keep the same array sorting2 = load(sorting.to_dict()) assert np.shares_memory(sorting2._cached_spike_vector, sorting._cached_spike_vector) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 8cc069ce1e..527e053184 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -1,5 +1,6 @@ import pytest from pathlib import Path +import warnings import numpy as np import shutil @@ -20,10 +21,14 @@ _sort_extensions_by_dependency, ) from spikeinterface.core.analyzer_extension_core import BaseSpikeVectorExtension +from spikeinterface.core.base import minimum_spike_dtype # to test basespikevectorextension with node pipeline from spikeinterface.core.node_pipeline import SpikeRetriever from spikeinterface.core.tests.test_node_pipeline import AmplitudeExtractionNode +from spikeinterface.core.core_tools import _is_zarr_write_supported + +analyzer_formats = ("memory", "binary_folder", "zarr") if _is_zarr_write_supported() else ("memory", "binary_folder") def get_dataset(): @@ -92,6 +97,7 @@ def test_SortingAnalyzer_memory(tmp_path, dataset): def test_SortingAnalyzer_binary_folder(tmp_path, dataset): recording, sorting = dataset + recording = recording.save(folder=tmp_path / "recording_binary") folder = tmp_path / "test_SortingAnalyzer_binary_folder" if folder.exists(): @@ -136,8 +142,10 @@ def test_SortingAnalyzer_binary_folder(tmp_path, dataset): assert "number" in sorting_analyzer_reloded.sorting.get_property_keys() +@pytest.mark.requires_zarr_write def test_SortingAnalyzer_zarr(tmp_path, dataset): recording, sorting = dataset + recording = recording.save(folder=tmp_path / "recording_zarr") folder = tmp_path / "test_SortingAnalyzer_zarr.zarr" @@ -306,20 +314,22 @@ def test_load_without_runtime_info(tmp_path, dataset): with pytest.warns(UserWarning): sorting_analyzer = load_sorting_analyzer(folder, format="auto") - # zarr - folder = tmp_path / "test_SortingAnalyzer_run_info.zarr" - sorting_analyzer = create_sorting_analyzer( - sorting, recording, format="zarr", folder=folder, sparse=False, sparsity=None - ) - sorting_analyzer.compute(extensions) - # remove run_info from attrs to mimic a previous version of spikeinterface - root = sorting_analyzer._get_zarr_root(mode="r+") - for ext in extensions: - del root["extensions"][ext].attrs["run_info"] - zarr.consolidate_metadata(root.store) - # should raise a warning for missing run_info - with pytest.warns(UserWarning): - sorting_analyzer = load_sorting_analyzer(folder, format="auto") + # TODO: remove once writing to zarr is supported with zarr>=3 + if _is_zarr_write_supported(): + # zarr + folder = tmp_path / "test_SortingAnalyzer_run_info.zarr" + sorting_analyzer = create_sorting_analyzer( + sorting, recording, format="zarr", folder=folder, sparse=False, sparsity=None + ) + sorting_analyzer.compute(extensions) + # remove run_info from attrs to mimic a previous version of spikeinterface + root = sorting_analyzer._get_zarr_root(mode="r+") + for ext in extensions: + del root["extensions"][ext].attrs["run_info"] + zarr.consolidate_metadata(root.store) + # should raise a warning for missing run_info + with pytest.warns(UserWarning): + sorting_analyzer = load_sorting_analyzer(folder, format="auto") def test_SortingAnalyzer_tmp_recording(dataset): @@ -363,7 +373,7 @@ def test_SortingAnalyzer_interleaved_probegroup(dataset): assert np.array_equal(recording.get_channel_locations(), sorting_analyzer.get_channel_locations()) -@pytest.mark.parametrize("format", ["binary_folder", "zarr"]) +@pytest.mark.parametrize("format", ["binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) def test_load_in_lazy_mode(tmp_path, dataset, format): recording, sorting = dataset @@ -444,7 +454,7 @@ def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder): assert sorting_analyzer.has_recording() # save to several format - for format in ("memory", "binary_folder", "zarr"): + for format in analyzer_formats: if format != "memory": if format == "zarr": folder = cache_folder / f"test_SortingAnalyzer_save_as_{format}.zarr" @@ -474,7 +484,7 @@ def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder): assert sorting_analyzer2.sparsity == sorting_analyzer.sparsity # select unit_ids to several format - for format in ("memory", "binary_folder", "zarr"): + for format in analyzer_formats: if format != "memory": if format == "zarr": folder = cache_folder / f"test_SortingAnalyzer_select_units_with_{format}.zarr" @@ -737,6 +747,32 @@ def test_extension(): register_result_extension(DummyAnalyzerExtension2) +def test_select_units_reordered_keeps_extension_alignment(): + """Extensions slice per-spike data with a mask on the old spike vector, which keeps the old + order. The selected sorting's spike vector must keep that same order, even for cotemporal + spikes and even when the selection reorders the units.""" + register_result_extension(DummyAnalyzerExtension) + + rng = np.random.default_rng(0) + num_spikes, num_units = 2000, 5 + spikes = np.empty(num_spikes, dtype=minimum_spike_dtype) + spikes["sample_index"] = np.sort(rng.integers(0, 200, size=num_spikes)) + spikes["unit_index"] = rng.integers(0, num_units, size=num_spikes) + spikes["segment_index"] = 0 + unit_ids = np.array(["u0", "u1", "u2", "u3", "u4"]) + sorting = NumpySorting(spikes, 30_000.0, unit_ids) + recording = generate_recording(num_channels=4, durations=[1.0], sampling_frequency=30_000.0, seed=0) + + analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=False) + analyzer.compute("dummy") + + selected = analyzer.select_units(unit_ids[::-1]) + # "result_two" is a copy of unit_index array + old_unit_id_per_spike = unit_ids[selected.get_extension("dummy").data["result_two"]] + new_unit_id_per_spike = selected.unit_ids[selected.sorting.to_spike_vector()["unit_index"]] + assert np.array_equal(old_unit_id_per_spike, new_unit_id_per_spike) + + def test_excess_spikes(dataset): """ If there are spikes that occur after the recording end time, @@ -855,7 +891,9 @@ def _compute_reference_pipeline_data(dataset): return analyzer.get_extension("dummy_pipeline").get_data() -@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +@pytest.mark.parametrize( + "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) @pytest.mark.parametrize("lazy", [True, False]) def test_compute_pipeline_extension_gather_to_disk_lazy(tmp_path, dataset, format, lazy): """ @@ -920,7 +958,7 @@ def test_compute_pipeline_extension_gather_to_disk_lazy(tmp_path, dataset, forma assert np.array_equal(load_sorting_analyzer(folder).get_extension("dummy_pipeline").get_data(), amp_ref) -@pytest.mark.parametrize("format", ["binary_folder", "zarr"]) +@pytest.mark.parametrize("format", ["binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) def test_compute_pipeline_extension_save_false(tmp_path, dataset, format): """ With save=False on a disk-backed analyzer, node-pipeline extensions are computed in memory @@ -942,7 +980,9 @@ def test_compute_pipeline_extension_save_false(tmp_path, dataset, format): assert not analyzer_reloaded.has_extension("dummy_pipeline") -@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +@pytest.mark.parametrize( + "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) @pytest.mark.parametrize("lazy", [True, False]) def test_compute_one_pipeline_extension_gather_to_disk(tmp_path, dataset, format, lazy): """ @@ -1272,6 +1312,28 @@ def test_merge_units_main_channel_id_disagreement(): assert merged_main_channel_id == "chB" +def test_no_mergeable_units_with_external_unit_ids(dataset): + """ + Test that merge_units is robust against no mergeable units left + """ + recording, sorting = dataset + sorting_analyzer = create_sorting_analyzer( + sorting, recording, format="memory", sparse=True, sparsity_kwargs=dict(method="best_channels", num_channels=3) + ) + merge_unit_groups = [["0", "1"], ["2", "3"]] + new_unit_ids = ["8", "9"] + with pytest.warns(UserWarning, match="No mergeable unit groups found."): + _, new_unit_ids = sorting_analyzer.merge_units( + merge_unit_groups=merge_unit_groups, + new_unit_ids=new_unit_ids, + sparsity_overlap=1, + raise_error_if_overlap_fails=False, + return_new_unit_ids=True, + ) + # Assert that no merges were made + assert len(new_unit_ids) == 0 + + @pytest.mark.parametrize("unit_indices", [[4, 3, 2, 1, 0], [3, 1]]) def test_select_units_reordered_sparsity(dataset, unit_indices): recording, sorting = dataset diff --git a/src/spikeinterface/core/tests/test_template_class.py b/src/spikeinterface/core/tests/test_template_class.py index 953527c8d5..20f88817cb 100644 --- a/src/spikeinterface/core/tests/test_template_class.py +++ b/src/spikeinterface/core/tests/test_template_class.py @@ -115,6 +115,7 @@ def test_initialization_fail_with_dense_templates(): template = generate_test_template(template_type="sparse_with_dense_templates") +@pytest.mark.requires_zarr_write @pytest.mark.parametrize("is_in_uV", [True, False]) @pytest.mark.parametrize("template_type", ["dense", "sparse"]) def test_save_and_load_zarr(template_type, is_in_uV, tmp_path): diff --git a/src/spikeinterface/core/tests/test_time_handling.py b/src/spikeinterface/core/tests/test_time_handling.py index 936a9498da..4a08a1c7c8 100644 --- a/src/spikeinterface/core/tests/test_time_handling.py +++ b/src/spikeinterface/core/tests/test_time_handling.py @@ -114,7 +114,7 @@ def test_has_time_vector(self, time_vector_recording): assert raw_recording.has_time_vector(segment_idx) is False assert times_recording.has_time_vector(segment_idx) is True - @pytest.mark.parametrize("mode", ["binary", "zarr"]) + @pytest.mark.parametrize("mode", ["binary", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) @pytest.mark.parametrize("fixture_name", ["time_vector_recording", "t_start_recording"]) def test_times_propagated_to_save_folder(self, request, fixture_name, mode, tmp_path): """ @@ -375,7 +375,7 @@ def test_save_and_load_time_shift(self, request, fixture_name, tmp_path): times_recording.get_times(segment_index=idx), loaded_recording.get_times(segment_index=idx) ) - @pytest.mark.parametrize("save_format", ["binary", "zarr"]) + @pytest.mark.parametrize("save_format", ["binary", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) def test_shift_times_after_load(self, request, save_format, tmp_path): """ Shift times on a recording loaded from disk as a read-only np.memmap @@ -475,7 +475,7 @@ def test_get_times_with_time_vector_slicing(): sampling_frequency = 10_000.0 recording = generate_recording(durations=[1.0], num_channels=3, sampling_frequency=sampling_frequency) times = 1.0 + np.arange(0, 10_000) / sampling_frequency - recording.set_times(times=times, segment_index=0, with_warning=False) + recording.set_times(times=times, segment_index=0) # Full get_times should return the complete time vector times_full = recording.get_times(segment_index=0) @@ -639,7 +639,7 @@ def test_shift_times_with_time_vector(self): 100.0 + np.cumsum(np.random.RandomState(0).uniform(0.5, 1.5, num_samples)) / recording.get_sampling_frequency() ) - recording.set_times(times, segment_index=0, with_warning=False) + recording.set_times(times, segment_index=0) sorting = generate_sorting(num_units=5, durations=[1.0]) sorting.register_recording(recording) diff --git a/src/spikeinterface/core/tests/test_unitsselectionsorting.py b/src/spikeinterface/core/tests/test_unitsselectionsorting.py index 3aa7bc7577..309563df62 100644 --- a/src/spikeinterface/core/tests/test_unitsselectionsorting.py +++ b/src/spikeinterface/core/tests/test_unitsselectionsorting.py @@ -1,8 +1,10 @@ import pytest import numpy as np -from pathlib import Path -from spikeinterface.core import UnitsSelectionSorting +from spikeinterface.core import NumpySorting, UnitsSelectionSorting +from spikeinterface.core.base import minimum_spike_dtype +from spikeinterface.core.basesorting import LEXSORT_UNIT_COMPACT +from spikeinterface.core.testing import check_sortings_equal from spikeinterface.core.generate import generate_sorting @@ -52,5 +54,224 @@ def test_compute_and_cache_spike_vector(): assert np.all(cached_spike_vector == computed_spike_vector) +PARENT_UNIT_IDS = ["22", "45", "29", "7", "3"] + + +def _make_parent_with_shuffled_ties(unit_ids=PARENT_UNIT_IDS, num_segments=2, num_spikes=2000, seed=42, dtype=None): + """A sorting whose cotemporal spikes are in arbitrary unit_index order (see #4606), with + unit_ids that are deliberately not sorted.""" + rng = np.random.default_rng(seed) + num_units = len(unit_ids) + + # Far fewer samples than spikes, so cotemporal spikes are abundant. + spikes = np.empty(num_spikes, dtype=minimum_spike_dtype if dtype is None else dtype) + spikes["sample_index"] = rng.integers(0, 200, size=num_spikes) + spikes["unit_index"] = rng.integers(0, num_units, size=num_spikes) + spikes["segment_index"] = rng.integers(0, num_segments, size=num_spikes) + spikes = spikes[np.lexsort((rng.random(num_spikes), spikes["sample_index"], spikes["segment_index"]))] + + sorting = NumpySorting(spikes, 30_000.0, np.asarray(unit_ids)) + assert sorting.get_num_segments() == num_segments + return sorting + + +def _mask_and_remap(parent, selected_parent_ids): + """The parent's spike vector filtered to the selected units, + with unit_index remapped to the selection order. + (This is the same mask the SortingAnalyzer + extensions apply to their per-spike data.)""" + spikes = parent.to_spike_vector() + lut = np.full(parent.get_num_units(), -1, dtype=np.int64) + lut[parent.ids_to_indices(selected_parent_ids)] = np.arange(len(selected_parent_ids)) + new_unit_index = lut[spikes["unit_index"]] + keep = new_unit_index >= 0 + expected = spikes[keep].copy() + expected["unit_index"] = new_unit_index[keep] + return expected + + +def _assert_partial_invariant(spikes): + assert np.all(np.diff(spikes["segment_index"]) >= 0) + for segment_index in np.unique(spikes["segment_index"]): + assert np.all(np.diff(spikes["sample_index"][spikes["segment_index"] == segment_index]) >= 0) + + +def _assert_has_shuffled_ties(spikes): + full_lexsort = np.lexsort((spikes["unit_index"], spikes["sample_index"], spikes["segment_index"])) + assert not np.array_equal(spikes, spikes[full_lexsort]) + + +@pytest.mark.parametrize( + "unit_ids, renamed_unit_ids", + [ + (["29", "22", "3"], None), + (["3", "7", "29", "45", "22"], None), + (["22", "45"], ["b", "a"]), + (["29", "45"], None), + ], + ids=["reorder", "reverse", "renamed_order_preserving", "unsorted_parent_order_preserving"], +) +def test_selection_preserves_parent_order(unit_ids, renamed_unit_ids): + """A selection is the parent's spike vector filtered and remapped, nothing more: the parent's + (unspecified) order of cotemporal spikes must carry over untouched.""" + parent = _make_parent_with_shuffled_ties() + child = UnitsSelectionSorting(parent, unit_ids=unit_ids, renamed_unit_ids=renamed_unit_ids) + + expected = _mask_and_remap(parent, unit_ids) + _assert_has_shuffled_ties(expected) + + spikes = child.to_spike_vector() + assert np.array_equal(spikes, expected) + _assert_partial_invariant(spikes) + + spike_trains = [] + for segment_index in range(parent.get_num_segments()): + spike_trains.append({}) + for new_id, parent_id in zip(child.unit_ids, unit_ids): + parent_train = parent.get_unit_spike_train(parent_id, segment_index=segment_index, use_cache=False) + assert np.array_equal(child.get_unit_spike_train(new_id, segment_index=segment_index), parent_train) + spike_trains[segment_index][new_id] = parent_train + + parent_counts = parent.count_num_spikes_per_unit() + child_counts = child.count_num_spikes_per_unit() + for new_id, parent_id in zip(child.unit_ids, unit_ids): + assert child_counts[new_id] == parent_counts[parent_id] + + reference = NumpySorting.from_unit_dict(spike_trains, parent.sampling_frequency) + check_sortings_equal(child, reference, check_exact_lexsort=False) + + +def test_selection_keeps_extra_fields(): + """Make sure fields beyond `minimum_spike_dtype` + (e.g. the "channel_index" that `to_spike_vector(main_channel_indices=...)` adds) + stay with their spike through a selection.""" + wide_dtype = minimum_spike_dtype + [("channel_index", "int64")] + parent = _make_parent_with_shuffled_ties(dtype=wide_dtype) + parent._cached_spike_vector["channel_index"] = np.arange(parent._cached_spike_vector.size) + + unit_ids = ["3", "22", "29"] + child = UnitsSelectionSorting(parent, unit_ids=unit_ids) + spikes = child.to_spike_vector() + + expected = _mask_and_remap(parent, unit_ids) + assert spikes.dtype == wide_dtype + assert np.array_equal(spikes, expected) + assert np.array_equal(spikes["channel_index"], expected["channel_index"]) + + +def test_zero_units_and_zero_spikes(): + parent = _make_parent_with_shuffled_ties() + child = parent.select_units([]) + assert child.get_num_units() == 0 + assert child.to_spike_vector().size == 0 + assert child.to_spike_vector().dtype == minimum_spike_dtype + assert child.count_num_spikes_per_unit() == {} + assert len(child.to_spike_vector(concatenated=False)) == parent.get_num_segments() + + empty_parent = NumpySorting(np.zeros(0, dtype=minimum_spike_dtype), 30_000.0, np.array([1, 2, 3])) + child = empty_parent.select_units([3, 1]) + assert child.to_spike_vector().size == 0 + assert np.array_equal(child._get_spike_vector_segment_slices(), [[0, 0]]) + + +NEW_UNIT_IDS = ["a", "b", "c", "d", "e"] + +IDENTITY_SELECTIONS = { + "select_all": lambda parent: parent.select_units(parent.unit_ids), + "rename_units": lambda parent: parent.rename_units(NEW_UNIT_IDS), + "remove_none": lambda parent: parent.remove_units([]), + "default_args": lambda parent: UnitsSelectionSorting(parent), +} + + +def _assert_shares_parent_caches(child, parent): + assert child._is_identity_selection is True + assert child.to_spike_vector() is parent.to_spike_vector() + assert child._get_spike_vector_segment_slices() is parent._get_spike_vector_segment_slices() + assert child._cached_lexsorted_spike_vector is parent._cached_lexsorted_spike_vector + + +@pytest.mark.parametrize("make_child", IDENTITY_SELECTIONS.values(), ids=IDENTITY_SELECTIONS.keys()) +def test_identity_selection_shares_parent_cache(make_child): + """Same units in the same order, renamed or not: the parent's spike vector, segment slices and + reorderings are exactly the child's, so they are shared by reference rather than recomputed.""" + parent = _make_parent_with_shuffled_ties() + child = make_child(parent) + _assert_shares_parent_caches(child, parent) + + # The parent computes the reordered spike vector first; the child must find it in the shared cache. + parent_reordered = parent.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False, return_slices=False) + child_reordered = child.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False, return_slices=False) + assert child_reordered is parent_reordered + assert len(parent._cached_lexsorted_spike_vector) == 1 + + # The child computes the reordered spike vector first (via get_unit_spike_train); the parent must find it. + parent = _make_parent_with_shuffled_ties() + child = make_child(parent) + child.get_unit_spike_train(child.unit_ids[0], segment_index=0) + assert len(parent._cached_lexsorted_spike_vector) == 1 + parent_reordered = parent.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False, return_slices=False) + child_reordered = child.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False, return_slices=False) + assert child_reordered is parent_reordered + + # The shared vector and reorder cache are indexed by unit position, not id. Trains and counts + # the child serves from them under its (possibly renamed) ids must match what the parent serves + # for the same positions (without using the cache). + parent_counts = parent.count_num_spikes_per_unit() + child_counts = child.count_num_spikes_per_unit() + for new_id, parent_id in zip(child.unit_ids, parent.unit_ids): + assert child_counts[new_id] == parent_counts[parent_id] + for segment_index in range(parent.get_num_segments()): + assert np.array_equal( + child.get_unit_spike_train(new_id, segment_index=segment_index), + parent.get_unit_spike_train(parent_id, segment_index=segment_index, use_cache=False), + ) + + empty_parent = NumpySorting(np.zeros(0, dtype=minimum_spike_dtype), 30_000.0, np.asarray(PARENT_UNIT_IDS)) + _assert_shares_parent_caches(make_child(empty_parent), empty_parent) + + +@pytest.mark.parametrize( + "unit_ids", + [PARENT_UNIT_IDS[::-1], PARENT_UNIT_IDS[:2]], + ids=["reversed", "subset"], +) +def test_non_identity_selection_does_not_share(unit_ids): + parent = _make_parent_with_shuffled_ties() + child = parent.select_units(unit_ids) + assert child._is_identity_selection is False + assert child.to_spike_vector() is not parent.to_spike_vector() + assert child._cached_lexsorted_spike_vector is not parent._cached_lexsorted_spike_vector + + child.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT) + assert len(child._cached_lexsorted_spike_vector) == 1 + assert len(parent._cached_lexsorted_spike_vector) == 0 + + +@pytest.mark.requires_zarr_write +def test_identity_selection_keeps_lazy_zarr_vector(tmp_path): + """A lazy parent spike vector should stay lazy through an identity selection.""" + from spikeinterface.core.zarrextractors import ZarrSortingExtractor, ZarrSpikeVector + + folder = tmp_path / "sorting.zarr" + ZarrSortingExtractor.write_sorting(_make_parent_with_shuffled_ties(), folder) + lazy_parent = ZarrSortingExtractor(folder, lazy_spike_vector=True) + assert isinstance(lazy_parent.to_spike_vector(), ZarrSpikeVector) + + renamed = lazy_parent.rename_units(NEW_UNIT_IDS) + assert renamed.to_spike_vector() is lazy_parent.to_spike_vector() + assert renamed._get_spike_vector_segment_slices() is lazy_parent._get_spike_vector_segment_slices() + + eager_parent = ZarrSortingExtractor(folder) + frames = [-1, 0, 5, 7, 10_000] + np.testing.assert_array_equal( + renamed.search_cached_spikes_sorted(frames, segment_index=0), + eager_parent.search_cached_spikes_sorted(frames, segment_index=0), + ) + subset = lazy_parent.select_units(["29", "22"]) + assert isinstance(subset.to_spike_vector(), np.ndarray) + assert np.array_equal(subset.to_spike_vector(), _mask_and_remap(eager_parent, ["29", "22"])) + + if __name__ == "__main__": test_basic_functions() diff --git a/src/spikeinterface/core/tests/test_waveform_tools.py b/src/spikeinterface/core/tests/test_waveform_tools.py index 5e0350f833..c2750774b6 100644 --- a/src/spikeinterface/core/tests/test_waveform_tools.py +++ b/src/spikeinterface/core/tests/test_waveform_tools.py @@ -136,7 +136,6 @@ def test_waveform_tools(create_cache_folder): all_waveforms = extract_waveforms_to_single_buffer( recording, spikes, - unit_ids, nbefore, nafter, return_in_uV=False, @@ -158,6 +157,108 @@ def test_waveform_tools(create_cache_folder): _check_all_wf_equal(list_wfs_sparse) +@pytest.mark.requires_zarr_write +@pytest.mark.parametrize("sparse", [False, True]) +def test_extract_waveforms_to_single_buffer_zarr(tmp_path, sparse): + # the "zarr" mode writes waveforms directly to a zarr dataset. Workers return their block and + # the main process writes it (single writer), so parallel writes must match a reference and + # never race, even with n_jobs > 1. + import zarr + + recording, sorting = get_dataset() + sampling_frequency = recording.sampling_frequency + nbefore = ms_to_samples(3.0, sampling_frequency) + nafter = ms_to_samples(4.0, sampling_frequency) + spikes = sorting.to_spike_vector() + unit_ids = sorting.unit_ids + dtype = recording.get_dtype() + + if sparse: + sparsity_mask = np.random.RandomState(0).randint( + 0, 2, size=(unit_ids.size, recording.channel_ids.size), dtype="bool" + ) + else: + sparsity_mask = None + + common = dict(return_in_uV=False, dtype=dtype, sparsity_mask=sparsity_mask, copy=True, progress_bar=False) + + # reference computed in shared memory, single job + reference = extract_waveforms_to_single_buffer( + recording, spikes, nbefore, nafter, mode="shared_memory", n_jobs=1, **common + ) + + # zarr mode, single and parallel jobs, must match the reference and reload from disk + for n_jobs in (1, 2): + dataset_path = tmp_path / f"analyzer_{sparse}_{n_jobs}.zarr" / "extensions" / "waveforms" / "waveforms" + zarr_waveforms = extract_waveforms_to_single_buffer( + recording, + spikes, + nbefore, + nafter, + mode="zarr", + file_path=dataset_path, + n_jobs=n_jobs, + chunk_duration="0.3s", + **common, + ) + assert isinstance(zarr_waveforms, zarr.Array) + assert zarr_waveforms.shape == reference.shape + assert np.array_equal(reference, zarr_waveforms[:]) + + # reload from disk + store_path = tmp_path / f"analyzer_{sparse}_{n_jobs}.zarr" + reloaded = zarr.open(str(store_path), mode="r")["extensions"]["waveforms"]["waveforms"] + assert np.array_equal(reference, reloaded[:]) + + +@pytest.mark.requires_zarr_write +def test_waveforms_at_segment_borders(tmp_path): + # spikes near the segment borders are partially filled: samples outside of the segment are 0 + from spikeinterface.core import NumpySorting + + recording = generate_recording(num_channels=4, durations=[2.0, 3.0], seed=0) + nbefore, nafter = 30, 60 + samples_list, labels_list = [], [] + for segment_index in range(recording.get_num_segments()): + num_samples = recording.get_num_samples(segment_index=segment_index) + samples_list.append(np.array([5, 1000, num_samples - 10])) + labels_list.append(np.array([0, 1, 0])) + sorting = NumpySorting.from_samples_and_labels(samples_list, labels_list, recording.sampling_frequency) + spikes = sorting.to_spike_vector() + + expected = np.zeros((spikes.size, nbefore + nafter, recording.get_num_channels()), dtype="float32") + for i, spike in enumerate(spikes): + segment_index = spike["segment_index"] + num_samples = recording.get_num_samples(segment_index=segment_index) + start, end = spike["sample_index"] - nbefore, spike["sample_index"] + nafter + traces = recording.get_traces( + start_frame=max(start, 0), end_frame=min(end, num_samples), segment_index=segment_index + ) + expected[i, max(0, -start) : max(0, -start) + traces.shape[0]] = traces + + job_kwargs = dict(n_jobs=1, chunk_duration="0.5s", progress_bar=False) + common = dict(return_in_uV=False, dtype="float32", **job_kwargs) + shared = extract_waveforms_to_single_buffer(recording, spikes, nbefore, nafter, mode="shared_memory", **common) + zarr_waveforms = extract_waveforms_to_single_buffer( + recording, spikes, nbefore, nafter, mode="zarr", file_path=tmp_path / "wf.zarr" / "waveforms", **common + ) + by_units = extract_waveforms_to_buffers( + recording, spikes, sorting.unit_ids, nbefore, nafter, mode="shared_memory", copy=True, **common + ) + np.testing.assert_array_equal(shared, expected) + np.testing.assert_array_equal(zarr_waveforms[:], expected) + for unit_index, unit_id in enumerate(sorting.unit_ids): + np.testing.assert_array_equal(by_units[unit_id], expected[spikes["unit_index"] == unit_index]) + + templates = estimate_templates_with_accumulator( + recording, spikes, sorting.unit_ids, nbefore, nafter, return_in_uV=False, **job_kwargs + ) + for unit_index in range(sorting.unit_ids.size): + np.testing.assert_allclose( + templates[unit_index], expected[spikes["unit_index"] == unit_index].mean(axis=0), rtol=1e-5, atol=1e-5 + ) + + def test_estimate_templates_with_accumulator(): recording, sorting = get_dataset() diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index 6278541698..2b388839d1 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -1,6 +1,7 @@ import pytest from pathlib import Path +import numpy as np import zarr from spikeinterface.core import ( @@ -10,12 +11,14 @@ ) from spikeinterface.core.zarrextractors import ( ZarrRecordingExtractor, + ZarrSampleIndexSearch, ZarrSortingExtractor, add_sorting_to_zarr_group, get_default_zarr_compressor, ) +@pytest.mark.requires_zarr_write def test_zarr_compression_options(tmp_path): from numcodecs import Blosc, Delta, FixedScaleOffset @@ -59,6 +62,7 @@ def test_zarr_compression_options(tmp_path): assert rec_other._root["times_seg0"].filters == other_filters2 +@pytest.mark.requires_zarr_write def test_ZarrSortingExtractor(tmp_path): np_sorting = generate_sorting() @@ -78,7 +82,65 @@ def test_ZarrSortingExtractor(tmp_path): sorting = load(sorting.to_dict()) +def test_ZarrSampleIndexSearch(tmp_path): + rng = np.random.default_rng(0) + # two "segments", each sorted, with long runs of equal values so that runs cross + # the (tiny) chunk boundaries + segments = [np.sort(rng.integers(0, 40, size=101)), np.sort(rng.integers(0, 25, size=58))] + sample_index = np.concatenate(segments) + bounds = np.cumsum([0] + [len(s) for s in segments]) + z = zarr.open(tmp_path / "sample_index.zarr", mode="w", shape=sample_index.shape, chunks=(7,), dtype="int64") + z[:] = sample_index + + values = np.arange(-3, 45) + for chunk_firsts in (sample_index[::7], None): + search = ZarrSampleIndexSearch(z, chunk_firsts) + for start, stop in zip(bounds[:-1], bounds[1:]): + expected = np.searchsorted(sample_index[start:stop], values, side="left") + np.testing.assert_array_equal(search.searchsorted(values, start, stop), expected) + # empty range + np.testing.assert_array_equal(search.searchsorted([5], 10, 10), [0]) + + +@pytest.mark.requires_zarr_write +def test_ZarrSortingExtractor_lazy_search(tmp_path): + sorting = generate_sorting(num_units=10, durations=[5.0, 3.0, 4.0], firing_rates=40.0, seed=0) + folder = tmp_path / "sorting.zarr" + ZarrSortingExtractor.write_sorting(sorting, folder) + # re-store sample_index in small chunks, so that the search crosses chunk boundaries + # and segments start in the middle of a chunk + spikes_group = zarr.open(folder, mode="a")["spikes"] + sample_index = spikes_group["sample_index"][:] + del spikes_group["sample_index"], spikes_group["sample_index_chunk_firsts"] + spikes_group.create_dataset("sample_index", data=sample_index, chunks=(97,)) + spikes_group.create_dataset("sample_index_chunk_firsts", data=sample_index[::97], compressor=None) + assert spikes_group["sample_index"].nchunks > 3 + + in_ram = ZarrSortingExtractor(folder) + lazy = ZarrSortingExtractor(folder, lazy_spike_vector=True) + assert type(lazy.to_spike_vector()).__name__ == "ZarrSpikeVector" + + rng = np.random.default_rng(2308) + num_samples = int(sorting.to_spike_vector()["sample_index"].max()) + 100 + for segment_index in range(sorting.get_num_segments()): + frames = np.sort(rng.integers(-10, num_samples, size=200)) + expected = in_ram.search_cached_spikes_sorted(frames, segment_index=segment_index) + np.testing.assert_array_equal(lazy.search_cached_spikes_sorted(frames, segment_index=segment_index), expected) + + # stores written before the chunk index existed rebuild it from the data + del zarr.open(folder, mode="a")["spikes/sample_index_chunk_firsts"] + old = ZarrSortingExtractor(folder, lazy_spike_vector=True) + frames = np.arange(-5, num_samples, 37) + for segment_index in range(sorting.get_num_segments()): + np.testing.assert_array_equal( + old.search_cached_spikes_sorted(frames, segment_index=segment_index), + in_ram.search_cached_spikes_sorted(frames, segment_index=segment_index), + ) + + if __name__ == "__main__": tmp_path = Path("tmp") test_zarr_compression_options(tmp_path) test_ZarrSortingExtractor(tmp_path) + test_ZarrSampleIndexSearch(tmp_path) + test_ZarrSortingExtractor_lazy_search(tmp_path) diff --git a/src/spikeinterface/core/time_series.py b/src/spikeinterface/core/time_series.py index b3dfe89857..1d5fd0afe1 100644 --- a/src/spikeinterface/core/time_series.py +++ b/src/spikeinterface/core/time_series.py @@ -222,7 +222,7 @@ def has_time_vector(self, segment_index: Optional[int] = None): segment_index = self._check_segment_index(segment_index) return self.segments[segment_index].has_time_vector() - def set_times(self, times, segment_index=None, with_warning=True): + def set_times(self, times, segment_index=None): """Set times for a recording segment. Parameters @@ -232,8 +232,6 @@ def set_times(self, times, segment_index=None, with_warning=True): into memory and cast to float64 before being stored on the segment. segment_index : int or None, default: None The segment index (required for multi-segment) - with_warning : bool, default: True - If True, a warning is printed """ segment_index = self._check_segment_index(segment_index) rs = self.segments[segment_index] @@ -243,13 +241,6 @@ def set_times(self, times, segment_index=None, with_warning=True): rs._t_start = None rs._time_vector = times.astype("float64", copy=False) - - if with_warning: - warnings.warn( - "Setting times with Recording.set_times() is not recommended because " - "times are not always propagated across preprocessing" - "Use this carefully!" - ) self._time_info_modified = True def reset_times(self): diff --git a/src/spikeinterface/core/unitsselectionsorting.py b/src/spikeinterface/core/unitsselectionsorting.py index fdba9b361a..33fec6fa94 100644 --- a/src/spikeinterface/core/unitsselectionsorting.py +++ b/src/spikeinterface/core/unitsselectionsorting.py @@ -33,6 +33,11 @@ def __init__(self, parent_sorting, unit_ids=None, renamed_unit_ids=None): BaseSorting.__init__(self, sampling_frequency, self._renamed_unit_ids) + self._is_identity_selection = bool(np.array_equal(self._unit_ids, parents_unit_ids)) + if self._is_identity_selection: + # Same units (possibly renamed), same order => we can use the parent's cached spike vector + self._cached_lexsorted_spike_vector = parent_sorting._cached_lexsorted_spike_vector + for parent_segment in self._parent_sorting.segments: sub_segment = UnitsSelectionSortingSegment(parent_segment, ids_conversion) self.add_sorting_segment(sub_segment) @@ -54,29 +59,28 @@ def _compute_and_cache_spike_vector(self) -> None: if self._parent_sorting._cached_spike_vector is None: return + if self._is_identity_selection: + self._cached_spike_vector = self._parent_sorting._cached_spike_vector + self._cached_spike_vector_segment_slices = self._parent_sorting._get_spike_vector_segment_slices() + return + spike_vector, _ = remap_unit_indices_in_vector( vector=self._parent_sorting._cached_spike_vector, all_old_unit_ids=self._parent_sorting.unit_ids, all_new_unit_ids=self._unit_ids, ) - - # check if order is preserved - pos = np.searchsorted(self._parent_sorting.unit_ids, self.unit_ids) - order_is_preserved = np.all(np.diff(pos) > 0) - - if not order_is_preserved: - # Note: this can be a very high cost and make big dataset very slow - # the only goal of this is to ensure the unit_index order when the sample is the same - # TODO: maybe we can remove it, if we don't guarantee the order of unit_index - # when sample_index is the same, but it can be a problem for some downstream analysis - - # lexsort by segment_index, sample_index, unit_index - sort_indices = np.lexsort( - (spike_vector["unit_index"], spike_vector["sample_index"], spike_vector["segment_index"]) - ) - spike_vector = spike_vector[sort_indices] self._cached_spike_vector = spike_vector + def search_cached_spikes_sorted( + self, + indices: list[int], + segment_index: int | None = None, + ): + if self._is_identity_selection: + # same spikes in the same positions, and the parent knows how to search a lazy vector + return self._parent_sorting.search_cached_spikes_sorted(indices=indices, segment_index=segment_index) + return super().search_cached_spikes_sorted(indices=indices, segment_index=segment_index) + class UnitsSelectionSortingSegment(BaseSortingSegment): def __init__(self, parent_segment, ids_conversion): diff --git a/src/spikeinterface/core/waveform_tools.py b/src/spikeinterface/core/waveform_tools.py index 9fedc89324..bcbdaf7e14 100644 --- a/src/spikeinterface/core/waveform_tools.py +++ b/src/spikeinterface/core/waveform_tools.py @@ -350,8 +350,6 @@ def _worker_distribute_buffers(segment_index, start_frame, end_frame, worker_dic inds_by_unit = worker_dict["inds_by_unit"] sparsity_mask = worker_dict["sparsity_mask"] - seg_size = recording.get_num_samples(segment_index=segment_index) - # take only spikes with the correct segment_index # this is a slice so no copy!! s0, s1 = np.searchsorted(spikes["segment_index"], [segment_index, segment_index + 1]) @@ -359,27 +357,31 @@ def _worker_distribute_buffers(segment_index, start_frame, end_frame, worker_dic # take only spikes in range [start_frame, end_frame] # this is a slice so no copy!! - # the border of segment are protected by nbefore on left an nafter on the right - i0, i1 = np.searchsorted( - in_seg_spikes["sample_index"], [max(start_frame, nbefore), min(end_frame, seg_size - nafter)] - ) + # spikes near the segment borders are kept: their waveforms are partially filled (zero padded) + i0, i1 = np.searchsorted(in_seg_spikes["sample_index"], [start_frame, end_frame]) # slice in absolut in spikes vector l0 = i0 + s0 l1 = i1 + s0 + num_samples = recording.get_num_samples(segment_index=segment_index) + if l1 > l0: sub_spikes = in_seg_spikes[i0:i1] start = sub_spikes[0]["sample_index"] - nbefore end = sub_spikes[-1]["sample_index"] + nafter - # load trace in memory + # load trace in memory, clipped to the segment: spikes near the segment borders are partially filled + traces_start = max(start, 0) traces = recording.get_traces( - start_frame=start, end_frame=end, segment_index=segment_index, return_in_uV=return_in_uV + start_frame=traces_start, + end_frame=min(end, num_samples), + segment_index=segment_index, + return_in_uV=return_in_uV, ) - onset = start + nbefore + onset = traces_start + nbefore offset = nbefore + nafter for unit_ind, unit_id in enumerate(unit_ids): @@ -398,18 +400,20 @@ def _worker_distribute_buffers(segment_index, start_frame, end_frame, worker_dic for pos in in_chunk_pos: sample_index = spikes["sample_index"][inds[pos]] - onset - wf = traces[sample_index : sample_index + offset, :] + # sample_index < 0 only near the segment start, near the segment end wf is shorter + wf_start = max(0, -sample_index) + wf = traces[sample_index + wf_start : sample_index + offset, :] + wf_end = wf_start + wf.shape[0] if sparsity_mask is None: - wfs[pos, :, :] = wf + wfs[pos, wf_start:wf_end, :] = wf else: - wfs[pos, :, :] = wf[:, sparsity_mask[unit_ind]] + wfs[pos, wf_start:wf_end, :] = wf[:, sparsity_mask[unit_ind]] def extract_waveforms_to_single_buffer( recording, spikes, - unit_ids, nbefore, nafter, mode="memmap", @@ -420,6 +424,7 @@ def extract_waveforms_to_single_buffer( copy=True, job_name=None, verbose=False, + zarr_target_chunk_bytes: int = 10 * 1024 * 1024, **job_kwargs, ): """ @@ -429,7 +434,8 @@ def extract_waveforms_to_single_buffer( needed to recover waveforms unit by unit. Importantly in case of sparsity, the channels are not aligned across units. - Note: spikes near borders (nbefore/nafter) are not extracted and 0 are put the output buffer. + Note: spikes near the segment borders (nbefore/nafter) are partially filled: the samples outside of the + segment are set to 0. This ensures that spikes.shape[0] == all_waveforms.shape[0]. Important note: for the "shared_memory" mode wf_array_info contains reference to @@ -446,8 +452,6 @@ def extract_waveforms_to_single_buffer( spikes: 1d numpy array with several fields Spikes handled as a unique vector. This vector can be obtained with: `spikes = Sorting.to_spike_vector()` - unit_ids: list ot numpy - List of unit_ids nbefore: int N samples before spike nafter: int @@ -458,7 +462,10 @@ def extract_waveforms_to_single_buffer( If True and the recording has scaling (gain_to_uV and offset_to_uV properties), traces are scaled to uV file_path: str or path or None, default: None - In case of memmap mode, file to save npy file + In "memmap" mode, the npy file to save the waveforms to. + In "zarr" mode, a path pointing inside a zarr store, e.g. + "my-analyzer.zarr/extensions/waveforms/waveforms" (the ".zarr" part identifies the store and + the dataset is created on the fly). dtype: numpy.dtype, default: None dtype for waveforms buffer sparsity_mask: None or array of bool, default: None @@ -470,6 +477,8 @@ def extract_waveforms_to_single_buffer( need to be referenced as long as all_waveforms will be used otherwise it might produce segmentation faults which are hard to debug. Also when copy=False the SharedMemory will need to be unlink manually if proper cleanup of resources is desired. + zarr_target_chunk_bytes: int, default: 10 * 1024 * 1024 + Target chunk size in bytes for zarr storage. {} @@ -491,8 +500,9 @@ def extract_waveforms_to_single_buffer( if mode == "shared_memory": assert file_path is None - else: + elif mode == "memmap": file_path = Path(file_path) + # for mode == "zarr", file_path is a path pointing inside a zarr store (handled below) num_spikes = spikes.size if sparsity_mask is None: @@ -502,6 +512,7 @@ def extract_waveforms_to_single_buffer( num_chans = int(np.max(np.sum(sparsity_mask, axis=1), initial=0)) # This is a numpy scalar, so we cast to int shape = (int(num_spikes), int(n_samples), int(num_chans)) + gather_func = None if mode == "memmap": all_waveforms = np.lib.format.open_memmap(file_path, mode="w+", dtype=dtype, shape=shape) # wf_array_info = str(file_path) @@ -516,6 +527,18 @@ def extract_waveforms_to_single_buffer( shm_name = shm.name # wf_array_info = (shm, shm_name, dtype.str, shape) wf_array_info = dict(shm=shm, shm_name=shm_name, dtype=dtype.str, shape=shape) + elif mode == "zarr": + # Because zarr's write unit is a whole (compressed) chunk, parallel workers cannot safely write + # directly (two workers may touch the same boundary chunk). Instead workers return the block of + # waveforms of their chunk and the main process appends it to the zarr dataset (single writer). + # Blocks arrive in chunk order and every spike is in exactly one block, so rows match `spikes`. + from .node_pipeline import GatherToZarr + + assert file_path is not None, "zarr mode requires a `file_path` pointing inside a .zarr store" + gather_func = GatherToZarr( + dest=[file_path], names=["waveforms"], zarr_target_chunk_bytes=zarr_target_chunk_bytes + ) + wf_array_info = None else: raise ValueError("allocate_waveforms_buffers bad mode") @@ -523,28 +546,62 @@ def extract_waveforms_to_single_buffer( if num_spikes > 0 and num_chans > 0: # and run - func = _worker_distribute_single_buffer - init_func = _init_worker_distribute_single_buffer + if mode == "zarr": + # workers return their block, the main process appends it (single writer) + func = _worker_return_single_buffer + init_func = _init_worker_return_single_buffer + init_args = ( + recording, + spikes, + nbefore, + nafter, + return_in_uV, + sparsity_mask, + dtype.str, + int(n_samples), + int(num_chans), + ) + else: + func = _worker_distribute_single_buffer + init_func = _init_worker_distribute_single_buffer + + init_args = ( + recording, + spikes, + wf_array_info, + nbefore, + nafter, + return_in_uV, + mode, + sparsity_mask, + ) - init_args = ( - recording, - spikes, - wf_array_info, - nbefore, - nafter, - return_in_uV, - mode, - sparsity_mask, - ) if job_name is None: - job_name = f"extract waveforms {mode} mono buffer" - + job_name = f"extract waveforms zarr {mode} buffer" processor = TimeSeriesChunkExecutor( - recording, func, init_func, init_args, job_name=job_name, verbose=verbose, **job_kwargs + recording, + func, + init_func, + init_args, + gather_func=gather_func, + job_name=job_name, + verbose=verbose, + **job_kwargs, ) processor.run() - if mode == "memmap": + if mode == "zarr": + all_waveforms = gather_func.finalize_buffers(squeeze_output=True) + if all_waveforms is None: + # no block was gathered (no spikes or no channels): create the empty dataset + all_waveforms = _create_empty_zarr_waveforms(file_path, shape, dtype, zarr_target_chunk_bytes) + if all_waveforms.shape[0] != num_spikes: + raise ValueError( + f"Extracted {all_waveforms.shape[0]} waveforms for {num_spikes} spikes: some spikes are " + "outside of the recording segments" + ) + + if mode in ("memmap", "zarr"): return all_waveforms elif mode == "shared_memory": if copy: @@ -586,7 +643,7 @@ def _init_worker_distribute_single_buffer( worker_dict["shm"] = shm worker_dict["all_waveforms"] = all_waveforms - # prepare segment slices + # prepare segment slices: since this is a subset of spikes, it's cheap to compute segment slices segment_slices = [] for segment_index in range(recording.get_num_segments()): s0, s1 = np.searchsorted(spikes["segment_index"], [segment_index, segment_index + 1]) @@ -608,48 +665,159 @@ def _worker_distribute_single_buffer(segment_index, start_frame, end_frame, work sparsity_mask = worker_dict["sparsity_mask"] all_waveforms = worker_dict["all_waveforms"] - seg_size = recording.get_num_samples(segment_index=segment_index) - s0, s1 = segment_slices[segment_index] in_seg_spikes = spikes[s0:s1] # take only spikes in range [start_frame, end_frame] # this is a slice so no copy!! - # the border of segment are protected by nbefore on left an nafter on the right - i0, i1 = np.searchsorted( - in_seg_spikes["sample_index"], [max(start_frame, nbefore), min(end_frame, seg_size - nafter)] - ) + # spikes near the segment borders are kept: their waveforms are partially filled (zero padded) + i0, i1 = np.searchsorted(in_seg_spikes["sample_index"], [start_frame, end_frame]) + num_samples = recording.get_num_samples(segment_index=segment_index) if i1 > i0: sub_spikes = in_seg_spikes[i0:i1] start = sub_spikes[0]["sample_index"] - nbefore end = sub_spikes[-1]["sample_index"] + nafter - # load trace in memory + # load trace in memory, clipped to the segment: spikes near the segment borders are partially filled + traces_start = max(start, 0) traces = recording.get_traces( - start_frame=start, end_frame=end, segment_index=segment_index, return_in_uV=return_in_uV + start_frame=traces_start, + end_frame=min(end, num_samples), + segment_index=segment_index, + return_in_uV=return_in_uV, ) - onset = start + nbefore + onset = traces_start + nbefore offset = nbefore + nafter sample_indices = sub_spikes["sample_index"] - onset unit_indices = sub_spikes["unit_index"] spike_indices = s0 + np.arange(i0, i1) for sample_index, unit_index, spike_index in zip(sample_indices, unit_indices, spike_indices): - wf = traces[sample_index : sample_index + offset, :] + # sample_index < 0 only near the segment start, near the segment end wf is shorter + wf_start = max(0, -sample_index) + wf = traces[sample_index + wf_start : sample_index + offset, :] + wf_end = wf_start + wf.shape[0] if sparsity_mask is None: - all_waveforms[spike_index, :, :] = wf + all_waveforms[spike_index, wf_start:wf_end, :] = wf else: mask = sparsity_mask[unit_index, :] wf = wf[:, mask] - all_waveforms[spike_index, :, : wf.shape[1]] = wf + all_waveforms[spike_index, wf_start:wf_end, : wf.shape[1]] = wf if worker_dict["mode"] == "memmap": all_waveforms.flush() +def _create_empty_zarr_waveforms(file_path, shape, dtype, zarr_target_chunk_bytes): + """ + Create the waveforms zarr dataset when no block is gathered by `GatherToZarr` in + `extract_waveforms_to_single_buffer` (no spikes or no channels). + """ + import zarr + + from .zarrextractors import get_default_zarr_compressor + from .node_pipeline import _split_zarr_store_path + + store_path, dataset_path = _split_zarr_store_path(file_path) + row_nbytes = int(np.prod(shape[1:])) * dtype.itemsize + chunk0 = max(1, zarr_target_chunk_bytes // max(1, row_nbytes)) + zarr_root = zarr.open(str(store_path), mode="a") + return zarr_root.create_dataset( + name=dataset_path, + shape=shape, + chunks=(chunk0,) + shape[1:], + dtype=dtype, + compressor=get_default_zarr_compressor(), + overwrite=True, + ) + + +def _init_worker_return_single_buffer( + recording, spikes, nbefore, nafter, return_in_uV, sparsity_mask, dtype, n_samples, num_chans +): + worker_dict = {} + worker_dict["recording"] = recording + worker_dict["spikes"] = spikes + worker_dict["nbefore"] = nbefore + worker_dict["nafter"] = nafter + worker_dict["return_in_uV"] = return_in_uV + worker_dict["sparsity_mask"] = sparsity_mask + worker_dict["dtype"] = np.dtype(dtype) + worker_dict["n_samples"] = n_samples + worker_dict["num_chans"] = num_chans + + # prepare segment slices + segment_slices = [] + for segment_index in range(recording.get_num_segments()): + s0, s1 = np.searchsorted(spikes["segment_index"], [segment_index, segment_index + 1]) + segment_slices.append((s0, s1)) + worker_dict["segment_slices"] = segment_slices + + return worker_dict + + +# used by TimeSeriesChunkExecutor for mode="zarr": build and return a contiguous block of waveforms +# (rather than writing to a shared buffer), so the main process can append it to the zarr array. +def _worker_return_single_buffer(segment_index, start_frame, end_frame, worker_dict): + recording = worker_dict["recording"] + segment_slices = worker_dict["segment_slices"] + spikes = worker_dict["spikes"] + nbefore = worker_dict["nbefore"] + nafter = worker_dict["nafter"] + return_in_uV = worker_dict["return_in_uV"] + sparsity_mask = worker_dict["sparsity_mask"] + dtype = worker_dict["dtype"] + n_samples = worker_dict["n_samples"] + num_chans = worker_dict["num_chans"] + + s0, s1 = segment_slices[segment_index] + in_seg_spikes = spikes[s0:s1] + + # take only spikes in range [start_frame, end_frame]; spikes near the segment borders are kept and + # their waveforms are partially filled (zero padded) + i0, i1 = np.searchsorted(in_seg_spikes["sample_index"], [start_frame, end_frame]) + num_samples = recording.get_num_samples(segment_index=segment_index) + + if i1 <= i0: + return None + + sub_spikes = in_seg_spikes[i0:i1] + start = sub_spikes[0]["sample_index"] - nbefore + end = sub_spikes[-1]["sample_index"] + nafter + + # load trace in memory, clipped to the segment: spikes near the segment borders are partially filled + traces_start = max(start, 0) + traces = recording.get_traces( + start_frame=traces_start, + end_frame=min(end, num_samples), + segment_index=segment_index, + return_in_uV=return_in_uV, + ) + + onset = traces_start + nbefore + offset = nbefore + nafter + sample_indices = sub_spikes["sample_index"] - onset + unit_indices = sub_spikes["unit_index"] + + block = np.zeros((i1 - i0, n_samples, num_chans), dtype=dtype) + for local_index, (sample_index, unit_index) in enumerate(zip(sample_indices, unit_indices)): + # sample_index < 0 only near the segment start, near the segment end wf is shorter + wf_start = max(0, -sample_index) + wf = traces[sample_index + wf_start : sample_index + offset, :] + wf_end = wf_start + wf.shape[0] + if sparsity_mask is None: + block[local_index, wf_start:wf_end, :] = wf + else: + mask = sparsity_mask[unit_index, :] + wf = wf[:, mask] + block[local_index, wf_start:wf_end, : wf.shape[1]] = wf + + return block + + def split_waveforms_by_units(unit_ids, spikes, all_waveforms, sparsity_mask=None, folder=None): """ Split a single buffer waveforms into waveforms by units (multi buffers or multi files). @@ -784,7 +952,6 @@ def estimate_templates( all_waveforms, wf_array_info = extract_waveforms_to_single_buffer( recording, spikes, - unit_ids, nbefore, nafter, mode="shared_memory", @@ -1007,17 +1174,14 @@ def _worker_estimate_templates(segment_index, start_frame, end_frame, worker_dic return_in_uV = worker_dict["return_in_uV"] sparsity_mask = worker_dict["sparsity_mask"] - seg_size = recording.get_num_samples(segment_index=segment_index) - s0, s1 = segment_slices[segment_index] in_seg_spikes = spikes[s0:s1] # take only spikes in range [start_frame, end_frame] # this is a slice so no copy!! - # the border of segment are protected by nbefore on left an nafter on the right - i0, i1 = np.searchsorted( - in_seg_spikes["sample_index"], [max(start_frame, nbefore), min(end_frame, seg_size - nafter)] - ) + # spikes near the segment borders are kept: their waveforms are partially filled (zero padded) + i0, i1 = np.searchsorted(in_seg_spikes["sample_index"], [start_frame, end_frame]) + num_samples = recording.get_num_samples(segment_index=segment_index) if i1 > i0: sub_spikes = in_seg_spikes[i0:i1] @@ -1025,28 +1189,37 @@ def _worker_estimate_templates(segment_index, start_frame, end_frame, worker_dic start = sub_spikes[0]["sample_index"] - nbefore end = sub_spikes[-1]["sample_index"] + nafter - # load trace in memory + # load trace in memory, clipped to the segment: spikes near the segment borders are partially filled + traces_start = max(start, 0) traces = recording.get_traces( - start_frame=start, end_frame=end, segment_index=segment_index, return_in_uV=return_in_uV + start_frame=traces_start, + end_frame=min(end, num_samples), + segment_index=segment_index, + return_in_uV=return_in_uV, ) - onset = start + nbefore + onset = traces_start + nbefore offset = nbefore + nafter sample_indices = sub_spikes["sample_index"] - onset unit_indices = sub_spikes["unit_index"] for sample_index, unit_index in zip(sample_indices, unit_indices): - wf = traces[sample_index : sample_index + offset, :] + # sample_index < 0 only near the segment start, near the segment end wf is shorter + wf_start = max(0, -sample_index) + wf = traces[sample_index + wf_start : sample_index + offset, :] + wf_end = wf_start + wf.shape[0] if sparsity_mask is None: - waveform_accumulator_per_worker[worker_index, unit_index, :, :] += wf + waveform_accumulator_per_worker[worker_index, unit_index, wf_start:wf_end, :] += wf if waveform_squared_accumulator_per_worker is not None: - waveform_squared_accumulator_per_worker[worker_index, unit_index, :, :] += wf**2 + waveform_squared_accumulator_per_worker[worker_index, unit_index, wf_start:wf_end, :] += wf**2 else: mask = sparsity_mask[unit_index, :] wf = wf[:, mask] - waveform_accumulator_per_worker[worker_index, unit_index, :, : wf.shape[1]] += wf + waveform_accumulator_per_worker[worker_index, unit_index, wf_start:wf_end, : wf.shape[1]] += wf if waveform_squared_accumulator_per_worker is not None: - waveform_squared_accumulator_per_worker[worker_index, unit_index, :, : wf.shape[1]] += wf**2 + waveform_squared_accumulator_per_worker[ + worker_index, unit_index, wf_start:wf_end, : wf.shape[1] + ] += (wf**2) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 448dbd94a7..d2a604e38b 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -11,7 +11,13 @@ from .baserecording import BaseRecording, BaseRecordingSegment from .basesorting import BaseSorting, SpikeVectorSortingSegment from .job_tools import split_job_kwargs -from .core_tools import define_function_from_class, check_json, retrieve_importing_provenance, is_path_remote +from .core_tools import ( + define_function_from_class, + check_json, + retrieve_importing_provenance, + is_path_remote, + _check_zarr_write_is_supported, +) from .time_series_tools import _write_time_series_to_zarr @@ -78,9 +84,11 @@ def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: d else: if not Path(folder_path).is_dir(): raise ValueError(f"Folder {folder_path} does not exist") + # zarr>=3 refuses storage_options for a local path, even an empty dict + local_kwargs = dict(storage_options=storage_options) if storage_options else dict() for open_func in open_funcs: try: - root = open_func(str(folder_path), mode=mode, storage_options=storage_options) + root = open_func(str(folder_path), mode=mode, **local_kwargs) break except Exception as e: exception = e @@ -453,6 +461,66 @@ def __array__(self, dtype=None): arr["segment_index"][s0:s1] = seg_idx return arr if dtype is None else arr.astype(dtype) + def copy(self): + return np.copy(np.asarray(self)) + + +class ZarrSampleIndexSearch: + """ + `np.searchsorted` on a zarr `sample_index` array without materialising it. + + The array is sorted within each segment and stored in fixed-size zarr chunks. The + search keeps the first value of every chunk (8 B per chunk), picks the one chunk that + can hold the answer, and decodes only that chunk — once per unique chunk per call. + + Parameters + ---------- + sample_index : zarr.Array + The 1D `sample_index` array of a spike vector. + chunk_firsts : np.ndarray | None, default: None + `sample_index[::chunk_length]`, as written by `add_sorting_to_zarr_group`. + If None (stores written before this index existed), it is read from the array, + which decodes every chunk once. + """ + + def __init__(self, sample_index, chunk_firsts=None): + self._sample_index = sample_index + self._num_spikes = sample_index.shape[0] + self._chunk_length = sample_index.chunks[0] + if chunk_firsts is None: + starts = np.arange(0, self._num_spikes, self._chunk_length) + chunk_firsts = sample_index.get_orthogonal_selection(starts) if starts.size else [] + self._chunk_firsts = np.asarray(chunk_firsts, dtype="int64") + + def searchsorted(self, values, start, stop): + """ + Equivalent to `np.searchsorted(sample_index[start:stop], values, side="left")`. + + Returns positions relative to `start`. `start:stop` must be a range over which + `sample_index` is sorted, i.e. one segment. + """ + values = np.atleast_1d(np.asarray(values)) + out = np.zeros(values.size, dtype="int64") + if stop <= start: + return out + first_chunk = start // self._chunk_length + last_chunk = (stop - 1) // self._chunk_length + # first values of the chunks after `first_chunk`, all of which lie inside start:stop + later_firsts = self._chunk_firsts[first_chunk + 1 : last_chunk + 1] + # assign every value to its chunk in one vectorised call: side="left" keeps runs of + # equal values that cross a chunk boundary in the earlier chunk + chunk_indices = first_chunk + np.searchsorted(later_firsts, values, side="left") + # decode each touched chunk exactly once and answer all values that fall in it + unique_chunks, inverse = np.unique(chunk_indices, return_inverse=True) + for k, chunk_index in enumerate(unique_chunks): + mask = inverse == k + chunk_start = chunk_index * self._chunk_length + lo = max(chunk_start, start) + hi = min(chunk_start + self._chunk_length, stop) + block = self._sample_index[lo:hi] + out[mask] = lo + np.searchsorted(block, values[mask], side="left") - start + return out + class ZarrSortingExtractor(BaseSorting): """ @@ -524,7 +592,9 @@ def __init__( # we do not need to lexsort at init (very high cost) because there already sorted by frame before to be saved. # In version 0.104.X this was fully lexsorted, but we don't need it anymore because it's only important in the context of SpikeVectorBased extensions in the SortingAnalyzer, which stores its own copy of the Sorting object. This makes the extension data and the spike vector always matching their order. # spikes = spikes[np.lexsort((spikes["unit_index"], spikes["sample_index"], spikes["segment_index"]))] - + self._lazy_spike_vector = lazy_spike_vector + self._spikes_group = spikes_group + self._sample_index_search = None self._cached_spike_vector = spikes # pre-populate segment slices so _get_spike_vector_segment_slices() never # needs to materialise the full segment_index array @@ -553,12 +623,35 @@ def __init__( "lazy_spike_vector": lazy_spike_vector, } + def search_cached_spikes_sorted( + self, + indices: list[int], + segment_index: int | None = None, + ): + if not self._lazy_spike_vector: + return super().search_cached_spikes_sorted( + indices=indices, + segment_index=segment_index, + ) + if segment_index is None: + assert self.get_num_segments() == 1, "segment_index is required for multi-segment sortings" + segment_index = 0 + if self._sample_index_search is None: + # the chunk index is written with the sorting; older stores rebuild it from the data + chunk_firsts = None + if "sample_index_chunk_firsts" in self._spikes_group: + chunk_firsts = self._spikes_group["sample_index_chunk_firsts"][:] + self._sample_index_search = ZarrSampleIndexSearch(self._spikes_group["sample_index"], chunk_firsts) + start, stop = self._cached_spike_vector_segment_slices[segment_index] + return self._sample_index_search.searchsorted(indices, int(start), int(stop)) + @staticmethod def write_sorting( sorting: BaseSorting, folder_path: str | Path, overwrite: bool = False, storage_options: dict | None = None, + relative_to: str | Path | None = None, **kwargs, ): """ @@ -567,7 +660,7 @@ def write_sorting( folder_path = create_zarr_path_for_write(folder_path, overwrite=overwrite) zarr_root = zarr.open(str(folder_path), mode="w", storage_options=storage_options) zarr_root.attrs["zarr_class_info"] = retrieve_importing_provenance(ZarrSortingExtractor) - add_sorting_to_zarr_group(sorting, zarr_root, **kwargs) + add_sorting_to_zarr_group(sorting, zarr_root, relative_to=relative_to, **kwargs) return ZarrSortingExtractor(folder_path, storage_options=storage_options) @@ -625,7 +718,7 @@ def resolve_zarr_path(folder_path: str | Path): return folder_path, folder_path else: folder_path = Path(folder_path) - folder_path_kwarg = str(Path(folder_path).resolve()) + folder_path_kwarg = str(Path(folder_path).absolute()) return folder_path, folder_path_kwarg @@ -638,6 +731,7 @@ def create_zarr_path_for_write(folder_path: str | Path, overwrite: bool = False) folder_path : str or Path Path to the zarr root file """ + _check_zarr_write_is_supported() if not is_path_remote(folder_path): folder_path = Path(folder_path) folder_path = folder_path.with_suffix(".zarr") @@ -734,7 +828,9 @@ def add_properties_and_annotations(zarr_group: zarr.Group, recording_or_sorting: zarr_group.attrs["annotations"] = check_json(recording_or_sorting._annotations) -def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kwargs): +def add_sorting_to_zarr_group( + sorting: BaseSorting, zarr_group: zarr.Group, relative_to: str | Path | None = None, **kwargs +): """ Add a sorting extractor to a zarr group. @@ -749,6 +845,11 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw """ from numcodecs import Delta + if sorting.check_serializability("json"): + zarr_group.attrs["provenance"] = check_json(sorting.to_dict(recursive=True, relative_to=relative_to)) + else: + zarr_group.attrs["provenance"] = None + num_segments = sorting.get_num_segments() zarr_group.attrs["sampling_frequency"] = float(sorting.sampling_frequency) zarr_group.attrs["num_segments"] = int(num_segments) @@ -774,14 +875,30 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw segment_slices.append([i0, i1]) spikes_group.create_dataset(name="segment_slices", data=segment_slices, compressor=None) + # first sample_index of every zarr chunk: lets a lazy reader search sample_index + # one chunk at a time (see ZarrSampleIndexSearch) instead of materialising it + chunk_length = spikes_group["sample_index"].chunks[0] + spikes_group.create_dataset( + name="sample_index_chunk_firsts", + data=np.asarray(spikes["sample_index"][::chunk_length], dtype="int64"), + compressor=None, + ) + add_properties_and_annotations(zarr_group, sorting) -def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group, verbose=False, dtype=None, **kwargs): +def add_recording_to_zarr_group( + recording: BaseRecording, + zarr_group: zarr.Group, + verbose=False, + dtype=None, + relative_to: str | Path | None = None, + **kwargs, +): zarr_kwargs, job_kwargs = split_job_kwargs(kwargs) if recording.check_serializability("json"): - zarr_group.attrs["provenance"] = check_json(recording.to_dict(recursive=True)) + zarr_group.attrs["provenance"] = check_json(recording.to_dict(recursive=True, relative_to=relative_to)) else: zarr_group.attrs["provenance"] = None diff --git a/src/spikeinterface/curation/train_manual_curation.py b/src/spikeinterface/curation/train_manual_curation.py index 5e714fda17..e34951820a 100644 --- a/src/spikeinterface/curation/train_manual_curation.py +++ b/src/spikeinterface/curation/train_manual_curation.py @@ -832,7 +832,9 @@ def check_metric_names_are_the_same(metrics_for_each_analyzer): def _format_metric_dataframe(input_data): - input_data = input_data.map(lambda x: np.nan if np.isinf(x) else x) - input_data = input_data.astype("float32") + # cast before replacing infs: finite float64 values beyond float32 range overflow to inf + with np.errstate(over="ignore"): + input_data = input_data.astype("float32") + input_data = input_data.replace([np.inf, -np.inf], np.nan) return input_data diff --git a/src/spikeinterface/extractors/alfsortingextractor.py b/src/spikeinterface/extractors/alfsortingextractor.py index ba08bca2e8..880158a64f 100644 --- a/src/spikeinterface/extractors/alfsortingextractor.py +++ b/src/spikeinterface/extractors/alfsortingextractor.py @@ -11,6 +11,7 @@ class ALFSortingExtractor(BaseSorting): Parameters ---------- + folder_path : str or Path Path to the ALF folder. sampling_frequency : int, default: 30000 @@ -40,7 +41,7 @@ def __init__(self, folder_path, sampling_frequency=30000): sorting_segment = ALFSortingSegment(spikes["clusters"], spikes["samples"]) self.add_sorting_segment(sorting_segment) self.extra_requirements.append("ONE-api") - self._kwargs = {"folder_path": str(Path(folder_path).resolve()), "sampling_frequency": sampling_frequency} + self._kwargs = {"folder_path": str(Path(folder_path).absolute()), "sampling_frequency": sampling_frequency} class ALFSortingSegment(BaseSortingSegment): diff --git a/src/spikeinterface/extractors/cbin_ibl.py b/src/spikeinterface/extractors/cbin_ibl.py index d2c8afe23e..adf04f366a 100644 --- a/src/spikeinterface/extractors/cbin_ibl.py +++ b/src/spikeinterface/extractors/cbin_ibl.py @@ -194,9 +194,9 @@ def __init__(self, folder_path=None, load_sync_channel=False, stream_name="ap", self.set_property("inter_sample_shift", sample_shifts) self._kwargs = { - "folder_path": str(Path(folder_path).resolve()), + "folder_path": str(Path(folder_path).absolute()), "load_sync_channel": load_sync_channel, - "cbin_file_path": str(Path(cbin_file_path).resolve()), + "cbin_file_path": str(Path(cbin_file_path).absolute()), } diff --git a/src/spikeinterface/extractors/matlabhelpers.py b/src/spikeinterface/extractors/matlabhelpers.py index f76952fd06..f7ffb45065 100644 --- a/src/spikeinterface/extractors/matlabhelpers.py +++ b/src/spikeinterface/extractors/matlabhelpers.py @@ -32,7 +32,7 @@ def __init__(self, file_path): if not isinstance(file_path, Path): raise TypeError(f"Expected a str or Path file_path but got '{type(file_path).__name__}'") - file_path = file_path.resolve() # get absolute path to this file + file_path = file_path.absolute() # get absolute path to this file if not file_path.is_file(): raise ValueError(f"Specified file path '{file_path}' is not a file.") diff --git a/src/spikeinterface/extractors/neoextractors/__init__.py b/src/spikeinterface/extractors/neoextractors/__init__.py index a90e0954f3..4f04d9db0b 100644 --- a/src/spikeinterface/extractors/neoextractors/__init__.py +++ b/src/spikeinterface/extractors/neoextractors/__init__.py @@ -18,7 +18,12 @@ read_neuroscope_sorting, read_neuroscope, ) -from .neuroexplorer import NeuroExplorerRecordingExtractor, read_neuroexplorer +from .neuroexplorer import ( + NeuroExplorerRecordingExtractor, + NeuroExplorerSortingExtractor, + read_neuroexplorer, + read_neuroexplorer_sorting, +) from .nix import NixRecordingExtractor, read_nix from .openephys import ( OpenEphysLegacyRecordingExtractor, @@ -79,6 +84,9 @@ PlexonSortingExtractor: dict(wrapper_string="read_plexon_sorting", wrapper_class=read_plexon_sorting), Plexon2SortingExtractor: dict(wrapper_string="read_plexon2_sorting", wrapper_class=read_plexon2_sorting), NeuroScopeSortingExtractor: dict(wrapper_string="read_neuroscope_sorting", wrapper_class=read_neuroscope_sorting), + NeuroExplorerSortingExtractor: dict( + wrapper_string="read_neuroexplorer_sorting", wrapper_class=read_neuroexplorer_sorting + ), } neo_event_extractors_dict = { diff --git a/src/spikeinterface/extractors/neoextractors/intan.py b/src/spikeinterface/extractors/neoextractors/intan.py index 793026ea94..877f7f33b2 100644 --- a/src/spikeinterface/extractors/neoextractors/intan.py +++ b/src/spikeinterface/extractors/neoextractors/intan.py @@ -70,7 +70,7 @@ def __init__( self._add_channel_groups() self._kwargs.update( - dict(file_path=str(Path(file_path).resolve()), ignore_integrity_checks=ignore_integrity_checks), + dict(file_path=str(Path(file_path).absolute()), ignore_integrity_checks=ignore_integrity_checks), ) @classmethod @@ -195,7 +195,7 @@ def __init__( # Update kwargs to include our specific parameters self._kwargs = dict( - folder_path=str(Path(folder_path).resolve()), + folder_path=str(Path(folder_path).absolute()), mode=mode, stream_id=stream_id, stream_name=stream_name, diff --git a/src/spikeinterface/extractors/neoextractors/neuroexplorer.py b/src/spikeinterface/extractors/neoextractors/neuroexplorer.py index 8a076264ec..d00995aeb7 100644 --- a/src/spikeinterface/extractors/neoextractors/neuroexplorer.py +++ b/src/spikeinterface/extractors/neoextractors/neuroexplorer.py @@ -1,8 +1,14 @@ from pathlib import Path +from spikeinterface.core import BaseSorting from spikeinterface.core.core_tools import define_function_from_class -from .neobaseextractor import NeoBaseRecordingExtractor +from .neobaseextractor import ( + NeoBaseRecordingExtractor, + NeoBaseSortingExtractor, + NeoSortingSegment, + _NeoBaseExtractor, +) class NeuroExplorerRecordingExtractor(NeoBaseRecordingExtractor): @@ -72,3 +78,63 @@ def map_to_neo_kwargs(cls, file_path): read_neuroexplorer = define_function_from_class(source_class=NeuroExplorerRecordingExtractor, name="read_neuroexplorer") + + +class NeuroExplorerSortingExtractor(NeoBaseSortingExtractor): + """ + Class for reading the sorted units of NEX (NeuroExplorer data format) files. + + Based on :py:class:`neo.rawio.NeuroExplorerRawIO` + + Units are the neuron variables of the file and the unit ids are their names. Waveform variables + (for example the ``_wf`` and ``_template`` companions of each unit in Plexon exports) repeat + the spike times of a neuron variable or hold a mean waveform, so they are not units. They + remain available through ``neo_reader``. + + Parameters + ---------- + file_path : str | Path + The file path to load the sorting from. + """ + + NeoRawIOClass = "NeuroExplorerRawIO" + neo_returns_frames = True + + def __init__(self, file_path): + neo_kwargs = self.map_to_neo_kwargs(file_path) + _NeoBaseExtractor.__init__(self, block_index=None, **neo_kwargs) + + # Only neuron variables are units: they are the spike channels without waveforms + spike_channels = self.neo_reader.header["spike_channels"] + unit_ids = spike_channels["name"][spike_channels["wf_sampling_rate"] == 0] + # Spike timestamps are ticks of the file's global clock + sampling_frequency = self.neo_reader.global_header["freq"] + BaseSorting.__init__(self, sampling_frequency, unit_ids) + + sorting_segment = NeuroExplorerSortingSegment( + neo_reader=self.neo_reader, + block_index=self.block_index, + segment_index=0, + t_start=None, + sampling_frequency=sampling_frequency, + neo_returns_frames=self.neo_returns_frames, + ) + self.add_sorting_segment(sorting_segment) + self._kwargs = {"file_path": str(Path(file_path).absolute())} + + @classmethod + def map_to_neo_kwargs(cls, file_path): + neo_kwargs = {"filename": str(file_path)} + return neo_kwargs + + +class NeuroExplorerSortingSegment(NeoSortingSegment): + def map_from_unit_id_to_spike_channel_index(self, unit_id): + # Units are a subset of the spike channels, so the position of a unit id is not its spike channel index + spike_channel_names = list(self.neo_reader.header["spike_channels"]["name"]) + return spike_channel_names.index(unit_id) + + +read_neuroexplorer_sorting = define_function_from_class( + source_class=NeuroExplorerSortingExtractor, name="read_neuroexplorer_sorting" +) diff --git a/src/spikeinterface/extractors/neoextractors/neuronexus.py b/src/spikeinterface/extractors/neoextractors/neuronexus.py index 92bb90a08b..c1366e3e06 100644 --- a/src/spikeinterface/extractors/neoextractors/neuronexus.py +++ b/src/spikeinterface/extractors/neoextractors/neuronexus.py @@ -51,7 +51,7 @@ def __init__( **neo_kwargs, ) - self._kwargs.update(dict(file_path=str(Path(file_path).resolve()))) + self._kwargs.update(dict(file_path=str(Path(file_path).absolute()))) @classmethod def map_to_neo_kwargs(cls, file_path): diff --git a/src/spikeinterface/extractors/neoextractors/openephys.py b/src/spikeinterface/extractors/neoextractors/openephys.py index b03f731909..969a5ea899 100644 --- a/src/spikeinterface/extractors/neoextractors/openephys.py +++ b/src/spikeinterface/extractors/neoextractors/openephys.py @@ -355,7 +355,7 @@ def __init__( else: sync_times = None try: - self.set_times(times=sync_times, segment_index=segment_index, with_warning=False) + self.set_times(times=sync_times, segment_index=segment_index) except: warnings.warn(f"Could not load synchronized timestamps for {stream_name}") diff --git a/src/spikeinterface/extractors/neoextractors/plexon.py b/src/spikeinterface/extractors/neoextractors/plexon.py index f8131a816a..2a96385dd0 100644 --- a/src/spikeinterface/extractors/neoextractors/plexon.py +++ b/src/spikeinterface/extractors/neoextractors/plexon.py @@ -54,7 +54,7 @@ def __init__( use_names_as_ids=use_names_as_ids, **neo_kwargs, ) - self._kwargs.update({"file_path": str(Path(file_path).resolve())}) + self._kwargs.update({"file_path": str(Path(file_path).absolute())}) @classmethod def map_to_neo_kwargs(cls, file_path): @@ -82,7 +82,7 @@ def __init__(self, file_path): self.neo_reader = NeoBaseSortingExtractor.get_neo_io_reader(self.NeoRawIOClass, **neo_kwargs) sampling_frequency = self.neo_reader._global_ssampling_rate NeoBaseSortingExtractor.__init__(self, sampling_frequency=sampling_frequency, **neo_kwargs) - self._kwargs = {"file_path": str(Path(file_path).resolve())} + self._kwargs = {"file_path": str(Path(file_path).absolute())} @classmethod def map_to_neo_kwargs(cls, file_path): diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index e6a30e9fb0..0d1946d39a 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -5,6 +5,7 @@ import importlib.util import numpy as np +from probeinterface import Probe, ProbeGroup from spikeinterface import get_global_tmp_folder from spikeinterface.core import BaseRecording, BaseRecordingSegment, BaseSorting, BaseSortingSegment @@ -332,7 +333,9 @@ def _find_neurodata_type_from_backend(group, path="", result=None, neurodata_typ if result is None: result = [] - for neurodata_name, value in group.items(): + # zarr>=3 groups have no `items()`, `keys()` works for h5py and both zarr versions + for neurodata_name in group.keys(): + value = group[neurodata_name] # Check if it's a group and if it has the neurodata_type if isinstance(value, group_class): current_path = f"{path}/{neurodata_name}" if path else neurodata_name @@ -353,8 +356,8 @@ def _retrieve_electrodes_indices_from_electrical_series_backend(open_file, elect if backend == "zarr": import zarr - # links must be resolved - zarr_links = electrical_series.attrs["zarr_link"] + # links must be resolved, hdmf-zarr>=0.14 stores them under "_LINKS" instead of "zarr_link" + zarr_links = electrical_series.attrs.get("zarr_link", electrical_series.attrs.get("_LINKS")) electrodes_path = None for zarr_link in zarr_links: if zarr_link["name"] == "electrodes": @@ -1040,9 +1043,21 @@ def __init__( self.set_channel_gains(gains_to_uV) self.set_channel_offsets(offsets_to_uV) locations = self._reader.locations() - if locations is not None: - self.set_channel_locations(locations) groups = self._reader.groups() + if locations is not None: + # rel_x, rel_y and rel_z are relative to each electrode group, so each group is its own + # probe. Two groups can then share positions (e.g. two identical probes in one + # ElectricalSeries), so the overlap check across probes is skipped. + probegroup = ProbeGroup() + for group_name in np.unique(groups): + channel_indices = np.flatnonzero(groups == group_name) + probe = Probe(ndim=2) + probe.set_contacts(locations[channel_indices, :2], shapes="circle", shape_params={"radius": 1}) + probe.set_device_channel_indices(channel_indices) + if locations.shape[1] == 3: + probe = probe.to_3d(axes="xy") + probegroup.add_probe(probe) + self.set_probegroup(probegroup, check_overlap=False) if groups is not None: self.set_channel_groups(groups) @@ -1070,7 +1085,7 @@ def __init__( self.set_property(property_name, values) if stream_mode is None and file_path is not None: - file_path = str(Path(file_path).resolve()) + file_path = str(Path(file_path).absolute()) if stream_mode == "fsspec" and stream_cache_path is not None: stream_cache_path = str(Path(self.stream_cache_path).absolute()) @@ -1404,7 +1419,7 @@ def __init__( self.extra_requirements.append(stream_mode) if stream_mode is None and file_path is not None: - file_path = str(Path(file_path).resolve()) + file_path = str(Path(file_path).absolute()) if storage_options is not None and stream_mode == "zarr": warnings.warn( @@ -1456,6 +1471,23 @@ def _compute_and_cache_spike_vector(self) -> None: self._cached_spike_vector = spikes self._cached_spike_vector_segment_slices = np.array([[0, num_spikes]], dtype="int64") + def get_last_spike_frame(self, segment_index: int | None = None) -> int: + # Performance: the generic BaseSorting version builds the full spike vector to take one maximum. + # spike_times_index holds the end offset of each unit's spikes, so reading only the last spike of + # each unit is enough. This assumes spike times are ordered within each unit, an NWB best practice + # checked by nwbinspector, so each unit's last stored time is its latest. + segment_index = self._check_segment_index(segment_index) + segment = self.segments[segment_index] + unit_ends = np.asarray(segment.spike_times_index_data[:], dtype="int64") + unit_counts = np.diff(unit_ends, prepend=0) + last_spike_positions = unit_ends[unit_counts > 0] - 1 + if last_spike_positions.size == 0: + return 0 + last_spike_times = np.asarray(segment.spike_times_data[last_spike_positions], dtype="float64") + last_spike_time = last_spike_times.max() + last_spike_frame = segment._times_to_samples(last_spike_time) + return int(last_spike_frame) + @staticmethod def fetch_available_units_tables( file_path: str | Path, @@ -1612,7 +1644,9 @@ def _find_timeseries_from_backend(group, path="", result=None, backend="hdf5"): if result is None: result = [] - for name, value in group.items(): + # zarr>=3 groups have no `items()`, `keys()` works for h5py and both zarr versions + for name in group.keys(): + value = group[name] if isinstance(value, group_class): current_path = f"{path}/{name}" if path else name if value.attrs.get("neurodata_type") == "TimeSeries": diff --git a/src/spikeinterface/extractors/tests/test_neoextractors.py b/src/spikeinterface/extractors/tests/test_neoextractors.py index 7522218d52..64b597d68a 100644 --- a/src/spikeinterface/extractors/tests/test_neoextractors.py +++ b/src/spikeinterface/extractors/tests/test_neoextractors.py @@ -19,6 +19,7 @@ IntanRecordingExtractor, NeuroScopeRecordingExtractor, NeuroExplorerRecordingExtractor, + NeuroExplorerSortingExtractor, NeuroScopeSortingExtractor, NeuroNexusRecordingExtractor, PlexonRecordingExtractor, @@ -270,6 +271,15 @@ class NeuroExplorerRecordingTest(RecordingCommonTestSuite, unittest.TestCase): ] +class NeuroExplorerSortingTest(SortingCommonTestSuite, unittest.TestCase): + ExtractorClass = NeuroExplorerSortingExtractor + downloads = ["neuroexplorer"] + entities = [ + "neuroexplorer/File_neuroexplorer_1.nex", + "neuroexplorer/File_neuroexplorer_2.nex", + ] + + class NeuroScopeSortingTest(SortingCommonTestSuite, unittest.TestCase): ExtractorClass = NeuroScopeSortingExtractor downloads = ["neuroscope"] diff --git a/src/spikeinterface/extractors/tests/test_nwbextractors.py b/src/spikeinterface/extractors/tests/test_nwbextractors.py index d3f26f84db..381c4f6161 100644 --- a/src/spikeinterface/extractors/tests/test_nwbextractors.py +++ b/src/spikeinterface/extractors/tests/test_nwbextractors.py @@ -311,6 +311,48 @@ def test_nwb_extractor_electrodes_region_out_of_order(tmp_path, use_pynwb): assert np.array_equal(recording.get_channel_offsets(), np.array([float(i) for i in region]) * 1e6) +@pytest.mark.parametrize("use_pynwb", [True, False]) +def test_nwb_extractor_one_probe_per_electrode_group(tmp_path, use_pynwb): + """rel_x and rel_y are relative to each electrode group, so two identical probes written to one + ElectricalSeries share positions. Building a single probe from all channels failed on the + duplicate positions (GH-3999).""" + from pynwb import NWBHDF5IO + from pynwb.testing.mock.file import mock_NWBFile + from pynwb.testing.mock.device import mock_Device + from pynwb.testing.mock.ecephys import mock_ElectrodeGroup, mock_ElectricalSeries + + nwbfile = mock_NWBFile() + device = mock_Device(nwbfile=nwbfile) + + # Both groups get the same contact positions, so the two probes overlap. + group_names = ["probe_A", "probe_B"] + probe_locations = np.array([[0.0, 0.0], [0.0, 20.0], [0.0, 40.0], [0.0, 60.0]]) + for group_name in group_names: + electrode_group = mock_ElectrodeGroup(name=group_name, device=device, nwbfile=nwbfile) + for rel_x, rel_y in probe_locations: + nwbfile.add_electrode(group=electrode_group, location="brain", rel_x=rel_x, rel_y=rel_y) + + channels_per_group = len(probe_locations) + num_electrodes = len(group_names) * channels_per_group + electrode_region = nwbfile.create_electrode_table_region(region=list(range(num_electrodes)), description="all") + mock_ElectricalSeries( + name="ElectricalSeries", data=np.ones((10, num_electrodes)), electrodes=electrode_region, nwbfile=nwbfile + ) + + nwbfile_path = tmp_path / "two_electrode_groups.nwb" + with NWBHDF5IO(str(nwbfile_path), mode="w") as io: + io.write(nwbfile) + + recording = NwbRecordingExtractor( + nwbfile_path, electrical_series_path="acquisition/ElectricalSeries", use_pynwb=use_pynwb + ) + + assert len(recording.get_probegroup().probes) == len(group_names) + recordings_by_group = recording.split_by("group") + for group_name in group_names: + assert np.array_equal(recordings_by_group[group_name].get_channel_locations(), probe_locations) + + @pytest.mark.parametrize("use_pynwb", [True, False]) def test_nwb_extractor_offset_from_series(generate_nwbfile, use_pynwb): """Test that the offset is retrieved from the ElectricalSeries if it is present.""" diff --git a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py index 161040cbaf..9a7b2c0c67 100644 --- a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py +++ b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py @@ -3,12 +3,18 @@ import pytest import numpy as np +import zarr from spikeinterface import load from spikeinterface.core.testing import check_recordings_equal from spikeinterface.core.testing import check_recordings_equal, check_sortings_equal from spikeinterface.extractors.extractor_classes import NwbRecordingExtractor, NwbSortingExtractor +# TODO: remove once streaming zarr is supported with zarr>=3 +skip_zarr_3 = pytest.mark.skipif( + int(zarr.__version__.split(".")[0]) >= 3, reason="Streaming zarr is not supported yet with zarr>=3" +) + @pytest.mark.streaming_extractors @pytest.mark.parametrize("cache", [True, False]) # Test with and without cache @@ -185,6 +191,7 @@ def test_sorting_s3_nwb_remfile(tmp_path): @pytest.mark.streaming_extractors +@skip_zarr_3 def test_sorting_s3_nwb_zarr(tmp_path): file_path = ( "s3://aind-open-data/ecephys_625749_2022-08-03_15-15-06_nwb_2023-05-16_16-34-55/" diff --git a/src/spikeinterface/generation/tests/test_template_database.py b/src/spikeinterface/generation/tests/test_template_database.py index 78dfb7a6ec..3a3034acf3 100644 --- a/src/spikeinterface/generation/tests/test_template_database.py +++ b/src/spikeinterface/generation/tests/test_template_database.py @@ -1,3 +1,6 @@ +import pytest +import zarr + from spikeinterface.core.template import Templates from spikeinterface.generation import ( @@ -7,7 +10,13 @@ query_templates_from_database, ) +# TODO: remove once probeinterface can read zarr>=3 groups (Probe.from_zarr_group calls len() on a zarr array) +skip_zarr_3 = pytest.mark.skipif( + int(zarr.__version__.split(".")[0]) >= 3, reason="Reading templates from zarr is not supported yet with zarr>=3" +) + +@skip_zarr_3 def test_fetch_template_object_from_database(): available_datasets = list_available_datasets_in_template_database() @@ -31,6 +40,7 @@ def test_fetch_templates_database_info(): assert "dataset" in templates_info.columns +@skip_zarr_3 def test_query_templates_from_database(): templates_info = fetch_templates_database_info() diff --git a/src/spikeinterface/metrics/conftest.py b/src/spikeinterface/metrics/conftest.py index aefe2e6cfe..9c495cd36f 100644 --- a/src/spikeinterface/metrics/conftest.py +++ b/src/spikeinterface/metrics/conftest.py @@ -8,12 +8,13 @@ job_kwargs = dict(n_jobs=2, progress_bar=True, chunk_duration="1s") -def make_small_analyzer(): +def make_small_analyzer(folder): recording, sorting = generate_ground_truth_recording( durations=[10.0], num_units=10, seed=1205, ) + recording = recording.save(folder=folder / "rec_for_small_analyzer", format="binary") sorting = sorting.select_units(["2", "7", "0"], ["#3", "#9", "#4"]) @@ -36,12 +37,12 @@ def make_small_analyzer(): @pytest.fixture(scope="module") -def small_sorting_analyzer(): - return make_small_analyzer() +def small_sorting_analyzer(create_cache_folder): + return make_small_analyzer(create_cache_folder) @pytest.fixture(scope="module") -def sorting_analyzer_simple(): +def sorting_analyzer_simple(create_cache_folder): # we need high firing rate for amplitude_cutoff recording, sorting = generate_ground_truth_recording( durations=[ @@ -64,6 +65,7 @@ def sorting_analyzer_simple(): noise_kwargs=dict(noise_levels=5.0, strategy="tile_pregenerated"), seed=1205, ) + recording = recording.save(folder=create_cache_folder / "rec_for_sorting_analyzer_simple", format="binary") sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=True) diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py index be5967ad1b..69a1420638 100644 --- a/src/spikeinterface/postprocessing/amplitude_scalings.py +++ b/src/spikeinterface/postprocessing/amplitude_scalings.py @@ -178,13 +178,12 @@ def __init__( PipelineNode.__init__(self, recording, parents=parents, return_output=return_output) self.return_in_uV = return_in_uV if return_in_uV and recording.has_scaleable_traces(): - self._dtype = np.float32 self._gains = recording.get_channel_gains() self._offsets = recording.get_channel_offsets() else: - self._dtype = recording.get_dtype() self._gains = None self._offsets = None + self._dtype = np.float32 spike_retriever = find_parent_of_type(parents, SpikeRetriever) assert isinstance( spike_retriever, SpikeRetriever @@ -192,17 +191,22 @@ def __init__( assert spike_retriever.include_spikes_in_margin, "Need SpikeRetriever with include_spikes_in_margin=True" if not handle_collisions: self._margin = max(nbefore, nafter) + overlap_matrix = None else: # in this case we extend the margin to be able to get with collisions outside the chunk margin_waveforms = max(nbefore, nafter) max_margin_collisions = delta_collision_samples + margin_waveforms self._margin = max_margin_collisions + # sparsity_mask is fixed for the node's lifetime, so the unit-pair overlap matrix + # is computed once here instead of once per chunk inside compute()/find_collisions. + overlap_matrix = _unit_pair_overlap_matrix(sparsity_mask) # for some edge cases a template can be zero, leading to problems later template_is_zero = [np.all(template == 0) for template in all_templates] self._all_templates = all_templates self._sparsity_mask = sparsity_mask + self._overlap_matrix = overlap_matrix self._nbefore = nbefore self._nafter = nafter self._cut_out_before = cut_out_before @@ -232,6 +236,7 @@ def compute(self, traces, peaks): offsets = self._offsets all_templates = self._all_templates sparsity_mask = self._sparsity_mask + overlap_matrix = self._overlap_matrix nbefore = self._nbefore cut_out_before = self._cut_out_before cut_out_after = self._cut_out_after @@ -255,14 +260,14 @@ def compute(self, traces, peaks): local_spikes, local_spikes_within_margin, delta_collision_samples, - sparsity_mask, + overlap_matrix, local_spike_indices, ) else: collisions = {} # compute the scaling for each spike - scalings = np.zeros(len(local_spikes), dtype=float) + scalings = np.zeros(len(local_spikes), dtype=self._dtype) spike_collision_mask = np.zeros(len(local_spikes), dtype=bool) for spike_index, spike in enumerate(local_spikes): @@ -324,30 +329,34 @@ def get_margin(self): ### Collision handling ### -def _are_units_spatially_overlapping(sparsity_mask, i, j): +def _unit_pair_overlap_matrix(sparsity_mask): """ - Returns True if the unit indices i and j are - spatially overlapping, False otherwise + Precompute, once, whether every pair of units shares at least one channel. + + Unit-pair spatial overlap is a fixed fact of `sparsity_mask` alone: there are only + `num_units ** 2` possible (i, j) answers, independent of which spikes are being compared. + `find_collisions` looks this up once per temporally-overlapping spike-pair candidate + (millions of times on a realistic recording), so computing it here with one matrix + multiplication instead of a fresh `np.any(sparsity_mask[i] & sparsity_mask[j])` per lookup + removes an O(num_channels) recomputation of an already-known answer. Parameters ---------- - sparsity_mask: boolean mask + sparsity_mask : boolean mask A num_units x num_channels boolean array indicating whether the unit is represented on the channel. - i: int - The first unit index - j: int - The second unit index Returns ------- - bool - True if the units i and j are spatially overlapping, False otherwise + np.ndarray + A num_units x num_units boolean array where entry (i, j) is True if units i and j + are spatially overlapping, False otherwise. """ - if np.any(sparsity_mask[i] & sparsity_mask[j]): - return True - else: - return False + # int32 avoids any risk of the dot-product overflowing (it accumulates at most num_channels + # per entry, far below the int32 range) while staying far cheaper than a bool broadcast that + # would materialize a full num_units x num_units x num_channels intermediate array. + sparsity_mask_int = np.asarray(sparsity_mask, dtype=np.int32) + return (sparsity_mask_int @ sparsity_mask_int.T) > 0 def _ordinary_scaling_slope(template, local_waveform): @@ -387,7 +396,7 @@ def _ordinary_scaling_slope(template, local_waveform): return covariance / template_variance -def find_collisions(spikes, spikes_within_margin, delta_collision_samples, sparsity_mask, spike_indices): +def find_collisions(spikes, spikes_within_margin, delta_collision_samples, overlap_matrix, spike_indices): """ Finds the collisions between spikes. @@ -415,9 +424,12 @@ def find_collisions(spikes, spikes_within_margin, delta_collision_samples, spars another spike within a given margin delta_collision_samples: int The maximum number of samples between two spikes to consider them as overlapping - sparsity_mask: boolean mask - A num_units x num_channels boolean array indicating whether - the unit is represented on the channel. + overlap_matrix : np.ndarray + A num_units x num_units boolean array, as returned by `_unit_pair_overlap_matrix`, + where entry (i, j) is True if units i and j are spatially overlapping. Callers that + run this once per chunk (like `AmplitudeScalingNode.compute`) should precompute it + once from `sparsity_mask` and reuse it, since spatial overlap is a fixed fact of the + sparsity mask, not of the spikes being compared. spike_indices : np.ndarray The indices of `spikes` in `spikes_within_margin`. Providing these indices avoids searching `spikes_within_margin` once for every spike. @@ -428,7 +440,8 @@ def find_collisions(spikes, spikes_within_margin, delta_collision_samples, spars A dictionary with collisions. The key is the index of the spike with collision, the value is an array of overlapping spikes, including the spike itself at position 0. """ - # TODO: refactor to speed-up + spikes_within_margin_unit_index = spikes_within_margin["unit_index"] + collision_spikes_dict = {} for spike_index, (spike, spike_index_within_margin) in enumerate(zip(spikes, spike_indices, strict=True)): # find the spikes that fall within a temporal window around the spike peak @@ -451,19 +464,16 @@ def find_collisions(spikes, spikes_within_margin, delta_collision_samples, spars (pre_possible_consecutive_spike_indices, post_possible_consecutive_spike_indices) ) - # Build the collusion_spikes_dict including only - # spikes that overlap spatially - collision_spikes = [] - for possible_overlapping_spike_index in possible_overlapping_spike_indices: - - if _are_units_spatially_overlapping( - sparsity_mask, - spike["unit_index"], - spikes_within_margin[possible_overlapping_spike_index]["unit_index"], - ): - collision_spikes.append(spikes_within_margin[possible_overlapping_spike_index]) - if collision_spikes: - collision_spikes_dict[spike_index] = np.array([spike, *collision_spikes], dtype=spikes.dtype) + # Keep only the candidates that overlap spatially, looked up from the precomputed matrix + # instead of recomputed per candidate. + other_unit_indices = spikes_within_margin_unit_index[possible_overlapping_spike_indices] + overlaps = overlap_matrix[spike["unit_index"], other_unit_indices] + collision_spike_indices = possible_overlapping_spike_indices[overlaps] + + if collision_spike_indices.size: + collision_spikes_dict[spike_index] = np.array( + [spike, *spikes_within_margin[collision_spike_indices]], dtype=spikes.dtype + ) return collision_spikes_dict diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 7767104484..092344a512 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -171,6 +171,8 @@ def get_pca_model(self): * if mode is "by_channel_local", "pca_model" is a list of PCA model by channel * if mode is "by_channel_global" or "concatenated", "pca_model" is a single PCA model """ + if self.sorting_analyzer._lazy: + raise RuntimeError("PCA models are not loaded in lazy mode. Reload the sorting analyzer in non-lazy mode.") mode = self.params["mode"] if mode == "by_channel_local": pca_models = [] @@ -280,17 +282,20 @@ def get_some_projections(self, channel_ids=None, unit_ids=None): else: # need re-alignement some_projections = np.zeros((selected_inds.size, num_components, channel_indices.size), dtype=dtype) + # read all requested rows at once: per-unit reads on a remote zarr array re-fetch the same chunks + selected_projections = slice_rows(all_projections, selected_inds) for unit_id in unit_ids: unit_index = sorting.id_to_index(unit_id) - sparse_projection, local_chan_inds = self.get_projections_one_unit(unit_id, sparse=True) + spike_mask = np.flatnonzero(spike_unit_indices == unit_index) + local_chan_inds = sparsity.unit_id_to_channel_indices[unit_id] + sparse_projection = selected_projections[spike_mask][:, :, : local_chan_inds.size] # keep only requested channels channel_mask = np.isin(local_chan_inds, channel_indices) sparse_projection = sparse_projection[:, :, channel_mask] local_chan_inds = local_chan_inds[channel_mask] - spike_mask = np.flatnonzero(spike_unit_indices == unit_index) proj = np.zeros((spike_mask.size, num_components, channel_indices.size), dtype=dtype) # inject in requested channels channel_mask = np.isin(channel_indices, local_chan_inds) @@ -325,6 +330,8 @@ def _run(self, verbose=False, **job_kwargs): Compute the PCs on waveforms extacted within the by ComputeWaveforms. Projections are computed only on the waveforms sampled by the SortingAnalyzer. """ + if self.sorting_analyzer._lazy: + raise RuntimeError("PCA models are not loaded in lazy mode. Reload the sorting analyzer in non-lazy mode.") p = self.params mode = p["mode"] @@ -651,16 +658,14 @@ def _all_pc_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx): nsamples = nbefore + nafter - # Extract all waveforms in the chunk at once + # Prepare waveform sample offsets for all spikes in the chunk spike_times_in_chunk = spike_times[i0:i1] # Offset spike times to be relative to the start of the traces buffer spike_times_offset = spike_times_in_chunk - start - nbefore spike_indices = np.arange(i0, i1) - # Build waveform array: (n_spikes, nsamples, n_channels) - # Use fancy indexing to extract all snippets at once - sample_indices = spike_times_offset[:, None] + np.arange(nsamples)[None, :] # (n_spikes, nsamples) - all_wfs = traces[sample_indices] # (n_spikes, nsamples, n_channels) + # This adds a sample axis using strides and does not allocate waveform data. + waveform_view = np.lib.stride_tricks.sliding_window_view(traces, nsamples, axis=0) # Vectorized PCA: batch by channel across all spikes in the chunk. # For each unique channel, find all spikes that use it (via their unit's @@ -685,7 +690,7 @@ def _all_pc_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx): global_idxs = spike_indices[all_local_idxs] # Batch waveforms for this channel: (n_spikes, nsamples) - wfs_batch = all_wfs[all_local_idxs, :, chan_ind] + wfs_batch = waveform_view[spike_times_offset[all_local_idxs], chan_ind] if wfs_batch.size == 0: continue diff --git a/src/spikeinterface/postprocessing/spike_locations.py b/src/spikeinterface/postprocessing/spike_locations.py index 32334e6b3e..d864a0d73e 100644 --- a/src/spikeinterface/postprocessing/spike_locations.py +++ b/src/spikeinterface/postprocessing/spike_locations.py @@ -69,19 +69,24 @@ def _get_pipeline_nodes(self): recording = self.sorting_analyzer.recording sorting = self.sorting_analyzer.sorting + main_channel_indices = self.sorting_analyzer.get_main_channels(outputs="index", with_dict=False) retriever_kwargs = { "channel_from_template": True, **self.params["spike_retriever_kwargs"], } retriever = SpikeRetriever(sorting, recording, **retriever_kwargs) + spikes = sorting.to_spike_vector() + channel_index_array = main_channel_indices[spikes["unit_index"]] nodes = get_localization_pipeline_nodes( recording, retriever, + peaks_or_spikes=spikes, method=self.params["method"], method_kwargs=self.params["method_kwargs"], ms_before=self.params["ms_before"], ms_after=self.params["ms_after"], + channel_index_array=channel_index_array, seed=self.params.get("seed"), ) return nodes diff --git a/src/spikeinterface/postprocessing/template_similarity.py b/src/spikeinterface/postprocessing/template_similarity.py index 2ada11a7e7..ef40671b1f 100644 --- a/src/spikeinterface/postprocessing/template_similarity.py +++ b/src/spikeinterface/postprocessing/template_similarity.py @@ -297,13 +297,11 @@ def _compute_similarity_matrix_numba( elif method == "cosine": metric = 2 + overlapping_i_list = typed.List() overlapping_j_list = typed.List() active_channels_list = typed.List() for src_unit in range(num_templates): - overlapping_ids = typed.List() - overlapping_chs = typed.List() - start = src_unit if same_array else 0 for tgt_unit in range(start, other_num_templates): @@ -329,11 +327,9 @@ def _compute_similarity_matrix_numba( ch = np.arange(num_channels, dtype=np.uint16) if len(ch) > 0: - overlapping_ids.append(np.uint16(tgt_unit)) - overlapping_chs.append(ch) - - overlapping_j_list.append(overlapping_ids) - active_channels_list.append(overlapping_chs) + overlapping_i_list.append(np.int64(src_unit)) + overlapping_j_list.append(np.int64(tgt_unit)) + active_channels_list.append(ch) for count in range(len(shift_loop)): shift = shift_loop[count] @@ -341,44 +337,40 @@ def _compute_similarity_matrix_numba( src_sliced = templates_array[:, num_shifts : num_samples - num_shifts] tgt_sliced = other_templates_array[:, num_shifts + shift : num_samples - num_shifts + shift] - for i in prange(num_templates): - i_ = np.int64(i) - src_template = src_sliced[i_] - overlapping_ids = overlapping_j_list[i_] - overlapping_chs = active_channels_list[i_] - - for pair_idx in range(len(overlapping_ids)): - j = np.uint16(overlapping_ids[pair_idx]) - ch = overlapping_chs[pair_idx] - - src_ch = src_template[:, ch] - tgt_ch = tgt_sliced[j][:, ch] - - if metric == 0: - # l1 - norm_i = np.sum(np.abs(src_ch)) - norm_j = np.sum(np.abs(tgt_ch)) - dist = np.sum(np.abs(src_ch - tgt_ch)) - distances[count, i, j] = dist / (norm_i + norm_j) - - elif metric == 1: - # l2 - norm_i = sqrt(np.sum(src_ch**2)) - norm_j = sqrt(np.sum(tgt_ch**2)) - dist = sqrt(np.sum((src_ch - tgt_ch) ** 2)) - distances[count, i, j] = dist / (norm_i + norm_j) - - elif metric == 2: - # cosine - dot = np.sum(src_ch * tgt_ch) - norm_i = sqrt(np.sum(src_ch**2)) - norm_j = sqrt(np.sum(tgt_ch**2)) - denom = norm_i * norm_j - if denom > 0.0: - distances[count, i, j] = 1.0 - dot / denom - - if same_array: - distances[count, j, i] = distances[count, i, j] + for pair_idx in prange(len(overlapping_j_list)): + pair_idx_ = np.int64(pair_idx) + i = overlapping_i_list[pair_idx_] + j = overlapping_j_list[pair_idx_] + active_channels = active_channels_list[pair_idx_] + + src_ch = src_sliced[i][:, active_channels] + tgt_ch = tgt_sliced[j][:, active_channels] + + if metric == 0: + # l1 + norm_i = np.sum(np.abs(src_ch)) + norm_j = np.sum(np.abs(tgt_ch)) + dist = np.sum(np.abs(src_ch - tgt_ch)) + distances[count, i, j] = dist / (norm_i + norm_j) + + elif metric == 1: + # l2 + norm_i = sqrt(np.sum(src_ch**2)) + norm_j = sqrt(np.sum(tgt_ch**2)) + dist = sqrt(np.sum((src_ch - tgt_ch) ** 2)) + distances[count, i, j] = dist / (norm_i + norm_j) + + elif metric == 2: + # cosine + dot = np.sum(src_ch * tgt_ch) + norm_i = sqrt(np.sum(src_ch**2)) + norm_j = sqrt(np.sum(tgt_ch**2)) + denom = norm_i * norm_j + if denom > 0.0: + distances[count, i, j] = 1.0 - dot / denom + + if same_array: + distances[count, j, i] = distances[count, i, j] if same_array and shift != 0: distances[num_shifts_both_sides - count - 1] = distances[count].T diff --git a/src/spikeinterface/postprocessing/tests/common_extension_tests.py b/src/spikeinterface/postprocessing/tests/common_extension_tests.py index 58e2092149..fb7770befe 100644 --- a/src/spikeinterface/postprocessing/tests/common_extension_tests.py +++ b/src/spikeinterface/postprocessing/tests/common_extension_tests.py @@ -9,6 +9,7 @@ estimate_sparsity, ) from spikeinterface.core.sortinganalyzer import get_extension_class +from spikeinterface.core.core_tools import _is_zarr_write_supported extensions_which_allow_unit_ids = ["unit_locations"] extensions_with_unit_by_unit_data = ["correlograms", "template_similarity"] @@ -241,8 +242,10 @@ def run_extension_tests(self, extension_class, params): of interest with the passed parameters. Will perform tests for sparsity and format. """ + # TODO: remove once writing to zarr is supported with zarr>=3 + formats = ("memory", "binary_folder", "zarr") if _is_zarr_write_supported() else ("memory", "binary_folder") for sparse in (True, False): - for format in ("memory", "binary_folder", "zarr"): + for format in formats: print("sparse", sparse, format) sorting_analyzer = self._prepare_sorting_analyzer( format, sparse, extension_class, extension_params=params diff --git a/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py b/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py index d6faabb616..0d6488e6b4 100644 --- a/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py +++ b/src/spikeinterface/postprocessing/tests/test_amplitude_scalings.py @@ -5,7 +5,12 @@ from spikeinterface.postprocessing import ComputeAmplitudeScalings from spikeinterface.core.base import spike_peak_dtype -from spikeinterface.postprocessing.amplitude_scalings import _ordinary_scaling_slope, find_collisions, fit_collision +from spikeinterface.postprocessing.amplitude_scalings import ( + _ordinary_scaling_slope, + _unit_pair_overlap_matrix, + find_collisions, + fit_collision, +) def test_ordinary_scaling_slope_float32_precision(): @@ -86,7 +91,7 @@ def test_find_collisions_with_margin_indices(monkeypatch): spikes, spikes_within_margin, delta_collision_samples=4, - sparsity_mask=sparsity_mask, + overlap_matrix=_unit_pair_overlap_matrix(sparsity_mask), spike_indices=spike_indices, ) @@ -96,6 +101,40 @@ def test_find_collisions_with_margin_indices(monkeypatch): np.testing.assert_array_equal(collisions[2], spikes_within_margin[[3, 2]]) +def test_unit_pair_overlap_matrix_matches_naive_reference(): + """ + `_unit_pair_overlap_matrix` replaces a per-pair `np.any(sparsity_mask[i] & sparsity_mask[j])` + lookup (previously recomputed on every temporally-overlapping spike-pair candidate inside + `find_collisions`, millions of times on a realistic recording) with one matrix multiplication + computed once. Verify it agrees, entry by entry, with the direct naive reference on + NON-CONTIGUOUS per-unit channel subsets (a bounded/contiguous slice would pass even a mutant + that only checks a channel range), and on a unit with an entirely empty sparsity row (no + channels at all -- that unit must not spuriously overlap with anything, including itself). + """ + rng = np.random.default_rng(99) + num_units, num_channels = 15, 24 + + sparsity_mask = np.zeros((num_units, num_channels), dtype=bool) + for unit_index in range(num_units): + if unit_index == 0: + continue # unit 0: entirely empty sparsity row + num_active = rng.integers(1, num_channels // 2) + active_channels = rng.choice(num_channels, size=num_active, replace=False) # non-contiguous + sparsity_mask[unit_index, active_channels] = True + + overlap_matrix = _unit_pair_overlap_matrix(sparsity_mask) + + assert overlap_matrix.shape == (num_units, num_units) + assert overlap_matrix.dtype == np.bool_ + assert not overlap_matrix[0, :].any() # the empty-row unit overlaps with nothing + assert not overlap_matrix[:, 0].any() + + for i in range(num_units): + for j in range(num_units): + expected = np.any(sparsity_mask[i] & sparsity_mask[j]) + assert overlap_matrix[i, j] == expected, f"mismatch at unit pair ({i}, {j})" + + class TestAmplitudeScalingsExtension(AnalyzerExtensionCommonTestSuite): @pytest.mark.parametrize("params", [dict(handle_collisions=True), dict(handle_collisions=False)]) diff --git a/src/spikeinterface/postprocessing/tests/test_empty_extensions.py b/src/spikeinterface/postprocessing/tests/test_empty_extensions.py new file mode 100644 index 0000000000..3102015997 --- /dev/null +++ b/src/spikeinterface/postprocessing/tests/test_empty_extensions.py @@ -0,0 +1,26 @@ +import numpy as np +import pytest + +from spikeinterface import create_sorting_analyzer, generate_ground_truth_recording + + +@pytest.mark.parametrize("sparse", [False, True]) +def test_postprocessing_metrics_with_no_units(sparse): + recording, sorting = generate_ground_truth_recording(durations=[0.5], num_channels=4, num_units=2, seed=42) + analyzer = create_sorting_analyzer(sorting.select_units([]), recording, sparse=sparse) + analyzer.compute(["random_spikes", "templates", "noise_levels"], n_jobs=1, progress_bar=False) + for name in ("correlograms", "isi_histograms"): + result, bins = analyzer.compute(name, method="numpy").get_data() + expected_shape = (0, 0, len(bins) - 1) if name == "correlograms" else (0, len(bins) - 1) + assert result.shape == expected_shape + assert np.all(np.diff(bins) > 0) + similarity = analyzer.compute("template_similarity").get_data() + assert similarity.shape == (0, 0) + for name, metrics in ( + ("quality_metrics", ["firing_rate", "snr"]), + ("template_metrics", ["peak_to_trough_duration"]), + ): + result = analyzer.compute(name, metric_names=metrics).get_data() + assert result.shape == (0, len(metrics)) + assert list(result.columns) == metrics + assert list(result.index) == [] diff --git a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py index 1bb3d383a2..6ae72e374d 100644 --- a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py +++ b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py @@ -69,7 +69,7 @@ } -def get_dataset_to_merge(): +def get_dataset_to_merge(cache_folder): # generate a dataset with some split units to minimize merge errors recording, sorting = generate_ground_truth_recording( durations=[30.0], @@ -81,6 +81,7 @@ def get_dataset_to_merge(): generate_unit_locations_kwargs=dict(margin_um=10.0, minimum_z=2.0, maximum_z=15.0, minimum_distance=20), seed=2205, ) + recording = recording.save(folder=cache_folder / "recording_merge") # since templates are going to be averaged and this might be a problem for amplitude scaling # we select the 3 units with the largest templates to split @@ -96,7 +97,7 @@ def get_dataset_to_merge(): return recording, sorting_with_splits, split_unit_ids -def get_dataset_to_split(): +def get_dataset_to_split(cache_folder): # generate a dataset and return large unit to split to minimize split errors recording, sorting = generate_ground_truth_recording( durations=[30.0], @@ -107,6 +108,7 @@ def get_dataset_to_split(): noise_kwargs=dict(noise_levels=5.0, strategy="tile_pregenerated"), seed=2205, ) + recording = recording.save(folder=cache_folder / "recording_split") # since templates are going to be averaged and this might be a problem for amplitude scaling # we select the 3 units with the largest templates to split @@ -120,18 +122,20 @@ def get_dataset_to_split(): @pytest.fixture(scope="module") -def dataset_to_merge(): - return get_dataset_to_merge() +def dataset_to_merge(create_cache_folder): + return get_dataset_to_merge(create_cache_folder) @pytest.fixture(scope="module") -def dataset_to_split(): - return get_dataset_to_split() +def dataset_to_split(create_cache_folder): + return get_dataset_to_split(create_cache_folder) @pytest.mark.parametrize("lazy", [False, True]) @pytest.mark.parametrize("sparse", [False, True]) -@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +@pytest.mark.parametrize( + "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, lazy, sparse, format, tmp_path): if format == "memory" and lazy: pytest.skip("lazy has no effect for format='memory' (nothing on disk to load lazily)") @@ -263,7 +267,9 @@ def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, lazy, sparse, fo @pytest.mark.parametrize("lazy", [False, True]) @pytest.mark.parametrize("sparse", [False, True]) -@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +@pytest.mark.parametrize( + "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] +) def test_SortingAnalyzer_split_all_extensions(dataset_to_split, lazy, sparse, format, tmp_path): if format == "memory" and lazy: pytest.skip("lazy has no effect for format='memory' (nothing on disk to load lazily)") diff --git a/src/spikeinterface/postprocessing/tests/test_template_similarity.py b/src/spikeinterface/postprocessing/tests/test_template_similarity.py index 9fa7a73fec..fd963fa013 100644 --- a/src/spikeinterface/postprocessing/tests/test_template_similarity.py +++ b/src/spikeinterface/postprocessing/tests/test_template_similarity.py @@ -129,6 +129,74 @@ def test_equal_results_numba(params): assert np.allclose(result_numpy, result_numba, 1e-3) +def test_equal_results_numba_asymmetric_sparse(): + rng = np.random.default_rng(seed=2205) + templates_array = rng.random(size=(2, 20, 8), dtype=np.float32) + other_templates_array = rng.random(size=(7, 20, 8), dtype=np.float32) + sparsity_mask = rng.random((2, 8)) > 0.4 + other_sparsity_mask = rng.random((7, 8)) > 0.4 + sparsity_mask[0] = np.array([True, True, False, False, False, False, False, False]) + other_sparsity_mask[0] = np.array([False, False, True, True, False, False, False, False]) + + for method in ("cosine", "l1", "l2"): + for support in ("dense", "union", "intersection"): + result_numba = _compute_similarity_matrix_numba( + templates_array, + other_templates_array, + num_shifts=3, + method=method, + sparsity_mask=sparsity_mask, + other_sparsity_mask=other_sparsity_mask, + support=support, + ) + result_numpy = _compute_similarity_matrix_numpy( + templates_array, + other_templates_array, + num_shifts=3, + method=method, + sparsity_mask=sparsity_mask, + other_sparsity_mask=other_sparsity_mask, + support=support, + ) + + np.testing.assert_allclose(result_numba, result_numpy, rtol=1e-5, atol=1e-5) + + +def test_equal_results_numba_same_array_sparse(): + # templates_array is passed as both sides so same_array=True, exercising the + # pair-list restriction to tgt_unit >= src_unit and the diagonal mirroring, + # which the asymmetric test above cannot reach. + rng = np.random.default_rng(seed=2205) + templates_array = rng.random(size=(6, 20, 8), dtype=np.float32) + sparsity_mask = rng.random((6, 8)) > 0.4 + sparsity_mask[0] = np.array([True, True, False, False, False, False, False, False]) + sparsity_mask[1] = np.array([False, False, True, True, False, False, False, False]) + + for method in ("cosine", "l1", "l2"): + for support in ("dense", "union", "intersection"): + for num_shifts in (0, 3): + result_numba = _compute_similarity_matrix_numba( + templates_array, + templates_array, + num_shifts=num_shifts, + method=method, + sparsity_mask=sparsity_mask, + other_sparsity_mask=sparsity_mask, + support=support, + ) + result_numpy = _compute_similarity_matrix_numpy( + templates_array, + templates_array, + num_shifts=num_shifts, + method=method, + sparsity_mask=sparsity_mask, + other_sparsity_mask=sparsity_mask, + support=support, + ) + + np.testing.assert_allclose(result_numba, result_numpy, rtol=1e-5, atol=1e-5) + + if __name__ == "__main__": from spikeinterface.postprocessing.tests.common_extension_tests import get_dataset from spikeinterface.core import estimate_sparsity diff --git a/src/spikeinterface/postprocessing/unit_locations.py b/src/spikeinterface/postprocessing/unit_locations.py index 4d913ab4f9..7cd94d7ee9 100644 --- a/src/spikeinterface/postprocessing/unit_locations.py +++ b/src/spikeinterface/postprocessing/unit_locations.py @@ -7,10 +7,10 @@ # this dict is for peak location dtype_localize_by_method = { - "center_of_mass": [("x", "float64"), ("y", "float64")], - "grid_convolution": [("x", "float64"), ("y", "float64"), ("z", "float64")], - "peak_channel": [("x", "float64"), ("y", "float64")], - "monopolar_triangulation": [("x", "float64"), ("y", "float64"), ("z", "float64"), ("alpha", "float64")], + "center_of_mass": [("x", "float32"), ("y", "float32")], + "grid_convolution": [("x", "float32"), ("y", "float32"), ("z", "float32")], + "peak_channel": [("x", "float32"), ("y", "float32")], + "monopolar_triangulation": [("x", "float32"), ("y", "float32"), ("z", "float32"), ("alpha", "float32")], } possible_localization_methods = list(dtype_localize_by_method.keys()) diff --git a/src/spikeinterface/postprocessing/valid_unit_periods.py b/src/spikeinterface/postprocessing/valid_unit_periods.py index 214ebb5efc..da28b20e00 100644 --- a/src/spikeinterface/postprocessing/valid_unit_periods.py +++ b/src/spikeinterface/postprocessing/valid_unit_periods.py @@ -546,7 +546,7 @@ def _get_data(self, outputs: str = "by_unit"): # by_unit unit_ids = self.sorting_analyzer.unit_ids good_periods = [] - good_periods_array = self.data["valid_unit_periods"] + good_periods_array = materialize_array(self.data["valid_unit_periods"]) for segment_index in range(self.sorting_analyzer.get_num_segments()): segment_mask = good_periods_array["segment_index"] == segment_index periods_dict = {} diff --git a/src/spikeinterface/preprocessing/detect_artifacts.py b/src/spikeinterface/preprocessing/detect_artifacts.py index 95f286a85a..ab19c49edc 100644 --- a/src/spikeinterface/preprocessing/detect_artifacts.py +++ b/src/spikeinterface/preprocessing/detect_artifacts.py @@ -126,6 +126,15 @@ def __init__( else: self.diff_threshold_unscaled = None + self._kwargs.update( + dict( + saturation_threshold_uV=saturation_threshold_uV, + diff_threshold_uV=diff_threshold_uV, + proportion=proportion, + signed=signed, + ) + ) + def get_margin(self) -> int: """Return the number of margin samples required on each side of a chunk.""" return 0 @@ -424,9 +433,13 @@ def compute( ``"front"`` (``True`` for rising edge, ``False`` for falling edge). """ z = np.median((traces - self.medians) / self.abs_thresholds, axis=1) - threshold_mask = np.diff((z > 1) != 0, axis=0) + above_threshold = z > 1 - indices = np.flatnonzero(threshold_mask) + # the chunk is treated as self-contained: the signal is considered below threshold before the + # first and after the last sample. This guarantees that crossings always come in + # (rising, falling) pairs, even when a period straddles a chunk boundary. Such a period is + # emitted as two adjacent periods, which are merged back by `_collapse_events`. + indices = np.flatnonzero(np.diff(above_threshold, prepend=False, append=False)) threshold_crossings = np.zeros(indices.size, dtype=self._dtype) threshold_crossings["sample_index"] = indices threshold_crossings["segment_index"] = segment_index diff --git a/src/spikeinterface/preprocessing/detect_bad_channels.py b/src/spikeinterface/preprocessing/detect_bad_channels.py index 7b865bd594..968c159666 100644 --- a/src/spikeinterface/preprocessing/detect_bad_channels.py +++ b/src/spikeinterface/preprocessing/detect_bad_channels.py @@ -39,6 +39,9 @@ The standard deviation/mad multiplier threshold psd_hf_threshold : float, default: 0.02 For coherence+psd - an absolute threshold (uV^2/Hz) used as a cutoff for noise channels. + Traces are retrieved with ``return_in_uV=True`` before computing the PSD, so this + threshold refers to microvolts regardless of the stored sample units. The recording + must have correct gains and offsets for conversion to microvolts. Channels with average power at >80% Nyquist larger than this threshold will be labeled as noise. IBL suggests 0.02 for AP data and 1.4 for LFP data. dead_channel_threshold : float, default: -0.5 @@ -475,6 +478,11 @@ def detect_bad_channels_ibl( psd_hf_threshold : float Threshold for high frequency PSD. If mean PSD above `nyquist_threshold` * fn is greater than this value, channels are flagged as noisy (together with channel coherence condition). + Units are the units of ``raw`` squared per Hz. This low-level function does not + apply recording gains or offsets. For ``raw`` in microvolts, 0.02 means + 0.02 uV^2/Hz; the equivalent threshold for the same signal in volts is 2e-14 V^2/Hz. + The public ``detect_bad_channels`` function converts recording traces to microvolts + before calling this function. dead_channel_thr : float, default: -0.5 Threshold for channel coherence below which channels are labeled as dead noisy_channel_thr : float, default: 1 diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index 3000d0c143..c01a9acf59 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -9,10 +9,7 @@ class PhaseShiftRecording(BasePreprocessor): """ - This apply a phase shift to a recording to cancel the small sampling - delay across for some recording system. - - This is particularly relevant for neuropixel recording. + Apply phase shifts to compensate for the staggered ADC sampling times in Neuropixels recordings. This code is inspired from from IBL lib. https://github.com/int-brain-lab/ibllib/blob/master/ibllib/dsp/fourier.py @@ -145,24 +142,21 @@ def apply_frequency_shift(signal, shift_samples, axis=0): from scipy.fft import rfft, irfft signal_length = signal.shape[axis] - num_channels = shift_samples.size - fourier_signal_size = signal_length // 2 + 1 - frequency_domain_signal = rfft(signal, n=signal_length, axis=axis, overwrite_x=True) - fourier_signal_size = frequency_domain_signal.shape[0] if axis == 0: - frequency_grid = np.empty(shape=(fourier_signal_size, num_channels)) - # Note that np.fft.rfttfreq handles both even and odd signal lengths - frequency_grid[:, :] = 2 * np.pi * np.fft.rfftfreq(signal_length)[:, np.newaxis] - shifts = np.multiply(frequency_grid, shift_samples[np.newaxis, :], out=frequency_grid) + angular_frequencies = 2 * np.pi * np.fft.rfftfreq(signal_length) + # Neuropixels channels reuse a small set of ADC sampling delays, so compute each rotation once. + unique_shifts, shift_indices = np.unique(shift_samples, return_inverse=True) + unique_angles = (angular_frequencies[:, np.newaxis] * unique_shifts[np.newaxis, :]).astype( + np.float64, copy=False + ) + unique_rotations = np.exp(-1j * unique_angles) + rotations = unique_rotations[:, shift_indices] + phase_shifted_signal = np.multiply(frequency_domain_signal, rotations, out=rotations) else: raise NotImplementedError("Axis != 0 is not implemented yet") - # Rotate the signal in the frequency domain - rotations = np.exp(-1j * shifts) - phase_shifted_signal = np.multiply(frequency_domain_signal, rotations, out=rotations) - # Inverse FFT to get the translated signal shifted_signal = irfft(phase_shifted_signal, n=signal_length, axis=axis, overwrite_x=True) return shifted_signal diff --git a/src/spikeinterface/preprocessing/tests/test_detect_artifacts.py b/src/spikeinterface/preprocessing/tests/test_detect_artifacts.py index 76347f56b8..2374927416 100644 --- a/src/spikeinterface/preprocessing/tests/test_detect_artifacts.py +++ b/src/spikeinterface/preprocessing/tests/test_detect_artifacts.py @@ -26,10 +26,18 @@ def test_detect_artifact_by_envelope(debug_plots): for start, stop in zip(artifact_starts, artifact_stops): data[start:stop, :] = sat_value + # one long artifact straddling a chunk boundary, to check that periods split across chunks + # are recovered as a single period + data[chunk_size - 500 : chunk_size + 500, :] = sat_value + recording = NumpyRecording(data, sampling_frequency) artifacts, envelope = detect_artifact_periods_by_envelope( - recording, apply_envelope_common_reference=False, return_envelope=True, random_slices_kwargs={"seed": 2308} + recording, + apply_envelope_common_reference=False, + return_envelope=True, + random_slices_kwargs={"seed": 2308}, + job_kwargs={"chunk_size": chunk_size, "n_jobs": 1}, ) if debug_plots: @@ -43,6 +51,24 @@ def test_detect_artifact_by_envelope(debug_plots): # it finds some artifacts assert len(artifacts) > 0 + # chunking must not change the result: the same periods are found when the whole + # recording is processed as a single chunk (edges can differ by a few samples because + # the gaussian filter is applied chunk-wise) + artifacts_single_chunk = detect_artifact_periods_by_envelope( + recording, + apply_envelope_common_reference=False, + random_slices_kwargs={"seed": 2308}, + job_kwargs={"chunk_size": data.shape[0], "n_jobs": 1}, + ) + assert artifacts.size == artifacts_single_chunk.size + for field in ("start_sample_index", "end_sample_index"): + assert np.allclose(artifacts[field], artifacts_single_chunk[field], atol=100) + + # the artifact straddling the chunk boundary is detected as one single period + straddling = artifacts[artifacts["start_sample_index"] < chunk_size] + straddling = straddling[straddling["end_sample_index"] > chunk_size] + assert straddling.size == 1 + def test_detect_saturation_periods(debug_plots): """ diff --git a/src/spikeinterface/preprocessing/tests/test_phase_shift.py b/src/spikeinterface/preprocessing/tests/test_phase_shift.py index 6bf7d50dd4..23718f6ac6 100644 --- a/src/spikeinterface/preprocessing/tests/test_phase_shift.py +++ b/src/spikeinterface/preprocessing/tests/test_phase_shift.py @@ -93,5 +93,24 @@ def test_phase_shift(): # ~ plt.show() +def test_phase_shift_recovers_delayed_signal(): + sampling_frequency = 1000.0 + times = np.arange(5000) / sampling_frequency + inter_sample_shift = np.tile([0.0, 0.25, 0.5, 0.75], 8) + + def signal(time): + return np.sin(2 * np.pi * 2.5 * time) + np.sin(2 * np.pi * 8.5 * time) + + traces = signal(times[:, np.newaxis] + inter_sample_shift / sampling_frequency) + recording = NumpyRecording([traces], sampling_frequency) + recording.set_property("inter_sample_shift", inter_sample_shift) + + shifted_traces = phase_shift(recording).get_traces() + + interior = slice(100, -100) + expected = np.broadcast_to(signal(times)[:, np.newaxis], traces.shape) + np.testing.assert_allclose(shifted_traces[interior], expected[interior], atol=1e-3) + + if __name__ == "__main__": test_phase_shift() diff --git a/src/spikeinterface/preprocessing/tests/test_pipeline.py b/src/spikeinterface/preprocessing/tests/test_pipeline.py index 5f5c298430..7b53b0ab92 100644 --- a/src/spikeinterface/preprocessing/tests/test_pipeline.py +++ b/src/spikeinterface/preprocessing/tests/test_pipeline.py @@ -2,6 +2,7 @@ from spikeinterface.core.testing import check_recordings_equal from spikeinterface.core import create_sorting_analyzer +from spikeinterface.core.core_tools import _is_zarr_write_supported from spikeinterface.generation import generate_recording, generate_ground_truth_recording from spikeinterface.preprocessing import ( apply_preprocessing_pipeline, @@ -169,6 +170,7 @@ def test_loading_provenance(create_cache_folder): cache_folder = create_cache_folder / "preprocessed_rec_for_pipeline" rec, _ = generate_ground_truth_recording(seed=0, num_channels=6) + rec = rec.save(folder=cache_folder / "recording_provenance") pp_rec = detect_and_remove_bad_channels( bandpass_filter(common_reference(rec, operator="average")), noisy_channel_threshold=0.3, @@ -176,9 +178,9 @@ def test_loading_provenance(create_cache_folder): # when several run seed=2205, ) - pp_rec.save(folder=cache_folder) + pp_rec.save(folder=cache_folder / "preprocessed_recording") - loaded_pp_list = get_preprocessing_list_from_file(cache_folder / "provenance.pkl") + loaded_pp_list = get_preprocessing_list_from_file(cache_folder / "preprocessed_recording" / "provenance.json") pipeline_rec_applying_precomputed_kwargs = apply_preprocessing_pipeline( rec, @@ -204,6 +206,7 @@ def test_loading_from_analyzer(create_cache_folder): cache_folder = create_cache_folder recording, sorting = generate_ground_truth_recording() + recording = recording.save(folder=cache_folder / "recording") preprocessing_dict = {"common_reference": {}, "highpass_filter": {"freq_min": 301.0}} pp_recording = apply_preprocessing_pipeline(recording, preprocessing_dict) @@ -216,11 +219,13 @@ def test_loading_from_analyzer(create_cache_folder): pp_recording_from_binary = apply_preprocessing_pipeline(recording, pp_list_from_binary) check_recordings_equal(pp_recording, pp_recording_from_binary) - analyzer_zarr_folder = cache_folder / "zarr_format.zarr" - _ = create_sorting_analyzer(sorting=sorting, recording=pp_recording, format="zarr", folder=analyzer_zarr_folder) - pp_list_from_zarr = get_preprocessing_list_from_analyzer(analyzer_zarr_folder) - pp_recording_from_zarr = apply_preprocessing_pipeline(recording, pp_list_from_zarr) - check_recordings_equal(pp_recording, pp_recording_from_zarr) + # TODO: remove once writing to zarr is supported with zarr>=3 + if _is_zarr_write_supported(): + analyzer_zarr_folder = cache_folder / "zarr_format.zarr" + _ = create_sorting_analyzer(sorting=sorting, recording=pp_recording, format="zarr", folder=analyzer_zarr_folder) + pp_list_from_zarr = get_preprocessing_list_from_analyzer(analyzer_zarr_folder) + pp_recording_from_zarr = apply_preprocessing_pipeline(recording, pp_list_from_zarr) + check_recordings_equal(pp_recording, pp_recording_from_zarr) def test_pipeline_recording_arg_substitution(create_cache_folder): diff --git a/src/spikeinterface/preprocessing/tests/test_resample.py b/src/spikeinterface/preprocessing/tests/test_resample.py index 9d745fa1dc..de5f19def4 100644 --- a/src/spikeinterface/preprocessing/tests/test_resample.py +++ b/src/spikeinterface/preprocessing/tests/test_resample.py @@ -188,7 +188,7 @@ def test_resample_rational_grid(): approximate = not np.isclose(achieved_rate, requested_rate, rtol=1e-12, atol=0) for explicit_times in [False, True]: if explicit_times: - parent.set_times(parent.get_times(), with_warning=False) + parent.set_times(parent.get_times()) with pytest.warns(UserWarning, match="achieves") if approximate else nullcontext(): processed = resample(parent, requested_rate, max_denominator=max_denominator) assert processed.get_sampling_frequency() == pytest.approx(achieved_rate, rel=1e-15) @@ -207,7 +207,7 @@ def test_resample_rational_grid(): def test_resample_short_sections(): for rate in [100, 1500]: parent = NumpyRecording(np.array([[1.0], [9.0]]), 1000) - parent.set_times(np.array([10.0, 20.0]), with_warning=False) + parent.set_times(np.array([10.0, 20.0])) processed = resample(parent, rate, gap_tolerance_ms=0) sections = [resample(parent.frame_slice(i, i + 1), rate) for i in range(2)] expected = np.concatenate([section.get_traces() for section in sections]) diff --git a/src/spikeinterface/sorters/container_tools.py b/src/spikeinterface/sorters/container_tools.py index b08cd8d60e..a75ca57be7 100644 --- a/src/spikeinterface/sorters/container_tools.py +++ b/src/spikeinterface/sorters/container_tools.py @@ -8,14 +8,19 @@ # TODO move this inside functions -from spikeinterface.core.core_tools import recursive_path_modifier, _get_paths_list +from spikeinterface.core.core_tools import recursive_path_modifier, _get_paths_list, is_path_remote def find_recording_folders(d): """Finds all recording folders 'paths' in a dict""" path_list = _get_paths_list(d=d) - folders_to_mount = [Path(p).resolve().parent for p in path_list] + # Remote paths (URLs) are read over the network from inside the container, so there is nothing + # to mount for them. + local_path_list = [p for p in path_list if not is_path_remote(p)] + folders_to_mount = [Path(p).resolve().parent for p in local_path_list] + if len(folders_to_mount) == 0: + return [] try: # this will fail if on different drives (Windows) base_folders_to_mount = [Path(os.path.commonpath(folders_to_mount))] diff --git a/src/spikeinterface/sorters/internal/lupin.py b/src/spikeinterface/sorters/internal/lupin.py index f6a055c5c8..6d79bb036d 100644 --- a/src/spikeinterface/sorters/internal/lupin.py +++ b/src/spikeinterface/sorters/internal/lupin.py @@ -11,6 +11,7 @@ ) from spikeinterface.core.job_tools import fix_job_kwargs +from spikeinterface.core.core_tools import _check_zarr_write_is_supported from spikeinterface.preprocessing import bandpass_filter, common_reference, zscore, whiten from spikeinterface.core.base import minimum_spike_dtype @@ -124,6 +125,10 @@ def get_sorter_version(cls): @classmethod def _run_from_folder(cls, sorter_output_folder, params, verbose): + # TODO: remove once writing to zarr is supported with zarr>=3 + if params["save_array"]: + _check_zarr_write_is_supported() + from spikeinterface.sortingcomponents.tools import get_prototype_and_waveforms_from_recording from spikeinterface.sortingcomponents.matching import find_spikes_from_templates from spikeinterface.sortingcomponents.peak_detection import detect_peaks diff --git a/src/spikeinterface/sorters/internal/tests/test_lupin.py b/src/spikeinterface/sorters/internal/tests/test_lupin.py index df2666be1d..c4eb788a6d 100644 --- a/src/spikeinterface/sorters/internal/tests/test_lupin.py +++ b/src/spikeinterface/sorters/internal/tests/test_lupin.py @@ -1,5 +1,7 @@ import unittest +import pytest + from spikeinterface.sorters.tests.common_tests import SorterCommonTestSuite from spikeinterface.sorters import LupinSorter, run_sorter @@ -10,6 +12,11 @@ class LupinSorterCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): SorterClass = LupinSorter + # TODO: remove once writing to zarr is supported with zarr>=3 (save_array writes the templates to zarr) + @pytest.mark.requires_zarr_write + def test_with_run(self): + super().test_with_run() + if __name__ == "__main__": test = LupinSorterCommonTestSuite() diff --git a/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py b/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py index 01284d7ee6..cb6ed09ce4 100644 --- a/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py +++ b/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py @@ -1,5 +1,7 @@ import unittest +import pytest + from spikeinterface.sorters.tests.common_tests import SorterCommonTestSuite from spikeinterface.sorters import Tridesclous2Sorter, run_sorter @@ -10,6 +12,11 @@ class Tridesclous2SorterCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): SorterClass = Tridesclous2Sorter + # TODO: remove once writing to zarr is supported with zarr>=3 (save_array writes the templates to zarr) + @pytest.mark.requires_zarr_write + def test_with_run(self): + super().test_with_run() + @unittest.skip("performance reason") def test_with_numpy_gather(self): recording = self.recording diff --git a/src/spikeinterface/sorters/internal/tridesclous2.py b/src/spikeinterface/sorters/internal/tridesclous2.py index d5c66f35be..55d377d4db 100644 --- a/src/spikeinterface/sorters/internal/tridesclous2.py +++ b/src/spikeinterface/sorters/internal/tridesclous2.py @@ -15,6 +15,7 @@ ) from spikeinterface.core.job_tools import fix_job_kwargs +from spikeinterface.core.core_tools import _check_zarr_write_is_supported from spikeinterface.preprocessing import bandpass_filter, common_reference, zscore, whiten from spikeinterface.core.base import minimum_spike_dtype @@ -118,6 +119,10 @@ def get_sorter_version(cls): @classmethod def _run_from_folder(cls, sorter_output_folder, params, verbose): + # TODO: remove once writing to zarr is supported with zarr>=3 + if params["save_array"]: + _check_zarr_write_is_supported() + from spikeinterface.sortingcomponents.matching import find_spikes_from_templates from spikeinterface.sortingcomponents.peak_detection import detect_peaks from spikeinterface.sortingcomponents.peak_selection import select_peaks diff --git a/src/spikeinterface/sorters/tests/test_container_tools.py b/src/spikeinterface/sorters/tests/test_container_tools.py index 606fe9940e..2f3b176cc5 100644 --- a/src/spikeinterface/sorters/tests/test_container_tools.py +++ b/src/spikeinterface/sorters/tests/test_container_tools.py @@ -64,6 +64,12 @@ def test_find_recording_folders(setup_module): assert len(f3) == 3 +def test_find_recording_folders_skips_remote_paths(tmp_path): + """A remote path is read from inside the container, so it is not mounted (GH-2370).""" + recording_dict = {"kwargs": {"file_paths": [str(tmp_path / "recording.bin"), "s3://bucket/recording.zarr"]}} + assert find_recording_folders(recording_dict) == [tmp_path.resolve()] + + @pytest.mark.skipif(ON_GITHUB, reason="Docker tests don't run on github: test locally") def test_ContainerClient(): mode = "docker" diff --git a/src/spikeinterface/sorters/tests/test_launcher.py b/src/spikeinterface/sorters/tests/test_launcher.py index 2cd25f119a..35db70e677 100644 --- a/src/spikeinterface/sorters/tests/test_launcher.py +++ b/src/spikeinterface/sorters/tests/test_launcher.py @@ -13,10 +13,12 @@ SORTERS = ["tridesclous2"] -def create_recordings(NUM_RECORDINGS=2, base_seed=42): +def create_recordings(folder, NUM_RECORDINGS=2, base_seed=42): recordings = [] for i in range(NUM_RECORDINGS): + recording, _ = generate_ground_truth_recording(num_channels=8, durations=[10.0], seed=base_seed + i) + recording = recording.save(folder=folder / f"rec_for_generate_recording_{i}", format="binary", overwrite=True) if i % 2 == 0: recording.set_channel_groups(["0"] * 4 + ["1"] * 4) @@ -28,7 +30,7 @@ def create_recordings(NUM_RECORDINGS=2, base_seed=42): def get_job_list(base_folder): jobs = [] - recordings = create_recordings(NUM_RECORDINGS) + recordings = create_recordings(base_folder, NUM_RECORDINGS) for i, recording in enumerate(recordings): for sorter_name in SORTERS: kwargs = dict( @@ -50,6 +52,7 @@ def job_list(create_cache_folder): return get_job_list(folder) +@pytest.mark.requires_zarr_write def test_run_sorter_jobs_loop(job_list): sortings = run_sorter_jobs(job_list, engine="loop", return_output=True) print(sortings) @@ -199,6 +202,7 @@ def test_run_sorter_jobs_slurm_kwargs(mocker, tmp_path, job_list): assert str(tmp_script_folder) in mock_subprocess_run.call_args_list[-1].args[0][5] +@pytest.mark.requires_zarr_write def test_run_sorter_by_property(create_cache_folder): cache_folder = create_cache_folder working_folder1 = cache_folder / "test_run_sorter_by_property_1" @@ -209,7 +213,7 @@ def test_run_sorter_by_property(create_cache_folder): if working_folder2.is_dir(): shutil.rmtree(working_folder2) - recordings = create_recordings(NUM_RECORDINGS) + recordings = create_recordings(create_cache_folder, NUM_RECORDINGS) rec0 = recordings[0] rec0_by = rec0.split_by("group") diff --git a/src/spikeinterface/sorters/tests/test_runsorter.py b/src/spikeinterface/sorters/tests/test_runsorter.py index 332e6e857e..1316430e55 100644 --- a/src/spikeinterface/sorters/tests/test_runsorter.py +++ b/src/spikeinterface/sorters/tests/test_runsorter.py @@ -13,16 +13,18 @@ ON_GITHUB = bool(os.getenv("GITHUB_ACTIONS")) -def _generate_recording(): +def _generate_recording(folder): recording, _ = generate_ground_truth_recording(num_channels=8, durations=[10.0], seed=2205) + recording = recording.save(folder=folder / "rec_for_generate_recording", format="binary") return recording @pytest.fixture(scope="module") -def generate_recording(): - return _generate_recording() +def generate_recording(create_cache_folder): + return _generate_recording(create_cache_folder) +@pytest.mark.requires_zarr_write @pytest.mark.xfail( platform.system() == "Windows" and parse(platform.python_version()) > parse("3.12"), reason="3rd parth threadpoolctl issue: OSError('GetModuleFileNameEx failed')", diff --git a/src/spikeinterface/sortingcomponents/matching/tests/test_wobble.py b/src/spikeinterface/sortingcomponents/matching/tests/test_wobble.py index 0d46b790ad..36d5452933 100644 --- a/src/spikeinterface/sortingcomponents/matching/tests/test_wobble.py +++ b/src/spikeinterface/sortingcomponents/matching/tests/test_wobble.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import pytest import numpy as np @@ -261,6 +263,63 @@ def test_compute_objective(): assert np.allclose(objective, expected_objective) +def test_find_strict_peaks_matches_argrelmax(): + from scipy.signal import argrelmax + + rng = np.random.default_rng(0) + nan_objective = rng.random(200).astype(np.float32) + nan_objective[50] = np.nan + cases = [ + (rng.random(2000).astype(np.float32), 5, 0.3), + (rng.random(50), 60, -1.0), + (rng.random(50), 61, -1.0), + (np.array([0.0, 1.0, 1.0, 1.0, 0.5, 2.0, 2.0, 0.0], dtype=np.float32), 2, -1.0), + (np.array([0.0, 3.0, 0.0, -np.inf, -np.inf, 0.0, 2.0, 0.0], dtype=np.float32), 2, -1.0), + (nan_objective, 4, 0.2), + (rng.random(2).astype(np.float32), 1, -1.0), + (rng.random(500).astype(np.float32), 3, 10.0), + ] + + for objective, order, threshold in cases: + expected = argrelmax(objective, order=order)[0] + expected = expected[objective[expected] > threshold] + actual = wobble._find_strict_peaks(objective, order, threshold) + assert np.array_equal(actual, expected) + + +def test_find_peaks_matches_argrelmax_windowing(monkeypatch): + from scipy.signal import argrelmax + + cases = [ + ([0.0, 0.0, 1.0, 4.0, 1.0, 3.0, 1.0, 0.0, 0.0], 3, 2.0), + ([0.0, 0.0, 0.0, 5.0, 1.0, 2.0, 1.0, 4.0, 0.0, 0.0], 4, 1.5), + ([0.0, 0.0, 2.0, 1.0, 2.0, 0.0, 0.0], 3, 3.0), + ] + + def no_high_res_shift(cls, spike_time_indices, *args): + zeros = np.zeros(spike_time_indices.size, dtype=np.int64) + return zeros, zeros, np.zeros(0, dtype=np.int64), np.zeros(0) + + monkeypatch.setattr(wobble.WobbleMatch, "calculate_high_res_shift", classmethod(no_high_res_shift)) + for objective_maximum, num_samples, threshold in cases: + objective_normalized = np.asarray([objective_maximum], dtype=np.float32) + window = objective_normalized[0, num_samples - 1 : -num_samples] + expected_indices = argrelmax(window, order=num_samples - 1)[0] + expected_indices = expected_indices[window[expected_indices] > threshold] + spike_train, scalings, distance_metric = wobble.WobbleMatch.find_peaks( + objective=objective_normalized, + objective_normalized=objective_normalized, + spike_trains=np.zeros((0, 2), dtype=np.int64), + params=SimpleNamespace(threshold=threshold, jitter_factor=1), + template_data=None, + template_meta=SimpleNamespace(num_samples=num_samples), + ) + + assert np.array_equal(spike_train[:, 0], expected_indices) + assert np.array_equal(scalings, np.ones(expected_indices.size, dtype=np.float32)) + assert np.array_equal(distance_metric, window[expected_indices]) + + def test_compute_scale_amplitudes(): # Arrange: Generate random 'data' seed = 0 diff --git a/src/spikeinterface/sortingcomponents/matching/wobble.py b/src/spikeinterface/sortingcomponents/matching/wobble.py index 92bf814f72..e9ee961c3e 100644 --- a/src/spikeinterface/sortingcomponents/matching/wobble.py +++ b/src/spikeinterface/sortingcomponents/matching/wobble.py @@ -326,6 +326,30 @@ def __post_init__(self): self.temporal, self.singular, self.spatial, self.temporal_jittered = self.compressed_templates +def _find_strict_peaks(objective, order, threshold): + """Find values strictly greater than their neighbors within ``order`` samples.""" + from scipy.ndimage import maximum_filter1d + + if objective.size < 3: + return np.zeros(0, dtype=np.intp) + if np.isnan(objective).any(): + from scipy.signal import argrelmax + + peak_indices = argrelmax(objective, order=order)[0] + return peak_indices[objective[peak_indices] > threshold] + + trailing_maximum = maximum_filter1d(objective, size=order, origin=(order - 1) // 2, mode="nearest") + left_maximum = np.empty_like(objective) + left_maximum[0] = objective[0] + left_maximum[1:] = trailing_maximum[:-1] + + trailing_maximum = maximum_filter1d(objective[::-1], size=order, origin=(order - 1) // 2, mode="nearest") + right_maximum = np.empty_like(objective) + right_maximum[-1] = objective[-1] + right_maximum[:-1] = trailing_maximum[-2::-1] + return np.flatnonzero((objective > threshold) & (objective > left_maximum) & (objective > right_maximum)) + + class WobbleMatch(BaseTemplateMatching): """Template matching method from the Paninski lab. @@ -618,16 +642,16 @@ def find_peaks( Finally, it generates a new spike train from the spike times, and returns it along with additional metrics about each spike. """ - from scipy import signal - # Get spike times (indices) using peaks in the objective objective_template_max = np.max(objective_normalized, axis=0) spike_window = (template_meta.num_samples - 1, objective_normalized.shape[1] - template_meta.num_samples) objective_windowed = objective_template_max[spike_window[0] : spike_window[1]] - spike_time_indices = signal.argrelmax(objective_windowed, order=template_meta.num_samples - 1)[0] + spike_time_indices = _find_strict_peaks( + objective_windowed, + order=template_meta.num_samples - 1, + threshold=params.threshold, + ) spike_time_indices += template_meta.num_samples - 1 - objective_spikes = objective_template_max[spike_time_indices] - spike_time_indices = spike_time_indices[objective_spikes > params.threshold] if len(spike_time_indices) == 0: # No new spikes found return np.zeros((0, 2), dtype=np.int32), np.zeros(0), np.zeros(0) diff --git a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py index 9635bad48b..97558448f4 100644 --- a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py +++ b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py @@ -224,7 +224,7 @@ def test_cross_band_interpolation(): assert np.array_equal(traces_corrected, target) -def test_InterpolateMotionRecording(): +def test_InterpolateMotionRecording(create_cache_folder): # rec, sorting = make_dataset() # 2 segments @@ -244,6 +244,7 @@ def test_InterpolateMotionRecording(): noise_kwargs=dict(noise_levels=5.0, strategy="on_the_fly"), seed=2205, ) + rec = rec.save(folder=create_cache_folder / "rec_for_test_interpolate_motion", format="binary") motion = make_fake_motion(rec) diff --git a/src/spikeinterface/sortingcomponents/peak_localization/main.py b/src/spikeinterface/sortingcomponents/peak_localization/main.py index c788f8b7fa..ffbdb5973d 100644 --- a/src/spikeinterface/sortingcomponents/peak_localization/main.py +++ b/src/spikeinterface/sortingcomponents/peak_localization/main.py @@ -20,10 +20,12 @@ def get_localization_pipeline_nodes( recording, peak_source, + peaks_or_spikes=None, method="center_of_mass", method_kwargs=None, ms_before=0.5, ms_after=0.5, + channel_index_array=None, job_kwargs=None, seed=None, ): @@ -44,16 +46,28 @@ def get_localization_pipeline_nodes( if method == "grid_convolution" and "prototype" not in method_kwargs: assert isinstance(peak_source, (PeakRetriever, SpikeRetriever)) + assert peaks_or_spikes is not None, "peaks_or_spikes must be provided for grid_convolution method." # extract prototypes silently from ..tools import get_prototype_and_waveforms_from_peaks + if isinstance(peak_source, SpikeRetriever): + assert channel_index_array is not None, "For SpikeRetriever, channel_index_array must be provided." + job_kwargs = fix_job_kwargs(job_kwargs) job_kwargs["progress_bar"] = False + # How do you get peaks + method_kwargs = method_kwargs.copy() method_kwargs["prototype"], _, _ = get_prototype_and_waveforms_from_peaks( - recording, peaks=peak_source.peaks, ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs, seed=seed + recording, + peaks=peaks_or_spikes, + ms_before=ms_before, + ms_after=ms_after, + channel_index_array=channel_index_array, + job_kwargs=job_kwargs, + seed=seed, ) localization_nodes = method_class(recording, parents=[peak_source, extract_dense_waveforms], **method_kwargs) @@ -141,6 +155,7 @@ def localize_peaks( pipeline_nodes = get_localization_pipeline_nodes( recording, peak_source, + peaks_or_spikes=peaks, method=method, method_kwargs=method_kwargs, ms_before=ms_before, diff --git a/src/spikeinterface/sortingcomponents/peak_selection.py b/src/spikeinterface/sortingcomponents/peak_selection.py index ddec61c1ae..018e627217 100644 --- a/src/spikeinterface/sortingcomponents/peak_selection.py +++ b/src/spikeinterface/sortingcomponents/peak_selection.py @@ -108,6 +108,11 @@ def select_peak_indices(peaks, method, seed, **method_kwargs): rng = np.random.default_rng(seed=seed) + # raise error if "channel_index" is not in peaks dtype and select_channel + channel_index_array = method_kwargs.get("channel_index_array", None) + if channel_index_array is None and "channel_index" in peaks.dtype.names: + channel_index_array = peaks["channel_index"] + if method == "uniform": params = {"select_per_channel": False, "n_peaks": None} @@ -116,9 +121,11 @@ def select_peak_indices(peaks, method, seed, **method_kwargs): assert params["n_peaks"] is not None, "n_peaks should be defined!" if params["select_per_channel"]: + if channel_index_array is None: + raise ValueError("Peaks do not have 'channel_index' field and channel_index_array is not provided.") ## This method will randomly select max_peaks_per_channel peaks per channels - for channel in np.unique(peaks["channel_index"]): - peaks_indices = np.where(peaks["channel_index"] == channel)[0] + for channel in np.unique(channel_index_array): + peaks_indices = np.where(channel_index_array == channel)[0] max_peaks = min(peaks_indices.size, params["n_peaks"]) selected_indices += [rng.choice(peaks_indices, size=max_peaks, replace=False)] else: @@ -144,8 +151,10 @@ def select_peak_indices(peaks, method, seed, **method_kwargs): assert params["noise_levels"] is not None, "Noise levels should be provided" if params["select_per_channel"]: - for channel in np.unique(peaks["channel_index"]): - peaks_indices = np.where(peaks["channel_index"] == channel)[0] + if channel_index_array is None: + raise ValueError("Peaks do not have 'channel_index' field and channel_index_array is not provided.") + for channel in np.unique(channel_index_array): + peaks_indices = np.where(channel_index_array == channel)[0] if params["n_peaks"] > peaks_indices.size: selected_indices += [peaks_indices] else: @@ -170,6 +179,10 @@ def select_peak_indices(peaks, method, seed, **method_kwargs): if params["n_peaks"] > peaks.size: selected_indices += [np.arange(peaks.size)] else: + if channel_index_array is None: + raise ValueError( + "Peaks do not have 'channel_index' field and channel_index_array is not provided." + ) snrs = peaks["amplitude"] / params["noise_levels"][peaks["channel_index"]] preprocessing = QuantileTransformer(output_distribution="uniform", n_quantiles=min(100, len(snrs))) snrs = preprocessing.fit_transform(snrs[:, np.newaxis]) diff --git a/src/spikeinterface/sortingcomponents/tools.py b/src/spikeinterface/sortingcomponents/tools.py index b445513a6c..75161725a5 100644 --- a/src/spikeinterface/sortingcomponents/tools.py +++ b/src/spikeinterface/sortingcomponents/tools.py @@ -41,7 +41,9 @@ def make_multi_method_doc(methods, indent=" "): return doc -def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, job_name=None, job_kwargs=None): +def extract_waveform_at_max_channel( + rec, peaks, ms_before=0.5, ms_after=1.5, channel_index_array=None, job_name=None, job_kwargs=None +): """ Helper function to extract waveforms at the max channel from a peak list @@ -50,15 +52,20 @@ def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, job job_kwargs = fix_job_kwargs(job_kwargs) n = rec.get_num_channels() - unit_ids = np.arange(n, dtype="int64") sparsity_mask = np.eye(n, dtype="bool") spikes = np.zeros( peaks.size, dtype=[("sample_index", "int64"), ("unit_index", "int64"), ("segment_index", "int64")] ) spikes["sample_index"] = peaks["sample_index"] - spikes["unit_index"] = peaks["channel_index"] spikes["segment_index"] = peaks["segment_index"] + if channel_index_array is not None: + assert len(channel_index_array) == peaks.size, "Length of channel_index_array must match number of peaks." + spikes["unit_index"] = channel_index_array + else: + if "channel_index" not in peaks.dtype.names: + raise ValueError("The peaks array must have a 'channel_index' field.") + spikes["unit_index"] = peaks["channel_index"] nbefore = ms_to_samples(ms_before, rec.sampling_frequency) nafter = ms_to_samples(ms_after, rec.sampling_frequency) @@ -66,7 +73,6 @@ def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, job all_wfs = extract_waveforms_to_single_buffer( rec, spikes, - unit_ids, nbefore, nafter, mode="shared_memory", @@ -83,7 +89,7 @@ def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, job def get_prototype_and_waveforms_from_peaks( - recording, peaks, n_peaks=5000, ms_before=0.5, ms_after=0.5, seed=None, job_kwargs=None + recording, peaks, n_peaks=5000, ms_before=0.5, ms_after=0.5, channel_index_array=None, seed=None, job_kwargs=None ): """ Function to extract a prototype waveform from peaks. @@ -93,13 +99,16 @@ def get_prototype_and_waveforms_from_peaks( recording : Recording The recording object containing the data. peaks : numpy.array, optional - Array of peaks, if None, peaks will be detected, by default None. + Array of peaks or spikes, if None, peaks will be detected, by default None. + If spikes, the channel_index_array must be provided to indicate the channel of each spike. n_peaks : int, optional Number of peaks to consider, by default 5000. ms_before : float, optional Time in milliseconds before the peak to extract the waveform, by default 0.5. ms_after : float, optional Time in milliseconds after the peak to extract the waveform, by default 0.5. + channel_index_array : numpy.array | None, optional + Array of channel indices corresponding to each peak, by default None. seed : int or None, optional Seed for random number generator, by default None. job_kwargs : dict @@ -122,7 +131,13 @@ def get_prototype_and_waveforms_from_peaks( nafter = ms_to_samples(ms_after, recording.sampling_frequency) few_peaks = select_peaks( - peaks, recording=recording, method="uniform", n_peaks=n_peaks, margin=(nbefore, nafter), seed=seed + peaks, + recording=recording, + method="uniform", + n_peaks=n_peaks, + margin=(nbefore, nafter), + channel_index_array=channel_index_array, + seed=seed, ) waveforms = extract_waveform_at_max_channel( recording, @@ -131,6 +146,7 @@ def get_prototype_and_waveforms_from_peaks( ms_after=ms_after, job_kwargs=job_kwargs, job_name="waveform prototype", + channel_index_array=channel_index_array, ) with np.errstate(divide="ignore", invalid="ignore"): diff --git a/src/spikeinterface/sortingcomponents/waveforms/features_from_peaks.py b/src/spikeinterface/sortingcomponents/waveforms/features_from_peaks.py index 1695f2cc59..0c7ffe7447 100644 --- a/src/spikeinterface/sortingcomponents/waveforms/features_from_peaks.py +++ b/src/spikeinterface/sortingcomponents/waveforms/features_from_peaks.py @@ -91,8 +91,8 @@ def __init__( self.all_channels = all_channels self.peak_sign = peak_sign - self._kwargs.update(dict(all_channels=all_channels, peak_sign=peak_sign)) self._dtype = recording.get_dtype() + self._kwargs.update(dict(all_channels=all_channels, peak_sign=peak_sign)) def get_dtype(self): return self._dtype @@ -125,8 +125,8 @@ def __init__( self.channel_distance = get_channel_distances(recording) self.neighbours_mask = self.channel_distance <= radius_um self.all_channels = all_channels - self._kwargs.update(dict(radius_um=radius_um, all_channels=all_channels)) self._dtype = recording.get_dtype() + self._kwargs.update(dict(radius_um=radius_um, all_channels=all_channels)) def get_dtype(self): return self._dtype @@ -168,6 +168,7 @@ def __init__( self.radius_um = radius_um self.sparse = sparse self.noise_threshold = noise_threshold + self._dtype = recording.get_dtype() self._kwargs.update( dict( projections=projections, @@ -177,7 +178,6 @@ def __init__( feature=feature, ) ) - self._dtype = recording.get_dtype() def get_dtype(self): return self._dtype