Fix: provenance handling in SortingAnalyzer

examples/tutorials/core/plot_4_sorting_analyzer.py

@@ -33,6 +33,8 @@
 # First let's generate a simulated recording and sorting
 
 recording, sorting = generate_ground_truth_recording()
+# Make the recording JSON serializable, so it can be reloaded later
+recording = recording.save(folder="recording_folder")
 print(recording)
 print(sorting)
 

src/spikeinterface/core/basesorting.py

@@ -535,6 +535,8 @@ def save(self, format="numpy_folder", **save_kwargs):
             * "numpy_folder" format:
                 - folder : str or Path
                     The folder where the files will be saved.
+                - mmap_mode : str or None, default: None
+                    The memory-mapping mode to use when saving numpy files. If None, no memory mapping is used.
                 - overwrite : bool, default: False
                     If True, existing files in the folder will be overwritten.
             * "zarr" format:

src/spikeinterface/core/binaryfolder.py

@@ -15,7 +15,7 @@
     make_paths_absolute,
     load_properties_from_folder,
     save_properties_to_folder,
-    save_extractor_provenance,
+    save_provenance_to_folder,
     save_annotations_to_folder,
     load_annotations_from_folder,
 )
@@ -160,7 +160,7 @@ def write_recording(
         )
 
         save_properties_to_folder(folder_path / "properties", recording)
-        save_extractor_provenance(folder_path, recording)
+        save_provenance_to_folder(folder_path, recording, relative_to=folder_path)
         # new in version 0.105.0, before that annotations were handle by "si_folder.json" file
         save_annotations_to_folder(folder_path, recording)
 

src/spikeinterface/core/core_tools.py

@@ -934,11 +934,11 @@ def load_annotations_from_folder(folder: str | Path, extractor: "BaseExtractor")
                 extractor._annotations.update(annotations)
 
 
-def save_extractor_provenance(folder: str | Path, extractor: "BaseExtractor"):
+def save_provenance_to_folder(folder: str | Path, extractor: "BaseExtractor", relative_to: str | Path | None = None):
     folder = Path(folder)
     if extractor.check_serializability("json"):
         provenance_file_path = folder / f"provenance.json"
-        extractor.dump_to_json(file_path=provenance_file_path, relative_to=folder)
+        extractor.dump_to_json(file_path=provenance_file_path, relative_to=relative_to)
     elif extractor.check_serializability("pickle"):
         provenance_file = folder / f"provenance.pkl"
         extractor.dump_to_pickle(provenance_file, relative_to=folder)

src/spikeinterface/core/numpyextractors.py

@@ -276,8 +276,7 @@ def __init__(self, spikes, sampling_frequency, unit_ids):
 
         self._serializability["memory"] = True
         self._serializability["json"] = False
-        # theorically this should be False but for simplicity make generators simples we still need this.
-        self._serializability["pickle"] = True
+        self._serializability["pickle"] = False
 
         if spikes.size == 0:
             nseg = 1

src/spikeinterface/core/sortinganalyzer.py

