Repository: SpikeInterface/spikeinterface
Pull request #2785 (2024-04-30 12:19:45)

Analyzer set recording
This is useful for example if the analyzer was set up and ran with a tmp cached recording, but later on one wants to set a lazy recording to perform additional operations and computations


Patch:
diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py
index c32baf9d5..024382dea 100644
--- a/src/spikeinterface/core/recording_tools.py
+++ b/src/spikeinterface/core/recording_tools.py
@@ -917,3 +917,34 @@ def get_rec_attributes(recording):
         dtype=recording.get_dtype(),
     )
     return rec_attributes
+
+
+def do_recording_attributes_match(recording1, recording2_attributes) -> bool:
+    """
+    Check if two recordings have the same attributes
+
+    Parameters
+    ----------
+    recording1 : BaseRecording
+        The first recording object
+    recording2_attributes : dict
+        The recording attributes to test against
+
+    Returns
+    -------
+    bool
+        True if the recordings have the same attributes
+    """
+    recording1_attributes = get_rec_attributes(recording1)
+    recording2_attributes = deepcopy(recording2_attributes)
+    recording1_attributes.pop("properties")
+    recording2_attributes.pop("properties")
+
+    return (
+        np.array_equal(recording1_attributes["channel_ids"], recording2_attributes["channel_ids"])
+        and recording1_attributes["sampling_frequency"] == recording2_attributes["sampling_frequency"]
+        and recording1_attributes["num_channels"] == recording2_attributes["num_channels"]
+        and recording1_attributes["num_samples"] == recording2_attributes["num_samples"]
+        and recording1_attributes["is_filtered"] == recording2_attributes["is_filtered"]
+        and recording1_attributes["dtype"] == recording2_attributes["dtype"]
+    )
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index c541634e9..46d02099d 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/src/spikeinterface/core/sortinganalyzer.py
@@ -22,7 +22,7 @@ from .baserecording import BaseRecording
 from .basesorting import BaseSorting
 
 from .base import load_extractor
-from .recording_tools import check_probe_do_not_overlap, get_rec_attributes
+from .recording_tools import check_probe_do_not_overlap, get_rec_attributes, do_recording_attributes_match
 from .core_tools import check_json, retrieve_importing_provenance
 from .job_tools import split_job_kwargs
 from .numpyextractors import NumpySorting
@@ -203,6 +203,8 @@ class SortingAnalyzer:
         self.format = format
         self.sparsity = sparsity
         self.return_scaled = return_scaled
+        # this is used to store temporary recording
+        self._temporary_recording = None
 
         # extensions are not loaded at init
         self.extensions = dict()
@@ -605,13 +607,37 @@ class SortingAnalyzer:
 
         return sorting_analyzer
 
+    def set_temporary_recording(self, recording: BaseRecording):
+        """
+        Sets a temporary recording object. This function can be useful to temporarily set
+        a "cached" recording object that is not saved in the SortingAnalyzer object to speed up
+        computations. Upon reloading, the SortingAnalyzer object will try to reload the recording
+        from the original location in a lazy way.
+
+
+        Parameters
+        ----------
+        recording : BaseRecording
+            The recording object to set as temporary recording.
+        """
+        # check that recording is compatible
+        assert do_recording_attributes_match(recording, self.rec_attributes), "Recording attributes do not match."
+        assert np.array_equal(
+            recording.get_channel_locations(), self.get_channel_locations()
+        ), "Recording channel locations do not match."
+        if self._recording is not None:
+            warnings.warn("SortingAnalyzer recording is already set. The current recording is temporarily replaced.")
+        self._temporary_recording = recording
+
     def _save_or_select(self, format="binary_folder", folder=None, unit_ids=None) -> "SortingAnalyzer":
         """
         Internal used by both save_as(), copy() and select_units() which are more or less the same.
         """
 
         if self.has_recording():
-            recording = self.recording
+            recording = self._recording
+        elif self.has_temporary_recording():
+            recording = self._temporary_recording
         else:
             recording = None
 
@@ -728,9 +754,9 @@ class SortingAnalyzer:
 
     @property
     def recording(self) -> BaseRecording:
