Extract waveforms to zarr dataset

src/spikeinterface/core/analyzer_extension_core.py

@@ -192,6 +192,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"
@@ -218,6 +222,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
 

src/spikeinterface/core/sortinganalyzer.py

@@ -2415,7 +2415,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)):

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):

src/spikeinterface/core/tests/test_waveform_tools.py

@@ -158,6 +158,60 @@ 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, unit_ids, 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,
+            unit_ids,
+            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_estimate_templates_with_accumulator():
     recording, sorting = get_dataset()