Repository: SpikeInterface/spikeinterface

Issue:
Storing channel index from MDA sorting in MdaSortingExtractor for downstream use
I am using spikeinterface and NeuroConv to convert a firings.mda sorting file from an older version of MountainSort to an NWB file. Currently this creates an NWB Units table with spike times and unit labels from MountainSort, but the table lacks information about which channel a unit was maximally detected on, which is available in the MDA file. NWB has an "electrodes" column in the Units table to store that.

`MdaSortingSegment._max_channels` (https://github.com/SpikeInterface/spikeinterface/blob/dd0c6d33f98f071362fa06abb631a7629b97de58/src/spikeinterface/extractors/mdaextractors.py) stores channel/electrode IDs for each spike event from the MDA file but those are not used anywhere.

In the MDA files that I have, for each unit, every spike assigned to that unit has the same max channel index. Is that always the case, @magland?

If so, I propose making `MdaSortingSegment._max_channels` public, as in drop the leading underscore in the variable name, so that downstream tools, e.g., NeuroConv, can use that information, e.g., when writing the sorting to an NWB file.

Does that sound reasonable? I am not an expert in spikeinterface and neuroconv. If there is a better approach, please let me know. Thanks!

Reference for firings.mda file format: https://mountainsort.readthedocs.io/en/latest/first_sort.html#format-of-the-firings-mda

Hints from issue comments:
@rly: Also, if I have only one `firings.mda` file for a session, and I am not concatenating across sessions, should there always be only one `MdaSortingSegment` per `MdaSortingExtractor`?@alejoe91: Hi @rly 

Second question first: there is always a Segment. That's the class in charge of retrieving data. Most Extractors are in fact mono-segment, but this organization allows us also to seamlessly deal with multi-segments.

First question:
Currently, there is no way in SI to attach spike vectors to spike trains. Since max channels are defined for each spike and might change over time, we are not attaching it to the object. 
If you need a "mean max channel", the easiest option is to get the max channel with the largest count for each unit an use the `set_property`. Properties are automatically saved as NWB unit columns by neuroconv. 
I wouldn't save them to the `electrodes` column though by default. We normally use that column when we also save waflceform means and stds to point to the electrodes table.@magland: > In the MDA files that I have, for each unit, every spike assigned to that unit has the same max channel index. Is that always the case, @magland?

Yes, that is always the case. @rly @alejoe91: Ok, then we can just add max channel as a property :)@rly: Great. Thanks @magland and @alejoe91 ! I can give this a try

Solution patch:
diff --git a/src/spikeinterface/extractors/mdaextractors.py b/src/spikeinterface/extractors/mdaextractors.py
index d2886d9e7..c1fc1144d 100644
--- a/src/spikeinterface/extractors/mdaextractors.py
+++ b/src/spikeinterface/extractors/mdaextractors.py
@@ -180,6 +180,13 @@ class MdaRecordingSegment(BaseRecordingSegment):
 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,6 +209,19 @@ class MdaSortingExtractor(BaseSorting):
         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,
@@ -209,7 +229,7 @@ class MdaSortingExtractor(BaseSorting):
 
     @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 @@ class MdaSortingExtractor(BaseSorting):
             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(


Test patch:
diff --git a/src/spikeinterface/extractors/tests/test_mdaextractors.py b/src/spikeinterface/extractors/tests/test_mdaextractors.py
index 78e6afb65..ef8c73ac6 100644
--- a/src/spikeinterface/extractors/tests/test_mdaextractors.py
+++ b/src/spikeinterface/extractors/tests/test_mdaextractors.py
@@ -21,6 +21,7 @@ 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()
@@ -28,6 +29,17 @@ def test_mda_extractors(create_cache_folder):
 
     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()