Moved previous changes to kilosortbase.py to kilosort4.py because they apply only to Kilosort4.

Author: Nikhil Chandra

src/spikeinterface/sorters/external/kilosort4.py

@@ -138,7 +138,7 @@ def check_sorter_version(cls):
 
     @classmethod
     def _setup_recording(cls, recording, sorter_output_folder, params, verbose):
-        KilosortBase._generate_channel_map_file(recording, sorter_output_folder, format="json")
+        cls._setup_json_probe_map(recording, sorter_output_folder)
 
         if params["use_binary_file"]:
             if not recording.binary_compatible_with(time_axis=0, file_paths_length=1):
@@ -382,3 +382,28 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
     @classmethod
     def _get_result_from_folder(cls, sorter_output_folder):
         return KilosortBase._get_result_from_folder(sorter_output_folder)
+
+    @classmethod
+    def _setup_json_probe_map(cls, recording, sorter_output_folder):
+        """Create a JSON probe map file for Kilosort4."""
+        from kilosort.io import save_probe
+        import numpy as np
+        
+        groups = recording.get_channel_groups()
+        positions = np.array(recording.get_channel_locations())
+        if positions.shape[1] != 2:
+            raise RuntimeError("3D 'location' are not supported. Set 2D locations instead.")
+        
+        nchan = recording.get_num_channels()
+        xcoords = ([p[0] for p in positions],)
+        ycoords = ([p[1] for p in positions],)
+        kcoords = (groups,)
+        
+        probe = {
+            'chanMap': np.arange(nchan),
+            'xc': np.array(xcoords[0]).astype(float).squeeze(),
+            'yc': np.array(ycoords[0]).astype(float).squeeze(),
+            'kcoords': np.array(kcoords).astype(float).squeeze(),
+            'n_chan': nchan,
+            }
+        save_probe(probe, str(sorter_output_folder / "chanMap.json"))

src/spikeinterface/sorters/external/kilosortbase.py

@@ -28,10 +28,9 @@ class KilosortBase:
     requires_binary_data = True
 
     @staticmethod
-    def _generate_channel_map_file(recording, sorter_output_folder, format="mat"):
+    def _generate_channel_map_file(recording, sorter_output_folder):
         """
         This function generates channel map data for kilosort and saves as `chanMap.mat`
-        or `chanMap.json`, depending on the format chosen by the caller.
 
         Loading example in Matlab (shouldn't be assigned to a variable):
         >> load('/path/to/sorter_output_folder/chanMap.mat');
@@ -42,15 +41,9 @@ def _generate_channel_map_file(recording, sorter_output_folder, format="mat"):
             The recording to generate the channel map file
         sorter_output_folder: pathlib.Path
             Path object to save `chanMap.mat` file
-        format: str
-            The format of the channel map file. Currently, 'mat'
-            and 'json' are supported ('mat' is the default).
         """
-        if format not in ("mat", "json"):
-            raise ValueError("Only 'mat' and 'json' formats are supported for Kilosort channel maps.")
-
         # prepare electrode positions for this group (only one group, the split is done in basesorter)
-        groups = recording.get_channel_groups()
+        groups = [1] * recording.get_num_channels()
         positions = np.array(recording.get_channel_locations())
         if positions.shape[1] != 2:
             raise RuntimeError("3D 'location' are not supported. Set 2D locations instead")
@@ -60,33 +53,21 @@ def _generate_channel_map_file(recording, sorter_output_folder, format="mat"):
         ycoords = ([p[1] for p in positions],)
         kcoords = (groups,)
 
-        if format == "mat":
-            channel_map = {}
-            channel_map["Nchannels"] = nchan
-            channel_map["connected"] = np.full((nchan, 1), True)
-            channel_map["chanMap0ind"] = np.arange(nchan)
-            channel_map["chanMap"] = channel_map["chanMap0ind"] + 1
-
-            channel_map["xcoords"] = np.array(xcoords).astype(float)
-            channel_map["ycoords"] = np.array(ycoords).astype(float)
-            channel_map["kcoords"] = np.array(kcoords).astype(float)
-
-            sample_rate = recording.get_sampling_frequency()
-            channel_map["fs"] = float(sample_rate)
-            import scipy.io
-
-            scipy.io.savemat(str(sorter_output_folder / "chanMap.mat"), channel_map)
-        elif format == "json":
-            from kilosort.io import save_probe
-
-            probe = {
-                "chanMap": np.arange(nchan),
-                "xc": np.array(xcoords[0]).astype(float).squeeze(),
-                "yc": np.array(ycoords[0]).astype(float).squeeze(),
-                "kcoords": np.array(kcoords).astype(float).squeeze(),
-                "n_chan": nchan,
-            }
-            save_probe(probe, str(sorter_output_folder / "chanMap.json"))
+        channel_map = {}
+        channel_map["Nchannels"] = nchan
+        channel_map["connected"] = np.full((nchan, 1), True)
+        channel_map["chanMap0ind"] = np.arange(nchan)
+        channel_map["chanMap"] = channel_map["chanMap0ind"] + 1
+
+        channel_map["xcoords"] = np.array(xcoords).astype(float)
+        channel_map["ycoords"] = np.array(ycoords).astype(float)
+        channel_map["kcoords"] = np.array(kcoords).astype(float)
+
+        sample_rate = recording.get_sampling_frequency()
+        channel_map["fs"] = float(sample_rate)
+        import scipy.io
+
+        scipy.io.savemat(str(sorter_output_folder / "chanMap.mat"), channel_map)
 
     @classmethod
     def _generate_ops_file(cls, recording, params, sorter_output_folder, binary_file_path):