From 626f70aaef7c1b65204a90ab83b25ba04c2e1fed Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 21 Sep 2026 17:48:59 +0200 Subject: [PATCH 01/38] After release (#4800) --- pyproject.toml | 10 +++++----- src/spikeinterface/__init__.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index fc44c7030d..a23e7b1d65 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" }, @@ -25,10 +25,10 @@ dependencies = [ "tqdm", "zarr>=2.18,<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", - "neo>=0.14.5", - "probeinterface>=0.4.0", + "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", + "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", + # "neo>=0.14.5", + # "probeinterface>=0.4.0", "packaging", "pydantic", "numcodecs<0.16.0", # For supporting zarr < 3 diff --git a/src/spikeinterface/__init__.py b/src/spikeinterface/__init__.py index f7c43a845f..a184e9f487 100644 --- a/src/spikeinterface/__init__.py +++ b/src/spikeinterface/__init__.py @@ -27,5 +27,5 @@ # This flag must be set to False for release # This avoids using versioning that contains ".dev0" (and this is a better choice) # This is mainly useful when using run_sorter in a container and spikeinterface install -# DEV_MODE = True -DEV_MODE = False +DEV_MODE = True +# DEV_MODE = False From b17c60f8f620f40cc224a805e4a639f799b93980 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= <42451234+JESUSROYETH@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:54:25 -0300 Subject: [PATCH 02/38] Speed up Wobble peak search with a linear maximum filter (#4782) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jesús Royeth Co-authored-by: Pierre Yger --- .../matching/tests/test_wobble.py | 59 +++++++++++++++++++ .../sortingcomponents/matching/wobble.py | 34 +++++++++-- 2 files changed, 88 insertions(+), 5 deletions(-) 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) From 80aa6436dfc203305859aef8430fb0c63f5cbc22 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 23 Sep 2026 16:05:46 +0200 Subject: [PATCH 03/38] Fix: zarr pickling in `SortingAnalyzer` with `try`/`except` (#4801) --- src/spikeinterface/core/sortinganalyzer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 3e80b9fab4..953f75259d 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -1103,7 +1103,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: From 959e1248d3db7b2e72f160ae75eb3fa81d764f01 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 28 Sep 2026 14:32:28 +0200 Subject: [PATCH 04/38] Fix: artifact detection crossing chunks (#4803) --- .../preprocessing/detect_artifacts.py | 8 ++++-- .../tests/test_detect_artifacts.py | 28 ++++++++++++++++++- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/preprocessing/detect_artifacts.py b/src/spikeinterface/preprocessing/detect_artifacts.py index 95f286a85a..c29199ae3b 100644 --- a/src/spikeinterface/preprocessing/detect_artifacts.py +++ b/src/spikeinterface/preprocessing/detect_artifacts.py @@ -424,9 +424,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/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): """ From e094b5e3656c5f94b58d73750cb7eda39406c9b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= <42451234+JESUSROYETH@users.noreply.github.com> Date: Mon, 28 Sep 2026 09:42:31 -0300 Subject: [PATCH 05/38] Speed up template similarity with pair-level parallelism (#4785) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jesús Royeth --- .../postprocessing/template_similarity.py | 84 +++++++++---------- .../tests/test_template_similarity.py | 68 +++++++++++++++ 2 files changed, 106 insertions(+), 46 deletions(-) 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/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 From d7cdc97268d6624e55cbf01176cfd35af85ea665 Mon Sep 17 00:00:00 2001 From: Amirali Moradniaei Date: Mon, 28 Sep 2026 16:41:04 +0330 Subject: [PATCH 06/38] Clarify PSD threshold units and recording scaling (#4808) Co-authored-by: Alessio Buccino --- src/spikeinterface/preprocessing/detect_bad_channels.py | 8 ++++++++ 1 file changed, 8 insertions(+) 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 From ed12c03e74fee7827a0f8850c68af4f2df064ebc Mon Sep 17 00:00:00 2001 From: Amirali Moradniaei Date: Mon, 28 Sep 2026 16:43:22 +0330 Subject: [PATCH 07/38] Test empty postprocessing and metric extension outputs (#4807) --- .../tests/test_empty_extensions.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 src/spikeinterface/postprocessing/tests/test_empty_extensions.py 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) == [] From 398222dbe567af1021609e5c9ea2d56b34320494 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 28 Sep 2026 15:46:29 +0200 Subject: [PATCH 08/38] Fix: speed up loading (and compute sparse PCA projections) in lazy mode (#4806) --- src/spikeinterface/core/sortinganalyzer.py | 2 ++ .../postprocessing/principal_component.py | 11 +++++++++-- .../postprocessing/valid_unit_periods.py | 2 +- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 953f75259d..1a648d1c37 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -3408,6 +3408,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/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 7767104484..b3bb845b84 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"] 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 = {} From 4e0397b1b19d9171eeec8d4587719edf09c0a98a Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 28 Sep 2026 15:46:57 +0200 Subject: [PATCH 09/38] Fix: drop new_unit_ids of non-mergeable units (#4804) --- src/spikeinterface/core/sortinganalyzer.py | 12 ++++++++-- .../core/tests/test_sortinganalyzer.py | 23 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 1a648d1c37..4f8c3cd001 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -2006,9 +2006,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) @@ -2021,9 +2024,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) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 8cc069ce1e..2028c53b0e 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 @@ -1272,6 +1273,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 From 93c612f51cca09ecafe2c23a283132cf0ad57869 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= <42451234+JESUSROYETH@users.noreply.github.com> Date: Mon, 28 Sep 2026 11:33:22 -0300 Subject: [PATCH 10/38] Precompute unit-pair spatial overlap once in find_collisions (#4797) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jesús Royeth --- .../postprocessing/amplitude_scalings.py | 77 +++++++++++-------- .../tests/test_amplitude_scalings.py | 43 ++++++++++- 2 files changed, 85 insertions(+), 35 deletions(-) diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py index be5967ad1b..d28d2ab1c0 100644 --- a/src/spikeinterface/postprocessing/amplitude_scalings.py +++ b/src/spikeinterface/postprocessing/amplitude_scalings.py @@ -192,17 +192,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 +237,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,7 +261,7 @@ def compute(self, traces, peaks): local_spikes, local_spikes_within_margin, delta_collision_samples, - sparsity_mask, + overlap_matrix, local_spike_indices, ) else: @@ -324,30 +330,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 +397,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 +425,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 +441,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 +465,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/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)]) From 1666a0eabde747ba62c0f41b15b0a423b5201cb7 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Fri, 2 Oct 2026 02:37:49 -0600 Subject: [PATCH 11/38] Don't follow symlinks when storing extractor paths (#4812) --- src/spikeinterface/core/zarrextractors.py | 2 +- src/spikeinterface/extractors/alfsortingextractor.py | 3 ++- src/spikeinterface/extractors/cbin_ibl.py | 4 ++-- src/spikeinterface/extractors/matlabhelpers.py | 2 +- src/spikeinterface/extractors/neoextractors/intan.py | 4 ++-- src/spikeinterface/extractors/neoextractors/neuronexus.py | 2 +- src/spikeinterface/extractors/neoextractors/plexon.py | 4 ++-- src/spikeinterface/extractors/nwbextractors.py | 4 ++-- 8 files changed, 13 insertions(+), 12 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 448dbd94a7..78c10df8d4 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -625,7 +625,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 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/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/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/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..5f7f88e996 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -1070,7 +1070,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 +1404,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( From 7d7bd2140b3266c7c4205d22a5c46f9440bf9951 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Fri, 2 Oct 2026 09:22:59 -0600 Subject: [PATCH 12/38] Enable case omission on `get_unit_spike_train_in_seconds` (#4811) --- src/spikeinterface/core/basesorting.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index 1e136f2dff..cccfb4a443 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -200,6 +200,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 +238,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 +263,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 +288,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 +327,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 From 0c9b6b9155056aa18394805bf639d659c3c9b77b Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 5 Oct 2026 12:38:06 +0200 Subject: [PATCH 13/38] Fix: provenance handling in `SortingAnalyzer` (#4805) --- .../tutorials/core/plot_4_sorting_analyzer.py | 2 + src/spikeinterface/core/basesorting.py | 2 + src/spikeinterface/core/binaryfolder.py | 4 +- src/spikeinterface/core/core_tools.py | 4 +- src/spikeinterface/core/numpyextractors.py | 3 +- src/spikeinterface/core/sortinganalyzer.py | 73 ++++++------------- src/spikeinterface/core/sortingfolder.py | 15 +++- .../core/tests/test_sortinganalyzer.py | 2 + src/spikeinterface/core/zarrextractors.py | 26 ++++++- src/spikeinterface/metrics/conftest.py | 10 ++- .../tests/test_multi_extensions.py | 14 ++-- .../preprocessing/tests/test_pipeline.py | 6 +- .../sorters/tests/test_launcher.py | 8 +- .../sorters/tests/test_runsorter.py | 7 +- .../motion/tests/test_motion_interpolation.py | 3 +- 15 files changed, 96 insertions(+), 83 deletions(-) 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/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index cccfb4a443..36a53e51dc 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -540,6 +540,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: 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/core_tools.py b/src/spikeinterface/core/core_tools.py index b5ee518208..25e1da5e93 100644 --- a/src/spikeinterface/core/core_tools.py +++ b/src/spikeinterface/core/core_tools.py @@ -934,11 +934,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/numpyextractors.py b/src/spikeinterface/core/numpyextractors.py index b58a98eddb..e2c2df6280 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 diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 4f8c3cd001..1c6604516b 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -714,7 +714,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 +730,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 +766,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", @@ -1042,9 +1032,7 @@ 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 is_remote = is_path_remote(folder) if not is_remote: @@ -1071,27 +1059,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: @@ -1135,7 +1106,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", @@ -2183,34 +2154,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 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/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 2028c53b0e..856612a6cf 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -93,6 +93,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(): @@ -139,6 +140,7 @@ def test_SortingAnalyzer_binary_folder(tmp_path, dataset): 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" diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 78c10df8d4..06f887871b 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -453,6 +453,9 @@ 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 ZarrSortingExtractor(BaseSorting): """ @@ -559,6 +562,7 @@ def write_sorting( folder_path: str | Path, overwrite: bool = False, storage_options: dict | None = None, + relative_to: str | Path | None = None, **kwargs, ): """ @@ -567,7 +571,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) @@ -734,7 +738,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 +755,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) @@ -777,11 +788,18 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw 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/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/tests/test_multi_extensions.py b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py index 1bb3d383a2..7b53f24f29 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,13 +122,13 @@ 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]) diff --git a/src/spikeinterface/preprocessing/tests/test_pipeline.py b/src/spikeinterface/preprocessing/tests/test_pipeline.py index 5f5c298430..4a62bc0088 100644 --- a/src/spikeinterface/preprocessing/tests/test_pipeline.py +++ b/src/spikeinterface/preprocessing/tests/test_pipeline.py @@ -169,6 +169,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 +177,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 +205,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) diff --git a/src/spikeinterface/sorters/tests/test_launcher.py b/src/spikeinterface/sorters/tests/test_launcher.py index 2cd25f119a..6619431866 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( @@ -209,7 +211,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..c07cd94c9b 100644 --- a/src/spikeinterface/sorters/tests/test_runsorter.py +++ b/src/spikeinterface/sorters/tests/test_runsorter.py @@ -13,14 +13,15 @@ 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.xfail( 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) From 3d3ee44ebe345326d1780adcbca4af44239db465 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Mon, 5 Oct 2026 07:49:57 -0600 Subject: [PATCH 14/38] Drop the separate core test workflow (#4822) --- .github/workflows/all-tests.yml | 16 +++++------ .github/workflows/core-test.yml | 48 --------------------------------- 2 files changed, 8 insertions(+), 56 deletions(-) delete mode 100644 .github/workflows/core-test.yml diff --git a/.github/workflows/all-tests.yml b/.github/workflows/all-tests.yml index 25b220a8a6..033ee7492c 100644 --- a/.github/workflows/all-tests.yml +++ b/.github/workflows/all-tests.yml @@ -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 @@ -137,10 +141,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 From b0aba191892202be9243db8160e4dbe75110fc86 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Mon, 5 Oct 2026 09:54:28 -0600 Subject: [PATCH 15/38] Run the tests with `forkserver` on Linux (#4818) --- conftest.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/conftest.py b/conftest.py index e326c60d71..9b1296bb0b 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): @@ -32,6 +51,9 @@ def pytest_collection_modifyitems(config, items): rootdir = Path(config.rootdir) modules_location = rootdir / "src" / "spikeinterface" for item in items: + 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: From 6bec8d9a417dc534c99b1f90a4f874c20906528e Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Mon, 5 Oct 2026 10:29:54 -0600 Subject: [PATCH 16/38] Add a `dev` dependency group so uv installs pre-commit by default (#4815) --- pyproject.toml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index a23e7b1d65..f9a956a844 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -340,6 +340,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 = [ From d0c9339d3df3af7437976571ac83fb066c7f09e2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 6 Oct 2026 10:42:56 +0200 Subject: [PATCH 17/38] [pre-commit.ci] pre-commit autoupdate (#4828) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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/ From 6fd8630bf66983066ef1b42c0f65bff48e7d5d3a Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Tue, 6 Oct 2026 02:44:20 -0600 Subject: [PATCH 18/38] Don't mount remote paths when running a sorter in a container (#4821) --- src/spikeinterface/core/core_tools.py | 3 ++- src/spikeinterface/sorters/container_tools.py | 9 +++++++-- src/spikeinterface/sorters/tests/test_container_tools.py | 6 ++++++ 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/core_tools.py b/src/spikeinterface/core/core_tools.py index 25e1da5e93..519771144c 100644 --- a/src/spikeinterface/core/core_tools.py +++ b/src/spikeinterface/core/core_tools.py @@ -777,7 +777,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: 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/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" From 04cae98fcf056ae74fcf72ae6a0cc848b6f797e4 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Tue, 6 Oct 2026 02:45:17 -0600 Subject: [PATCH 19/38] Read the last spike from the unit end offsets in `NwbSortingExtractor` (#4819) --- src/spikeinterface/extractors/nwbextractors.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index 5f7f88e996..7ba2dd02d1 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -1456,6 +1456,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, From 2750986478b23a5a4dac721dc550358ebe098813 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 6 Oct 2026 12:18:08 +0200 Subject: [PATCH 20/38] fix: cast infinites to NaN is overflow error (#4830) --- src/spikeinterface/curation/train_manual_curation.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 From 11e5ca4bb4558ecbc422de98856cf853d9038414 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= <42451234+JESUSROYETH@users.noreply.github.com> Date: Wed, 7 Oct 2026 04:33:37 -0300 Subject: [PATCH 21/38] Speed up FFT phase shifts for Neuropixels recordings (#4802) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jesús Royeth --- .../preprocessing/phase_shift.py | 26 +++++++------------ .../preprocessing/tests/test_phase_shift.py | 19 ++++++++++++++ 2 files changed, 29 insertions(+), 16 deletions(-) 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_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() From d89765e9b55219ba7d760422b4c0564cb6403454 Mon Sep 17 00:00:00 2001 From: Joe Ziminski <55797454+JoeZiminski@users.noreply.github.com> Date: Wed, 7 Oct 2026 08:38:07 +0100 Subject: [PATCH 22/38] Fix sorting times segment propagation (#4796) Co-authored-by: Alessio Buccino --- src/spikeinterface/core/numpyextractors.py | 3 +++ .../core/tests/test_numpy_extractors.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/src/spikeinterface/core/numpyextractors.py b/src/spikeinterface/core/numpyextractors.py index e2c2df6280..7f5f5321b9 100644 --- a/src/spikeinterface/core/numpyextractors.py +++ b/src/spikeinterface/core/numpyextractors.py @@ -303,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/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) From 10ffd5285e4aa0d7fd4bc0748a152a263246f2bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= <42451234+JESUSROYETH@users.noreply.github.com> Date: Wed, 7 Oct 2026 04:52:50 -0300 Subject: [PATCH 23/38] Avoid dense waveform allocation when exporting PC features (#4789) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jesús Royeth Co-authored-by: Alessio Buccino --- .../postprocessing/principal_component.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index b3bb845b84..092344a512 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -658,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 @@ -692,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 From 6ef71956f7b3d75d11a7c62aa1cb4512aba29346 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 7 Oct 2026 09:54:00 +0200 Subject: [PATCH 24/38] fix: remove `set_times` warning (#4834) --- src/spikeinterface/core/tests/test_baserecording.py | 4 ++-- src/spikeinterface/core/tests/test_time_handling.py | 4 ++-- src/spikeinterface/core/time_series.py | 11 +---------- .../extractors/neoextractors/openephys.py | 2 +- .../preprocessing/tests/test_resample.py | 4 ++-- 5 files changed, 8 insertions(+), 17 deletions(-) diff --git a/src/spikeinterface/core/tests/test_baserecording.py b/src/spikeinterface/core/tests/test_baserecording.py index 63fe18f899..f2ddc9c5a7 100644 --- a/src/spikeinterface/core/tests/test_baserecording.py +++ b/src/spikeinterface/core/tests/test_baserecording.py @@ -477,7 +477,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 +504,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_time_handling.py b/src/spikeinterface/core/tests/test_time_handling.py index 936a9498da..2f3cc1a6bf 100644 --- a/src/spikeinterface/core/tests/test_time_handling.py +++ b/src/spikeinterface/core/tests/test_time_handling.py @@ -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/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/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/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]) From 0421d0b2e57cb5b21870d0efd35066c05bd7a277 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Wed, 7 Oct 2026 01:54:18 -0600 Subject: [PATCH 25/38] Add `NeuroExplorerSortingExtractor` (#4825) --- doc/api.rst | 1 + doc/modules/extractors.rst | 1 + .../extractors/neoextractors/__init__.py | 10 ++- .../extractors/neoextractors/neuroexplorer.py | 68 ++++++++++++++++++- .../extractors/tests/test_neoextractors.py | 10 +++ 5 files changed, 88 insertions(+), 2 deletions(-) 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/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/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/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/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"] From 67083114edb7202d29d6c9e8dd3e57eeeb33dcb3 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 7 Oct 2026 10:41:14 +0200 Subject: [PATCH 26/38] Extract waveforms to zarr dataset (#4708) Co-authored-by: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com> --- .../core/analyzer_extension_core.py | 8 +- src/spikeinterface/core/sortinganalyzer.py | 1 - .../tests/test_analyzer_extension_core.py | 55 ++++ .../core/tests/test_waveform_tools.py | 101 +++++- src/spikeinterface/core/waveform_tools.py | 297 ++++++++++++++---- src/spikeinterface/sortingcomponents/tools.py | 2 - 6 files changed, 396 insertions(+), 68 deletions(-) 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/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 1c6604516b..cc24463117 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -2402,7 +2402,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)): diff --git a/src/spikeinterface/core/tests/test_analyzer_extension_core.py b/src/spikeinterface/core/tests/test_analyzer_extension_core.py index f606ca34ac..d96ccf0e37 100644 --- a/src/spikeinterface/core/tests/test_analyzer_extension_core.py +++ b/src/spikeinterface/core/tests/test_analyzer_extension_core.py @@ -115,6 +115,61 @@ def test_ComputeWaveforms(format, sparse, create_cache_folder): _check_result_extension(sorting_analyzer, "waveforms", cache_folder) +@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", "zarr"]) @pytest.mark.parametrize("sparse", [True, False]) def test_ComputeTemplates(format, sparse, create_cache_folder): diff --git a/src/spikeinterface/core/tests/test_waveform_tools.py b/src/spikeinterface/core/tests/test_waveform_tools.py index 5e0350f833..6623c92046 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,106 @@ def test_waveform_tools(create_cache_folder): _check_all_wf_equal(list_wfs_sparse) +@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[:]) + + +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/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/sortingcomponents/tools.py b/src/spikeinterface/sortingcomponents/tools.py index b445513a6c..c3782975f1 100644 --- a/src/spikeinterface/sortingcomponents/tools.py +++ b/src/spikeinterface/sortingcomponents/tools.py @@ -50,7 +50,6 @@ 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( @@ -66,7 +65,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", From 5cacee5421d41485acfa41eefac082f76a53b888 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Wed, 7 Oct 2026 03:06:15 -0600 Subject: [PATCH 27/38] Build one probe per electrode group in NwbRecordingExtractor (#4820) --- .../extractors/nwbextractors.py | 17 +++++++- .../extractors/tests/test_nwbextractors.py | 42 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index 7ba2dd02d1..13f1d2d605 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 @@ -1040,9 +1041,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) 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.""" From a633c22bb3e2b0ac8f63e97d2806d6915abcc41f Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 7 Oct 2026 14:19:25 +0200 Subject: [PATCH 28/38] Fix group (and other property) propagation in channel slice (#4838) --- src/spikeinterface/core/channelslice.py | 4 +++- .../core/tests/test_channelslicerecording.py | 13 ++++++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) 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/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() From 6ab884d9c9d2c865b88d42e94b29019d821439c5 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Thu, 8 Oct 2026 02:57:14 -0600 Subject: [PATCH 29/38] Support python 3.14 (#4173) Co-authored-by: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com> Co-authored-by: chrishalcrow Co-authored-by: Alessio Buccino --- .github/workflows/all-tests.yml | 5 +- conftest.py | 6 + pyproject.toml | 10 +- src/spikeinterface/core/core_tools.py | 21 +++ src/spikeinterface/core/node_pipeline.py | 3 +- src/spikeinterface/core/sortinganalyzer.py | 4 + src/spikeinterface/core/template.py | 7 + .../tests/test_analyzer_extension_core.py | 22 +++- .../core/tests/test_baserecording.py | 55 ++++---- .../core/tests/test_basesorting.py | 23 ++-- src/spikeinterface/core/tests/test_loading.py | 12 +- .../core/tests/test_node_pipeline.py | 123 +++++++++--------- .../core/tests/test_sortinganalyzer.py | 50 ++++--- .../core/tests/test_template_class.py | 1 + .../core/tests/test_time_handling.py | 4 +- .../core/tests/test_waveform_tools.py | 2 + .../core/tests/test_zarrextractors.py | 2 + src/spikeinterface/core/zarrextractors.py | 13 +- .../extractors/nwbextractors.py | 12 +- .../tests/test_nwbextractors_streaming.py | 7 + .../tests/test_template_database.py | 10 ++ .../tests/common_extension_tests.py | 5 +- .../tests/test_multi_extensions.py | 8 +- .../preprocessing/tests/test_pipeline.py | 13 +- src/spikeinterface/sorters/internal/lupin.py | 5 + .../sorters/internal/tests/test_lupin.py | 7 + .../internal/tests/test_tridesclous2.py | 7 + .../sorters/internal/tridesclous2.py | 5 + .../sorters/tests/test_launcher.py | 2 + .../sorters/tests/test_runsorter.py | 1 + 30 files changed, 301 insertions(+), 144 deletions(-) diff --git a/.github/workflows/all-tests.yml b/.github/workflows/all-tests.yml index 033ee7492c..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 @@ -129,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 diff --git a/conftest.py b/conftest.py index 9b1296bb0b..e4a08c170a 100644 --- a/conftest.py +++ b/conftest.py @@ -48,9 +48,15 @@ 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")) diff --git a/pyproject.toml b/pyproject.toml index f9a956a844..32442a7e72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] @@ -364,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/core_tools.py b/src/spikeinterface/core/core_tools.py index 519771144c..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://). diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py index 6ecf66918d..13796cb8d9 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: @@ -1044,6 +1044,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/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index cc24463117..37def0c435 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, @@ -1013,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) @@ -1034,6 +1037,7 @@ def create_zarr( ) -> "SortingAnalyzer": 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) 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 d96ccf0e37..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,6 +122,7 @@ def test_ComputeWaveforms(format, sparse, create_cache_folder): _check_result_extension(sorting_analyzer, "waveforms", cache_folder) +@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 @@ -170,7 +178,9 @@ def test_ComputeWaveforms_consistent_across_formats(create_cache_folder, sparse) assert np.array_equal(zarr_waveforms, wfs_mem) -@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_ComputeTemplates(format, sparse, create_cache_folder): cache_folder = create_cache_folder @@ -259,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 f2ddc9c5a7..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): 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_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_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 856612a6cf..84440c42a1 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -25,6 +25,9 @@ # 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(): @@ -138,6 +141,7 @@ 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") @@ -309,20 +313,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): @@ -366,7 +372,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 @@ -447,7 +453,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" @@ -477,7 +483,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" @@ -858,7 +864,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): """ @@ -923,7 +931,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 @@ -945,7 +953,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): """ 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 2f3cc1a6bf..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 diff --git a/src/spikeinterface/core/tests/test_waveform_tools.py b/src/spikeinterface/core/tests/test_waveform_tools.py index 6623c92046..c2750774b6 100644 --- a/src/spikeinterface/core/tests/test_waveform_tools.py +++ b/src/spikeinterface/core/tests/test_waveform_tools.py @@ -157,6 +157,7 @@ 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 @@ -210,6 +211,7 @@ def test_extract_waveforms_to_single_buffer_zarr(tmp_path, sparse): 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 diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index 6278541698..4d546b4c4a 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -16,6 +16,7 @@ ) +@pytest.mark.requires_zarr_write def test_zarr_compression_options(tmp_path): from numcodecs import Blosc, Delta, FixedScaleOffset @@ -59,6 +60,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() diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 06f887871b..7a90089edd 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 @@ -642,6 +650,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") diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index 13f1d2d605..0d1946d39a 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -333,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 @@ -354,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": @@ -1642,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_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/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_multi_extensions.py b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py index 7b53f24f29..6ae72e374d 100644 --- a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py +++ b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py @@ -133,7 +133,9 @@ def 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)") @@ -265,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/preprocessing/tests/test_pipeline.py b/src/spikeinterface/preprocessing/tests/test_pipeline.py index 4a62bc0088..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, @@ -218,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/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_launcher.py b/src/spikeinterface/sorters/tests/test_launcher.py index 6619431866..35db70e677 100644 --- a/src/spikeinterface/sorters/tests/test_launcher.py +++ b/src/spikeinterface/sorters/tests/test_launcher.py @@ -52,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) @@ -201,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" diff --git a/src/spikeinterface/sorters/tests/test_runsorter.py b/src/spikeinterface/sorters/tests/test_runsorter.py index c07cd94c9b..1316430e55 100644 --- a/src/spikeinterface/sorters/tests/test_runsorter.py +++ b/src/spikeinterface/sorters/tests/test_runsorter.py @@ -24,6 +24,7 @@ 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')", From 2182c28ebf0e8bb12c3443b1657df09ead162048 Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Thu, 8 Oct 2026 08:54:17 -0500 Subject: [PATCH 30/38] UnitSelectionSorting improvements: remove lexsort, fix bug, save memory by sharing parent cache when possible (#4783) Co-authored-by: Alessio Buccino --- .../core/tests/test_sortinganalyzer.py | 27 +++ .../core/tests/test_unitsselectionsorting.py | 219 +++++++++++++++++- .../core/unitsselectionsorting.py | 26 +-- 3 files changed, 254 insertions(+), 18 deletions(-) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 84440c42a1..527e053184 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -21,6 +21,7 @@ _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 @@ -746,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, diff --git a/src/spikeinterface/core/tests/test_unitsselectionsorting.py b/src/spikeinterface/core/tests/test_unitsselectionsorting.py index 3aa7bc7577..0b9cf6c813 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,218 @@ 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 + + +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) + 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/unitsselectionsorting.py b/src/spikeinterface/core/unitsselectionsorting.py index fdba9b361a..31fbe2098f 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,27 +59,16 @@ 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 From 9d016014fbe2e6fb748799cd2fdcadf2c4f1248b Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 8 Oct 2026 15:55:27 +0200 Subject: [PATCH 31/38] Optimize `SpikeRetriever` and RAM usage for extensions (#4809) --- src/spikeinterface/core/basesorting.py | 53 ++++++++- src/spikeinterface/core/node_pipeline.py | 112 +++++++++++++----- src/spikeinterface/core/sortinganalyzer.py | 19 +-- .../core/tests/test_zarrextractors.py | 59 +++++++++ src/spikeinterface/core/zarrextractors.py | 92 +++++++++++++- .../postprocessing/amplitude_scalings.py | 5 +- .../postprocessing/spike_locations.py | 5 + .../postprocessing/unit_locations.py | 8 +- .../preprocessing/detect_artifacts.py | 9 ++ .../peak_localization/main.py | 17 ++- .../sortingcomponents/peak_selection.py | 21 +++- src/spikeinterface/sortingcomponents/tools.py | 28 ++++- .../waveforms/features_from_peaks.py | 6 +- 13 files changed, 365 insertions(+), 69 deletions(-) diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index 36a53e51dc..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, @@ -1083,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/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py index 13796cb8d9..c66d68a2e4 100644 --- a/src/spikeinterface/core/node_pipeline.py +++ b/src/spikeinterface/core/node_pipeline.py @@ -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) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 37def0c435..b702237445 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -1175,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: diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index 4d546b4c4a..e1b315e592 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,6 +11,7 @@ ) from spikeinterface.core.zarrextractors import ( ZarrRecordingExtractor, + ZarrSampleIndexSearch, ZarrSortingExtractor, add_sorting_to_zarr_group, get_default_zarr_compressor, @@ -80,7 +82,64 @@ 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]) + + +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/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 7a90089edd..d2a604e38b 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -465,6 +465,63 @@ 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): """ SortingExtractor for a zarr format @@ -535,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 @@ -564,6 +623,28 @@ 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, @@ -794,6 +875,15 @@ def add_sorting_to_zarr_group( 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) diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py index d28d2ab1c0..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 @@ -268,7 +267,7 @@ def compute(self, traces, peaks): 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): 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/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/preprocessing/detect_artifacts.py b/src/spikeinterface/preprocessing/detect_artifacts.py index c29199ae3b..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 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 c3782975f1..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 @@ -56,8 +58,14 @@ def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, job 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) @@ -81,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. @@ -91,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 @@ -120,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, @@ -129,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 From 8462838a133e17053c298e97b526900bda74a7fa Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 8 Oct 2026 16:22:54 +0200 Subject: [PATCH 32/38] doc: add 0.105.1 release notes --- doc/releases/0.105.1.rst | 83 ++++++++++++++++++++++++++++++++++++++++ doc/whatisnew.rst | 6 +++ 2 files changed, 89 insertions(+) create mode 100644 doc/releases/0.105.1.rst diff --git a/doc/releases/0.105.1.rst b/doc/releases/0.105.1.rst new file mode 100644 index 0000000000..e7f4e68c9b --- /dev/null +++ b/doc/releases/0.105.1.rst @@ -0,0 +1,83 @@ +.. _release0.105.1: + +SpikeInterface 0.105.1 release notes +------------------------------------ + +October 8th 2026 + +Minor release with bug fixes and performance improvements + + +core: + +* 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: + +* Filter one channel at a time to reduce peak memory usage (#4840) +* 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: + +* 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: + +* @AtomicGlance +* @JESUSROYETH +* @JoeZiminski +* @alejoe91 +* @grahamfindlay +* @h-mayorquin 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 =============== From 74522cc113f1b991751019d15458e8aa4b7a0b38 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 8 Oct 2026 16:24:49 +0200 Subject: [PATCH 33/38] fix: pyproject.toml --- pyproject.toml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 32442a7e72..5fb13c9edd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,10 +27,10 @@ dependencies = [ "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", - # "neo>=0.14.5", - # "probeinterface>=0.4.0", + # "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", + # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", + "neo>=0.14.5", + "probeinterface>=0.4.0", "packaging", "pydantic", "numcodecs<0.16.0;python_version<'3.14'", # For supporting zarr < 3 From c7cac346f771afd3c70eec049fbd63482930c76c Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 8 Oct 2026 16:25:10 +0200 Subject: [PATCH 34/38] fix: __init__ --- src/spikeinterface/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/__init__.py b/src/spikeinterface/__init__.py index a184e9f487..f7c43a845f 100644 --- a/src/spikeinterface/__init__.py +++ b/src/spikeinterface/__init__.py @@ -27,5 +27,5 @@ # This flag must be set to False for release # This avoids using versioning that contains ".dev0" (and this is a better choice) # This is mainly useful when using run_sorter in a container and spikeinterface install -DEV_MODE = True -# DEV_MODE = False +# DEV_MODE = True +DEV_MODE = False From 58ebbaf9cb1ddae838afbbb7e06bee1813c7ac3b Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Fri, 9 Oct 2026 03:09:29 -0500 Subject: [PATCH 35/38] Override search_cached_spikes_sorted for identity selections (#4845) Co-authored-by: Alessio Buccino --- .../core/tests/test_unitsselectionsorting.py | 6 ++++++ src/spikeinterface/core/tests/test_zarrextractors.py | 1 + src/spikeinterface/core/unitsselectionsorting.py | 10 ++++++++++ 3 files changed, 17 insertions(+) diff --git a/src/spikeinterface/core/tests/test_unitsselectionsorting.py b/src/spikeinterface/core/tests/test_unitsselectionsorting.py index 0b9cf6c813..309563df62 100644 --- a/src/spikeinterface/core/tests/test_unitsselectionsorting.py +++ b/src/spikeinterface/core/tests/test_unitsselectionsorting.py @@ -248,6 +248,7 @@ def test_non_identity_selection_does_not_share(unit_ids): 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 @@ -262,6 +263,11 @@ def test_identity_selection_keeps_lazy_zarr_vector(tmp_path): 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"])) diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index e1b315e592..2b388839d1 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -102,6 +102,7 @@ def test_ZarrSampleIndexSearch(tmp_path): 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" diff --git a/src/spikeinterface/core/unitsselectionsorting.py b/src/spikeinterface/core/unitsselectionsorting.py index 31fbe2098f..33fec6fa94 100644 --- a/src/spikeinterface/core/unitsselectionsorting.py +++ b/src/spikeinterface/core/unitsselectionsorting.py @@ -71,6 +71,16 @@ def _compute_and_cache_spike_vector(self) -> None: ) 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): From 1d0dd8b1761a99a4002a3cbe1c064772e1faf0a7 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 9 Oct 2026 10:21:57 +0200 Subject: [PATCH 36/38] doc: update PR list --- doc/releases/0.105.1.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/releases/0.105.1.rst b/doc/releases/0.105.1.rst index e7f4e68c9b..3ffb4faa5b 100644 --- a/doc/releases/0.105.1.rst +++ b/doc/releases/0.105.1.rst @@ -3,13 +3,14 @@ SpikeInterface 0.105.1 release notes ------------------------------------ -October 8th 2026 +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) @@ -32,7 +33,6 @@ extractors: preprocessing: -* Filter one channel at a time to reduce peak memory usage (#4840) * Fix: artifact detection crossing chunks (#4803) * Speed up FFT phase shifts for Neuropixels recordings (#4802) * Fix: frame slicing after resample/decimation (#4799) From 64b126b059b7947b0b6bcbf9c7bd1fc3185e0966 Mon Sep 17 00:00:00 2001 From: Ishaan <101312037+guptaishaan@users.noreply.github.com> Date: Fri, 9 Oct 2026 01:31:12 -0700 Subject: [PATCH 37/38] Fix manual split example in process_by_channel_group docs (#4846) Co-authored-by: guptaishaan Co-authored-by: Alessio Buccino --- doc/how_to/process_by_channel_group.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From 8e8b5489d6dd152e4989cb34f4540ad183c7c286 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 9 Oct 2026 10:35:02 +0200 Subject: [PATCH 38/38] doc: add last PR --- doc/releases/0.105.1.rst | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/doc/releases/0.105.1.rst b/doc/releases/0.105.1.rst index 3ffb4faa5b..fa72a6f9b2 100644 --- a/doc/releases/0.105.1.rst +++ b/doc/releases/0.105.1.rst @@ -58,6 +58,7 @@ sortingcomponents: documentation: +* Fix manual split example in process_by_channel_group docs (#4846) * Clarify PSD threshold units and recording scaling (#4808) continuous integration: @@ -75,9 +76,10 @@ testing: Contributors: -* @AtomicGlance -* @JESUSROYETH -* @JoeZiminski * @alejoe91 +* @AtomicGlance * @grahamfindlay +* @guptaishaan * @h-mayorquin +* @JESUSROYETH +* @JoeZiminski