diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py index dde1ba8620..85dcf4b890 100644 --- a/src/spikeinterface/core/analyzer_extension_core.py +++ b/src/spikeinterface/core/analyzer_extension_core.py @@ -1532,13 +1532,23 @@ def _set_params(self, **kwargs): def _run(self, verbose=False, **job_kwargs): from spikeinterface.core.node_pipeline import run_node_pipeline - # TODO: should we save directly to npy in binary_folder format / or to zarr? - # if self.sorting_analyzer.format == "binary_folder": - # gather_mode = "npy" - # extension_folder = self.sorting_analyzer.folder / "extenstions" / self.extension_name - # gather_kwargs = {"folder": extension_folder} - gather_mode = "memory" - gather_kwargs = {} + # gather results directly to the final on-disk location (one npy file / zarr dataset per + # nodepipeline variable) to avoid an extra in-memory copy. This is only done when we are + # actually saving to a disk format (see AnalyzerExtension.run()); otherwise gather in memory. + gather_to_disk = getattr(self, "_save_to_disk", False) and self.format in ("binary_folder", "zarr") + if gather_to_disk: + extension_folder = self.sorting_analyzer.folder / "extensions" / self.extension_name + names = self.nodepipeline_variables + if self.format == "binary_folder": + gather_mode = "npy" + folder = [extension_folder / f"{name}.npy" for name in names] + else: + gather_mode = "zarr" + folder = [extension_folder / name for name in names] + else: + gather_mode = "memory" + folder = None + names = None job_kwargs = fix_job_kwargs(job_kwargs) nodes = self.get_pipeline_nodes() @@ -1548,7 +1558,8 @@ def _run(self, verbose=False, **job_kwargs): job_kwargs=job_kwargs, job_name=self.extension_name, gather_mode=gather_mode, - gather_kwargs=gather_kwargs, + folder=folder, + names=names, verbose=False, ) if isinstance(data, tuple): @@ -1599,6 +1610,11 @@ def _get_data(self, outputs="numpy", concatenated=False, return_data_name=None, ), f"return_data_name {return_data_name} not in nodepipeline_variables {self.nodepipeline_variables}" all_data = self.data[return_data_name] + # data gathered directly into a zarr store (e.g. by the node pipeline) is kept as a + # zarr.Array handle. On a non-lazy analyzer we materialize it to a numpy array (this mirrors + # the non-lazy load convention). A memmap is an np.ndarray subclass so it is left untouched. + if not self.sorting_analyzer._lazy and not isinstance(all_data, np.ndarray): + all_data = np.asarray(all_data) keep_mask = None if periods is not None: keep_mask = select_sorting_periods_mask( @@ -1659,7 +1675,9 @@ def _select_units_extension_data(self, unit_ids): new_data = dict() for data_name in self.nodepipeline_variables: if self.data.get(data_name) is not None: - new_data[data_name] = self.data[data_name][keep_spike_mask] + # np.asarray materializes a zarr.Array to numpy (needed for boolean-mask indexing) + # while leaving a numpy array / memmap untouched (the mask then reads only the subset) + new_data[data_name] = np.asarray(self.data[data_name])[keep_spike_mask] return new_data @@ -1670,15 +1688,18 @@ def _merge_extension_data( for data_name in self.nodepipeline_variables: if self.data.get(data_name) is not None: if keep_mask is None: - new_data[data_name] = self.data[data_name].copy() + # return an independent copy (also materializes a zarr.Array to numpy) + new_data[data_name] = np.array(self.data[data_name]) else: - new_data[data_name] = self.data[data_name][keep_mask] + # np.asarray leaves a numpy array / memmap untouched; the mask reads only the subset + new_data[data_name] = np.asarray(self.data[data_name])[keep_mask] return new_data def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs): - # splitting only changes random spikes assignments - return self.data.copy() + # splitting only changes random spikes assignments, the per-spike data is unchanged. + # materialize to numpy in case data is a zarr.Array so it can be saved to the new store. + return {data_name: np.array(data) for data_name, data in self.data.items()} def _update_data_after_merge_or_split(old_analyzer, new_analyzer, old_arr, new_sub_arr, new_unit_ids): diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py index 556370f1ed..4fea3bf0c8 100644 --- a/src/spikeinterface/core/node_pipeline.py +++ b/src/spikeinterface/core/node_pipeline.py @@ -1,6 +1,6 @@ -from typing import Type -import struct +from typing import Type, Literal import copy +import struct import warnings from pathlib import Path @@ -534,10 +534,10 @@ def run_node_pipeline( nodes: list[PipelineNode], job_kwargs: dict, job_name: str = "pipeline", - gather_mode: str = "memory", + gather_mode: Literal["memory", "npy", "zarr"] = "memory", gather_kwargs: dict = {}, squeeze_output: bool = True, - folder: str | None = None, + folder: str | Path | list | None = None, names: list[str] | None = None, verbose: bool = False, skip_after_n_peaks: int | None = None, @@ -580,14 +580,16 @@ def run_node_pipeline( The classical job_kwargs job_name : str The name of the pipeline used for the progress_bar - gather_mode : "memory" | "npy" + gather_mode : "memory" | "npy" | "zarr" How to gather the output of the nodes. gather_kwargs : dict - Options to control the "gather engine". See GatherToMemory or GatherToNpy. + Options to control the "gather engine". See GatherToMemory, GatherToNpy or GatherToZarr. squeeze_output : bool, default True If only one output node then squeeze the tuple - folder : str | Path | None - Used for gather_mode="npy" + folder : str | Path | list | None + Used for gather_mode="npy" or gather_mode="zarr". Either a single folder (one file/array + per name is created inside it) or a list of explicit per-output destinations + (see GatherToNpy and GatherToZarr). names : list of str Names of outputs. verbose : bool, default False @@ -622,6 +624,8 @@ def run_node_pipeline( gather_func = GatherToMemory() elif gather_mode == "npy": gather_func = GatherToNpy(folder, names, **gather_kwargs) + elif gather_mode == "zarr": + gather_func = GatherToZarr(folder, names, **gather_kwargs) else: raise ValueError(f"wrong gather_mode : {gather_mode}") @@ -807,24 +811,54 @@ class GatherToNpy: * speculate on a header length (1024) * accumulate in C order the buffer * create the npy v1.0 header at the end with the correct shape and dtype + + Parameters + ---------- + folder : str | Path | list of (str | Path) | None + Where to write the npy files. Two modes: + + * a single folder (str | Path) : one ``.npy`` file is created per name inside it. + * a list of file paths : explicit destination ``.npy`` file per output. This is useful + to gather directly to a final location (e.g. an extension folder). When `names` is + None, they are derived from the file stems. + names : list of str | None + Names of the outputs. Can be None when `folder` is a list of file paths. + npy_header_size : int, default: 1024 + The reserved header size for the npy files. + exist_ok : bool, default: False + Whether the `folder` is allowed to already exist. Only used when `folder` is a single folder. """ - def __init__(self, folder, names, npy_header_size=1024, exist_ok=False): - self.folder = Path(folder) - self.folder.mkdir(parents=True, exist_ok=exist_ok) - assert names is not None - self.names = names + def __init__(self, folder=None, names=None, npy_header_size=1024, exist_ok=False): self.npy_header_size = npy_header_size + if isinstance(folder, (list, tuple)): + # explicit destination file per output + self.file_paths = [Path(p) for p in folder] + if names is None: + names = [file_path.stem for file_path in self.file_paths] + assert len(self.file_paths) == len(names), "`folder` (list of files) must have the same length as `names`" + self.names = names + self.folder = None + # make sure parent folders exist + for file_path in self.file_paths: + file_path.parent.mkdir(parents=True, exist_ok=True) + else: + assert folder is not None, "`folder` must be given" + assert names is not None, "`names` must be given when `folder` is a single folder" + self.folder = Path(folder) + self.folder.mkdir(parents=True, exist_ok=exist_ok) + self.names = names + self.file_paths = [self.folder / (name + ".npy") for name in names] + self.tuple_mode = None self.files = [] self.dtypes = [] self.shapes0 = [] self.final_shapes = [] - for name in names: - filename = self.folder / (name + ".npy") - f = open(filename, "wb+") + for file_path in self.file_paths: + f = open(file_path, "wb+") f.seek(npy_header_size) self.files.append(f) self.dtypes.append(None) @@ -862,7 +896,7 @@ def finalize_buffers(self, squeeze_output=False): f.close() for i, name in enumerate(self.names): - filename = self.folder / (name + ".npy") + filename = self.file_paths[i] shape = (self.shapes0[i],) if self.final_shapes[i] is not None: @@ -892,7 +926,7 @@ def finalize_buffers(self, squeeze_output=False): if self.tuple_mode: outs = () for i, name in enumerate(self.names): - filename = self.folder / (name + ".npy") + filename = self.file_paths[i] outs += (np.load(filename, mmap_mode="r"),) if len(outs) == 1 and squeeze_output: @@ -903,10 +937,205 @@ def finalize_buffers(self, squeeze_output=False): return outs else: # only one file - filename = self.folder / (self.names[0] + ".npy") + filename = self.file_paths[0] return np.load(filename, mmap_mode="r") class GatherToZarr: - pass - # Fot me (sam) this is not necessary unless someone realy really want to use + """ + Gather output of nodes into a zarr folder and then open them back as zarr arrays. + + Each returned output ("name") is stored as a zarr array in the root group. Buffers + are appended chunk by chunk along the first axis as the pipeline runs, so the full + result never has to be held in memory. This is the on-disk equivalent of + GatherToNpy but with chunking and compression. + + Parameters + ---------- + folder : str | Path | list of (str | Path | zarr.Array) | None + Where to write the zarr arrays. Two modes: + + * a single folder (str | Path) : a fresh zarr store is created there and one zarr + array is created per name in its root. + * a list of explicit destinations : buffers are appended directly to these datasets + instead of creating a fresh store. This is useful to gather directly to a final + location (e.g. a SortingAnalyzer extension group). Each item can be: + + - a path pointing inside a zarr store, e.g. + ``"my-analyzer.zarr/extensions/spike_amplitudes/amplitudes"``. The store root + (the ``.zarr`` part) is opened in append mode and the dataset is created on the + fly (with the right dtype and trailing shape) once the first buffer arrives, or + - a pre-created (and resizable) ``zarr.Array``, empty on the first axis + (shape ``(0, *trailing)``) with the correct trailing shape and dtype. + + When `names` is None, they are derived from the datasets' basenames. + names : list of str | None + Names of the outputs. Can be None when `folder` is a list of explicit destinations. + compressor : numcodecs codec | "default" | None, default: "default" + The compressor used for every array. If "default", the SpikeInterface default + zarr compressor is used (Blosc-zstd, level 5, bitshuffle). If None, no compression. + Ignored for destinations that are already a ``zarr.Array`` (they keep their own compressor). + zarr_chunk_size : int | None, default: None + Number of rows (first axis) per zarr chunk. If None, it is computed automatically + so that each chunk is about `zarr_target_chunk_bytes` (see below), which gives a + sensible chunk size regardless of the dtype and trailing shape. Ignored for + destinations that are already a ``zarr.Array``. + zarr_target_chunk_bytes : int, default: 10485760 (10 MiB) + Target (uncompressed) size in bytes of one zarr chunk, used to compute the number of + rows per chunk when `zarr_chunk_size` is None. Ignored for destinations that are + already a ``zarr.Array``. + """ + + def __init__( + self, + folder=None, + names=None, + compressor="default", + zarr_chunk_size=None, + zarr_target_chunk_bytes=10 * 1024 * 1024, + ): + import zarr + + from spikeinterface.core.zarrextractors import get_default_zarr_compressor + + self.zarr_chunk_size = zarr_chunk_size + self.zarr_target_chunk_bytes = zarr_target_chunk_bytes + + if compressor == "default": + compressor = get_default_zarr_compressor() + self.compressor = compressor + + if isinstance(folder, (list, tuple)): + # append to explicit destinations (paths inside a store or pre-created zarr.Arrays) + num_datasets = len(folder) + self.arrays = [None] * num_datasets + # per output : (root_group, internal_path) used to lazily create the dataset, + # or None when the dataset is already a zarr.Array + self._create_specs = [None] * num_datasets + derived_names = [] + root_cache = {} + for i, dataset in enumerate(folder): + if isinstance(dataset, (str, Path)): + store_path, internal_path = _split_zarr_store_path(dataset) + store_path = str(store_path) + if store_path not in root_cache: + # append mode : do not wipe an existing store (e.g. an analyzer) + root_cache[store_path] = zarr.open(store_path, mode="a") + self._create_specs[i] = (root_cache[store_path], internal_path) + derived_names.append(internal_path.split("/")[-1]) + else: + # already a zarr.Array + self.arrays[i] = dataset + derived_names.append(dataset.basename) + if names is None: + names = derived_names + assert num_datasets == len(names), "`folder` (list of datasets) must have the same length as `names`" + self.names = names + self.folder = None + self.zarr_root = None + # we do not own the store so we must not consolidate/close it + self._owns_store = False + else: + assert folder is not None, "`folder` must be given" + assert names is not None, "`names` must be given when `folder` is a single folder" + self.names = names + self.folder = Path(folder) + self.zarr_root = zarr.open(str(self.folder), mode="w") + # arrays are created lazily on the first buffer so we know dtype and trailing shape + self.arrays = [None] * len(names) + self._create_specs = [(self.zarr_root, name) for name in names] + self._owns_store = True + + self.tuple_mode = None + + def __call__(self, res): + if res is None: + return + + if self.tuple_mode is None: + # first loop only + self.tuple_mode = isinstance(res, tuple) + if self.tuple_mode: + assert len(self.names) == len(res) + else: + assert len(self.names) == 1 + + if not self.tuple_mode: + res = (res,) + + # distribute buffers to zarr arrays + for i, name in enumerate(self.names): + buf = np.require(res[i], requirements="C") + if self.arrays[i] is None: + # first loop only : create the array with the right dtype and trailing shape + root, internal_path = self._create_specs[i] + trailing_shape = buf.shape[1:] + if self.zarr_chunk_size is not None: + chunk0 = self.zarr_chunk_size + else: + # pick the number of rows per chunk to target ~zarr_target_chunk_bytes per chunk + row_nbytes = int(np.prod(trailing_shape, dtype="int64")) * buf.dtype.itemsize + chunk0 = max(1, self.zarr_target_chunk_bytes // max(1, row_nbytes)) + self.arrays[i] = root.create_dataset( + name=internal_path, + shape=(0,) + trailing_shape, + chunks=(chunk0,) + trailing_shape, + dtype=buf.dtype, + compressor=self.compressor, + overwrite=True, + ) + self.arrays[i].append(buf, axis=0) + + def finalize_buffers(self, squeeze_output=False): + import zarr + + if self._owns_store: + # consolidate metadata for faster/cleaner re-opening + zarr.consolidate_metadata(self.zarr_root.store) + + if self.tuple_mode: + outs = tuple(self.arrays) + + if len(outs) == 1 and squeeze_output: + # when tuple size == 1 then remove the tuple + return outs[0] + else: + # always a tuple even of size 1 + return outs + else: + # only one array + return self.arrays[0] + + +def _split_zarr_store_path(dataset_path): + """ + Split a path pointing inside a zarr store into (store_path, internal_path). + + The store is the portion of the path up to and including the component ending with + ``.zarr``. The remaining components form the internal (group/dataset) path. + + Example + ------- + >>> _split_zarr_store_path("a/b/my-analyzer.zarr/extensions/spike_amplitudes/amplitudes") + (PosixPath('a/b/my-analyzer.zarr'), 'extensions/spike_amplitudes/amplitudes') + """ + parts = Path(dataset_path).parts + zarr_index = None + for index, part in enumerate(parts): + if part.endswith(".zarr"): + zarr_index = index + break + if zarr_index is None: + raise ValueError( + f"Could not find a '.zarr' store in path '{dataset_path}'. " + "Provide a path like '.zarr/'." + ) + internal_parts = parts[zarr_index + 1 :] + if len(internal_parts) == 0: + raise ValueError( + f"No dataset path inside the zarr store found in '{dataset_path}'. " + "Provide a path like '.zarr/'." + ) + store_path = Path(*parts[: zarr_index + 1]) + internal_path = "/".join(internal_parts) + return store_path, internal_path diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 62d5369b0d..0e4f94ebc3 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -2349,7 +2349,9 @@ def compute_one_extension(self, extension_name, save=True, verbose=False, **kwar self.extensions[extension_name] = extension_instance return extension_instance - def compute_several_extensions(self, extensions, save=True, verbose=False, **job_kwargs): + def compute_several_extensions( + self, extensions, save=True, verbose=False, gather_mode=None, gather_kwargs=None, **job_kwargs + ): """ Compute several extensions @@ -2365,6 +2367,11 @@ def compute_several_extensions(self, extensions, save=True, verbose=False, **job It the extension can be saved then it is saved. If not then the extension will only live in memory as long as the object is deleted. save=False is convenient to try some parameters without changing an already saved extension. + gather_mode : "memory" | "numpy" | "zarr" | None, default: None + Gather mode for node_pipeline extensions. If None, the results are gathered using the same format + as the SortingAnalyzer ("memory" for "memory", "npy" for "binary_folder", "zarr" for "zarr"). + gather_kwargs : dict | None, default: None + Additional keyword arguments for the gather function. If None, default gather kwargs are used. Returns ------- @@ -2430,14 +2437,40 @@ def compute_several_extensions(self, extensions, save=True, verbose=False, **job result_routage = [] extension_instances = {} + # decide how the pipeline results are gathered. + # when we save to a disk format we gather directly to the extension final location + # (npy files or zarr datasets) to avoid an extra in-memory copy. Otherwise (memory + # format, save=False or read-only analyzer) we gather in memory. + save_to_disk = save and not self.is_read_only() + if gather_mode is None: + if save_to_disk and self.format == "binary_folder": + gather_mode = "npy" + elif save_to_disk and self.format == "zarr": + gather_mode = "zarr" + else: + gather_mode = "memory" + if gather_kwargs is None: + gather_kwargs = {} + + # for disk gather modes we build one destination per output variable + gather_folder = [] if gather_mode in ("npy", "zarr") else None + gather_names = [] if gather_mode in ("npy", "zarr") else None + for extension_name, extension_params in extensions_with_pipeline.items(): extension_class = get_extension_class(extension_name) assert ( self.has_recording() or self.has_temporary_recording() ), f"Extension {extension_name} requires the recording" + extension_folder = self.folder / "extensions" / extension_name if gather_folder is not None else None for variable_name in extension_class.nodepipeline_variables: result_routage.append((extension_name, variable_name)) + if gather_mode == "npy": + gather_folder.append(extension_folder / f"{variable_name}.npy") + gather_names.append(variable_name) + elif gather_mode == "zarr": + gather_folder.append(extension_folder / variable_name) + gather_names.append(variable_name) extension_instance = extension_class(self) extension_instance.set_params(save=save, **extension_params) @@ -2446,6 +2479,13 @@ def compute_several_extensions(self, extensions, save=True, verbose=False, **job nodes = extension_instance.get_pipeline_nodes() all_nodes.extend(nodes) + # reset and save params before running so the pipeline can gather directly into the + # (freshly created) extension folders/groups (mirrors AnalyzerExtension.run()) + if save_to_disk: + for extension_instance in extension_instances.values(): + extension_instance._save_params() + extension_instance._save_importing_provenance() + job_name = "Compute : " + " + ".join(extensions_with_pipeline.keys()) t_start = perf_counter() @@ -2454,7 +2494,10 @@ def compute_several_extensions(self, extensions, save=True, verbose=False, **job all_nodes, job_kwargs=job_kwargs, job_name=job_name, - gather_mode="memory", + gather_mode=gather_mode, + folder=gather_folder, + names=gather_names, + gather_kwargs=gather_kwargs, squeeze_output=False, verbose=verbose, ) @@ -2470,7 +2513,18 @@ def compute_several_extensions(self, extensions, save=True, verbose=False, **job for extension_name, extension_instance in extension_instances.items(): self.extensions[extension_name] = extension_instance - if save: + if save_to_disk: + # params/provenance already saved above (before the run). Here we only persist + # the run info and the data. For disk gather modes the data is already written + # to its final location so _save_data() is a no-op for those variables. + extension_instance._save_run_info() + extension_instance._save_data() + if self.format == "zarr": + import zarr + + zarr.consolidate_metadata(self._get_zarr_root().store) + elif save: + # memory format or read-only : keep the previous behavior extension_instance.save() for extension_name, extension_params in extensions_post_pipeline.items(): @@ -2914,6 +2968,22 @@ def __init__(self, sorting_analyzer): self.run_info = self._default_run_info_dict() self.data = dict() + def __del__(self): + # Close any open memmap file handles held in `data` (e.g. when an extension gathers or + # loads its data as a memmap). On Windows an open memmap prevents deleting the underlying + # file, so releasing the handles here allows the extension folder to be removed when the + # extension is dropped (e.g. on recompute). Best-effort: __del__ must never raise. + data = self.__dict__.get("data", None) + if not data: + return + for value in data.values(): + mmap = getattr(value, "_mmap", None) if isinstance(value, np.memmap) else None + if mmap is not None: + try: + mmap.close() + except Exception: + pass + def _default_run_info_dict(self): return dict(run_completed=False, runtime_s=None) @@ -3281,18 +3351,24 @@ def split( return new_extension def run(self, save=True, **kwargs): - if save and not self.sorting_analyzer.is_read_only(): + save_to_disk = save and not self.sorting_analyzer.is_read_only() + if save_to_disk: # NB: this call to _save_params() also resets the folder or zarr group self._save_params() self._save_importing_provenance() + # let _run() know whether it may gather results directly to the final on-disk location. + # This is only valid when we actually save to a disk format (the folder/group has just + # been reset above). Otherwise _run() must gather in memory. + self._save_to_disk = save_to_disk + t_start = perf_counter() self._run(**kwargs) t_end = perf_counter() self.run_info["runtime_s"] = t_end - t_start self.run_info["run_completed"] = True - if save and not self.sorting_analyzer.is_read_only(): + if save_to_disk: self._save_run_info() self._save_data() if self.format == "zarr": @@ -3351,6 +3427,8 @@ def _save_data(self): except: raise Exception(f"Could not save {ext_data_name} as extension data") elif self.format == "zarr": + import zarr + saving_options = self.sorting_analyzer._backend_options.get("saving_options", {}) extension_group = self._get_zarr_extension_group(mode="r+") @@ -3359,6 +3437,10 @@ def _save_data(self): saving_options["compressor"] = get_default_zarr_compressor() for ext_data_name, ext_data in self.data.items(): + if isinstance(ext_data, zarr.Array): + # the data was gathered directly into the extension group (e.g. by the node + # pipeline), so it is already in its final location : nothing to copy + continue if ext_data_name in extension_group: del extension_group[ext_data_name] if isinstance(ext_data, (dict, list)): @@ -3366,7 +3448,10 @@ def _save_data(self): _write_object_array(extension_group, ext_data_name, ext_data_, codec="json") extension_group[ext_data_name].attrs["dict"] = True elif isinstance(ext_data, np.ndarray): - extension_group.create_dataset(name=ext_data_name, data=ext_data, **saving_options) + # only save the array if the dataset does not already exist, since it is created directly + # by the run_node_pipeline() function in the case of nodepipeline extensions + if ext_data_name not in extension_group: + extension_group.create_dataset(name=ext_data_name, data=ext_data, **saving_options) elif HAS_PANDAS and isinstance(ext_data, pd.DataFrame): df_group = extension_group.create_group(ext_data_name) # first we save the index @@ -3393,6 +3478,13 @@ def _reset_extension_folder(self): Delete the extension in a folder (binary or zarr) and create an empty one. """ if self.format == "binary_folder": + # Drop the analyzer's reference to a previously computed extension of the same name so + # it can be garbage collected. Its __del__ releases any open file handles (e.g. a memmap + # kept in `data` when gathering directly to npy) which, on Windows, would otherwise + # prevent deleting the folder below. + old_extension = self.sorting_analyzer.extensions.pop(self.extension_name, None) + del old_extension + extension_folder = self._get_binary_extension_folder() if extension_folder.is_dir(): shutil.rmtree(extension_folder) diff --git a/src/spikeinterface/core/tests/test_analyzer_extension_core.py b/src/spikeinterface/core/tests/test_analyzer_extension_core.py index fd550c729e..070172682d 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_node_pipeline.py b/src/spikeinterface/core/tests/test_node_pipeline.py index ab809b5d4d..831d481005 100644 --- a/src/spikeinterface/core/tests/test_node_pipeline.py +++ b/src/spikeinterface/core/tests/test_node_pipeline.py @@ -180,6 +180,89 @@ 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", + folder=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}" + if npy_files_folder.is_dir(): + shutil.rmtree(npy_files_folder) + npy_files = [ + npy_files_folder / "amp.npy", + npy_files_folder / "sub" / "rms.npy", + npy_files_folder / "denoised_rms.npy", + ] + output = run_node_pipeline( + recording, + nodes, + job_kwargs, + gather_mode="npy", + folder=npy_files, + ) + amplitudes_f, waveforms_rms_f, denoised_waveforms_rms_f = output + for npy_file in npy_files: + assert npy_file.is_file() + assert np.array_equal(amplitudes, amplitudes_f) + 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", + folder=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: import pickle @@ -188,6 +271,54 @@ def test_run_node_pipeline(cache_folder_creation): unpickled_node = pickle.loads(pickled_node) +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 + # collapses to 1 row per chunk on a quiet first chunk. + recording, sorting = generate_ground_truth_recording(num_channels=8, num_units=5, durations=[20.0], seed=7) + + # small chunks + n_jobs>1 so the first non-empty buffer is small (would give chunk0==1 with the + # old first-buffer heuristic) + job_kwargs = dict(chunk_duration="0.1s", n_jobs=2, progress_bar=False) + + spikes = sorting.to_spike_vector() + peaks = np.zeros(spikes.size, dtype=spike_peak_dtype) + peaks["sample_index"] = spikes["sample_index"] + peaks["segment_index"] = spikes["segment_index"] + + peak_retriever = PeakRetriever(recording, peaks) + ms_before, ms_after = 0.5, 1.0 + dense_waveforms = ExtractDenseWaveforms( + recording, parents=[peak_retriever], ms_before=ms_before, ms_after=ms_after, return_output=True + ) + nodes = [peak_retriever, dense_waveforms] + + # default byte target (10 MiB) + target_bytes = 10 * 1024 * 1024 + waveforms = run_node_pipeline( + recording, nodes, job_kwargs, gather_mode="zarr", folder=tmp_path / "wfs.zarr", names=["waveforms"] + ) + nbefore_after = dense_waveforms.nbefore + dense_waveforms.nafter + row_nbytes = nbefore_after * recording.get_num_channels() * np.dtype("float32").itemsize + expected_chunk0 = max(1, target_bytes // row_nbytes) + assert waveforms.chunks[0] == expected_chunk0 + assert waveforms.chunks[0] > 1 + assert waveforms.chunks[1:] == waveforms.shape[1:] + + # explicit override is respected + waveforms2 = run_node_pipeline( + recording, + nodes, + job_kwargs, + gather_mode="zarr", + folder=tmp_path / "wfs2.zarr", + names=["waveforms"], + gather_kwargs={"zarr_chunk_size": 1234}, + ) + assert waveforms2.chunks[0] == 1234 + assert np.array_equal(waveforms[:], waveforms2[:]) + + def test_skip_after_n_peaks_and_recording_slices(): recording, sorting = generate_ground_truth_recording(num_channels=10, num_units=10, durations=[10.0], seed=2205) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 25aeb78c1a..e8f3ad1213 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -792,6 +792,146 @@ def test_runtime_dependencies(dataset): assert not sorting_analyzer.has_extension("dummy_pipeline") +def _compute_reference_pipeline_data(dataset): + """Compute the dummy_pipeline extension in memory to use as a reference.""" + recording, sorting = dataset + analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=False, sparsity=None) + analyzer.compute(["random_spikes", "templates"]) + analyzer.compute({"dummy_pipeline": {"param0": 5.5}}) + return analyzer.get_extension("dummy_pipeline").get_data() + + +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +def test_compute_pipeline_extension_gather_to_disk(tmp_path, dataset, format): + """ + When computing node-pipeline extensions on a disk-backed analyzer, the results are gathered + directly to their final location (npy files for binary_folder, zarr datasets for zarr) instead + of being accumulated in memory and copied afterwards. This test checks that: + * the auto gather_mode selection matches the analyzer format + * the data is written in place and kept as a memmap / zarr.Array (no extra copy) + * the values match a plain in-memory computation and survive a reload + * recomputing (overwriting) works + """ + import zarr + + register_result_extension(DummyPipelineAnalyzerExtension) + recording, sorting = dataset + + amp_ref = _compute_reference_pipeline_data(dataset) + + if format == "memory": + folder = None + elif format == "binary_folder": + folder = tmp_path / "analyzer" + else: + folder = tmp_path / "analyzer.zarr" + + analyzer = create_sorting_analyzer(sorting, recording, format=format, folder=folder, sparse=False, sparsity=None) + analyzer.compute(["random_spikes", "templates"]) + analyzer.compute({"dummy_pipeline": {"param0": 5.5}}) + + # NB: do not keep a local reference to the extension (or to `ext.data["amp"]`) across the + # recompute below: on Windows an open memmap on amp.npy would prevent deleting the folder. + assert np.array_equal(analyzer.get_extension("dummy_pipeline").get_data(), amp_ref) + + if format == "binary_folder": + # written directly to the final npy file and kept as a memmap (not re-copied by _save_data) + amp_file = folder / "extensions" / "dummy_pipeline" / "amp.npy" + assert amp_file.is_file() + assert isinstance(analyzer.get_extension("dummy_pipeline").data["amp"], np.memmap) + elif format == "zarr": + # written directly as a zarr dataset in the extension group + root = analyzer._get_zarr_root(mode="r") + assert "amp" in root["extensions"]["dummy_pipeline"] + assert isinstance(analyzer.get_extension("dummy_pipeline").data["amp"], zarr.Array) + + if format != "memory": + # data must survive a reload from disk + assert np.array_equal(load_sorting_analyzer(folder).get_extension("dummy_pipeline").get_data(), amp_ref) + + # recompute (overwrite) must not corrupt or leave stale data behind + analyzer.compute({"dummy_pipeline": {"param0": 5.5}}) + assert np.array_equal(load_sorting_analyzer(folder).get_extension("dummy_pipeline").get_data(), amp_ref) + + +@pytest.mark.parametrize("format", ["binary_folder", "zarr"]) +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 + and nothing is written to disk. + """ + register_result_extension(DummyPipelineAnalyzerExtension) + recording, sorting = dataset + + folder = tmp_path / ("analyzer" if format == "binary_folder" else "analyzer.zarr") + analyzer = create_sorting_analyzer(sorting, recording, format=format, folder=folder, sparse=False, sparsity=None) + analyzer.compute(["random_spikes", "templates"]) + analyzer.compute({"dummy_pipeline": {"param0": 5.5}}, save=False) + + # in memory the extension is available + assert analyzer.has_extension("dummy_pipeline") + + # but nothing was written to disk + analyzer_reloaded = load_sorting_analyzer(folder) + assert not analyzer_reloaded.has_extension("dummy_pipeline") + + +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) +def test_compute_one_pipeline_extension_gather_to_disk(tmp_path, dataset, format): + """ + Same as test_compute_pipeline_extension_gather_to_disk but through compute_one_extension + (i.e. computing a single node-pipeline extension via a string input), which uses + BaseSpikeVectorExtension._run() to gather directly to disk. + """ + import zarr + + register_result_extension(DummyPipelineAnalyzerExtension) + recording, sorting = dataset + + amp_ref = _compute_reference_pipeline_data(dataset) + + if format == "memory": + folder = None + elif format == "binary_folder": + folder = tmp_path / "analyzer" + else: + folder = tmp_path / "analyzer.zarr" + + analyzer = create_sorting_analyzer(sorting, recording, format=format, folder=folder, sparse=False, sparsity=None) + analyzer.compute(["random_spikes", "templates"]) + # single string -> compute_one_extension -> BaseSpikeVectorExtension._run + analyzer.compute("dummy_pipeline", param0=5.5) + + # NB: do not keep a local reference to the extension (or to `ext.data["amp"]`) across the + # recompute below: on Windows an open memmap on amp.npy would prevent deleting the folder. + assert np.array_equal(analyzer.get_extension("dummy_pipeline").get_data(), amp_ref) + + if format == "binary_folder": + assert (folder / "extensions" / "dummy_pipeline" / "amp.npy").is_file() + assert isinstance(analyzer.get_extension("dummy_pipeline").data["amp"], np.memmap) + elif format == "zarr": + root = analyzer._get_zarr_root(mode="r") + assert "amp" in root["extensions"]["dummy_pipeline"] + assert isinstance(analyzer.get_extension("dummy_pipeline").data["amp"], zarr.Array) + + if format != "memory": + # data must survive a reload and recompute (overwrite) must work + assert np.array_equal(load_sorting_analyzer(folder).get_extension("dummy_pipeline").get_data(), amp_ref) + + analyzer.compute("dummy_pipeline", param0=5.5) + assert np.array_equal(load_sorting_analyzer(folder).get_extension("dummy_pipeline").get_data(), amp_ref) + + # save=False on a disk analyzer: computed in memory, nothing written to disk + folder2 = tmp_path / ("analyzer_nosave" + (".zarr" if format == "zarr" else "")) + analyzer2 = create_sorting_analyzer( + sorting, recording, format=format, folder=folder2, sparse=False, sparsity=None + ) + analyzer2.compute(["random_spikes", "templates"]) + analyzer2.compute("dummy_pipeline", param0=5.5, save=False) + assert analyzer2.has_extension("dummy_pipeline") + assert not load_sorting_analyzer(folder2).has_extension("dummy_pipeline") + + def test_select_channels(dataset): recording, sorting = dataset sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=False, sparsity=None) diff --git a/src/spikeinterface/core/tests/test_waveform_tools.py b/src/spikeinterface/core/tests/test_waveform_tools.py index 5e0350f833..1b52f63172 100644 --- a/src/spikeinterface/core/tests/test_waveform_tools.py +++ b/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() diff --git a/src/spikeinterface/core/waveform_tools.py b/src/spikeinterface/core/waveform_tools.py index 70f29d94b0..d25213ab58 100644 --- a/src/spikeinterface/core/waveform_tools.py +++ b/src/spikeinterface/core/waveform_tools.py @@ -465,15 +465,19 @@ def extract_waveforms_to_single_buffer( N samples before spike nafter: int N samples after spike - mode: "memmap" | "shared_memory", default: "memmap" - The mode to use for the buffer + mode: "memmap" | "shared_memory" | "zarr", default: "memmap" + The mode to use for the buffer. In "zarr" mode the waveforms are written directly to a zarr + dataset (workers return their block and the main process writes it, so parallel writes are safe). return_scaled : bool | None, default: None DEPRECATED. Use return_in_uV instead. return_in_uV : bool, default: False 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 @@ -516,8 +520,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: @@ -526,6 +531,7 @@ def extract_waveforms_to_single_buffer( num_chans = int(max(np.sum(sparsity_mask, axis=1))) # This is a numpy scalar, so we cast to int shape = (int(num_spikes), int(n_samples), int(num_chans)) + zarr_writer = 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) @@ -540,6 +546,33 @@ 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": + # Create the zarr dataset up front, then fill it. 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 their contiguous block and the main + # process writes it (single writer) via the `zarr_writer` gather function below. + import zarr + + from .zarrextractors import get_default_zarr_compressor + from .node_pipeline import _split_zarr_store_path + + assert file_path is not None, "zarr mode requires a `file_path` pointing inside a .zarr store" + store_path, dataset_path = _split_zarr_store_path(file_path) + # chunk along the first (spike) axis so that a chunk is about 10 MiB + row_nbytes = int(n_samples) * int(num_chans) * dtype.itemsize + chunk0 = max(1, (10 * 1024 * 1024) // max(1, row_nbytes)) + zarr_root = zarr.open(str(store_path), mode="a") + all_waveforms = zarr_root.create_dataset( + name=dataset_path, + shape=shape, + chunks=(chunk0, int(n_samples), int(num_chans)), + dtype=dtype, + fill_value=0, + compressor=get_default_zarr_compressor(), + overwrite=True, + ) + wf_array_info = None + zarr_writer = _ZarrSingleBufferWriter(all_waveforms) else: raise ValueError("allocate_waveforms_buffers bad mode") @@ -547,28 +580,57 @@ 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 - - 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" - - processor = TimeSeriesChunkExecutor( - recording, func, init_func, init_args, job_name=job_name, verbose=verbose, **job_kwargs - ) - processor.run() - - if mode == "memmap": + if mode == "zarr": + # workers return their contiguous block, the main process writes 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), + ) + if job_name is None: + job_name = "extract waveforms zarr mono buffer" + processor = TimeSeriesChunkExecutor( + recording, + func, + init_func, + init_args, + gather_func=zarr_writer, + job_name=job_name, + verbose=verbose, + **job_kwargs, + ) + processor.run() + 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, + ) + if job_name is None: + job_name = f"extract waveforms {mode} mono buffer" + + processor = TimeSeriesChunkExecutor( + recording, func, init_func, init_args, job_name=job_name, verbose=verbose, **job_kwargs + ) + processor.run() + + if mode in ("memmap", "zarr"): return all_waveforms elif mode == "shared_memory": if copy: @@ -674,6 +736,103 @@ def _worker_distribute_single_buffer(segment_index, start_frame, end_frame, work all_waveforms.flush() +class _ZarrSingleBufferWriter: + """ + Gather function used by `extract_waveforms_to_single_buffer` in "zarr" mode. + + Each worker returns a (start_row, block) tuple for a contiguous range of spikes. This is called + in the main process (single writer) so concurrent writes to the same zarr chunk cannot happen. + """ + + def __init__(self, zarr_array): + self.zarr_array = zarr_array + + def __call__(self, res): + if res is None: + return + start_row, block = res + self.zarr_array[start_row : start_row + block.shape[0]] = block + + +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 write 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"] + + 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]; borders are protected by nbefore/nafter + i0, i1 = np.searchsorted( + in_seg_spikes["sample_index"], [max(start_frame, nbefore), min(end_frame, seg_size - nafter)] + ) + + 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 + + traces = recording.get_traces( + start_frame=start, end_frame=end, segment_index=segment_index, return_in_uV=return_in_uV + ) + + onset = 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)): + wf = traces[sample_index : sample_index + offset, :] + if sparsity_mask is None: + block[local_index, :, :] = wf + else: + mask = sparsity_mask[unit_index, :] + wf = wf[:, mask] + block[local_index, :, : wf.shape[1]] = wf + + # spike_indices s0 + [i0, i1) are contiguous -> write as a single slice in the main process + start_row = s0 + i0 + return (start_row, 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). diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py index 7870c82be1..6befe8c9a9 100644 --- a/src/spikeinterface/postprocessing/amplitude_scalings.py +++ b/src/spikeinterface/postprocessing/amplitude_scalings.py @@ -1,6 +1,7 @@ import numpy as np from spikeinterface.core import ChannelSparsity +from spikeinterface.core.core_tools import ms_to_samples from spikeinterface.core.template_tools import get_dense_templates_array, _get_nbefore from spikeinterface.core.sortinganalyzer import register_result_extension from spikeinterface.core.analyzer_extension_core import BaseSpikeVectorExtension