Merge with main

Author: alejoe91

doc/how_to/benchmark_with_hybrid_recordings.rst

@@ -2531,9 +2531,8 @@ Although non of the sorters find all units perfectly, ``Kilosort2.5``,
 ``Kilosort4``, and ``SpyKING CIRCUS 2`` all find around 10-12 hybrid
 units with accuracy greater than 80%. ``Kilosort4`` has a better overall
 curve, being able to find almost all units with an accuracy above 50%.
-``Kilosort2.5`` performs well when looking at precision (finding all
-spikes in a hybrid unit), at the cost of lower recall (finding spikes
-when it shouldn’t).
+``Kilosort2.5`` performs well when looking at precision (not finding spikes
+when it shouldn’t), but it has a lower recall (finding all spikes in the ground truth).
 
 In this example, we showed how to:
 

examples/how_to/benchmark_with_hybrid_recordings.py

@@ -276,7 +276,8 @@
 # From the performance plots, we can see that there is no clear "winner", but `Kilosort3` definitely performs worse than the other options.
 #
 # Although non of the sorters find all units perfectly, `Kilosort2.5`, `Kilosort4`, and `SpyKING CIRCUS 2` all find around 10-12 hybrid units with accuracy greater than 80%.
-# `Kilosort4` has a better overall curve, being able to find almost all units with an accuracy above 50%. `Kilosort2.5` performs well when looking at precision (finding all spikes in a hybrid unit), at the cost of lower recall (finding spikes when it shouldn't).
+# `Kilosort4` has a better overall curve, being able to find almost all units with an accuracy above 50%. `Kilosort2.5` performs well when looking at precision (not finding spikes
+# when it shouldn’t), but it has a lower recall (finding all spikes in the ground truth).
 #
 #
 # In this example, we showed how to:

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
     "numpy>=1.20",
     "threadpoolctl>=3.0.0",
     "tqdm",
-    "zarr>=2.16,<2.18",
+    "zarr>=2.18,<3",
     "neo>=0.14.0",
     "probeinterface>=0.2.23",
     "packaging",

src/spikeinterface/benchmark/benchmark_base.py

@@ -11,7 +11,7 @@
 
 from spikeinterface.core import SortingAnalyzer
 
-from spikeinterface import load_extractor, create_sorting_analyzer, load_sorting_analyzer
+from spikeinterface import load, create_sorting_analyzer, load_sorting_analyzer
 from spikeinterface.widgets import get_some_colors
 
 
@@ -150,13 +150,13 @@ def scan_folder(self):
             analyzer = load_sorting_analyzer(folder)
             self.analyzers[key] = analyzer
             # the sorting is in memory here we take the saved one because comparisons need to pickle it later
-            sorting = load_extractor(analyzer.folder / "sorting")
+            sorting = load(analyzer.folder / "sorting")
             self.datasets[key] = analyzer.recording, sorting
 
         # for rec_file in (self.folder / "datasets" / "recordings").glob("*.pickle"):
         #     key = rec_file.stem
-        #     rec = load_extractor(rec_file)
-        #     gt_sorting = load_extractor(self.folder / f"datasets" / "gt_sortings" / key)
+        #     rec = load(rec_file)
+        #     gt_sorting = load(self.folder / f"datasets" / "gt_sortings" / key)
         #     self.datasets[key] = (rec, gt_sorting)
 
         with open(self.folder / "cases.pickle", "rb") as f:
@@ -428,7 +428,7 @@ def load_folder(cls, folder):
             elif format == "sorting":
                 from spikeinterface.core import load_extractor
 
-                result[k] = load_extractor(folder / k)
+                result[k] = load(folder / k)
             elif format == "Motion":
                 from spikeinterface.sortingcomponents.motion import Motion
 

src/spikeinterface/comparison/multicomparisons.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from spikeinterface.core import load_extractor, BaseSorting, BaseSortingSegment
+from spikeinterface.core import load, BaseSorting, BaseSortingSegment
 from spikeinterface.core.core_tools import define_function_from_class
 from .basecomparison import BaseMultiComparison, MixinSpikeTrainComparison, MixinTemplateComparison
 from .paircomparisons import SymmetricSortingComparison, TemplateComparison
@@ -230,7 +230,7 @@ def load_from_folder(folder_path):
         with (folder_path / "sortings.json").open() as f:
             dict_sortings = json.load(f)
         name_list = list(dict_sortings.keys())
-        sorting_list = [load_extractor(v, base_folder=folder_path) for v in dict_sortings.values()]
+        sorting_list = [load(v, base_folder=folder_path) for v in dict_sortings.values()]
         mcmp = MultiSortingComparison(sorting_list=sorting_list, name_list=list(name_list), do_matching=False, **kwargs)
         filename = str(folder_path / "multicomparison.gpickle")
         with open(filename, "rb") as f:

src/spikeinterface/comparison/tests/test_hybrid.py

@@ -1,7 +1,7 @@
 import pytest
 import shutil
 from pathlib import Path
