add max_channel column for phy interface

CHANGELOG.md

@@ -9,6 +9,7 @@
 * Support roiextractors 0.5.11 [PR #1236](https://github.com/catalystneuro/neuroconv/pull/1236)
 * Added stub_test option to TDTFiberPhotometryInterface [PR #1242](https://github.com/catalystneuro/neuroconv/pull/1242)
 * Added ThorImagingInterface for Thor TIFF files with OME metadata [PR #1238](https://github.com/catalystneuro/neuroconv/pull/1238)
+* For `PhySortingInterface`, automatically calculate `max_channel` for each unit and adds it to units table. [PR #961](https://github.com/catalystneuro/neuroconv/pull/961)
 
 ## Improvements
 * Filter out warnings for missing timezone information in continuous integration [PR #1240](https://github.com/catalystneuro/neuroconv/pull/1240)

src/neuroconv/datainterfaces/ecephys/phy/phydatainterface.py

@@ -1,8 +1,12 @@
-from typing import Optional
+from pathlib import Path
+from typing import Literal, Optional
 
+import numpy as np
 from pydantic import DirectoryPath, validate_call
+from pynwb.file import NWBFile
 
 from ..basesortingextractorinterface import BaseSortingExtractorInterface
+from ....utils import DeepDict
 
 
 class PhySortingInterface(BaseSortingExtractorInterface):
@@ -24,6 +28,23 @@ def get_source_schema(cls) -> dict:
         ] = "Path to the output Phy folder (containing the params.py)."
         return source_schema
 
+    def get_max_channel(self):
+        folder_path = Path(self.source_data["folder_path"])
+
+        templates = np.load(str(folder_path / "templates.npy"))
+        channel_map = np.load(str(folder_path / "channel_map.npy")).T
+        whitening_mat_inv = np.load(str(folder_path / "whitening_mat_inv.npy"))
+        templates_unwh = templates @ whitening_mat_inv
+
+        cluster_ids = self.sorting_extractor.get_property("original_cluster_id")
+        templates = templates_unwh[cluster_ids]
+
+        max_over_time = np.max(templates, axis=1)
+        idx_max_channel = np.argmax(max_over_time, axis=1)
+        max_channel = channel_map[idx_max_channel].ravel()
+
+        return max_channel
+
     @validate_call
     def __init__(
         self,
@@ -44,6 +65,33 @@ def __init__(
         """
         super().__init__(folder_path=folder_path, exclude_cluster_groups=exclude_cluster_groups, verbose=verbose)
 
+    def add_to_nwbfile(
+        self,
+        nwbfile: NWBFile,
+        metadata: Optional[DeepDict] = None,
+        stub_test: bool = False,
+        write_ecephys_metadata: bool = False,
+        write_as: Literal["units", "processing"] = "units",
+        units_name: str = "units",
+        units_description: str = "Imported from Phy",
+        include_max_channel: bool = True,
+    ):
+        if include_max_channel and "max_channel" not in self.sorting_extractor.get_property_keys():
+            max_channels = self.get_max_channel()
+            self.sorting_extractor.set_property("max_channel", max_channels)
+
+        super().add_to_nwbfile(
+            nwbfile=nwbfile,
+            metadata=metadata,
+            stub_test=stub_test,
+            write_ecephys_metadata=write_ecephys_metadata,
+            write_as=write_as,
+            units_name=units_name,
+            units_description=units_description,
+        )
+
+        return nwbfile
+
     def get_metadata(self):
         metadata = super().get_metadata()
         # See Kilosort save_to_phy() docstring for more info on these fields: https://github.com/MouseLand/Kilosort/blob/main/kilosort/io.py

tests/test_on_data/ecephys/test_sorting_interfaces.py

@@ -1,7 +1,9 @@
 from datetime import datetime
 
 import numpy as np
+from numpy.testing import assert_array_equal
 from pynwb import NWBHDF5IO
+from spikeinterface.extractors.nwbextractors import read_nwbfile
 
 from neuroconv.datainterfaces import (
     BlackrockRecordingInterface,
@@ -198,6 +200,21 @@ class TestPhySortingInterface(SortingExtractorInterfaceTestMixin):
     interface_kwargs = dict(folder_path=str(DATA_PATH / "phy" / "phy_example_0"))
     save_directory = OUTPUT_PATH
 
+    def check_read_nwb(self, nwbfile_path: str):
+        # Test that the max channel is correctly extracted
+        super().check_read_nwb(nwbfile_path)
+
+        # check that the max channel is correctly extracted
+        max_channel = self.interface.get_max_channel()
+        assert_array_equal(max_channel, [1, 2, 5, 5, 6, 21, 13, 13, 21, 21, 22, 22, 24])
+
+        # check that max channel was properly added to sorting extractor
+        assert_array_equal(self.interface.sorting_extractor.get_property("max_channel"), max_channel)
+
+        # check that max channels were properly added to the NWB file
+        nwbfile = read_nwbfile(file_path=nwbfile_path, backend="hdf5")
+        assert_array_equal(nwbfile.units["max_channel"].data[:], max_channel)
+
     def check_extracted_metadata(self, metadata: dict):
         assert metadata["Ecephys"]["UnitProperties"] == [
             dict(name="n_spikes", description="Number of spikes recorded from each unit."),