Merge branch 'main' into curation-pot-merges

Author: zm711

src/spikeinterface/comparison/paircomparisons.py

@@ -456,12 +456,12 @@ def print_summary(self, well_detected_score=None, redundant_score=None, overmerg
             num_gt=len(self.unit1_ids),
             num_tested=len(self.unit2_ids),
             num_well_detected=self.count_well_detected_units(well_detected_score),
-            num_redundant=self.count_redundant_units(redundant_score),
-            num_overmerged=self.count_overmerged_units(overmerged_score),
         )
 
         if self.exhaustive_gt:
             txt = txt + _template_summary_part2
+            d["num_redundant"] = self.count_redundant_units(redundant_score)
+            d["num_overmerged"] = self.count_overmerged_units(overmerged_score)
             d["num_false_positive_units"] = self.count_false_positive_units()
             d["num_bad"] = self.count_bad_units()
 
@@ -676,11 +676,11 @@ def count_units_categories(
 GT num_units: {num_gt}
 TESTED num_units: {num_tested}
 num_well_detected: {num_well_detected}
-num_redundant: {num_redundant}
-num_overmerged: {num_overmerged}
 """
 
-_template_summary_part2 = """num_false_positive_units {num_false_positive_units}
+_template_summary_part2 = """num_redundant: {num_redundant}
+num_overmerged: {num_overmerged}
+num_false_positive_units {num_false_positive_units}
 num_bad: {num_bad}
 """
 

src/spikeinterface/core/base.py

@@ -17,6 +17,7 @@
 from .globals import get_global_tmp_folder, is_set_global_tmp_folder
 from .core_tools import (
     check_json,
+    clean_zarr_folder_name,
     is_dict_extractor,
     SIJsonEncoder,
     make_paths_relative,
@@ -1061,9 +1062,7 @@ def save_to_zarr(
                         print(f"Use zarr_path={zarr_path}")
         else:
             if storage_options is None:
-                folder = Path(folder)
-                if folder.suffix != ".zarr":
-                    folder = folder.parent / f"{folder.stem}.zarr"
+                folder = clean_zarr_folder_name(folder)
                 if folder.is_dir() and overwrite:
                     shutil.rmtree(folder)
                 zarr_path = folder

src/spikeinterface/core/core_tools.py

@@ -153,6 +153,13 @@ def check_json(dictionary: dict) -> dict:
     return json.loads(json_string)
 
 
+def clean_zarr_folder_name(folder):
+    folder = Path(folder)
+    if folder.suffix != ".zarr":
+        folder = folder.parent / f"{folder.stem}.zarr"
+    return folder
+
+
 def add_suffix(file_path, possible_suffix):
     file_path = Path(file_path)
     if isinstance(possible_suffix, str):

src/spikeinterface/core/recording_tools.py

@@ -929,7 +929,9 @@ def get_rec_attributes(recording):
     return rec_attributes
 
 
-def do_recording_attributes_match(recording1, recording2_attributes) -> bool:
+def do_recording_attributes_match(
+    recording1: "BaseRecording", recording2_attributes: bool, check_dtype: bool = True
+) -> tuple[bool, str]:
     """
     Check if two recordings have the same attributes
 
@@ -939,22 +941,43 @@ def do_recording_attributes_match(recording1, recording2_attributes) -> bool:
         The first recording object
     recording2_attributes : dict
         The recording attributes to test against
+    check_dtype : bool, default: True
+        If True, check if the recordings have the same dtype
 
     Returns
     -------
     bool
         True if the recordings have the same attributes
+    str
+        A string with the exception message with the attributes that do not match
     """
     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"]
-    )
+    attributes_match = True
+    non_matching_attrs = []
+
+    if not np.array_equal(recording1_attributes["channel_ids"], recording2_attributes["channel_ids"]):
+        non_matching_attrs.append("channel_ids")
+    if not recording1_attributes["sampling_frequency"] == recording2_attributes["sampling_frequency"]:
+        non_matching_attrs.append("sampling_frequency")
+    if not recording1_attributes["num_channels"] == recording2_attributes["num_channels"]:
+        non_matching_attrs.append("num_channels")
+    if not recording1_attributes["num_samples"] == recording2_attributes["num_samples"]:
+        non_matching_attrs.append("num_samples")
+    # dtype is optional
+    if "dtype" in recording1_attributes and "dtype" in recording2_attributes:
+        if check_dtype:
+            if not recording1_attributes["dtype"] == recording2_attributes["dtype"]:
+                non_matching_attrs.append("dtype")
+
+    if len(non_matching_attrs) > 0:
+        attributes_match = False
+        exception_str = f"Recordings do not match in the following attributes: {non_matching_attrs}"
+    else:
+        attributes_match = True
+        exception_str = ""
+
+    return attributes_match, exception_str

src/spikeinterface/core/sortinganalyzer.py

@@ -23,7 +23,7 @@
 
 from .base import load_extractor
 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, is_path_remote
+from .core_tools import check_json, retrieve_importing_provenance, is_path_remote, clean_zarr_folder_name
 from .sorting_tools import generate_unit_ids_for_merge_group, _get_ids_after_merging
 from .job_tools import split_job_kwargs
 from .numpyextractors import NumpySorting
@@ -111,6 +111,8 @@ def create_sorting_analyzer(
     sparsity off (or give external sparsity) like this.
     """
     if format != "memory":
+        if format == "zarr":
+            folder = clean_zarr_folder_name(folder)
         if Path(folder).is_dir():
             if not overwrite:
                 raise ValueError(f"Folder already exists {folder}! Use overwrite=True to overwrite it.")
@@ -162,6 +164,8 @@ def load_sorting_analyzer(folder, load_extensions=True, format="auto"):
         The loaded SortingAnalyzer
 
     """
+    if format == "zarr":
+        folder = clean_zarr_folder_name(folder)
     return SortingAnalyzer.load(folder, load_extensions=load_extensions, format=format)
 
 
@@ -269,6 +273,8 @@ def create(
             sorting_analyzer = cls.load_from_binary_folder(folder, recording=recording)
             sorting_analyzer.folder = Path(folder)
         elif format == "zarr":
+            assert folder is not None, "For format='zarr' folder must be provided"
+            folder = clean_zarr_folder_name(folder)
             cls.create_zarr(folder, sorting, recording, sparsity, return_scaled, rec_attributes=None)
             sorting_analyzer = cls.load_from_zarr(folder, recording=recording)
             sorting_analyzer.folder = Path(folder)
@@ -487,10 +493,7 @@ def create_zarr(cls, folder, sorting, recording, sparsity, return_scaled, rec_at
         import zarr
         import numcodecs
 
-        folder = Path(folder)
-        # force zarr sufix
-        if folder.suffix != ".zarr":
-            folder = folder.parent / f"{folder.stem}.zarr"
+        folder = clean_zarr_folder_name(folder)
 
         if folder.is_dir():
             raise ValueError(f"Folder already exists {folder}")
@@ -613,7 +616,7 @@ def load_from_zarr(cls, folder, recording=None, storage_options=None):
 
         return sorting_analyzer
 
-    def set_temporary_recording(self, recording: BaseRecording):
+    def set_temporary_recording(self, recording: BaseRecording, check_dtype: bool = True):
         """
         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
@@ -625,12 +628,17 @@ def set_temporary_recording(self, recording: BaseRecording):
         ----------
         recording : BaseRecording
             The recording object to set as temporary recording.
+        check_dtype : bool, default: True
+            If True, check that the dtype of the temporary recording is the same as the original 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."
+        attributes_match, exception_str = do_recording_attributes_match(
+            recording, self.rec_attributes, check_dtype=check_dtype
+        )
+        if not attributes_match:
+            raise ValueError(exception_str)
+        if not np.array_equal(recording.get_channel_locations(), self.get_channel_locations()):
+            raise ValueError("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
@@ -768,9 +776,7 @@ def _save_or_select_or_merge(
 
         elif format == "zarr":
             assert folder is not None, "For format='zarr' folder must be provided"
-            folder = Path(folder)
-            if folder.suffix != ".zarr":
-                folder = folder.parent / f"{folder.stem}.zarr"
+            folder = clean_zarr_folder_name(folder)
             SortingAnalyzer.create_zarr(
                 folder, sorting_provenance, recording, sparsity, self.return_scaled, self.rec_attributes
             )
@@ -829,6 +835,8 @@ def save_as(self, format="memory", folder=None) -> "SortingAnalyzer":
         format : "memory" | "binary_folder" | "zarr", default: "memory"
             The new backend format to use
         """
+        if format == "zarr":
+            folder = clean_zarr_folder_name(folder)
         return self._save_or_select_or_merge(format=format, folder=folder)
 
     def select_units(self, unit_ids, format="memory", folder=None) -> "SortingAnalyzer":
@@ -854,6 +862,8 @@ def select_units(self, unit_ids, format="memory", folder=None) -> "SortingAnalyz
             The newly create sorting_analyzer with the selected units
         """
         # TODO check that unit_ids are in same order otherwise many extension do handle it properly!!!!
+        if format == "zarr":
+            folder = clean_zarr_folder_name(folder)
         return self._save_or_select_or_merge(format=format, folder=folder, unit_ids=unit_ids)
 
     def remove_units(self, remove_unit_ids, format="memory", folder=None) -> "SortingAnalyzer":
@@ -880,6 +890,8 @@ def remove_units(self, remove_unit_ids, format="memory", folder=None) -> "Sortin
         """
         # TODO check that unit_ids are in same order otherwise many extension do handle it properly!!!!
         unit_ids = self.unit_ids[~np.isin(self.unit_ids, remove_unit_ids)]
+        if format == "zarr":
+            folder = clean_zarr_folder_name(folder)
         return self._save_or_select_or_merge(format=format, folder=folder, unit_ids=unit_ids)
 
     def merge_units(
@@ -938,6 +950,9 @@ def merge_units(
             The newly create `SortingAnalyzer` with the selected units
         """
 
+        if format == "zarr":
+            folder = clean_zarr_folder_name(folder)
+
         assert merging_mode in ["soft", "hard"], "Merging mode should be either soft or hard"
 
         if len(merge_unit_groups) == 0:
@@ -1016,6 +1031,9 @@ def has_temporary_recording(self) -> bool:
     def is_sparse(self) -> bool:
         return self.sparsity is not None
 
+    def is_filtered(self) -> bool:
+        return self.rec_attributes["is_filtered"]
+
     def get_sorting_provenance(self):
         """
         Get the original sorting if possible otherwise return None

src/spikeinterface/core/tests/test_recording_tools.py

@@ -17,6 +17,8 @@
     get_channel_distances,
     get_noise_levels,
     order_channels_by_depth,
+    do_recording_attributes_match,
+    get_rec_attributes,
 )
 
 
@@ -300,6 +302,35 @@ def test_order_channels_by_depth():
     assert np.array_equal(order_2d[::-1], order_2d_fliped)
 
 
+def test_do_recording_attributes_match():
+    recording = NoiseGeneratorRecording(
+        num_channels=2, durations=[10.325, 3.5], sampling_frequency=30_000, strategy="tile_pregenerated"
+    )
+    rec_attributes = get_rec_attributes(recording)
+    do_match, _ = do_recording_attributes_match(recording, rec_attributes)
+    assert do_match
+
+    rec_attributes = get_rec_attributes(recording)
+    rec_attributes["sampling_frequency"] = 1.0
+    do_match, exc = do_recording_attributes_match(recording, rec_attributes)
+    assert not do_match
+    assert "sampling_frequency" in exc
+
+    # check dtype options
+    rec_attributes = get_rec_attributes(recording)
+    rec_attributes["dtype"] = "int16"
+    do_match, exc = do_recording_attributes_match(recording, rec_attributes)
+    assert not do_match
+    assert "dtype" in exc
+    do_match, exc = do_recording_attributes_match(recording, rec_attributes, check_dtype=False)
+    assert do_match
+
+    # check missing dtype
+    rec_attributes.pop("dtype")
+    do_match, exc = do_recording_attributes_match(recording, rec_attributes)
+    assert do_match
+
+
 if __name__ == "__main__":
     # Create a temporary folder using the standard library
     import tempfile

src/spikeinterface/core/tests/test_sortinganalyzer.py

@@ -141,7 +141,7 @@ def test_SortingAnalyzer_tmp_recording(dataset):
     recording_sliced = recording.channel_slice(recording.channel_ids[:-1])
 
     # wrong channels
-    with pytest.raises(AssertionError):
+    with pytest.raises(ValueError):
         sorting_analyzer.set_temporary_recording(recording_sliced)
 
 

src/spikeinterface/exporters/to_phy.py

@@ -167,7 +167,7 @@ def export_to_phy(
         f.write(f"dtype = '{dtype_str}'\n")
         f.write(f"offset = 0\n")
         f.write(f"sample_rate = {fs}\n")
-        f.write(f"hp_filtered = {sorting_analyzer.recording.is_filtered()}")
+        f.write(f"hp_filtered = {sorting_analyzer.is_filtered()}")
 
     # export spike_times/spike_templates/spike_clusters
     # here spike_labels is a remapping to unit_index