Handle case where pc_features does not exist

Author: kabilar

element_array_ephys/ephys_precluster.py

@@ -667,7 +667,8 @@ def make(self, key):
         # -- Spike-sites and Spike-depths --
         spike_sites = np.array([channel2electrodes[s]['electrode']
                                 for s in kilosort_dataset.data['spike_sites']])
-        spike_depths = kilosort_dataset.data['spike_depths']
+        spike_depths = kilosort_dataset.data['spike_depths'] \
+                        if 'pc_features' in kilosort_dataset.data else None
 
         # -- Insert unit, label, peak-chn
         units = []
@@ -685,7 +686,7 @@ def make(self, key):
                     'spike_times': unit_spike_times,
                     'spike_count': spike_count,
                     'spike_sites': spike_sites[kilosort_dataset.data['spike_clusters'] == unit],
-                    'spike_depths': spike_depths[kilosort_dataset.data['spike_clusters'] == unit]})
+                    'spike_depths': spike_depths[kilosort_dataset.data['spike_clusters'] == unit] if spike_depths else None})
 
         self.insert1(key)
         self.Unit.insert([{**key, **u} for u in units])

element_array_ephys/readers/kilosort.py

@@ -127,19 +127,20 @@ def get_best_channel(self, unit):
 
     def extract_spike_depths(self):
         """ Reimplemented from https://github.com/cortex-lab/spikes/blob/master/analysis/ksDriftmap.m """
-        ycoords = self.data['channel_positions'][:, 1]
-        pc_features = self.data['pc_features'][:, 0, :]  # 1st PC only
-        pc_features = np.where(pc_features < 0, 0, pc_features)
-
-        # ---- compute center of mass of these features (spike depths) ----
-
-        # which channels for each spike?
-        spk_feature_ind = self.data['pc_feature_ind'][self.data['spike_templates'], :]
-        # ycoords of those channels?
-        spk_feature_ycoord = ycoords[spk_feature_ind]
-        # center of mass is sum(coords.*features)/sum(features)
-        self._data['spike_depths'] = (np.sum(spk_feature_ycoord * pc_features**2, axis=1)
-                                      / np.sum(pc_features**2, axis=1))
+        if 'pc_features' in self.data:
+            ycoords = self.data['channel_positions'][:, 1]
+            pc_features = self.data['pc_features'][:, 0, :]  # 1st PC only
+            pc_features = np.where(pc_features < 0, 0, pc_features)
+
+            # ---- compute center of mass of these features (spike depths) ----
+
+            # which channels for each spike?
+            spk_feature_ind = self.data['pc_feature_ind'][self.data['spike_templates'], :]
+            # ycoords of those channels?
+            spk_feature_ycoord = ycoords[spk_feature_ind]
+            # center of mass is sum(coords.*features)/sum(features)
+            self._data['spike_depths'] = (np.sum(spk_feature_ycoord * pc_features**2, axis=1)
+                                        / np.sum(pc_features**2, axis=1))
 
         # ---- extract spike sites ----
         max_site_ind = np.argmax(np.abs(self.data['templates']).max(axis=1), axis=1)