Merge pull request #997 from rjurkus/fix/phyio_annotations

[phyio] Fixed spiketrain annotations

Author: samuelgarcia

neo/rawio/phyrawio.py

@@ -9,7 +9,7 @@
 """
 
 from .baserawio import (BaseRawIO, _signal_channel_dtype, _signal_stream_dtype,
-                _spike_channel_dtype, _event_channel_dtype)
+                        _spike_channel_dtype, _event_channel_dtype)
 
 import numpy as np
 from pathlib import Path
@@ -88,7 +88,8 @@ def _parse_header(self):
         signal_streams = np.array(signal_streams, dtype=_signal_stream_dtype)
 
         signal_channels = []
-        signal_channels = np.array(signal_channels, dtype=_signal_channel_dtype)
+        signal_channels = np.array(signal_channels,
+                                   dtype=_signal_channel_dtype)
 
         spike_channels = []
         for i, clust_id in enumerate(clust_ids):
@@ -132,6 +133,9 @@ def _parse_header(self):
         for index, clust_id in enumerate(clust_ids):
             spiketrain_an = seg_ann['spikes'][index]
 
+            # Add cluster_id annotation
+            spiketrain_an['cluster_id'] = clust_id
+
             # Loop over list of list of dict and annotate each st
             for annotation_list in annotation_lists:
                 clust_key, property_name = tuple(annotation_list[0].
@@ -172,8 +176,8 @@ def _spike_count(self, block_index, seg_index, spike_channel_index):
         nb_spikes = np.sum(mask)
         return nb_spikes
 
-    def _get_spike_timestamps(self, block_index, seg_index, spike_channel_index,
-                              t_start, t_stop):
+    def _get_spike_timestamps(self, block_index, seg_index,
+                              spike_channel_index, t_start, t_stop):
         assert block_index == 0
         assert seg_index == 0
 
@@ -183,7 +187,8 @@ def _get_spike_timestamps(self, block_index, seg_index, spike_channel_index,
 
         if t_start is not None:
             start_frame = int(t_start * self._sampling_frequency)
-            spike_timestamps = spike_timestamps[spike_timestamps >= start_frame]
+            spike_timestamps = \
+                spike_timestamps[spike_timestamps >= start_frame]
         if t_stop is not None:
             end_frame = int(t_stop * self._sampling_frequency)
             spike_timestamps = spike_timestamps[spike_timestamps < end_frame]
@@ -195,8 +200,8 @@ def _rescale_spike_timestamp(self, spike_timestamps, dtype):
         spike_times /= self._sampling_frequency
         return spike_times
 
-    def _get_spike_raw_waveforms(self, block_index, seg_index, spike_channel_index,
-                                 t_start, t_stop):
+    def _get_spike_raw_waveforms(self, block_index, seg_index,
+                                 spike_channel_index, t_start, t_stop):
         return None
 
     def _event_count(self, block_index, seg_index, event_channel_index):
@@ -215,14 +220,30 @@ def _rescale_epoch_duration(self, raw_duration, dtype):
     @staticmethod
     def _parse_tsv_or_csv_to_list_of_dict(filename):
         list_of_dict = list()
+        letter_pattern = re.compile('[a-zA-Z]')
+        float_pattern = re.compile(r'\d*\.')
         with open(filename) as csvfile:
             if filename.suffix == '.csv':
                 reader = csv.DictReader(csvfile, delimiter=',')
             elif filename.suffix == '.tsv':
                 reader = csv.DictReader(csvfile, delimiter='\t')
             else:
                 raise ValueError("Function parses only .csv or .tsv files")
+            line = 0
+
             for row in reader:
+                if line == 0:
+                    key1, key2 = tuple(row.keys())
+                # Convert cluster ID to int
+                row[key1] = int(row[key1])
+                # Convert strings without letters
+                if letter_pattern.match(row[key2]) is None:
+                    if float_pattern.match(row[key2]) is None:
+                        row[key2] = int(row[key2])
+                    else:
+                        row[key2] = float(row[key2])
+
                 list_of_dict.append(row)
+                line += 1
 
         return list_of_dict

neo/test/rawiotest/test_phyrawio.py

@@ -31,18 +31,26 @@ def test_csv_tsv_parser_with_csv(self):
         csv_tempfile = Path(tempfile.gettempdir()).joinpath('test.csv')
         with open(csv_tempfile, 'w') as csv_file:
             csv_writer = csv.writer(csv_file, delimiter=',')
-            csv_writer.writerow(['Header 1', 'Header 2'])
-            csv_writer.writerow(['Value 1', 'Value 2'])
+            csv_writer.writerow(['cluster_id', 'some_annotation'])
+            csv_writer.writerow([1, 'Good'])
+            csv_writer.writerow([2, 10])
+            csv_writer.writerow([3, 1.23])
 
         # the parser in PhyRawIO runs csv.DictReader to parse the file
         # csv.DictReader for python version 3.6+ returns list of OrderedDict
         if (3, 6) <= sys.version_info < (3, 8):
-            target = [OrderedDict({'Header 1': 'Value 1',
-                                   'Header 2': 'Value 2'})]
+            target = [OrderedDict({'cluster_id': 1,
+                                   'some_annotation': 'Good'}),
+                      OrderedDict({'cluster_id': 2,
+                                   'some_annotation': 10}),
+                      OrderedDict({'cluster_id': 3,
+                                   'some_annotation': 1.23})]
 
         # csv.DictReader for python version 3.8+ returns list of dict
         elif sys.version_info >= (3, 8):
-            target = [{'Header 1': 'Value 1', 'Header 2': 'Value 2'}]
+            target = [{'cluster_id': 1, 'some_annotation': 'Good'},
+                      {'cluster_id': 2, 'some_annotation': 10},
+                      {'cluster_id': 3, 'some_annotation': 1.23}]
 
         list_of_dict = PhyRawIO._parse_tsv_or_csv_to_list_of_dict(csv_tempfile)
 
@@ -52,18 +60,26 @@ def test_csv_tsv_parser_with_tsv(self):
         tsv_tempfile = Path(tempfile.gettempdir()).joinpath('test.tsv')
         with open(tsv_tempfile, 'w') as tsv_file:
             tsv_writer = csv.writer(tsv_file, delimiter='\t')
-            tsv_writer.writerow(['Header 1', 'Header 2'])
-            tsv_writer.writerow(['Value 1', 'Value 2'])
+            tsv_writer.writerow(['cluster_id', 'some_annotation'])
+            tsv_writer.writerow([1, 'Good'])
+            tsv_writer.writerow([2, 10])
+            tsv_writer.writerow([3, 1.23])
 
         # the parser in PhyRawIO runs csv.DictReader to parse the file
         # csv.DictReader for python version 3.6+ returns list of OrderedDict
         if (3, 6) <= sys.version_info < (3, 8):
-            target = [OrderedDict({'Header 1': 'Value 1',
-                                   'Header 2': 'Value 2'})]
+            target = [OrderedDict({'cluster_id': 1,
+                                   'some_annotation': 'Good'}),
+                      OrderedDict({'cluster_id': 2,
+                                   'some_annotation': 10}),
+                      OrderedDict({'cluster_id': 3,
+                                   'some_annotation': 1.23})]
 
         # csv.DictReader for python version 3.8+ returns list of dict
         elif sys.version_info >= (3, 8):
-            target = [{'Header 1': 'Value 1', 'Header 2': 'Value 2'}]
+            target = [{'cluster_id': 1, 'some_annotation': 'Good'},
+                      {'cluster_id': 2, 'some_annotation': 10},
+                      {'cluster_id': 3, 'some_annotation': 1.23}]
 
         list_of_dict = PhyRawIO._parse_tsv_or_csv_to_list_of_dict(tsv_tempfile)