Repository: SpikeInterface/spikeinterface
Pull request #578 (2022-03-30 09:22:54)

Add split recording class and function
New functions to split multi-segment recordings into a list of mono-segments:

- Return a list of mono-segment recordings
`recording_list = si.split_recording(rec_multi)`

- Return a single mono segment recording
`recording_mono = si.select_segment_recording(rec_multi, segment_index=1)`


Patch:
diff --git a/spikeinterface/comparison/basecomparison.py b/spikeinterface/comparison/basecomparison.py
index d761f6cba..4e6b9ddc3 100644
--- a/spikeinterface/comparison/basecomparison.py
+++ b/spikeinterface/comparison/basecomparison.py
@@ -175,8 +175,8 @@ class BaseMultiComparison(BaseComparison):
         self._new_units = {}
 
         # save new units
-        self.subgraphs = (self.clean_graph.subgraph(c).copy()
-                          for c in nx.connected_components(self.clean_graph))
+        self.subgraphs = [self.clean_graph.subgraph(c).copy()
+                          for c in nx.connected_components(self.clean_graph)]
         for new_unit, sg in enumerate(self.subgraphs):
             edges = list(sg.edges(data=True))
             if len(edges) > 0:
diff --git a/spikeinterface/core/__init__.py b/spikeinterface/core/__init__.py
index 81a203bac..f899f2552 100644
--- a/spikeinterface/core/__init__.py
+++ b/spikeinterface/core/__init__.py
@@ -32,6 +32,7 @@ from .unitsaggregationsorting import UnitsAggregationSorting, aggregate_units
 from .segmentutils import (
     append_recordings, AppendSegmentRecording,
     concatenate_recordings, ConcatenateSegmentRecording,
+    split_recording, select_segment_recording, SelectSegmentRecording,
     append_sortings, AppendSegmentSorting, 
     split_sorting, SplitSegmentSorting)
 
diff --git a/spikeinterface/core/baserecording.py b/spikeinterface/core/baserecording.py
index 27348afbe..153848dea 100644
--- a/spikeinterface/core/baserecording.py
+++ b/spikeinterface/core/baserecording.py
@@ -8,6 +8,7 @@ from probeinterface import Probe, ProbeGroup, write_probeinterface, read_probein
 
 from .base import BaseExtractor, BaseSegment
 from .core_tools import write_binary_recording, write_memory_recording, write_traces_to_zarr, check_json
+from .job_tools import job_keys
 
 from warnings import warn
 
@@ -165,8 +166,6 @@ class BaseRecording(BaseExtractor):
                           'times are not always propagated to across preprocessing'
                           'Use use this carefully!')
 
-    _job_keys = ['n_jobs', 'total_memory', 'chunk_size', 'chunk_memory', 'progress_bar', 'verbose']
-
     def _save(self, format='binary', **save_kwargs):
         """
         This function replaces the old CacheRecordingExtractor, but enables more engines
@@ -193,7 +192,7 @@ class BaseRecording(BaseExtractor):
             if dtype is None:
                 dtype = self.get_dtype()
 
-            job_kwargs = {k: save_kwargs[k] for k in self._job_keys if k in save_kwargs}
+            job_kwargs = {k: save_kwargs[k] for k in job_keys if k in save_kwargs}
             write_binary_recording(self, file_paths=file_paths, dtype=dtype, **job_kwargs)
 
             from .binaryrecordingextractor import BinaryRecordingExtractor
@@ -204,7 +203,7 @@ class BaseRecording(BaseExtractor):
                                               offset_to_uV=self.get_channel_offsets())
 
         elif format == 'memory':
-            job_kwargs = {k: save_kwargs[k] for k in self._job_keys if k in save_kwargs}
+            job_kwargs = {k: save_kwargs[k] for k in job_keys if k in save_kwargs}
             traces_list = write_memory_recording(self, dtype=None, **job_kwargs)
             from .numpyextractors import NumpyRecording
 
@@ -228,7 +227,7 @@ class BaseRecording(BaseExtractor):
             filters = save_kwargs.get('filters', None)
 
             job_kwargs = {k: save_kwargs[k]
-                          for k in self._job_keys if k in save_kwargs}
+                          for k in job_keys if k in save_kwargs}
             write_traces_to_zarr(self, zarr_root=zarr_root, zarr_path=zarr_path, dataset_paths=dataset_paths,
                                  dtype=dtype, compressor=compressor, filters=filters, **job_kwargs)
 
@@ -621,6 +620,23 @@ class BaseRecording(BaseExtractor):
             elif outputs == 'dict':
                 recordings[value] = subrec
         return recordings
+    
+    def select_segments(self, segment_indices):
+        """
+        Return a recording with the segments specified by 'segment_indices'
+
+        Parameters
+        ----------
+        segment_indices : list of int
+            List of segment indices to keep in the returned recording
+
+        Returns
+        -------
+        SelectSegmentRecording
+            The recording with the selected segments
+        """
+        from .segmentutils import SelectSegmentRecording
+        return SelectSegmentRecording(self, segment_indices=segment_indices)
 
     def planarize(self, axes: str = "xy"):
         """
