Added option to sort single shank to GUI

Author: Jacob Pennington

kilosort/gui/main.py

@@ -11,7 +11,7 @@
     DataConversionBox
 )
 from kilosort.gui.logger import setup_logger
-from kilosort.io import BinaryFiltered, remove_bad_channels
+from kilosort.io import BinaryFiltered, remove_bad_channels, select_shank
 from kilosort.utils import DOWNLOADS_DIR, download_probes
 from qtpy import QtCore, QtGui, QtWidgets
 
@@ -311,10 +311,14 @@ def load_data(self):
     def set_parameters(self):
         settings = self.settings_box.settings
         bad_channels = self.settings_box.bad_channels
+        shank_idx = self.settings_box.shank_idx
 
         self.data_path = settings["data_file_path"]
         self.results_directory = settings["results_dir"]
-        self.probe_layout = remove_bad_channels(settings["probe"], bad_channels)
+        probe = remove_bad_channels(settings["probe"], bad_channels)
+        if shank_idx is not None:
+            probe = select_shank(probe, shank_idx)
+        self.probe_layout = probe
         self.probe_name = settings["probe_name"]
         self.num_channels = settings["n_chan_bin"]
 
@@ -324,9 +328,6 @@ def set_parameters(self):
         params['do_CAR'] = self.run_box.do_CAR_check.isChecked()
         params['invert_sign'] = self.run_box.invert_sign_check.isChecked()
         params['verbose_log'] = self.run_box.verbose_check.isChecked()
-
-        assert params
-
         self.params = params
 
     def do_load(self):

kilosort/gui/probe_view_box.py

@@ -173,6 +173,13 @@ def generate_spots_list(self):
         self.center_spots = []
         self.number_spots = []
         bad_channels = self.gui.settings_box.get_bad_channels()
+        shank_idx = self.gui.settings_box.shank_idx
+        chan_map = self.active_layout['chanMap']
+        if shank_idx is None:
+            shank_channels = chan_map
+        else:
+            shank_map = (self.active_layout['kcoords'] == shank_idx).nonzero()[0]
+            shank_channels = chan_map[shank_map]
         channel_size = 10 * self.spot_scale.value()/4
         template_size = 5 * self.spot_scale.value()/4
         center_size = 20 * self.spot_scale.value()/4
@@ -182,7 +189,7 @@ def generate_spots_list(self):
             for x_pos, y_pos in zip(self.xc, self.yc):
                 index = self.channel_map_dict[(x_pos, y_pos)]
                 channel = self.channel_map[index]
-                if channel in bad_channels:
+                if (channel in bad_channels) or (channel not in shank_channels):
                     color = "b"
                 else:
                     color = "g"

kilosort/gui/settings_box.py

@@ -89,6 +89,21 @@ def __init__(self, parent):
         else:
             self.bad_channels = []
 
+        self.shank_idx_text = QtWidgets.QLabel("Shank Index:")
+        self.shank_idx_input = QtWidgets.QLineEdit()
+        if self.gui.qt_settings.contains('shank_idx'):
+            idx = self.gui.qt_settings.value('shank_idx')
+            if isinstance(idx, str):
+                self.shank_idx = float(idx)
+            elif isinstance(idx, list):
+                self.shank_idx = [float(s) for s in idx]
+            else:
+                self.shank_idx = idx
+        else:
+            self.shank_idx = None
+        self.shank_idx_input.setText(str(self.shank_idx))
+
+
         self.dtype_selector_text = QtWidgets.QLabel("Data dtype:")
         self.dtype_selector = QtWidgets.QComboBox()
         self.populate_dtype_selector()
@@ -227,6 +242,18 @@ def setup(self):
             "excluding them from the probe dictionary."
             )
 
+        row_count += rspan
+        layout.addWidget(self.shank_idx_text, row_count, col1, rspan, cspan1)
+        layout.addWidget(self.shank_idx_input, row_count, col2, rspan, cspan2)
+        self.shank_idx_input.editingFinished.connect(self.update_shank_idx)
+        self.shank_idx_text.setToolTip(
+            "If not None, only channels from the specified shank index will be used. "
+            "If a list is provided, each shank will be sorted sequentially and results "
+            "will be saved in separate subfolders. Note that the shank_idx value(s) "
+            "must match the actual value specified in `probe['kcoords']`. For example, "
+            "`probe_idx=0` will not work if `probe['kcoords']` uses 1,2,3,4."
+            )
+
 
         # Add small vertical space for visual grouping
         row_count += rspan
@@ -755,13 +782,23 @@ def update_bad_channels(self):
         # Remove brackets and white space if present, convert to list of ints.
         self.bad_channels = self.get_bad_channels()
         self.gui.qt_settings.setValue('bad_channels', self.bad_channels)
-
         if not self.pause_checks:
-            # Trigger update so that probe layout in main gets updated, then
-            # refresh probe view.
-            self.update_settings()
             self.previewProbe.emit()
 
+    @QtCore.Slot()
+    def update_shank_idx(self):
+        # Remove brackets and white space if present, convert to list of floats.
+        text = self.shank_idx_input.text()
+        text = text.replace(']','').replace('[','').replace(' ','')
+        if len(text) > 0 and text.lower() != 'none':
+            idx = float(text)
+        else:
+            idx = None
+        self.shank_idx = idx
+        self.gui.qt_settings.setValue('shank_idx', self.shank_idx)
+
+        if not self.pause_checks:
+            self.previewProbe.emit()
 
     def on_data_dtype_selected(self, data_dtype):
         self.data_dtype = data_dtype