Merge branch 'main' of github.com:SpikeInterface/spikeinterface

Author: alejoe91

src/spikeinterface/core/__init__.py

@@ -97,12 +97,14 @@
     get_best_job_kwargs,
     ensure_n_jobs,
     ensure_chunk_size,
-    ChunkExecutor,
+    TimeSeriesChunkExecutor,
     split_job_kwargs,
     fix_job_kwargs,
 )
 from .recording_tools import (
     write_binary_recording,
+    write_memory_recording,
+    write_recording_to_zarr,
     write_to_h5_dataset_format,
     get_random_data_chunks,
     get_channel_distances,

src/spikeinterface/core/baserecording.py

@@ -5,13 +5,13 @@
 import numpy as np
 from probeinterface import read_probeinterface, write_probeinterface
 
-from .chunkable import ChunkableSegment, ChunkableMixin
+from .time_series import TimeSeriesSegment, TimeSeries
 from .baserecordingsnippets import BaseRecordingSnippets
 from .core_tools import convert_bytes_to_str, convert_seconds_to_str
 from .job_tools import split_job_kwargs
 
 
-class BaseRecording(BaseRecordingSnippets, ChunkableMixin):
+class BaseRecording(BaseRecordingSnippets, TimeSeries):
     """
     Abstract class representing several a multichannel timeseries (or block of raw ephys traces).
     Internally handle list of RecordingSegment
@@ -305,7 +305,7 @@ def get_traces(
 
     def get_data(self, start_frame: int, end_frame: int, segment_index: int | None = None, **kwargs) -> np.ndarray:
         """
-        General retrieval function for chunkable objects
+        General retrieval function for time_series objects
         """
         return self.get_traces(segment_index=segment_index, start_frame=start_frame, end_frame=end_frame, **kwargs)
 
@@ -316,7 +316,7 @@ def _save(self, format="binary", verbose: bool = False, **save_kwargs):
         kwargs, job_kwargs = split_job_kwargs(save_kwargs)
 
         if format == "binary":
-            from .chunkable_tools import write_binary
+            from .time_series_tools import write_binary
 
             folder = kwargs["folder"]
             file_paths = [folder / f"traces_cached_seg{i}.raw" for i in range(self.get_num_segments())]
@@ -642,7 +642,7 @@ def astype(self, dtype, round: bool | None = None):
         return astype(self, dtype=dtype, round=round)
 
 
-class BaseRecordingSegment(ChunkableSegment):
+class BaseRecordingSegment(TimeSeriesSegment):
     """
     Abstract class representing a multichannel timeseries, or block of raw ephys traces
     """
@@ -677,6 +677,6 @@ def get_data(
         self, start_frame: int, end_frame: int, indices: list | np.ndarray | tuple | None = None
     ) -> np.ndarray:
         """
-        General retrieval function for chunkable objects
+        General retrieval function for time_series objects
         """
         return self.get_traces(start_frame=start_frame, end_frame=end_frame, channel_indices=indices)

src/spikeinterface/core/basesorting.py

@@ -327,6 +327,11 @@ def register_recording(self, recording, check_spike_frames: bool = True):
                     "Might be necessary for further postprocessing."
                 )
         self._recording = recording
+        # Copy the recording's start times into the sorting segments. This way,
+        # the sorting preserves the start time even if the recording is later
+        # detached (e.g. analyzer saved and reloaded without the recording).
+        for segment_index, segment in enumerate(self.segments):
+            segment._t_start = recording.get_start_time(segment_index=segment_index)
 
     @property
     def sorting_info(self):
@@ -352,6 +357,66 @@ def has_time_vector(self, segment_index: int | None = None) -> bool:
         else:
             return False
 
+    def get_start_time(self, segment_index: int | None = None) -> float:
+        """Get the start time of the sorting segment.
+
+        Parameters
+        ----------
+        segment_index : int or None, default: None
+            The segment index (required for multi-segment)
+
+        Returns
+        -------
+        float
+            The start time in seconds
+        """
+        segment_index = self._check_segment_index(segment_index)
+        segment = self.segments[segment_index]
+        return segment._t_start if segment._t_start is not None else 0.0
+
+    def get_end_time(self, segment_index: int | None = None) -> float:
+        """Get the end time of the sorting segment.
+
+        If a recording is registered, returns the recording's end time.
+        Otherwise returns the time of the last spike in the segment.
+
+        Parameters
+        ----------
+        segment_index : int or None, default: None
+            The segment index (required for multi-segment)
+
+        Returns
+        -------
+        float
+            The end time in seconds
+        """
+        segment_index = self._check_segment_index(segment_index)
+        if self.has_recording():
+            return self._recording.get_end_time(segment_index=segment_index)
+        else:
+            last_spike_frame = self.get_last_spike_frame(segment_index=segment_index)
+            return self.sample_index_to_time(last_spike_frame, segment_index=segment_index)
+
+    def get_last_spike_frame(self, segment_index: int | None = None) -> int:
+        """Get the frame index of the last spike in a segment across all units.
+
+        Parameters
+        ----------
+        segment_index : int or None, default: None
+            The segment index (required for multi-segment)
+
+        Returns
+        -------
+        int
+            The frame index of the last spike, or 0 if no spikes exist.
+        """
+        segment_index = self._check_segment_index(segment_index)
+        spike_vector = self.to_spike_vector(concatenated=False)
+        spikes_in_segment = spike_vector[segment_index]
+        if len(spikes_in_segment) == 0:
+            return 0
+        return int(np.max(spikes_in_segment["sample_index"]))
+
     def get_times(self, segment_index=None):
         """
         Get time vector for a registered recording segment.

src/spikeinterface/core/job_tools.py

@@ -205,7 +205,7 @@ def divide_segment_into_chunks(num_frames, chunk_size):
     return chunks
 
 
-def divide_chunkable_into_chunks(recording, chunk_size):
+def divide_time_series_into_chunks(recording, chunk_size):
     slices = []
     for segment_index in range(recording.get_num_segments()):
         num_frames = recording.get_num_samples(segment_index)
@@ -242,24 +242,24 @@ def ensure_n_jobs(extractor, n_jobs=1):
     return n_jobs
 
 
-def chunk_duration_to_chunk_size(chunk_duration, chunkable: "ChunkableMixin"):
+def chunk_duration_to_chunk_size(chunk_duration, time_series: "TimeSeries"):
     if isinstance(chunk_duration, float):
-        chunk_size = int(chunk_duration * chunkable.get_sampling_frequency())
+        chunk_size = int(chunk_duration * time_series.get_sampling_frequency())
     elif isinstance(chunk_duration, str):
         if chunk_duration.endswith("ms"):
             chunk_duration = float(chunk_duration.replace("ms", "")) / 1000.0
         elif chunk_duration.endswith("s"):
             chunk_duration = float(chunk_duration.replace("s", ""))
         else:
             raise ValueError("chunk_duration must ends with s or ms")
-        chunk_size = int(chunk_duration * chunkable.get_sampling_frequency())
+        chunk_size = int(chunk_duration * time_series.get_sampling_frequency())
     else:
         raise ValueError("chunk_duration must be str or float")
     return chunk_size
 
 
 def ensure_chunk_size(
-    chunkable: "ChunkableMixin",
+    time_series: "TimeSeries",
     total_memory=None,
     chunk_size=None,
     chunk_memory=None,
@@ -299,30 +299,30 @@ def ensure_chunk_size(
         assert total_memory is None
         # set by memory per worker size
         chunk_memory = convert_string_to_bytes(chunk_memory)
-        chunk_size = int(chunk_memory / chunkable.get_sample_size_in_bytes())
+        chunk_size = int(chunk_memory / time_series.get_sample_size_in_bytes())
     elif total_memory is not None:
         # clip by total memory size
-        n_jobs = ensure_n_jobs(chunkable, n_jobs=n_jobs)
+        n_jobs = ensure_n_jobs(time_series, n_jobs=n_jobs)
         total_memory = convert_string_to_bytes(total_memory)
-        chunk_size = int(total_memory / (chunkable.get_sample_size_in_bytes() * n_jobs))
+        chunk_size = int(total_memory / (time_series.get_sample_size_in_bytes() * n_jobs))
     elif chunk_duration is not None:
-        chunk_size = chunk_duration_to_chunk_size(chunk_duration, chunkable)
+        chunk_size = chunk_duration_to_chunk_size(chunk_duration, time_series)
     else:
         # Edge case to define single chunk per segment for n_jobs=1.
         # All chunking parameters equal None mean single chunk per segment
         if n_jobs == 1:
-            num_segments = chunkable.get_num_segments()
-            samples_in_larger_segment = max([chunkable.get_num_samples(segment) for segment in range(num_segments)])
+            num_segments = time_series.get_num_segments()
+            samples_in_larger_segment = max([time_series.get_num_samples(segment) for segment in range(num_segments)])
             chunk_size = samples_in_larger_segment
         else:
             raise ValueError("For n_jobs >1 you must specify total_memory or chunk_size or chunk_memory")
 
     return chunk_size
 
 
-class ChunkExecutor:
+class TimeSeriesChunkExecutor:
     """
-    Core class for parallel processing to run a "function" over chunks on a chunkable extractor.
+    Core class for parallel processing to run a "function" over chunks on a time_series extractor.
 
     It supports running a function:
         * in loop with chunk processing (low RAM usage)
@@ -334,8 +334,8 @@ class ChunkExecutor:
 
     Parameters
     ----------
-    chunkable : ChunkableMixin
-        The chunkable object to be processed.
+    time_series : TimeSeries
+        The time_series object to be processed.
     func : function
         Function that runs on each chunk
     init_func : function
@@ -383,7 +383,7 @@ class ChunkExecutor:
 
     def __init__(
         self,
-        chunkable: "ChunkableMixin",
+        time_series: "TimeSeries",
         func,
         init_func,
         init_args,
@@ -402,7 +402,7 @@ def __init__(
         max_threads_per_worker=1,
         need_worker_index=False,
     ):
-        self.chunkable = chunkable
+        self.time_series = time_series
         self.func = func
         self.init_func = init_func
         self.init_args = init_args
@@ -421,7 +421,7 @@ def __init__(
                 else:
                     mp_context = "spawn"
 
-            preferred_mp_context = chunkable.get_preferred_mp_context()
+            preferred_mp_context = time_series.get_preferred_mp_context()
             if preferred_mp_context is not None and preferred_mp_context != mp_context:
                 warnings.warn(
                     f"Your processing chain using pool_engine='process' and mp_context='{mp_context}' is not possible."
@@ -437,7 +437,7 @@ def __init__(
         self.handle_returns = handle_returns
         self.gather_func = gather_func
 
-        self.n_jobs = ensure_n_jobs(self.chunkable, n_jobs=n_jobs)
+        self.n_jobs = ensure_n_jobs(self.time_series, n_jobs=n_jobs)
         self.chunk_size = self.ensure_chunk_size(
             total_memory=total_memory,
             chunk_size=chunk_size,
@@ -455,7 +455,7 @@ def __init__(
         if verbose:
             chunk_memory = self.get_chunk_memory()
             total_memory = chunk_memory * self.n_jobs
-            chunk_duration = self.chunk_size / chunkable.sampling_frequency
+            chunk_duration = self.chunk_size / time_series.sampling_frequency
             chunk_memory_str = convert_bytes_to_str(chunk_memory)
             total_memory_str = convert_bytes_to_str(total_memory)
             chunk_duration_str = convert_seconds_to_str(chunk_duration)
@@ -471,13 +471,13 @@ def __init__(
             )
 
     def get_chunk_memory(self):
-        return self.chunk_size * self.chunkable.get_sample_size_in_bytes()
+        return self.chunk_size * self.time_series.get_sample_size_in_bytes()
 
     def ensure_chunk_size(
         self, total_memory=None, chunk_size=None, chunk_memory=None, chunk_duration=None, n_jobs=1, **other_kwargs
     ):
         return ensure_chunk_size(
-            self.chunkable, total_memory, chunk_size, chunk_memory, chunk_duration, n_jobs, **other_kwargs
+            self.time_series, total_memory, chunk_size, chunk_memory, chunk_duration, n_jobs, **other_kwargs
         )
 
     def run(self, slices=None):
@@ -487,7 +487,7 @@ def run(self, slices=None):
 
         if slices is None:
             # TODO: rename
-            slices = divide_chunkable_into_chunks(self.chunkable, self.chunk_size)
+            slices = divide_time_series_into_chunks(self.time_series, self.chunk_size)
 
         if self.handle_returns:
             returns = []