diff --git a/spikeinterface/core/job_tools.py b/spikeinterface/core/job_tools.py
index f591cdeb1..7d06f805b 100644
--- a/spikeinterface/core/job_tools.py
+++ b/spikeinterface/core/job_tools.py
@@ -27,6 +27,9 @@ _shared_job_kwargs_doc = \
             * progress_bar: bool
                 If True, a progress bar is printed
     """
+    
+job_keys = ['n_jobs', 'total_memory', 'chunk_size', 'chunk_memory', 'chunk_duration', 'progress_bar', 'verbose']
+
 
 
 def divide_segment_into_chunks(num_frames, chunk_size):
diff --git a/spikeinterface/core/segmentutils.py b/spikeinterface/core/segmentutils.py
index 29dc55c9d..07719c300 100644
--- a/spikeinterface/core/segmentutils.py
+++ b/spikeinterface/core/segmentutils.py
@@ -14,14 +14,16 @@ import numpy as np
 from .baserecording import BaseRecording, BaseRecordingSegment
 from .basesorting import BaseSorting, BaseSortingSegment
 
-from typing import List
+from typing import List, Union
+
 
 def _check_sampling_frequencies(sampling_frequency_list, sampling_frequency_max_diff):
     assert sampling_frequency_max_diff >= 0
     freq_0 = sampling_frequency_list[0]
     max_diff = max( abs(freq - freq_0) for freq in sampling_frequency_list)
     if max_diff > sampling_frequency_max_diff:
-        raise ValueError(f"Sampling frequencies across datasets differ by `{max_diff}`Hz which is more than `sampling_frequency_max_diff`={sampling_frequency_max_diff}Hz")
+        raise ValueError(f"Sampling frequencies across datasets differ by `{max_diff}`Hz which is more than "
+                         f"`sampling_frequency_max_diff`={sampling_frequency_max_diff}Hz")
     elif max_diff > 0:
         diff_sec = 24 * 3600 * max_diff / freq_0 
         import warnings
@@ -140,7 +142,8 @@ class ConcatenateSegmentRecording(BaseRecording):
                     assert d['t_start'] is None, ("ConcatenateSegmentRecording does not handle t_start. "
                                                   "Use ignore_times=True to ignore time information.")
                 parent_segments.append(parent_segment)
-        rec_seg = ProxyConcatenateRecordingSegment(parent_segments, one_rec.get_sampling_frequency(), ignore_times=ignore_times)
+        rec_seg = ProxyConcatenateRecordingSegment(parent_segments, one_rec.get_sampling_frequency(), 
+                                                   ignore_times=ignore_times)
         self.add_recording_segment(rec_seg)
 
         self._kwargs = {'recording_list': [rec.to_dict() for rec in recording_list],
@@ -218,6 +221,67 @@ def concatenate_recordings(*args, **kwargs):
 concatenate_recordings.__doc__ == ConcatenateSegmentRecording.__doc__
 
 
+class SelectSegmentRecording(BaseRecording):
+    """
+    Return a new recording with a single segment from a multi-segment recording.
+
+    Parameters
+    ----------
+    recording : BaseRecording
+        The multi-segment recording
+    segment_indices : list of int
+        The segment indices to select
+    """
+
+    def __init__(self, recording: BaseRecording, segment_indices: Union[int, List[int]]):
+        BaseRecording.__init__(self, recording.get_sampling_frequency(), 
+                               recording.channel_ids, recording.get_dtype())
+        recording.copy_metadata(self)
+        
+        if isinstance(segment_indices, int):
+            segment_indices = [segment_indices]
+        
+        num_segments = recording.get_num_segments()
+        assert all(0 <= s < num_segments for s in segment_indices), \
+            f"'segment_index' must be between 0 and {num_segments - 1}"
+
+        for segment_index in segment_indices:
+            rec_seg = recording._recording_segments[segment_index]
+            self.add_recording_segment(rec_seg)
+
+        self._kwargs = {'recording': recording.to_dict(),
+                        'segment_indices': [int(s) for s in segment_indices]}
+        
+
+def split_recording(recording: BaseRecording):
+    """
+    Return a list of mono-segment recordings from a multi-segment recording.
+
+    Parameters
+    ----------
+    recording : BaseRecording
+        The multi-segment recording
+
+    Returns
+    -------
+    recording_list
+        A list of mono-segment recordings
+    """
+    recording_list = []
+    for segment_index in range(recording.get_num_segments()):
+        rec_mono = SelectSegmentRecording(
+            recording=recording, segment_indices=[segment_index])
+        recording_list.append(rec_mono)
+    return recording_list
+
+
+def select_segment_recording(*args, **kwargs):
+    return SelectSegmentRecording(*args, **kwargs)
+
+
+select_segment_recording.__doc__ = SelectSegmentRecording.__doc__
+
+
 class AppendSegmentSorting(BaseSorting):
     """
     Return a sorting that "append" all segments from all sorting
