Merge pull request #3701 from rly/mda_channels

Set max_channel property of MdaSortingExtractor

Author: alejoe91

src/spikeinterface/extractors/mdaextractors.py

@@ -180,6 +180,13 @@ def get_traces(
 class MdaSortingExtractor(BaseSorting):
     """Load MDA format data as a sorting extractor.
 
+    NOTE: As in the MDA format, the max_channel property indexes the channels that are given as input
+    to the sorter.
+    If sorting was run on a subset of channels of the recording, then the max_channel values are
+    based on that subset, so care must be taken when associating these values with a recording.
+    If additional sorting segments are added to this sorting extractor after initialization,
+    then max_channel will not be updated. The max_channel indices begin at 1.
+
     Parameters
     ----------
     file_path : str or Path
@@ -202,14 +209,27 @@ def __init__(self, file_path, sampling_frequency):
         sorting_segment = MdaSortingSegment(firings)
         self.add_sorting_segment(sorting_segment)
 
+        # Store the max channel for each unit
+        # Every spike assigned to a unit (label) has the same max channel
+        # ref: https://github.com/SpikeInterface/spikeinterface/issues/3695#issuecomment-2663329006
+        max_channels = []
+        segment = self._sorting_segments[0]
+        for unit_id in self.unit_ids:
+            label_mask = segment._labels == unit_id
+            # since all max channels are the same, we can just grab the first occurrence for the unit
+            max_channel = segment._max_channels[label_mask][0]
+            max_channels.append(max_channel)
+
+        self.set_property(key="max_channel", values=max_channels)
+
         self._kwargs = {
             "file_path": str(Path(file_path).absolute()),
             "sampling_frequency": sampling_frequency,
         }
 
     @staticmethod
     def write_sorting(sorting, save_path, write_primary_channels=False):
-        assert sorting.get_num_segments() == 1, "MdaSorting.write_sorting() can only write a single segment " "sorting"
+        assert sorting.get_num_segments() == 1, "MdaSorting.write_sorting() can only write a single segment sorting"
         unit_ids = sorting.get_unit_ids()
         times_list = []
         labels_list = []
@@ -223,7 +243,7 @@ def write_sorting(sorting, save_path, write_primary_channels=False):
             else:
                 labels_list.append(np.ones(times.shape, dtype=int) * unit_index)
             if write_primary_channels:
-                if "max_channel" in sorting.get_unit_property_names(unit_id):
+                if "max_channel" in sorting.get_property_keys():
                     primary_channels_list.append([sorting.get_unit_property(unit_id, "max_channel")] * times.shape[0])
                 else:
                     raise ValueError(

src/spikeinterface/extractors/tests/test_mdaextractors.py

@@ -21,13 +21,25 @@ def test_mda_extractors(create_cache_folder):
 
     check_recordings_equal(rec, rec_mda, return_scaled=False)
 
+    # Write without setting max_channel
     MdaSortingExtractor.write_sorting(sort, cache_folder / "mdatest" / "firings.mda")
     sort_mda = MdaSortingExtractor(
         cache_folder / "mdatest" / "firings.mda", sampling_frequency=sort.get_sampling_frequency()
     )
 
     check_sortings_equal(sort, sort_mda)
 
+    # Set a fake max channel (1-indexed) for each unit
+    sort.set_property(key="max_channel", values=[i % rec.get_num_channels() + 1 for i in range(sort.get_num_units())])
+
+    # Write with setting max_channel
+    MdaSortingExtractor.write_sorting(sort, cache_folder / "mdatest" / "firings.mda", write_primary_channels=True)
+    sort_mda = MdaSortingExtractor(
+        cache_folder / "mdatest" / "firings.mda", sampling_frequency=sort.get_sampling_frequency()
+    )
+
+    check_sortings_equal(sort, sort_mda)
+
 
 if __name__ == "__main__":
     test_mda_extractors()