Load and annotate pcs and amplitudes

Author: kleinjohann

neo/io/phyio.py

@@ -7,6 +7,6 @@ class PhyIO(PhyRawIO, BaseFromRaw):
     description = "Phy IO"
     mode = 'dir'
 
-    def __init__(self, dirname):
-        PhyRawIO.__init__(self, dirname=dirname)
+    def __init__(self, dirname, load_pcs=False):
+        PhyRawIO.__init__(self, dirname=dirname, load_pcs=load_pcs)
         BaseFromRaw.__init__(self, dirname)

neo/rawio/phyrawio.py

@@ -16,6 +16,7 @@
 import re
 import csv
 import ast
+import warnings
 
 
 class PhyRawIO(BaseRawIO):
@@ -35,9 +36,10 @@ class PhyRawIO(BaseRawIO):
     extensions = []
     rawmode = 'one-dir'
 
-    def __init__(self, dirname=''):
+    def __init__(self, dirname='', load_pcs=False):
         BaseRawIO.__init__(self)
         self.dirname = dirname
+        self.load_pcs = load_pcs
 
     def _source_name(self):
         return self.dirname
@@ -53,16 +55,21 @@ def _parse_header(self):
         else:
             self._spike_clusters = self._spike_templates
 
-        # TODO: Add this when array_annotations are ready
-        # if (phy_folder / 'amplitudes.npy').is_file():
-        #     amplitudes = np.squeeze(np.load(phy_folder / 'amplitudes.npy'))
-        # else:
-        #     amplitudes = np.ones(len(spike_times))
-        #
-        # if (phy_folder / 'pc_features.npy').is_file():
-        #     pc_features = np.squeeze(np.load(phy_folder / 'pc_features.npy'))
-        # else:
-        #     pc_features = None
+        if (phy_folder / 'amplitudes.npy').is_file():
+            self._amplitudes = np.squeeze(np.load(phy_folder / 'amplitudes.npy'))
+        else:
+            self._amplitudes = None
+
+        self._pc_features = None
+        self._pc_feature_ind = None
+        if self.load_pcs:
+            if ((phy_folder / 'pc_features.npy').is_file()
+                    and (phy_folder / 'pc_feature_ind.npy').is_file()):
+                self._pc_features = np.squeeze(np.load(phy_folder / 'pc_features.npy'))
+                self._pc_feature_ind = np.squeeze(np.load(phy_folder / 'pc_feature_ind.npy'))
+            else:
+                warnings.warn('PCs requested but "pc_features.npy" and/or'
+                              '"pc_feature_ind.npy" not found in the data folder.')
 
         # SEE: https://stackoverflow.com/questions/4388626/
         #  python-safe-eval-string-to-bool-int-float-none-string
@@ -150,6 +157,25 @@ def _parse_header(self):
                             annotation_dict[property_name]
                         break
 
+            cluster_mask = (self._spike_clusters == clust_id).flatten()
+
+            if self._amplitudes is not None:
+                spiketrain_an['__array_annotations__']['amplitudes'] = \
+                    self._amplitudes[cluster_mask]
+
+            if self._pc_features is not None:
+                current_pc_features = self._pc_features[cluster_mask]
+                _, num_pcs, num_pc_channels = current_pc_features.shape
+                for pc_idx in range(num_pcs):
+                    for channel_idx in range(num_pc_channels):
+                        key = 'channel{channel_idx}_pc{pc_idx}'.format(channel_idx=channel_idx,
+                                                                       pc_idx=pc_idx)
+                        spiketrain_an['__array_annotations__'][key] = \
+                                current_pc_features[:, pc_idx, channel_idx]
+
+            if self._pc_feature_ind is not None:
+                spiketrain_an['pc_feature_ind'] = self._pc_feature_ind[index]
+
     def _segment_t_start(self, block_index, seg_index):
         assert block_index == 0
         return self._t_start