diff --git a/spikeinterface/toolkit/postprocessing/template_metrics.py b/spikeinterface/toolkit/postprocessing/template_metrics.py
index e97b2e178..2667537b5 100644
--- a/spikeinterface/toolkit/postprocessing/template_metrics.py
+++ b/spikeinterface/toolkit/postprocessing/template_metrics.py
@@ -78,7 +78,7 @@ def calculate_template_metrics(waveform_extractor, feature_names=None, peak_sign
                                                               outputs='id')
 
         template_metrics = pd.DataFrame(
-            index=unit_ids, columns=[feature_names])
+            index=unit_ids, columns=feature_names)
     else:
         extremum_channels_ids = get_template_channel_sparsity(
             waveform_extractor, **sparsity_dict)


Test patch:
diff --git a/spikeinterface/core/tests/test_segmentutils.py b/spikeinterface/core/tests/test_segmentutils.py
index 1dfdcafbb..a600516df 100644
--- a/spikeinterface/core/tests/test_segmentutils.py
+++ b/spikeinterface/core/tests/test_segmentutils.py
@@ -4,7 +4,8 @@ import numpy as np
 from spikeinterface.core import (
     append_recordings, AppendSegmentRecording,
     concatenate_recordings, ConcatenateSegmentRecording,
-    append_sortings, AppendSegmentSorting)
+    append_sortings, AppendSegmentSorting,
+    split_recording, select_segment_recording, split_sorting)
 
 from spikeinterface.core import NumpyRecording, NumpySorting
 
@@ -54,6 +55,29 @@ def test_append_concatenate_recordings():
     assert traces[-1, 0] == 499
 
 
+def test_split_recordings():
+    traces = np.zeros((1000, 5), dtype='float64')
+    traces[:] = np.arange(1000)[:, None]
+    sampling_frequency = 30000
+    rec0 = NumpyRecording([traces] * 3, sampling_frequency)
+    
+    rec_list = split_recording(rec0)
+    rec1 = select_segment_recording(rec0, segment_indices=[1])
+    rec2 = rec0.select_segments(segment_indices=[1, 2])
+    rec3 = select_segment_recording(rec0, segment_indices=2)
+
+    for i, rec in enumerate(rec_list):
+        assert rec.get_num_segments() == 1
+        assert np.allclose(rec.get_traces(), rec0.get_traces(segment_index=i))
+        
+    assert rec1.get_num_segments() == 1
+    assert np.allclose(rec1.get_traces(), rec0.get_traces(segment_index=1))
+    
+    assert np.allclose(rec2.get_traces(segment_index=0), rec0.get_traces(segment_index=1))
+    assert np.allclose(rec2.get_traces(segment_index=1), rec0.get_traces(segment_index=2))
+    assert np.allclose(rec3.get_traces(), rec0.get_traces(segment_index=2))
+
+
 def test_append_sortings():
     sampling_frequency = 30000.
     times = np.arange(0, 1000, 10)