add max_channel column for phy interface

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

@@ -1,7 +1,11 @@
-from typing import Optional
+from pathlib import Path
+from typing import Optional, Literal
+
+import numpy as np
+from pynwb.file import NWBFile
 
 from ..basesortingextractorinterface import BaseSortingExtractorInterface
-from ....utils import FolderPathType
+from ....utils import FolderPathType, DeepDict
 
 
 class PhySortingInterface(BaseSortingExtractorInterface):
@@ -23,6 +27,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'))
+        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
+
     def __init__(
         self,
         folder_path: FolderPathType,
@@ -41,3 +62,29 @@ def __init__(
         verbose : bool, default: True
         """
         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 = "Autogenerated by neuroconv.",
+    ):
+
+        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,
+        )
+
+        max_channel = self.get_max_channel()
+        nwbfile.units.add_column(name='max_channel', description='Channel with maximum amplitude', data=max_channel)
+
+        return nwbfile