-        if not self.has_recording():
+        if not self.has_recording() and not self.has_temporary_recording():
             raise ValueError("SortingAnalyzer could not load the recording")
-        return self._recording
+        return self._temporary_recording or self._recording
 
     @property
     def channel_ids(self) -> np.ndarray:
@@ -747,6 +773,9 @@ class SortingAnalyzer:
     def has_recording(self) -> bool:
         return self._recording is not None
 
+    def has_temporary_recording(self) -> bool:
+        return self._temporary_recording is not None
+
     def is_sparse(self) -> bool:
         return self.sparsity is not None
 


Test patch:
diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py
index 13e01c32d..d78093214 100644
--- a/src/spikeinterface/core/tests/test_sortinganalyzer.py
+++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py
@@ -32,8 +32,13 @@ def get_dataset():
     return recording, sorting
 
 
-def test_SortingAnalyzer_memory(tmp_path):
-    recording, sorting = get_dataset()
+@pytest.fixture(scope="module")
+def dataset():
+    return get_dataset()
+
+
+def test_SortingAnalyzer_memory(tmp_path, dataset):
+    recording, sorting = dataset
     sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=False, sparsity=None)
     _check_sorting_analyzers(sorting_analyzer, sorting, cache_folder=tmp_path)
 
@@ -52,8 +57,8 @@ def test_SortingAnalyzer_memory(tmp_path):
     assert not sorting_analyzer.return_scaled
 
 
-def test_SortingAnalyzer_binary_folder(tmp_path):
-    recording, sorting = get_dataset()
+def test_SortingAnalyzer_binary_folder(tmp_path, dataset):
+    recording, sorting = dataset
 
     folder = tmp_path / "test_SortingAnalyzer_binary_folder"
     if folder.exists():
@@ -82,8 +87,8 @@ def test_SortingAnalyzer_binary_folder(tmp_path):
     _check_sorting_analyzers(sorting_analyzer, sorting, cache_folder=tmp_path)
 
 
-def test_SortingAnalyzer_zarr(tmp_path):
-    recording, sorting = get_dataset()
+def test_SortingAnalyzer_zarr(tmp_path, dataset):
+    recording, sorting = dataset
 
     folder = tmp_path / "test_SortingAnalyzer_zarr.zarr"
     if folder.exists():
@@ -103,10 +108,27 @@ def test_SortingAnalyzer_zarr(tmp_path):
     )
 
 
-def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder):
+def test_SortingAnalyzer_tmp_recording(dataset):
+    recording, sorting = dataset
+    recording_cached = recording.save(mode="memory")
 
-    print()
-    print(sorting_analyzer)
+    sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=False, sparsity=None)
+    sorting_analyzer.set_temporary_recording(recording_cached)
+    assert sorting_analyzer.has_temporary_recording()
+    # check that saving as uses the original recording
+    sorting_analyzer_saved = sorting_analyzer.save_as(format="memory")
+    assert sorting_analyzer_saved.has_recording()
+    assert not sorting_analyzer_saved.has_temporary_recording()
+    assert isinstance(sorting_analyzer_saved.recording, type(recording))
+
+    recording_sliced = recording.channel_slice(recording.channel_ids[:-1])
+
+    # wrong channels
+    with pytest.raises(AssertionError):
+        sorting_analyzer.set_temporary_recording(recording_sliced)
+
+
+def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder):
 
     register_result_extension(DummyAnalyzerExtension)
 
@@ -290,8 +312,10 @@ def test_extensions_sorting():
 
 if __name__ == "__main__":
     tmp_path = Path("test_SortingAnalyzer")
-    test_SortingAnalyzer_memory(tmp_path)
-    test_SortingAnalyzer_binary_folder(tmp_path)
-    test_SortingAnalyzer_zarr(tmp_path)
+    dataset = _get_dataset()
+    test_SortingAnalyzer_memory(tmp_path, dataset)
+    test_SortingAnalyzer_binary_folder(tmp_path, dataset)
+    test_SortingAnalyzer_zarr(tmp_path, dataset)
+    test_SortingAnalyzer_tmp_recording(dataset)
     test_extension()
     test_extension_params()