diff --git a/.github/scripts/serialization/objects.py b/.github/scripts/serialization/objects.py index ce3d215260..78fb3aed0b 100644 --- a/.github/scripts/serialization/objects.py +++ b/.github/scripts/serialization/objects.py @@ -29,10 +29,13 @@ "json": ".json", "pickle": ".pkl", "binary": "_binary", + "binary_parallel": "_binary_parallel", "numpy_folder": "_numpy_folder", "zarr": ".zarr", + "zarr_parallel": "_parallel.zarr", } +DEFAULT_DURATION = 2.0 # seconds, for recordings and sortings # --- json (recipe) entries: moved class + a second class ------------------------- @@ -46,7 +49,12 @@ def _build_noise_generator_recording(): except ImportError: from spikeinterface.core.generate import NoiseGeneratorRecording - return NoiseGeneratorRecording(num_channels=4, sampling_frequency=30000.0, durations=[1.0, 1.5], seed=0) + return NoiseGeneratorRecording( + num_channels=4, + sampling_frequency=30000.0, + durations=[DEFAULT_DURATION, DEFAULT_DURATION + 0.5], + seed=0, + ) def _check_noise_generator_recording(rec): @@ -59,7 +67,7 @@ def _check_noise_generator_recording(rec): def _build_mock_recording(): from spikeinterface.core import generate_recording - return generate_recording(num_channels=4, durations=[1.0], sampling_frequency=30000.0, seed=0) + return generate_recording(num_channels=4, durations=[DEFAULT_DURATION], sampling_frequency=30000.0, seed=0) def _check_mock_recording(rec): @@ -77,7 +85,7 @@ def _build_recording_with_properties(): import numpy as np from spikeinterface.core import generate_recording - rec = generate_recording(num_channels=4, durations=[1.0], sampling_frequency=30000.0, seed=0) + rec = generate_recording(num_channels=4, durations=[DEFAULT_DURATION], sampling_frequency=30000.0, seed=0) rec.set_property("quality", np.array(["good", "good", "bad", "good"])) rec.annotate(experimenter="test") return rec @@ -94,7 +102,7 @@ def _build_recording_with_probe(): from probeinterface import generate_linear_probe from spikeinterface.core import generate_recording - rec = generate_recording(num_channels=8, durations=[1.0], sampling_frequency=30000.0, seed=0) + rec = generate_recording(num_channels=8, durations=[DEFAULT_DURATION], sampling_frequency=30000.0, seed=0) probe = generate_linear_probe(num_elec=8) probe.set_device_channel_indices(np.arange(8)) if parse(si_version) <= parse("0.105.0"): @@ -104,6 +112,20 @@ def _build_recording_with_probe(): return rec_with_probe +def _build_recording_with_timestamps(): + rec = _build_mock_recording() + times = rec.get_times(segment_index=0) + 100 + rec.set_times(times, segment_index=0) + return rec + + +def _check_recording_with_timestamps(rec): + import numpy as np + expected_times = np.arange(int(DEFAULT_DURATION * 30000)) / 30000.0 + 100 + times = rec.get_times(segment_index=0) + assert np.allclose(times, expected_times) + + def _check_recording_with_probe(rec): import numpy as np @@ -121,7 +143,7 @@ def _build_recording_with_interleaved_probes(): from probeinterface import ProbeGroup, generate_linear_probe from spikeinterface.core import generate_recording - rec = generate_recording(num_channels=8, durations=[1.0], sampling_frequency=30000.0, seed=0) + rec = generate_recording(num_channels=8, durations=[DEFAULT_DURATION], sampling_frequency=30000.0, seed=0) probe0 = generate_linear_probe(num_elec=4) probe1 = generate_linear_probe(num_elec=4) probe1.move([100.0, 0.0]) @@ -164,7 +186,7 @@ def _build_preprocessed_chain(): from spikeinterface.core import generate_recording from spikeinterface.preprocessing import common_reference, scale - rec = generate_recording(num_channels=4, durations=[1.0], sampling_frequency=30000.0, seed=0) + rec = generate_recording(num_channels=4, durations=[DEFAULT_DURATION], sampling_frequency=30000.0, seed=0) # Two nested scipy-free preprocessing wrappers (scale then common_reference): this # exercises recursive parent reload without pulling scipy into the environments. return common_reference(scale(rec, gain=2.0)) @@ -181,7 +203,7 @@ def _check_preprocessed_chain(rec): def _build_sorting(): from spikeinterface.core import generate_sorting - return generate_sorting(num_units=5, sampling_frequency=30000.0, durations=[1.0]) + return generate_sorting(num_units=5, sampling_frequency=30000.0, durations=[DEFAULT_DURATION]) def _check_sorting(sorting): @@ -195,7 +217,7 @@ def _build_sorting_with_properties(): import numpy as np from spikeinterface.core import generate_sorting - sorting = generate_sorting(num_units=4, sampling_frequency=30000.0, durations=[1.0]) + sorting = generate_sorting(num_units=4, sampling_frequency=30000.0, durations=[DEFAULT_DURATION]) sorting.set_property("quality", np.array(["good", "good", "bad", "good"])) sorting.annotate(experimenter="test") return sorting @@ -224,19 +246,25 @@ def _check_sorting_with_properties(sorting): "id": "recording_with_properties", "build": _build_recording_with_properties, "check": _check_recording_with_properties, - "formats": ["binary", "zarr"], + "formats": ["binary", "zarr", "binary_parallel", "zarr_parallel"], + }, + { + "id": "recording_with_timestamps", + "build": _build_recording_with_timestamps, + "check": _check_recording_with_timestamps, + "formats": ["binary", "zarr", "binary_parallel", "zarr_parallel"], }, { "id": "recording_with_probe", "build": _build_recording_with_probe, "check": _check_recording_with_probe, - "formats": ["binary", "zarr"], + "formats": ["binary", "zarr", "binary_parallel", "zarr_parallel"], }, { "id": "recording_with_interleaved_probes", "build": _build_recording_with_interleaved_probes, "check": _check_recording_with_interleaved_probes, - "formats": ["binary", "zarr"], + "formats": ["binary", "zarr", "binary_parallel", "zarr_parallel"], }, { "id": "preprocessed_chain", diff --git a/.github/scripts/serialization/serialize_objects.py b/.github/scripts/serialization/serialize_objects.py index 44a68ca26e..c4e9326f08 100644 --- a/.github/scripts/serialization/serialize_objects.py +++ b/.github/scripts/serialization/serialize_objects.py @@ -33,10 +33,14 @@ obj.dump_to_pickle(dest) elif fmt == "binary": obj.save(folder=dest, format="binary", overwrite=True) + elif fmt == "binary_parallel": + obj.save(folder=dest, format="binary", overwrite=True, n_jobs=2) elif fmt == "numpy_folder": obj.save(folder=dest, format="numpy_folder", overwrite=True) elif fmt == "zarr": obj.save(folder=dest, format="zarr", overwrite=True) + elif fmt == "zarr_parallel": + obj.save(folder=dest, format="zarr", overwrite=True, n_jobs=2) print(f" wrote {dest.name} ({fmt})") print(f"Fixtures written to: {out_dir.resolve()}") diff --git a/.github/workflows/cross_version_serialization.yml b/.github/workflows/cross_version_serialization.yml index 4712a3976d..4bafaed4e1 100644 --- a/.github/workflows/cross_version_serialization.yml +++ b/.github/workflows/cross_version_serialization.yml @@ -6,9 +6,8 @@ name: Cross-version serialization # # Delivery model: Live generation. Fixtures are regenerated from a real old install each # run in an isolated uv environment (nothing committed). The version list is computed -# from PyPI (latest patch of each minor at or above MIN_VERSION_TO_TEST): on a pull -# request only the latest released minor is tested; the full matrix runs weekly and on -# manual dispatch. +# from PyPI (latest patch of each minor at or above MIN_VERSION_TO_TEST), and the full +# matrix is tested on every pull request and on manual dispatch. on: pull_request: @@ -20,11 +19,9 @@ on: # watches core only; a break introduced in generation/ or preprocessing/ would not # trigger it. paths: - - "src/spikeinterface/core/**" - - ".github/scripts/serialization/**" - - ".github/workflows/cross_version_serialization.yml" - schedule: - - cron: "0 4 * * 0" # weekly, Sunday 04:00 UTC + - 'src/spikeinterface/core/**' + - '.github/scripts/serialization/**' + - '.github/workflows/cross_version_serialization.yml' workflow_dispatch: concurrency: @@ -33,8 +30,8 @@ concurrency: jobs: # Compute which released versions the matrix below tests loading from, exposed as the - # job output `list`. Latest patch of each minor at or above the floor on schedule / - # dispatch; only the latest minor on a pull request (keeps per-PR runs cheap). + # job output `list`: the latest patch of each minor at or above the floor, tested on + # every trigger. versions: name: Select versions runs-on: ubuntu-latest @@ -46,7 +43,6 @@ jobs: python-version: "3.10" - id: set env: - GITHUB_EVENT_NAME: ${{ github.event_name }} # Support floor: oldest minor to test. Below this, installs fail on the CI # Python (numcodecs dependency rot in 0.100/0.101). Raise when 0.102 rots. MIN_VERSION_TO_TEST: "0.102" # We will change this when making breaking changes (hopefully not too often) @@ -67,10 +63,6 @@ jobs: candidates = [v for v in available if not v.is_prerelease and (v.major, v.minor) >= (floor.major, floor.minor)] minor_releases = [max(group) for _, group in groupby(candidates, key=lambda v: (v.major, v.minor))] - # on a pull request, test only the latest minor - if os.environ.get('GITHUB_EVENT_NAME') == 'pull_request': - minor_releases = minor_releases[-1:] - print(json.dumps([str(v) for v in minor_releases])) ") diff --git a/src/spikeinterface/core/base.py b/src/spikeinterface/core/base.py index 9dd5d4d5b5..2c83bb1251 100644 --- a/src/spikeinterface/core/base.py +++ b/src/spikeinterface/core/base.py @@ -482,8 +482,8 @@ def to_dict( include_annotations: bool = False, include_properties: bool = False, relative_to: str | Path | None = None, - folder_metadata=None, recursive: bool = False, + _in_reduce: bool = False, ) -> dict: """ Construct a nested dictionary representation of the extractor. @@ -510,11 +510,11 @@ def to_dict( If provided, file and folder paths will be made relative to this path, enabling portability in folder formats such as the waveform extractor, by default None. - folder_metadata : str | Path | None, default: None - Path to a folder containing additional metadata files (e.g., probe information in BaseRecording) - in numpy `npy` format, by default None. recursive : bool, default: False If True, recursively apply `to_dict` to dictionaries within the kwargs, by default False. + _in_reduce : bool, default: False + Internal flag to indicate that the dictionary is being generated for pickle serialization purposes. + This can be used to handle special cases or optimizations during serialization, by default False. Raises ------ @@ -533,13 +533,11 @@ def to_dict( "relative_paths": , "annotations": , "properties": , - "folder_metadata": } Notes ----- - The `relative_to` argument only has an effect if `recursive` is set to True. - - The `folder_metadata` argument will be made relative to `relative_to` if both are specified. - The `version` field in the resulting dictionary reflects the version of the module from which the extractor class originates. - The full class attribute above is the full import of the class, e.g. @@ -558,7 +556,6 @@ def to_dict( include_properties=include_properties, # make_paths_relative() will make the recusrivity later: relative_to=None, - folder_metadata=folder_metadata, recursive=recursive, ) @@ -605,12 +602,7 @@ def to_dict( # warnings.warn("Try to BaseExtractor.to_dict() using relative_to but there is no common folder") dump_dict["relative_paths"] = False - if folder_metadata is not None: - if relative_to is not None: - folder_metadata = Path(folder_metadata).resolve().absolute().relative_to(relative_to) - dump_dict["folder_metadata"] = str(folder_metadata) - - self._extra_metadata_to_dict(dump_dict) + self._extra_metadata_to_dict(dump_dict, _in_reduce=_in_reduce) return dump_dict @@ -636,39 +628,8 @@ def from_dict(dictionary: dict, base_folder: Path | str | None = None) -> "BaseE assert base_folder is not None, "When relative_paths=True, need to provide base_folder" dictionary = make_paths_absolute(dictionary, base_folder) extractor = _load_extractor_from_dict(dictionary) - folder_metadata = dictionary.get("folder_metadata", None) - if folder_metadata is not None: - folder_metadata = Path(folder_metadata) - if dictionary.get("relative_paths", False): - folder_metadata = base_folder / folder_metadata - extractor.load_metadata_from_folder(folder_metadata) return extractor - def load_metadata_from_folder(self, folder_metadata: str | Path): - # hack to load probe for recording - folder_metadata = Path(folder_metadata) - - # load properties - prop_folder = folder_metadata / "properties" - if prop_folder.is_dir(): - for prop_file in prop_folder.iterdir(): - if prop_file.suffix == ".npy": - values = np.load(prop_file, allow_pickle=True) - key = prop_file.stem - self.set_property(key, values) - - self._extra_metadata_from_folder(folder_metadata) - - def save_metadata_to_folder(self, folder_metadata: str | Path): - self._extra_metadata_to_folder(folder_metadata) - - # save properties - prop_folder = Path(folder_metadata) / "properties" - prop_folder.mkdir(parents=True, exist_ok=False) - for key in self.get_property_keys(): - values = self.get_property(key) - np.save(prop_folder / (key + ".npy"), values) - def clone(self) -> "BaseExtractor": """ Clones an existing extractor into a new instance. @@ -727,7 +688,7 @@ def _get_file_path(file_path: str | Path, extensions: Sequence) -> Path: ) return file_path - def dump(self, file_path: str | Path, relative_to=None, folder_metadata=None) -> None: + def dump(self, file_path: str | Path, relative_to=None) -> None: """ Dumps extractor to json or pickle @@ -740,9 +701,9 @@ def dump(self, file_path: str | Path, relative_to=None, folder_metadata=None) -> This means that file and folder paths in extractor objects kwargs are changed to be relative rather than absolute. """ if str(file_path).endswith(".json"): - self.dump_to_json(file_path, relative_to=relative_to, folder_metadata=folder_metadata) + self.dump_to_json(file_path, relative_to=relative_to) elif str(file_path).endswith(".pkl") or str(file_path).endswith(".pickle"): - self.dump_to_pickle(file_path, relative_to=relative_to, folder_metadata=folder_metadata) + self.dump_to_pickle(file_path, relative_to=relative_to) else: raise ValueError("Dump: file must .json or .pkl") @@ -750,7 +711,6 @@ def dump_to_json( self, file_path: str | Path | None = None, relative_to: str | Path | bool | None = None, - folder_metadata: str | Path | None = None, ) -> None: """ Dump recording extractor to json file. @@ -763,8 +723,6 @@ def dump_to_json( relative_to: str, Path, True or None If not None, files and folders are serialized relative to this path. If True, the relative folder is the parent folder. This means that file and folder paths in extractor objects kwargs are changed to be relative rather than absolute. - folder_metadata: str, Path, or None - Folder with files containing additional information (e.g. probe in BaseRecording) and properties """ assert self.check_serializability("json"), "The extractor is not json serializable" @@ -777,7 +735,6 @@ def dump_to_json( include_annotations=True, include_properties=False, relative_to=relative_to, - folder_metadata=folder_metadata, recursive=True, ) file_path = self._get_file_path(file_path, [".json"]) @@ -792,7 +749,6 @@ def dump_to_pickle( file_path: str | Path | None = None, relative_to: str | Path | bool | None = None, include_properties: bool = True, - folder_metadata: str | Path | None = None, ): """ Dump recording extractor to a pickle file. @@ -807,8 +763,6 @@ def dump_to_pickle( This means that file and folder paths in extractor objects kwargs are changed to be relative rather than absolute. include_properties: bool If True, all properties are dumped - folder_metadata: str, Path, or None - Folder with files containing additional information (e.g. probe in BaseRecording) and properties. """ assert self.check_serializability("pickle"), "The extractor is not serializable to file with pickle" @@ -824,7 +778,6 @@ def dump_to_pickle( dump_dict = self.to_dict( include_annotations=True, include_properties=include_properties, - folder_metadata=folder_metadata, relative_to=relative_to, recursive=recursive, ) @@ -852,7 +805,7 @@ def __reduce__(self): This function is used by pickle to serialize the object. """ instance_constructor = self.from_dict - intialization_args = (self.to_dict(),) + intialization_args = (self.to_dict(_in_reduce=True),) return (instance_constructor, intialization_args) def _save(self, folder, **save_kwargs): @@ -860,19 +813,11 @@ def _save(self, folder, **save_kwargs): # this is internally call by cache(...) main function raise NotImplementedError - def _extra_metadata_from_folder(self, folder): - # This implemented in BaseRecording for probe - pass - - def _extra_metadata_to_folder(self, folder): - # This implemented in BaseRecording for probe - pass - def _extra_metadata_from_dict(self, dump_dict): # This implemented in BaseRecording for probe pass - def _extra_metadata_to_dict(self, dump_dict): + def _extra_metadata_to_dict(self, dump_dict, _in_reduce: bool = False): # This implemented in BaseRecording for probe pass @@ -1012,9 +957,7 @@ def save_to_folder( warnings.warn("The extractor is not serializable to file. The provenance will not be saved.") # save data (done the subclass) - self.save_metadata_to_folder(folder) cached = self._save(folder=folder, verbose=verbose, **save_kwargs) - cached.load_metadata_from_folder(folder) # copy properties/ self.copy_metadata(cached) diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index 5fc7d37505..e193ae668a 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -3,11 +3,11 @@ from pathlib import Path import numpy as np -from probeinterface import read_probeinterface, write_probeinterface +from probeinterface import write_probeinterface from .time_series import TimeSeriesSegment, TimeSeries from .baserecordingsnippets import BaseRecordingSnippets -from .core_tools import convert_bytes_to_str, convert_seconds_to_str +from .core_tools import convert_bytes_to_str, convert_seconds_to_str, save_properties_to_binary_folder from .job_tools import split_job_kwargs @@ -329,13 +329,24 @@ def _save(self, format="binary", verbose: bool = False, **save_kwargs): folder = kwargs["folder"] file_paths = [folder / f"traces_cached_seg{i}.raw" for i in range(self.get_num_segments())] dtype = kwargs.get("dtype", None) or self.get_dtype() - t_starts = self._get_t_starts() + t_starts = self.get_segment_t_starts() + # Check if there are any time vectors + if self.has_any_time_vector(): + file_timestamps_paths = [folder / f"times_cached_seg{i}.raw" for i in range(self.get_num_segments())] + else: + file_timestamps_paths = None - write_binary(self, file_paths=file_paths, dtype=dtype, verbose=verbose, **job_kwargs) + write_binary( + self, + file_paths=file_paths, + dtype=dtype, + verbose=verbose, + file_timestamps_paths=file_timestamps_paths, + **job_kwargs, + ) # This is created so it can be saved as json because the `BinaryFolderRecording` requires it loading # See the __init__ of `BinaryFolderRecording` - binary_rec = BinaryRecordingExtractor( file_paths=file_paths, sampling_frequency=self.get_sampling_frequency(), @@ -348,16 +359,20 @@ def _save(self, format="binary", verbose: bool = False, **save_kwargs): is_filtered=self.is_filtered(), gain_to_uV=self.get_channel_gains(), offset_to_uV=self.get_channel_offsets(), + file_timestamps_paths=file_timestamps_paths, ) binary_rec.dump(folder / "binary.json", relative_to=folder) - cached = BinaryFolderRecording(folder_path=folder) - # timestamps are not saved in binary, so we have to set them explicitly - for segment_index in range(self.get_num_segments()): - if self.has_time_vector(segment_index): - # the use of get_times is preferred since timestamps are converted to array - time_vector = self.get_times(segment_index=segment_index) - cached.set_times(time_vector, segment_index=segment_index) + # save properties + properties_folder = folder / "properties" + properties_folder.mkdir(parents=True, exist_ok=True) + save_properties_to_binary_folder(properties_folder, self) + # save probegroup + if self.has_probe(): + probegroup = self.get_probegroup() + write_probeinterface(folder / "probegroup.json", probegroup) + + cached = BinaryFolderRecording(folder_path=folder) elif format == "memory": if kwargs.get("sharedmem", True): @@ -396,26 +411,15 @@ def _save(self, format="binary", verbose: bool = False, **save_kwargs): return cached - def _extra_metadata_from_folder(self, folder): - # load probe - super()._extra_metadata_from_folder(folder) - - # load time vector if any - for segment_index, rs in enumerate(self.segments): - time_file = folder / f"times_cached_seg{segment_index}.npy" - if time_file.is_file(): - time_vector = np.load(time_file, mmap_mode="r") - rs._time_vector = time_vector - - def _extra_metadata_to_folder(self, folder): - super()._extra_metadata_to_folder(folder) - - # save time vector if any - for segment_index, rs in enumerate(self.segments): - d = rs.get_times_kwargs() - time_vector = d["time_vector"] - if time_vector is not None: - np.save(folder / f"times_cached_seg{segment_index}.npy", time_vector) + def _extra_metadata_to_dict(self, dump_dict, _in_reduce: bool = False): + super()._extra_metadata_to_dict(dump_dict, _in_reduce=_in_reduce) + + # Add times_kwargs if the recording has been modified in memory (e.g. by set_times / shift_times / reset_times) + if _in_reduce and self._time_info_modified: + dump_dict["times_kwargs"] = [] + for segment_index in range(self.get_num_segments()): + times_kwargs = self.segments[segment_index].get_times_kwargs() + dump_dict["times_kwargs"].append(times_kwargs) def select_channels(self, channel_ids: list | np.ndarray | tuple) -> "BaseRecording": """ diff --git a/src/spikeinterface/core/baserecordingsnippets.py b/src/spikeinterface/core/baserecordingsnippets.py index 66c6c5614b..777a48986a 100644 --- a/src/spikeinterface/core/baserecordingsnippets.py +++ b/src/spikeinterface/core/baserecordingsnippets.py @@ -320,31 +320,6 @@ def _extra_metadata_copy(self, other): if self._probegroup is not None: other._probegroup = self._probegroup.copy() - def _extra_metadata_from_folder(self, folder): - # load probe from folder - # Note: we don't need any fix for legacy probegroups, since the - # set_probegroup() method will handle the device_channel_indices - # sorting and global contact order - folder = Path(folder) - probe_file = folder / "probegroup.json" - legacy_probe_file = folder / "probe.json" - if probe_file.is_file(): - probegroup = read_probeinterface(probe_file) - self.set_probegroup(probegroup) - elif legacy_probe_file.is_file(): - probegroup = read_probeinterface(legacy_probe_file) - self.set_probegroup(probegroup) - - # remove "contact_vector" property if present as it is not needed anymore - if "contact_vector" in self.get_property_keys(): - self.delete_property("contact_vector") - - def _extra_metadata_to_folder(self, folder): - # save probe - if self.has_probe(): - probegroup = self.get_probegroup() - write_probeinterface(folder / "probegroup.json", probegroup) - def _extra_metadata_from_dict(self, dump_dict): # load probe and handle backward-compatibility with legacy "contact_vector"/"location" property if "probegroup" in dump_dict: @@ -352,7 +327,16 @@ def _extra_metadata_from_dict(self, dump_dict): probegroup = dump_dict["probegroup"] self._probegroup = ProbeGroup.from_dict(probegroup) - def _extra_metadata_to_dict(self, dump_dict): + if "times_kwargs" in dump_dict: + # When serializing, dump timestamps information because this could have been + # set in memory + times_kwargs_list = dump_dict["times_kwargs"] + for segment_index, times_kwargs in enumerate(times_kwargs_list): + self.segments[segment_index]._sampling_frequency = times_kwargs["sampling_frequency"] + self.segments[segment_index]._t_start = times_kwargs["t_start"] + self.segments[segment_index]._time_vector = times_kwargs["time_vector"] + + def _extra_metadata_to_dict(self, dump_dict, _in_reduce: bool = False): # save probe if self.has_probe(): probegroup = self.get_probegroup() diff --git a/src/spikeinterface/core/basesnippets.py b/src/spikeinterface/core/basesnippets.py index bb7704d7ec..b4b797ea36 100644 --- a/src/spikeinterface/core/basesnippets.py +++ b/src/spikeinterface/core/basesnippets.py @@ -3,6 +3,8 @@ import numpy as np from warnings import warn +from .core_tools import save_properties_to_binary_folder + # snippets segments? @@ -225,6 +227,10 @@ def _save(self, format="npy", **save_kwargs): from spikeinterface.core.npysnippetsextractor import NpySnippetsExtractor NpySnippetsExtractor.write_snippets(snippets=self, file_paths=file_paths, dtype=dtype) + # save properties + properties_folder = folder / "properties" + properties_folder.mkdir(parents=True, exist_ok=True) + save_properties_to_binary_folder(properties_folder, self) cached = NpySnippetsExtractor( file_paths=file_paths, sampling_frequency=self.get_sampling_frequency(), diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index c5c26c375d..f65370b14f 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -6,6 +6,7 @@ from .base import BaseExtractor, BaseSegment, minimum_spike_dtype from .waveform_tools import has_exceeding_spikes +from .core_tools import save_properties_to_binary_folder class BaseSorting(BaseExtractor): @@ -516,6 +517,10 @@ def _save(self, format: str = "numpy_folder", **save_kwargs): folder = save_kwargs.pop("folder") NumpyFolderSorting.write_sorting(self, folder) + # save properties + properties_folder = folder / "properties" + properties_folder.mkdir(parents=True, exist_ok=True) + save_properties_to_binary_folder(properties_folder, self) cached = NumpyFolderSorting(folder) elif format == "zarr": diff --git a/src/spikeinterface/core/binaryfolder.py b/src/spikeinterface/core/binaryfolder.py index ee86c4dc7c..51fe50904c 100644 --- a/src/spikeinterface/core/binaryfolder.py +++ b/src/spikeinterface/core/binaryfolder.py @@ -6,7 +6,11 @@ from probeinterface import read_probeinterface from .binaryrecordingextractor import BinaryRecordingExtractor -from .core_tools import define_function_from_class, make_paths_absolute +from .core_tools import ( + define_function_from_class, + make_paths_absolute, + load_properties_from_binary_folder, +) class BinaryFolderRecording(BinaryRecordingExtractor): @@ -28,6 +32,7 @@ class BinaryFolderRecording(BinaryRecordingExtractor): def __init__(self, folder_path): folder_path = Path(folder_path) + self.folder_path = folder_path with open(folder_path / "binary.json", "r") as f: d = json.load(f) @@ -42,15 +47,7 @@ def __init__(self, folder_path): BinaryRecordingExtractor.__init__(self, **d["kwargs"]) # Load properties - prop_folder = folder_path / "properties" - if prop_folder.is_dir(): - for prop_file in prop_folder.iterdir(): - if prop_file.suffix == ".npy": - values = np.load(prop_file, allow_pickle=True) - key = prop_file.stem - if key == "contact_vector": - continue - self.set_property(key, values) + load_properties_from_binary_folder(folder_path / "properties", self) # Load the probegroup probe_file = folder_path / "probegroup.json" @@ -87,17 +84,8 @@ def __init__(self, folder_path): if probegroup is not None: self._probegroup = probegroup - # Load time vectors if any - for segment_index, rs in enumerate(self.segments): - time_file = folder_path / f"times_cached_seg{segment_index}.npy" - if time_file.is_file(): - rs._time_vector = np.load(time_file, mmap_mode="r") - self._kwargs = dict(folder_path=str(Path(folder_path).absolute())) self._bin_kwargs = d["kwargs"] - if "num_channels" not in self._bin_kwargs: - assert "num_chan" in self._bin_kwargs, "Cannot find num_channels or num_chan in binary.json" - self._bin_kwargs["num_channels"] = self._bin_kwargs["num_chan"] def is_binary_compatible(self) -> bool: return True @@ -112,5 +100,21 @@ def get_binary_description(self): ) return d + def _handle_extractor_backward_compatibility(self): + """ + Handle backward compatibility for BinaryFolderRecording for loading timestamps. + In previous versions of spikeinterface (<0.105.0), the timestamps were saved in a + file called "times_cached_seg{i}.npy" for each segment by the _extra_metadata_to_folder method. + In the current version, the timestamps are saved in a file called "times_cached_seg{i}.raw" for each segment + by the _save method. This method checks for the existence of the old timestamp files and loads them if they + exist, ensuring that recordings saved with older versions of spikeinterface can still be loaded correctly. + """ + super()._handle_extractor_backward_compatibility() + # Load time vectors if any + for segment_index, rs in enumerate(self.segments): + time_file = self.folder_path / f"times_cached_seg{segment_index}.npy" + if time_file.is_file(): + rs._time_vector = np.load(time_file, mmap_mode="r") + read_binary_folder = define_function_from_class(source_class=BinaryFolderRecording, name="read_binary_folder") diff --git a/src/spikeinterface/core/binaryrecordingextractor.py b/src/spikeinterface/core/binaryrecordingextractor.py index c25f2fe9a4..7a3cb53355 100644 --- a/src/spikeinterface/core/binaryrecordingextractor.py +++ b/src/spikeinterface/core/binaryrecordingextractor.py @@ -38,6 +38,8 @@ class BinaryRecordingExtractor(BaseRecording): The offset to apply to the traces is_filtered : bool or None, default: None If True, the recording is assumed to be filtered. If None, is_filtered is not set. + file_timestamps_paths : str or Path or list, default: None + Path to the binary file containing timestamps for each segment. If None, timestamps are not loaded Notes ----- @@ -51,17 +53,18 @@ class BinaryRecordingExtractor(BaseRecording): def __init__( self, - file_paths, - sampling_frequency, - dtype, + file_paths: str | Path | list[str | Path], + sampling_frequency: float, + dtype: str | np.dtype, num_channels: int | None = None, - t_starts=None, - channel_ids=None, - time_axis=0, - file_offset=0, - gain_to_uV=None, - offset_to_uV=None, - is_filtered=None, + t_starts: list[float] | None = None, + channel_ids: list[str | int] | None = None, + time_axis: int = 0, + file_offset: int = 0, + gain_to_uV: float | np.ndarray | None = None, + offset_to_uV: float | np.ndarray | None = None, + is_filtered: bool | None = None, + file_timestamps_paths: str | Path | list[str | Path] | None = None, ): if channel_ids is None: @@ -82,6 +85,12 @@ def __init__( assert len(t_starts) == len(file_path_list), "t_starts must be a list of the same size as file_paths" t_starts = [float(t_start) for t_start in t_starts] + if file_timestamps_paths is not None: + if isinstance(file_timestamps_paths, list): + file_timestamps_paths = [Path(p) for p in file_timestamps_paths] + else: + file_timestamps_paths = [Path(file_timestamps_paths)] + dtype = np.dtype(dtype) for i, file_path in enumerate(file_path_list): @@ -89,8 +98,20 @@ def __init__( t_start = None else: t_start = t_starts[i] + if file_timestamps_paths is None: + file_timestamps_path = None + else: + file_timestamps_path = file_timestamps_paths[i] + rec_segment = BinaryRecordingSegment( - file_path, sampling_frequency, t_start, num_channels, dtype, time_axis, file_offset + file_path, + sampling_frequency, + t_start, + num_channels, + dtype, + time_axis, + file_offset, + file_timestamps_path, ) self.add_recording_segment(rec_segment) @@ -115,6 +136,9 @@ def __init__( "gain_to_uV": gain_to_uV, "offset_to_uV": offset_to_uV, "is_filtered": is_filtered, + "file_timestamps_paths": ( + [str(Path(e).absolute()) for e in file_timestamps_paths] if file_timestamps_paths is not None else None + ), } @classmethod @@ -174,7 +198,17 @@ def __del__(self): class BinaryRecordingSegment(BaseRecordingSegment): - def __init__(self, file_path, sampling_frequency, t_start, num_channels, dtype, time_axis, file_offset): + def __init__( + self, + file_path, + sampling_frequency, + t_start, + num_channels, + dtype, + time_axis, + file_offset, + file_timestamps_path=None, + ): BaseRecordingSegment.__init__(self, sampling_frequency=sampling_frequency, t_start=t_start) self.num_channels = num_channels self.dtype = np.dtype(dtype) @@ -185,6 +219,9 @@ def __init__(self, file_path, sampling_frequency, t_start, num_channels, dtype, self.bytes_per_sample = self.num_channels * self.dtype.itemsize self.data_size_in_bytes = Path(file_path).stat().st_size - file_offset self.num_samples = self.data_size_in_bytes // self.bytes_per_sample + self.file_timestamps_path = file_timestamps_path + if file_timestamps_path is not None: + self._time_vector = np.memmap(file_timestamps_path, dtype="float64", mode="r", shape=(self.num_samples,)) def get_num_samples(self) -> int: """Returns the number of samples in this signal block @@ -251,6 +288,13 @@ def __del__(self): warnings.warn(f"Error closing file handle in BinaryRecordingSegment: {e}") pass + if self._time_vector is not None: + try: + self._time_vector._mmap.close() # Close the underlying mmap object + del self._time_vector + except Exception as e: + pass + # For backward compatibility (old good time) BinDatRecordingExtractor = BinaryRecordingExtractor diff --git a/src/spikeinterface/core/core_tools.py b/src/spikeinterface/core/core_tools.py index ad64323909..aa5d919116 100644 --- a/src/spikeinterface/core/core_tools.py +++ b/src/spikeinterface/core/core_tools.py @@ -779,3 +779,42 @@ def is_path_remote(path: str | Path) -> bool: def ms_to_samples(ms: float, sampling_frequency: float) -> int: """Convert a duration in milliseconds to the nearest number of samples.""" return round(ms * sampling_frequency / 1000.0) + + +def load_properties_from_binary_folder(folder: str | Path, extractor: "BaseExtractor") -> dict: + """ + Load properties from a folder properties as .npy files and return sets them + as properties to the extractor. + + Parameters + ---------- + folder : str or Path + The folder containing the properties as .npy files. + extractor : BaseExtractor + The extractor to which the properties will be set. + """ + folder = Path(folder) + if folder.is_dir(): + for prop_file in folder.iterdir(): + if prop_file.suffix == ".npy": + values = np.load(prop_file, allow_pickle=True) + key = prop_file.stem + if key == "contact_vector": + continue + extractor.set_property(key, values) + + +def save_properties_to_binary_folder(folder: str | Path, extractor: "BaseExtractor"): + """ + Save properties from an extractor to a folder as .npy files. + + Parameters + ---------- + folder : str or Path + The folder where the properties will be saved as .npy files. + extractor : BaseExtractor + The extractor from which the properties will be saved. + """ + for key in extractor.get_property_keys(): + values = extractor.get_property(key) + np.save(folder / f"{key}.npy", values, allow_pickle=True) diff --git a/src/spikeinterface/core/frameslicerecording.py b/src/spikeinterface/core/frameslicerecording.py index 513c8b3dfb..d732f45936 100644 --- a/src/spikeinterface/core/frameslicerecording.py +++ b/src/spikeinterface/core/frameslicerecording.py @@ -68,10 +68,13 @@ def __init__(self, parent_recording_segment, start_frame, end_frame): d = d.copy() if d["time_vector"] is None: d["t_start"] = parent_recording_segment.sample_index_to_time(start_frame) + parent_has_time_vector = False else: - d["time_vector"] = d["time_vector"][start_frame:end_frame] + d["time_vector"] = None + parent_has_time_vector = True BaseRecordingSegment.__init__(self, **d) self._parent_recording_segment = parent_recording_segment + self._parent_has_time_vector = parent_has_time_vector self.start_frame = start_frame self.end_frame = end_frame @@ -85,3 +88,45 @@ def get_traces(self, start_frame, end_frame, channel_indices): start_frame=parent_start, end_frame=parent_end, channel_indices=channel_indices ) return traces + + # Times methods below mirror get_traces(): defer to the parent segment with an offset + # instead of slicing self._time_vector directly. Slicing here would force reading the + # whole window right away (e.g. a zarr.Array fetch+decompress), and since this segment + # is rebuilt from scratch once per worker process in a multiprocessing job, that cost + # would be paid again in every worker instead of being read lazily per chunk. + def get_times(self, start_frame=None, end_frame=None): + if self._parent_has_time_vector: + start_frame = int(start_frame) if start_frame is not None else 0 + end_frame = int(end_frame) if end_frame is not None else self.get_num_samples() + return self._parent_recording_segment.get_times( + start_frame=self.start_frame + start_frame, end_frame=self.start_frame + end_frame + ) + else: + return super().get_times(start_frame=start_frame, end_frame=end_frame) + + def get_start_time(self) -> float: + if self._parent_has_time_vector: + return self._parent_recording_segment.sample_index_to_time(self.start_frame) + else: + return super().get_start_time() + + def get_end_time(self) -> float: + if self._parent_has_time_vector: + return self._parent_recording_segment.sample_index_to_time(self.end_frame - 1) + else: + return super().get_end_time() + + def sample_index_to_time(self, sample_ind): + if self._parent_has_time_vector: + return self._parent_recording_segment.sample_index_to_time(self.start_frame + sample_ind) + else: + return super().sample_index_to_time(sample_ind) + + def time_to_sample_index(self, time_s): + if self._parent_has_time_vector: + return self._parent_recording_segment.time_to_sample_index(time_s) - self.start_frame + else: + return super().time_to_sample_index(time_s) + + def has_time_vector(self) -> bool: + return self._parent_has_time_vector diff --git a/src/spikeinterface/core/loading.py b/src/spikeinterface/core/loading.py index 4242473fd5..b24012606b 100644 --- a/src/spikeinterface/core/loading.py +++ b/src/spikeinterface/core/loading.py @@ -296,23 +296,17 @@ def _load_object_from_zarr(folder_or_url, object_type, **kwargs): elif object_type == "Recording": from .zarrextractors import read_zarr_recording - storage_options = kwargs.get("storage_options", None) - load_compression_ratio = kwargs.get("load_compression_ratio", False) - recording = read_zarr_recording( - folder_or_url, storage_options=storage_options, load_compression_ratio=load_compression_ratio - ) + recording = read_zarr_recording(folder_or_url, **kwargs) return recording elif object_type == "Sorting": from .zarrextractors import read_zarr_sorting - storage_options = kwargs.get("storage_options", None) - sorting = read_zarr_sorting(folder_or_url, storage_options=storage_options) + sorting = read_zarr_sorting(folder_or_url, **kwargs) return sorting elif object_type == "Recording|Sorting": # This case shoudl deprecated soon because the read_zarr is ultra ambiguous # just testing if the zarr contains unit_ids or channel_ids but many object also contains it (see template)!!!! from .zarrextractors import read_zarr - storage_options = kwargs.get("storage_options", None) - rec_or_sorting = read_zarr(folder_or_url, storage_options=storage_options) + rec_or_sorting = read_zarr(folder_or_url, **kwargs) return rec_or_sorting diff --git a/src/spikeinterface/core/npyfoldersnippets.py b/src/spikeinterface/core/npyfoldersnippets.py index 1e465d827a..7c29a4c730 100644 --- a/src/spikeinterface/core/npyfoldersnippets.py +++ b/src/spikeinterface/core/npyfoldersnippets.py @@ -2,7 +2,7 @@ import json from .npysnippetsextractor import NpySnippetsExtractor -from .core_tools import define_function_from_class, make_paths_absolute +from .core_tools import define_function_from_class, make_paths_absolute, load_properties_from_binary_folder class NpyFolderSnippets(NpySnippetsExtractor): @@ -41,8 +41,7 @@ def __init__(self, folder_path): NpySnippetsExtractor.__init__(self, **d["kwargs"]) - folder_metadata = folder_path - self.load_metadata_from_folder(folder_metadata) + load_properties_from_binary_folder(folder_path / "properties", self) self._kwargs = dict(folder_path=str(Path(folder_path).absolute())) self._bin_kwargs = d["kwargs"] diff --git a/src/spikeinterface/core/numpyextractors.py b/src/spikeinterface/core/numpyextractors.py index cc04d51b61..a83d8b87b8 100644 --- a/src/spikeinterface/core/numpyextractors.py +++ b/src/spikeinterface/core/numpyextractors.py @@ -80,7 +80,7 @@ def __init__(self, traces_list, sampling_frequency, t_starts=None, channel_ids=N def from_recording(source_recording, **job_kwargs): traces_list, shms = write_memory_recording(source_recording, dtype=None, **job_kwargs) - t_starts = source_recording._get_t_starts() + t_starts = source_recording.get_segment_t_starts() if shms[0] is not None: # if the computation was done in parallel then traces_list is shared array @@ -208,7 +208,7 @@ def __del__(self): def from_recording(source_recording, **job_kwargs): traces_list, shms = write_memory_recording(source_recording, buffer_type="sharedmem", **job_kwargs) - t_starts = source_recording._get_t_starts() + t_starts = source_recording.get_segment_t_starts() recording = SharedMemoryRecording( shm_names=[shm.name for shm in shms], diff --git a/src/spikeinterface/core/sortingfolder.py b/src/spikeinterface/core/sortingfolder.py index ad331cf835..91296114d4 100644 --- a/src/spikeinterface/core/sortingfolder.py +++ b/src/spikeinterface/core/sortingfolder.py @@ -5,7 +5,7 @@ from .basesorting import BaseSorting, SpikeVectorSortingSegment from .npzsortingextractor import NpzSortingExtractor -from .core_tools import define_function_from_class, make_paths_absolute +from .core_tools import define_function_from_class, load_properties_from_binary_folder, make_paths_absolute class NumpyFolderSorting(BaseSorting): @@ -44,8 +44,7 @@ def __init__(self, folder_path, mmap_mode: str | None = None): # important trick : the cache is already spikes vector self._cached_spike_vector = self.spikes - # Load metadata - self.load_metadata_from_folder(folder_path) + load_properties_from_binary_folder(folder_path / "properties", self) self._kwargs = dict(folder_path=str(folder_path.absolute()), mmap_mode=mmap_mode) @@ -107,8 +106,7 @@ def __init__(self, folder_path): NpzSortingExtractor.__init__(self, **d["kwargs"]) - folder_metadata = folder_path - self.load_metadata_from_folder(folder_metadata) + load_properties_from_binary_folder(folder_path / "properties", self) self._kwargs = dict(folder_path=str(folder_path.absolute())) self._npz_kwargs = d["kwargs"] diff --git a/src/spikeinterface/core/tests/test_baserecording.py b/src/spikeinterface/core/tests/test_baserecording.py index 7debd337c3..f20d174da9 100644 --- a/src/spikeinterface/core/tests/test_baserecording.py +++ b/src/spikeinterface/core/tests/test_baserecording.py @@ -5,6 +5,7 @@ import json import pickle +import platform from pathlib import Path import pytest import numpy as np @@ -286,9 +287,7 @@ def test_BaseRecording(create_cache_folder): traces_int16 = rec_int16.get_traces() assert traces_int16.dtype == "int16" - # Both return_scaled and return_in_uV raise error when no gain_to_uV/offset_to_uV properties - with pytest.raises(ValueError): - traces_float32 = rec_int16.get_traces(return_scaled=True) + # return_in_uV raise error when no gain_to_uV/offset_to_uV properties with pytest.raises(ValueError): traces_float32 = rec_int16.get_traces(return_in_uV=True) @@ -494,6 +493,57 @@ def test_time_slice_with_time_vector(): assert np.allclose(sliced_recording_times.get_traces(), sliced_recording_frames.get_traces()) +@pytest.mark.parametrize( + "mp_context", + [ + pytest.param( + "fork", marks=pytest.mark.skipif(platform.system() != "Linux", reason="fork only supported on Linux") + ), + pytest.param( + "forkserver", + marks=pytest.mark.skipif(platform.system() != "Linux", reason="forkserver only supported on Linux"), + ), + "spawn", + ], +) +def test_save_load_binary_with_time_vector(create_cache_folder, mp_context): + cache_folder = create_cache_folder + rec = generate_recording(durations=[5.0], num_channels=3, sampling_frequency=10_000.0) + times = rec.get_times(segment_index=0) + 100.0 + + # Set time vector + rec.set_times(times=times, segment_index=0, with_warning=False) + # 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)) + + # Save + rec_saved_par = rec.save( + folder=cache_folder / f"recording_with_time_vector_par_{mp_context}", + format="binary", + n_jobs=2, + mp_context=mp_context, + ) + assert np.allclose(rec.get_times(segment_index=0), rec_saved_par.get_times(segment_index=0)) + + # Now reset_times and save again, to check that the time vector is not saved + rec_saved.reset_times() + rec_saved_no_time_vector = rec_saved.save( + folder=cache_folder / f"recording_without_time_vector_{mp_context}", format="binary" + ) + assert not rec_saved_no_time_vector.has_time_vector(segment_index=0) + + # Now make sure the same happens if we save in parallel with multiple jobs, which requires pickling/unpickling + # the recording object + rec_saved_no_time_vector_par = rec_saved.save( + folder=cache_folder / f"recording_without_time_vector_par_{mp_context}", + format="binary", + n_jobs=2, + mp_context=mp_context, + ) + assert not rec_saved_no_time_vector_par.has_time_vector(segment_index=0) + + if __name__ == "__main__": import tempfile diff --git a/src/spikeinterface/core/time_series.py b/src/spikeinterface/core/time_series.py index 640a256d24..e17437a61c 100644 --- a/src/spikeinterface/core/time_series.py +++ b/src/spikeinterface/core/time_series.py @@ -33,6 +33,8 @@ class TimeSeries(ABC): """ _preferred_mp_context = None + # Flag to indicate whether time info has been modified in-memory (e.g. by set_times or shift_times). + _time_info_modified = False def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) @@ -216,9 +218,7 @@ def has_time_vector(self, segment_index: Optional[int] = None): True if the recording has time vectors, False otherwise """ segment_index = self._check_segment_index(segment_index) - rs = self.segments[segment_index] - d = rs.get_times_kwargs() - return d["time_vector"] is not None + return self.segments[segment_index].has_time_vector() def set_times(self, times, segment_index=None, with_warning=True): """Set times for a recording segment. @@ -248,6 +248,7 @@ def set_times(self, times, segment_index=None, with_warning=True): "times are not always propagated across preprocessing" "Use this carefully!" ) + self._time_info_modified = True def reset_times(self): """ @@ -262,6 +263,7 @@ def reset_times(self): rs._time_vector = None rs._t_start = None rs._sampling_frequency = self.sampling_frequency + self._time_info_modified = True def shift_times(self, shift: int | float, segment_index: int | None = None) -> None: """ @@ -297,6 +299,7 @@ def shift_times(self, shift: int | float, segment_index: int | None = None) -> N else: new_start_time = 0 + shift if rs._t_start is None else rs._t_start + shift rs._t_start = new_start_time + self._time_info_modified = True def sample_index_to_time(self, sample_ind, segment_index=None): """ @@ -361,7 +364,7 @@ def get_total_duration(self) -> float: duration = sum([self.get_duration(segment_index) for segment_index in range(self.get_num_segments())]) return duration - def _get_t_starts(self): + def get_segment_t_starts(self): # handle t_starts t_starts = [] for rs in self.segments: @@ -372,14 +375,11 @@ def _get_t_starts(self): t_starts = None return t_starts - def _get_time_vectors(self): - time_vectors = [] + def has_any_time_vector(self): for rs in self.segments: - d = rs.get_times_kwargs() - time_vectors.append(d["time_vector"]) - if all(time_vector is None for time_vector in time_vectors): - time_vectors = None - return time_vectors + if rs.has_time_vector(): + return True + return False def _searchsorted_right_lazy(time_vector: TimeVector, time_s: float | np.ndarray) -> np.int64 | np.ndarray: @@ -521,6 +521,17 @@ def time_to_sample_index(self, time_s): return sample_index + def has_time_vector(self) -> bool: + """ + Returns whether the segment has a time vector. + + Returns + ------- + bool + True if the segment has a time vector, False otherwise. + """ + return self._time_vector is not None + def get_num_samples(self) -> int: """Returns the number of samples in this signal segment diff --git a/src/spikeinterface/core/time_series_tools.py b/src/spikeinterface/core/time_series_tools.py index a5c64563ec..44b6925bb5 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -70,12 +70,20 @@ def write_binary( file_path_dict = {segment_index: file_path for segment_index, file_path in enumerate(file_path_list)} if file_timestamps_paths is not None: - file_timestamps_path_dict = { - segment_index: file_path for segment_index, file_path in enumerate(file_timestamps_paths) - } + file_timestamps_path_list = ( + [file_timestamps_paths] if not isinstance(file_timestamps_paths, list) else file_timestamps_paths + ) + if len(file_timestamps_path_list) != num_segments: + raise ValueError( + "'file_timestamps_paths' must be a list of the same size as the number of segments in the time_series" + ) else: - file_timestamps_path_dict = None - for segment_index, file_path in file_path_dict.items(): + file_timestamps_path_list = [None] * num_segments + + file_path_dict = {} + file_timestamps_path_dict = {} + for segment_index, file_path in enumerate(file_path_list): + file_path_dict[segment_index] = file_path num_samples = time_series.get_num_samples(segment_index=segment_index) data_size_bytes = sample_size_bytes * num_samples file_size_bytes = data_size_bytes + byte_offset @@ -86,14 +94,13 @@ def write_binary( file.seek(file_size_bytes - 1) file.write(b"\0") - if file_timestamps_path_dict is not None: - file_timestamps_path = file_timestamps_path_dict[segment_index] + file_timestamps_path = file_timestamps_path_list[segment_index] + if file_timestamps_path is not None and time_series.has_time_vector(segment_index=segment_index): + file_timestamps_path_dict[segment_index] = file_timestamps_path with open(file_timestamps_path, "wb+") as file: file.seek(num_samples * 8 - 1) # 8 bytes for float64 timestamps file.write(b"\0") - assert Path(file_path).is_file() - # use executor (loop or workers) func = _write_binary_chunk init_func = _init_binary_worker @@ -114,13 +121,22 @@ def _init_binary_worker(time_series, file_path_dict, dtype, byte_offset, file_ti file_dict = {segment_index: open(file_path, "rb+") for segment_index, file_path in file_path_dict.items()} worker_ctx["file_dict"] = file_dict - worker_ctx["file_timestamps_dict"] = file_timestamps_path_dict + if file_timestamps_path_dict is not None: + file_timestamps_dict = { + segment_index: open(file_timestamps_path, "rb+") + for segment_index, file_timestamps_path in file_timestamps_path_dict.items() + } + worker_ctx["file_timestamps_dict"] = file_timestamps_dict + else: + worker_ctx["file_timestamps_dict"] = None return worker_ctx # used by write_binary + TimeSeriesChunkExecutor def _write_binary_chunk(segment_index, start_frame, end_frame, worker_ctx): + import gc + # recover variables of the worker time_series = worker_ctx["time_series"] dtype = worker_ctx["dtype"] @@ -139,15 +155,26 @@ def _write_binary_chunk(segment_index, start_frame, end_frame, worker_ctx): file.write(data.data) # flush is important!! file.flush() + del data if file_timestamps_dict is not None: - file_timestamps = file_timestamps_dict[segment_index] - timestamps = time_series.get_times(start_frame=start_frame, end_frame=end_frame, segment_index=segment_index) - timestamps = timestamps.astype("float64", order="c", copy=False) - timestamp_byte_offset = start_frame * 8 # 8 bytes for float64 - file.seek(timestamp_byte_offset) - file.write(timestamps.data) - file.flush() + # Some segments might not have timestamps to save + if segment_index in file_timestamps_dict: + file_timestamps = file_timestamps_dict[segment_index] + timestamps = time_series.get_times( + start_frame=start_frame, end_frame=end_frame, segment_index=segment_index + ) + timestamps = timestamps.astype("float64", order="c", copy=False) + timestamp_byte_offset = start_frame * 8 # 8 bytes for float64 + file_timestamps.seek(timestamp_byte_offset) + file_timestamps.write(timestamps.data) + file_timestamps.flush() + del timestamps + + # fix memory leak by forcing garbage collection (same issue as _write_zarr_chunk, + # e.g. reading compressed zarr chunks leaves reference cycles that the generational + # GC doesn't clear promptly in a tight chunk loop) + gc.collect() write_binary.__doc__ = write_binary.__doc__.format(_shared_job_kwargs_doc) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 02ae5e9ecd..3923136c17 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -115,13 +115,16 @@ class ZarrRecordingExtractor(BaseRecording): """ def __init__( - self, folder_path: Path | str, storage_options: dict | None = None, load_compression_ratio: bool = False + self, + folder_path: Path | str, + storage_options: dict | None = None, + load_compression_ratio: bool = False, + load_times: bool = True, ): folder_path, folder_path_kwarg = resolve_zarr_path(folder_path) self._root = super_zarr_open(folder_path, mode="r", storage_options=storage_options) - sampling_frequency = self._root.attrs.get("sampling_frequency", None) num_segments = self._root.attrs.get("num_segments", None) assert "channel_ids" in self._root.keys(), "'channel_ids' dataset not found!" @@ -151,7 +154,7 @@ def __init__( time_kwargs = {} time_vector = self._root.get(f"times_seg{segment_index}", None) - if time_vector is not None: + if time_vector is not None and load_times: time_kwargs["time_vector"] = time_vector else: if t_starts is None: @@ -232,6 +235,7 @@ def __init__( "folder_path": folder_path_kwarg, "storage_options": storage_options, "load_compression_ratio": load_compression_ratio, + "load_times": load_times, } @staticmethod diff --git a/src/spikeinterface/preprocessing/basepreprocessor.py b/src/spikeinterface/preprocessing/basepreprocessor.py index 64d57d3637..d28de0b34f 100644 --- a/src/spikeinterface/preprocessing/basepreprocessor.py +++ b/src/spikeinterface/preprocessing/basepreprocessor.py @@ -34,3 +34,26 @@ def get_num_samples(self): def get_traces(self, start_frame, end_frame, channel_indices): raise NotImplementedError + + # Preprocessors never change the frame numbering (no offset, no resampling), so time + # handling is a pure pass-through to the parent segment. Delegating live (instead of + # relying on the time_vector/t_start copied into __init__ above) lets any lazy/offset-aware + # override further up the chain (e.g. FrameSliceRecordingSegment after a frame_slice) keep + # working without materializing a full time_vector every time this segment is reconstructed. + def get_times(self, start_frame=None, end_frame=None): + return self.parent_recording_segment.get_times(start_frame=start_frame, end_frame=end_frame) + + def get_start_time(self): + return self.parent_recording_segment.get_start_time() + + def get_end_time(self): + return self.parent_recording_segment.get_end_time() + + def sample_index_to_time(self, sample_ind): + return self.parent_recording_segment.sample_index_to_time(sample_ind) + + def time_to_sample_index(self, time_s): + return self.parent_recording_segment.time_to_sample_index(time_s) + + def get_times_kwargs(self): + return self.parent_recording_segment.get_times_kwargs()