-from spikeinterface.core import extract_waveforms, load_waveforms, load_extractor
+from spikeinterface.core import extract_waveforms, load_waveforms, load
 from spikeinterface.core.testing import check_recordings_equal
 from spikeinterface.comparison import (
     create_hybrid_units_recording,
@@ -52,7 +52,7 @@ def test_hybrid_units_recording(setup_module):
     )
 
     # Check dumpability
-    saved_loaded = load_extractor(hybrid_units_recording.to_dict())
+    saved_loaded = load(hybrid_units_recording.to_dict())
     check_recordings_equal(hybrid_units_recording, saved_loaded, return_scaled=False)
 
     saved_1job = hybrid_units_recording.save(folder=cache_folder / "units_1job")
@@ -81,7 +81,7 @@ def test_hybrid_spikes_recording(setup_module):
     )
 
     # Check dumpability
-    saved_loaded = load_extractor(hybrid_spikes_recording.to_dict())
+    saved_loaded = load(hybrid_spikes_recording.to_dict())
     check_recordings_equal(hybrid_spikes_recording, saved_loaded, return_scaled=False)
 
     saved_1job = hybrid_spikes_recording.save(folder=cache_folder / "spikes_1job")

src/spikeinterface/core/__init__.py

@@ -1,10 +1,11 @@
-from .base import load_extractor  # , load_extractor_from_dict, load_extractor_from_json, load_extractor_from_pickle
 from .baserecording import BaseRecording, BaseRecordingSegment
 from .basesorting import BaseSorting, BaseSortingSegment, SpikeVectorSortingSegment
 from .baseevent import BaseEvent, BaseEventSegment
 from .basesnippets import BaseSnippets, BaseSnippetsSegment
 from .baserecordingsnippets import BaseRecordingSnippets
 
+from .loading import load, load_extractor
+
 # main extractor from dump and cache
 from .binaryrecordingextractor import BinaryRecordingExtractor, read_binary
 from .npzsortingextractor import NpzSortingExtractor, read_npz_sorting

src/spikeinterface/core/base.py

