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 10e9829c65..6b6df17ab4 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: @@ -23,8 +22,6 @@ on: - 'src/spikeinterface/core/**' - '.github/scripts/serialization/**' - '.github/workflows/cross_version_serialization.yml' - schedule: - - cron: '0 4 * * 0' # weekly, Sunday 04:00 UTC 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.11' - 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 d6bdebff7b..f7b0f66d4f 100644 --- a/src/spikeinterface/core/base.py +++ b/src/spikeinterface/core/base.py @@ -481,7 +481,6 @@ def to_dict( include_annotations: bool = False, include_properties: bool = False, relative_to: str | Path | None = None, - folder_metadata=None, recursive: bool = False, ) -> dict: """ @@ -509,9 +508,6 @@ 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. @@ -532,13 +528,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. @@ -557,7 +551,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, ) @@ -604,11 +597,6 @@ 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) return dump_dict @@ -635,39 +623,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): - # 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): - self._extra_metadata_to_folder(folder_metadata) - - # save properties - prop_folder = 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. @@ -726,7 +683,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 @@ -739,9 +696,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") @@ -749,7 +706,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. @@ -762,8 +718,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" @@ -776,7 +730,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"]) @@ -791,7 +744,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. @@ -806,8 +758,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" @@ -823,7 +773,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, ) @@ -859,14 +808,6 @@ 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 @@ -1011,9 +952,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 e411362e6f..210ceecb56 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -7,7 +7,12 @@ 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, + load_properties_from_binary_folder, +) from .job_tools import split_job_kwargs @@ -330,8 +335,21 @@ def _save(self, format="binary", verbose: bool = False, **save_kwargs): 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() + time_vectors = self._get_time_vectors() + if time_vectors is not None: + file_timestamps_paths = [folder / f"times_cached_seg{i}.raw" for i in range(self.get_num_segments())] + kwargs["file_timestamps_paths"] = file_timestamps_paths + 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` @@ -348,16 +366,14 @@ 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 + save_properties_to_binary_folder(folder / "properties", self) + + cached = BinaryFolderRecording(folder_path=folder) elif format == "memory": if kwargs.get("sharedmem", True): @@ -396,27 +412,6 @@ 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 select_channels(self, channel_ids: list | np.ndarray | tuple) -> "BaseRecording": """ Returns a new recording object with a subset of channels. diff --git a/src/spikeinterface/core/baserecordingsnippets.py b/src/spikeinterface/core/baserecordingsnippets.py index 66c6c5614b..11eb909b59 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: diff --git a/src/spikeinterface/core/basesnippets.py b/src/spikeinterface/core/basesnippets.py index bb7704d7ec..212b833517 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,7 @@ 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_to_binary_folder(folder / "properties", 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 51ec90bf37..8c27b68619 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): @@ -514,6 +515,7 @@ def _save(self, format="numpy_folder", **save_kwargs): folder = save_kwargs.pop("folder") NumpyFolderSorting.write_sorting(self, folder) + save_properties_to_binary_folder(folder / "properties", self) cached = NumpyFolderSorting(folder) if self.has_recording(): diff --git a/src/spikeinterface/core/binaryfolder.py b/src/spikeinterface/core/binaryfolder.py index e9986193a3..a489db1457 100644 --- a/src/spikeinterface/core/binaryfolder.py +++ b/src/spikeinterface/core/binaryfolder.py @@ -6,7 +6,12 @@ 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, + save_properties_to_binary_folder, +) class BinaryFolderRecording(BinaryRecordingExtractor): @@ -28,6 +33,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 +48,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 +85,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 +101,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..6515b1fbf5 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 a6f33651d5..6f6bda675c 100644 --- a/src/spikeinterface/core/core_tools.py +++ b/src/spikeinterface/core/core_tools.py @@ -781,3 +781,44 @@ 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. + """ + folder = Path(folder) + folder.mkdir(parents=True, exist_ok=True) + 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..00344dbd6b 100644 --- a/src/spikeinterface/core/frameslicerecording.py +++ b/src/spikeinterface/core/frameslicerecording.py @@ -68,8 +68,6 @@ 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) - else: - d["time_vector"] = d["time_vector"][start_frame:end_frame] BaseRecordingSegment.__init__(self, **d) self._parent_recording_segment = parent_recording_segment self.start_frame = start_frame @@ -85,3 +83,37 @@ 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.time_vector is None: + return super().get_times(start_frame=start_frame, end_frame=end_frame) + 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 + ) + + def get_start_time(self) -> float: + if self.time_vector is None: + return super().get_start_time() + return self._parent_recording_segment.sample_index_to_time(self.start_frame) + + def get_end_time(self) -> float: + if self.time_vector is None: + return super().get_end_time() + return self._parent_recording_segment.sample_index_to_time(self.end_frame - 1) + + def sample_index_to_time(self, sample_ind): + if self.time_vector is None: + return super().sample_index_to_time(sample_ind) + return self._parent_recording_segment.sample_index_to_time(self.start_frame + sample_ind) + + def time_to_sample_index(self, time_s): + if self.time_vector is None: + return super().time_to_sample_index(time_s) + return self._parent_recording_segment.time_to_sample_index(time_s) - self.start_frame 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/sortingfolder.py b/src/spikeinterface/core/sortingfolder.py index c0d66393d2..50cc1729d9 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): # important trick : the cache is already spikes vector self._cached_spike_vector = self.spikes - 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())) @@ -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..ce29fb7b3b 100644 --- a/src/spikeinterface/core/tests/test_baserecording.py +++ b/src/spikeinterface/core/tests/test_baserecording.py @@ -286,9 +286,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 +492,34 @@ def test_time_slice_with_time_vector(): assert np.allclose(sliced_recording_times.get_traces(), sliced_recording_frames.get_traces()) +def test_save_load_binary_with_time_vector(create_cache_folder): + 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 / "recording_with_time_vector", 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 / "recording_with_time_vector_par", format="binary") + 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 / "recording_without_time_vector", 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 / "recording_without_time_vector_par", format="binary", n_jobs=2 + ) + 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_tools.py b/src/spikeinterface/core/time_series_tools.py index a5c64563ec..83259c729a 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -114,13 +114,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 +148,22 @@ 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() + 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 990ec3cca0..8695246fdd 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()