@@ -714,7 +714,6 @@ def create_binary_folder(
         info_file = folder / "spikeinterface_info.json"
         settings_file = folder / "settings.json"
         sorting_folder = folder / "sorting"
-        sorting_provenance_file = folder / "sorting_provenance.json"
         sparsity_file = folder / "sparsity_mask.npy"
         recording_provenance_file = folder / "recording"  # .json (default) or .pickle (if not json-serializable)
         rec_attributes_file = recording_info_folder / "recording_attributes.json"
@@ -731,17 +730,8 @@ def create_binary_folder(
             json.dump(check_json(settings), f, indent=4)
 
         # Save the sorting output
-        sorting = sorting.save(folder=sorting_folder)
-
-        # Dump sorting provenance
-        if sorting.check_serializability("json"):
-            sorting.dump(sorting_provenance_file, relative_to=folder)
-        elif sorting.check_serializability("pickle"):
-            sorting.dump(sorting_provenance_file.with_suffix(".pickle"), relative_to=folder)
-        else:
-            warnings.warn(
-                "The sorting provenance is not serializable! The sorting provenance link will be lost for future load"
-            )
+        mmap_mode = "r" if lazy else None
+        sorting_cached = sorting.save(folder=sorting_folder, relative_to=folder, mmap_mode=mmap_mode)
 
         # Save sparsity
         if sparsity is not None:
@@ -776,7 +766,7 @@ def create_binary_folder(
 
         # Create SortingAnalyzer
         sorting_analyzer = SortingAnalyzer(
-            sorting=sorting,
+            sorting=sorting_cached,
             recording=recording,
             rec_attributes={**rec_attributes_to_save, "probegroup": probegroup},
             format="binary_folder",
@@ -1042,9 +1032,7 @@ def create_zarr(
         backend_options: dict | None,
         lazy: bool = False,
     ) -> "SortingAnalyzer":
-        # used by create and save_as
-
-        from .zarrextractors import add_sorting_to_zarr_group
+        from .zarrextractors import add_sorting_to_zarr_group, ZarrSortingExtractor
 
         is_remote = is_path_remote(folder)
         if not is_remote:
@@ -1071,27 +1059,10 @@ def create_zarr(
         settings = {"return_in_uV": return_in_uV, "peak_sign": peak_sign, "peak_mode": peak_mode}
         zarr_root.attrs["settings"] = check_json(settings)
 
-        # Save the sorting output
-        add_sorting_to_zarr_group(sorting, sorting_group, **saving_options)
-
-        # Dump sorting provenance
         relative_to = None if is_remote else folder
-        sort_dict = sorting.to_dict(relative_to=relative_to, recursive=True)
-        if sorting.check_serializability("json"):
-            _write_object_array(zarr_root, "sorting_provenance", check_json(sort_dict), codec="json")
-        elif sorting.check_serializability("pickle"):
-            try:
-                _write_object_array(zarr_root, "sorting_provenance", sort_dict, codec="pickle")
-            except:
-                warnings.warn(
-                    "Failed to serialize sorting provenance with Pickle Codec! "
-                    "The sorting provenance link will be lost for future load"
-                )
-        else:
-            warnings.warn(
-                "The sorting provenance is not serializable! "
-                "The sorting provenance link will be lost for future load"
-            )
+        # Save the sorting output
+        add_sorting_to_zarr_group(sorting, sorting_group, relative_to=folder, **saving_options)
+        sorting_cached = ZarrSortingExtractor(folder / "sorting", lazy_spike_vector=lazy)
 
         # Save sparsity
         if sparsity is not None:
@@ -1135,7 +1106,7 @@ def create_zarr(
 
         # Create SortingAnalyzer
         sorting_analyzer = SortingAnalyzer(
-            sorting=NumpySorting.from_sorting(sorting, with_metadata=True, copy_spike_vector=True),
+            sorting=sorting_cached,
             recording=recording,
             rec_attributes={**rec_attributes_to_save, "probegroup": probegroup},
             format="zarr",
@@ -2175,34 +2146,36 @@ def get_sorting_provenance(self):
         """
         from .loading import load
 
+        sorting_provenance = None
         if self.format == "memory":
             # the original sorting provenance is not kept in that case
-            sorting_provenance = None
-
+            pass
         elif self.format == "binary_folder":
             for type in ("json", "pickle"):
-                filename = self.folder / f"sorting_provenance.{type}"
+                filename = self.folder / "sorting" / f"provenance.{type}"
                 sorting_provenance = None
+                if not filename.is_file():
+                    # Legacy location
+                    filename = self.folder / "sorting_provenance.json"
+                # try-except here is because it's not required to be able
+                # to load the sorting provenance, as the user might have deleted
+                # the original sorting folder
                 if filename.is_file():
-                    # try-except here is because it's not required to be able
-                    # to load the sorting provenance, as the user might have deleted
-                    # the original sorting folder
                     try:
                         sorting_provenance = load(filename, base_folder=self.folder)
                         break
                     except:
                         pass
-                        # sorting_provenance = None
-
         elif self.format == "zarr":
             zarr_root = self._get_zarr_root(mode="r")
             sorting_provenance = None
-            if "sorting_provenance" in zarr_root.keys():
-                # try-except here is because it's not required to be able
-                # to load the sorting provenance, as the user might have deleted
-                # the original sorting folder
-                try:
+            sort_dict = zarr_root["sorting"].attrs.get("provenance", None)
+            if sort_dict is None:
+                # Legacy
+                if "sorting_provenance" in zarr_root.keys():
                     sort_dict = zarr_root["sorting_provenance"][0]
+            if sort_dict is not None:
+                try:
                     sorting_provenance = load(sort_dict, base_folder=self.folder)
                 except:
                     pass

src/spikeinterface/core/sortingfolder.py

@@ -15,7 +15,7 @@
     save_properties_to_folder,
     save_annotations_to_folder,
     load_annotations_from_folder,
-    save_extractor_provenance,
+    save_provenance_to_folder,
 )
 
 
@@ -61,7 +61,13 @@ def __init__(self, folder_path, mmap_mode: str | None = None):
         self._kwargs = dict(folder_path=str(folder_path.absolute()), mmap_mode=mmap_mode)
 
     @staticmethod
-    def write_sorting(sorting, folder_path, overwrite: bool = False):
+    def write_sorting(
+        sorting,
+        folder_path,
+        overwrite: bool = False,
+        mmap_mode: str | None = None,
+        relative_to: str | Path | None = None,
+    ):
         # the folder can already exists but not contaning numpysorting_info.json
         folder_path = Path(folder_path)
         if folder_path.is_dir():
@@ -81,14 +87,15 @@ def write_sorting(sorting, folder_path, overwrite: bool = False):
         np.save(folder_path / "spikes.npy", sorting.to_spike_vector())
 
         save_properties_to_folder(folder_path / "properties", sorting)
-        save_extractor_provenance(folder_path, sorting)
+        relative_to = folder_path if relative_to is None else relative_to
+        save_provenance_to_folder(folder_path, sorting, relative_to=relative_to)
         # new in version 0.105.0, before that annotations were handle by "si_folder.json" file
         save_annotations_to_folder(folder_path, sorting)
 
         # Create the si_folder file to make the load() easier until version 0.105.0
         # All properties, annotations, and probe information are already saved in the folder,
         # so we don't need to include them in the si_folder.json
-        cached = NumpyFolderSorting(folder_path=folder_path)
+        cached = NumpyFolderSorting(folder_path=folder_path, mmap_mode=mmap_mode)
         si_folder_path = folder_path / f"si_folder.json"
         cached.dump_to_json(
             file_path=si_folder_path,

src/spikeinterface/core/tests/test_sortinganalyzer.py

@@ -92,6 +92,7 @@ def test_SortingAnalyzer_memory(tmp_path, dataset):
 
 def test_SortingAnalyzer_binary_folder(tmp_path, dataset):
     recording, sorting = dataset
+    recording = recording.save(folder=tmp_path / "recording_binary")
 
     folder = tmp_path / "test_SortingAnalyzer_binary_folder"
     if folder.exists():
@@ -138,6 +139,7 @@ def test_SortingAnalyzer_binary_folder(tmp_path, dataset):
 
 def test_SortingAnalyzer_zarr(tmp_path, dataset):
     recording, sorting = dataset
+    recording = recording.save(folder=tmp_path / "recording_zarr")
 
     folder = tmp_path / "test_SortingAnalyzer_zarr.zarr"
 

src/spikeinterface/core/zarrextractors.py

@@ -453,6 +453,9 @@ def __array__(self, dtype=None):
             arr["segment_index"][s0:s1] = seg_idx
         return arr if dtype is None else arr.astype(dtype)
 
+    def copy(self):
+        return np.copy(np.asarray(self))
+
 
 class ZarrSortingExtractor(BaseSorting):
     """
@@ -559,6 +562,7 @@ def write_sorting(
         folder_path: str | Path,
         overwrite: bool = False,
         storage_options: dict | None = None,
+        relative_to: str | Path | None = None,
         **kwargs,
     ):
         """
@@ -567,7 +571,7 @@ def write_sorting(
         folder_path = create_zarr_path_for_write(folder_path, overwrite=overwrite)
         zarr_root = zarr.open(str(folder_path), mode="w", storage_options=storage_options)
         zarr_root.attrs["zarr_class_info"] = retrieve_importing_provenance(ZarrSortingExtractor)
-        add_sorting_to_zarr_group(sorting, zarr_root, **kwargs)
+        add_sorting_to_zarr_group(sorting, zarr_root, relative_to=relative_to, **kwargs)
         return ZarrSortingExtractor(folder_path, storage_options=storage_options)
 
 
@@ -734,7 +738,9 @@ def add_properties_and_annotations(zarr_group: zarr.Group, recording_or_sorting:
     zarr_group.attrs["annotations"] = check_json(recording_or_sorting._annotations)
 
 
-def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kwargs):
+def add_sorting_to_zarr_group(
+    sorting: BaseSorting, zarr_group: zarr.Group, relative_to: str | Path | None = None, **kwargs
+):
     """
     Add a sorting extractor to a zarr group.
 
@@ -749,6 +755,11 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw
     """
     from numcodecs import Delta
 
+    if sorting.check_serializability("json"):
+        zarr_group.attrs["provenance"] = check_json(sorting.to_dict(recursive=True, relative_to=relative_to))
+    else:
+        zarr_group.attrs["provenance"] = None
+
     num_segments = sorting.get_num_segments()
     zarr_group.attrs["sampling_frequency"] = float(sorting.sampling_frequency)
     zarr_group.attrs["num_segments"] = int(num_segments)
@@ -777,11 +788,18 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw
     add_properties_and_annotations(zarr_group, sorting)
 
 
-def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group, verbose=False, dtype=None, **kwargs):
+def add_recording_to_zarr_group(
+    recording: BaseRecording,
+    zarr_group: zarr.Group,
+    verbose=False,
+    dtype=None,
+    relative_to: str | Path | None = None,
+    **kwargs,
+):
     zarr_kwargs, job_kwargs = split_job_kwargs(kwargs)
 
     if recording.check_serializability("json"):
-        zarr_group.attrs["provenance"] = check_json(recording.to_dict(recursive=True))
+        zarr_group.attrs["provenance"] = check_json(recording.to_dict(recursive=True, relative_to=relative_to))
     else:
         zarr_group.attrs["provenance"] = None
 

src/spikeinterface/metrics/conftest.py

@@ -8,12 +8,13 @@
 job_kwargs = dict(n_jobs=2, progress_bar=True, chunk_duration="1s")
 
 
-def make_small_analyzer():
+def make_small_analyzer(folder):
     recording, sorting = generate_ground_truth_recording(
         durations=[10.0],
         num_units=10,
         seed=1205,
     )
+    recording = recording.save(folder=folder / "rec_for_small_analyzer", format="binary")
 
     sorting = sorting.select_units(["2", "7", "0"], ["#3", "#9", "#4"])
 
@@ -36,12 +37,12 @@ def make_small_analyzer():
 
 
 @pytest.fixture(scope="module")
-def small_sorting_analyzer():
-    return make_small_analyzer()
+def small_sorting_analyzer(create_cache_folder):
+    return make_small_analyzer(create_cache_folder)
 
 
 @pytest.fixture(scope="module")
-def sorting_analyzer_simple():
+def sorting_analyzer_simple(create_cache_folder):
     # we need high firing rate for amplitude_cutoff
     recording, sorting = generate_ground_truth_recording(
         durations=[
@@ -64,6 +65,7 @@ def sorting_analyzer_simple():
         noise_kwargs=dict(noise_levels=5.0, strategy="tile_pregenerated"),
         seed=1205,
     )
+    recording = recording.save(folder=create_cache_folder / "rec_for_sorting_analyzer_simple", format="binary")
 
     sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=True)