@@ -16,7 +16,7 @@
 
 from .globals import get_global_tmp_folder, is_set_global_tmp_folder
 from .core_tools import (
-    check_json,
+    is_path_remote,
     clean_zarr_folder_name,
     is_dict_extractor,
     SIJsonEncoder,
@@ -673,7 +673,7 @@ def dump_to_json(
     ) -> None:
         """
         Dump recording extractor to json file.
-        The extractor can be re-loaded with load_extractor(json_file)
+        The extractor can be re-loaded with load(json_file)
 
         Parameters
         ----------
@@ -715,7 +715,7 @@ def dump_to_pickle(
     ):
         """
         Dump recording extractor to a pickle file.
-        The extractor can be re-loaded with load_extractor(pickle_file)
+        The extractor can be re-loaded with load(pickle_file)
 
         Parameters
         ----------
@@ -752,7 +752,9 @@ def dump_to_pickle(
         file_path.write_bytes(pickle.dumps(dump_dict))
 
     @staticmethod
-    def load(file_path: Union[str, Path], base_folder: Optional[Union[Path, str, bool]] = None) -> "BaseExtractor":
+    def load(
+        file_or_folder_path: Union[str, Path], base_folder: Optional[Union[Path, str, bool]] = None
+    ) -> "BaseExtractor":
         """
         Load extractor from file path (.json or .pkl)
 
@@ -761,62 +763,10 @@ def load(file_path: Union[str, Path], base_folder: Optional[Union[Path, str, boo
           * save (...)  a folder which contain data  + json (or pickle) + metadata.
 
         """
+        # use loading.py and keep backward compatibility
+        from .loading import load
 
-        file_path = Path(file_path)
-        if base_folder is True:
-            base_folder = file_path.parent
-
-        if file_path.is_file():
-            # standard case based on a file (json or pickle)
-            if str(file_path).endswith(".json"):
-                with open(file_path, "r") as f:
-                    d = json.load(f)
-            elif str(file_path).endswith(".pkl") or str(file_path).endswith(".pickle"):
-                with open(file_path, "rb") as f:
-                    d = pickle.load(f)
-            else:
-                raise ValueError(f"Impossible to load {file_path}")
-            if "warning" in d:
-                print("The extractor was not serializable to file")
-                return None
-
-            extractor = BaseExtractor.from_dict(d, base_folder=base_folder)
-            return extractor
-
-        elif file_path.is_dir():
-            # case from a folder after a calling extractor.save(...)
-            folder = file_path
-            file = None
-
-            if folder.suffix == ".zarr":
-                from .zarrextractors import read_zarr
-
-                extractor = read_zarr(folder)
-            else:
-                # the is spikeinterface<=0.94.0
-                # a folder came with 'cached.json'
-                for dump_ext in ("json", "pkl", "pickle"):
-                    f = folder / f"cached.{dump_ext}"
-                    if f.is_file():
-                        file = f
-
-                # spikeinterface>=0.95.0
-                f = folder / f"si_folder.json"
-                if f.is_file():
-                    file = f
-
-                if file is None:
-                    raise ValueError(f"This folder is not a cached folder {file_path}")
-                extractor = BaseExtractor.load(file, base_folder=folder)
-
-            return extractor
-
-        else:
-            error_msg = (
-                f"{file_path} is not a file or a folder. It should point to either a json, pickle file or a "
-                "folder that is the result of extractor.save(...)"
-            )
-            raise ValueError(error_msg)
+        return load(file_or_folder_path, base_folder=base_folder)
 
     def __reduce__(self):
         """
@@ -1167,50 +1117,6 @@ def _check_same_version(class_string, version):
         return "unknown"
 
 
-def load_extractor(file_or_folder_or_dict, base_folder=None) -> BaseExtractor:
-    """
-    Instantiate extractor from:
-      * a dict
-      * a json file
-      * a pickle file
-      * folder (after save)
-      * a zarr folder (after save)
-
-    Parameters
-    ----------
-    file_or_folder_or_dict : dictionary or folder or file (json, pickle)
-        The file path, folder path, or dictionary to load the extractor from
-    base_folder : str | Path | bool (optional)
-        The base folder to make relative paths absolute.
-        If True and file_or_folder_or_dict is a file, the parent folder of the file is used.
-
-    Returns
-    -------
-    extractor: Recording or Sorting
-        The loaded extractor object
-    """
-    if isinstance(file_or_folder_or_dict, dict):
-        assert not isinstance(base_folder, bool), "`base_folder` must be a string or Path when loading from dict"
-        return BaseExtractor.from_dict(file_or_folder_or_dict, base_folder=base_folder)
-    else:
-        return BaseExtractor.load(file_or_folder_or_dict, base_folder=base_folder)
-
-
-def load_extractor_from_dict(d, base_folder=None) -> BaseExtractor:
-    warnings.warn("Use load_extractor(..) instead")
-    return BaseExtractor.from_dict(d, base_folder=base_folder)
-
-
-def load_extractor_from_json(json_file, base_folder=None) -> "BaseExtractor":
-    warnings.warn("Use load_extractor(..) instead")
-    return BaseExtractor.load(json_file, base_folder=base_folder)
-
-
-def load_extractor_from_pickle(pkl_file, base_folder=None) -> "BaseExtractor":
-    warnings.warn("Use load_extractor(..) instead")
-    return BaseExtractor.load(pkl_file, base_folder=base_folder)
-
-
 class BaseSegment:
     def __init__(self):
         self._parent_extractor = None

src/spikeinterface/core/baserecording.py

@@ -232,15 +232,9 @@ def get_duration(self, segment_index=None) -> float:
         float
             The duration in seconds
         """
-        segment_index = self._check_segment_index(segment_index)
-
-        if self.has_time_vector(segment_index):
-            times = self.get_times(segment_index)
-            segment_duration = times[-1] - times[0] + (1 / self.get_sampling_frequency())
-        else:
-            segment_num_samples = self.get_num_samples(segment_index=segment_index)
-            segment_duration = segment_num_samples / self.get_sampling_frequency()
-
+        segment_duration = (
+            self.get_end_time(segment_index) - self.get_start_time(segment_index) + (1 / self.get_sampling_frequency())
+        )
         return segment_duration
 
     def get_total_duration(self) -> float:
@@ -252,7 +246,7 @@ def get_total_duration(self) -> float:
         float
             The duration in seconds
         """
-        duration = sum([self.get_duration(idx) for idx in range(self.get_num_segments())])
+        duration = sum([self.get_duration(segment_index) for segment_index in range(self.get_num_segments())])
         return duration
 
     def get_memory_size(self, segment_index=None) -> int:
@@ -445,6 +439,40 @@ def get_times(self, segment_index=None) -> np.ndarray:
         times = rs.get_times()
         return times
 
+    def get_start_time(self, segment_index=None) -> float:
+        """Get the start time of the recording 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)
+        rs = self._recording_segments[segment_index]
+        return rs.get_start_time()
+
+    def get_end_time(self, segment_index=None) -> float:
+        """Get the stop time of the recording segment.
+
+        Parameters
+        ----------
+        segment_index : int or None, default: None
+            The segment index (required for multi-segment)
+
+        Returns
+        -------
+        float
+            The stop time in seconds
+        """
+        segment_index = self._check_segment_index(segment_index)
+        rs = self._recording_segments[segment_index]
+        return rs.get_end_time()
+
     def has_time_vector(self, segment_index=None):
         """Check if the segment of the recording has a time vector.
 
@@ -903,6 +931,21 @@ def get_times(self) -> np.ndarray:
                 time_vector += self.t_start
             return time_vector
 
+    def get_start_time(self) -> float:
+        if self.time_vector is not None:
+            return self.time_vector[0]
+        else:
+            return self.t_start if self.t_start is not None else 0.0
+
+    def get_end_time(self) -> float:
+        if self.time_vector is not None:
+            return self.time_vector[-1]
+        else:
+            t_stop = (self.get_num_samples() - 1) / self.sampling_frequency
+            if self.t_start is not None:
+                t_stop += self.t_start
+            return t_stop
+
     def get_times_kwargs(self) -> dict:
         """
         Retrieves the timing attributes characterizing a RecordingSegment