Merge pull request #1057 from JuliaSprenger/enh/tdt-block

Extend tdtrawio to support single tdt block reading and empty streams

Author: samuelgarcia

neo/rawio/tdtrawio.py

@@ -28,49 +28,75 @@
 import numpy as np
 import os
 import re
+import warnings
 from collections import OrderedDict
+from pathlib import Path
 
 
 class TdtRawIO(BaseRawIO):
     rawmode = 'one-dir'
 
     def __init__(self, dirname='', sortname=''):
         """
-        'sortname' is used to specify the external sortcode generated by offline spike sorting.
-        if sortname=='PLX', there should be a ./sort/PLX/*.SortResult file in the tdt block,
-        which stores the sortcode for every spike; defaults to '',
-        which uses the original online sort.
+        Initialize reader for one or multiple TDT data blocks.
+
+        dirname (str, pathlib.Path):
+            tank-directory of a dataset to be read as multiple segments OR single file of dataset.
+            In the latter case only the corresponding segment will considered.
+        sortname (str):
+            'sortname' is used to specify the external sortcode generated by offline spike sorting.
+            if sortname=='PLX', there should be a ./sort/PLX/*.SortResult file in the tdt block,
+            which stores the sortcode for every spike
+            Default: '', uses the original online sort.
+
+
         """
         BaseRawIO.__init__(self)
-        dirname = str(dirname)
-        if dirname.endswith('/'):
-            dirname = dirname[:-1]
-        self.dirname = dirname
+        dirname = Path(dirname)
+        if dirname.is_dir():
+            self.dirname = Path(dirname)
+            self.tdt_block_mode = 'multi'
+        elif dirname.is_file():
+            # in single tdt block mode the dirname also contains the block prefix
+            self.dirname = dirname.with_suffix('')
+            self.tdt_block_mode = 'single'
+        else:
+            raise ValueError(f'No data folder or file found for {dirname}')
 
         self.sortname = sortname
 
     def _source_name(self):
         return self.dirname
 
-    def _parse_header(self):
-
-        tankname = os.path.basename(self.dirname)
+    def _get_filestem(self, segment_name=''):
+        if self.tdt_block_mode == 'multi':
+            return self.dirname / segment_name / f'{self.dirname.name}_{segment_name}'
+        else:
+            return self.dirname
 
+    def _parse_header(self):
         segment_names = []
-        for segment_name in os.listdir(self.dirname):
-            path = os.path.join(self.dirname, segment_name)
-            if is_tdtblock(path):
-                segment_names.append(segment_name)
+        if self.tdt_block_mode == 'multi':
+            tankname = self.dirname.stem
+            for path in self.dirname.iterdir():
+                if is_tdtblock(path):
+                    segment_names.append(path.stem)
+
+        # if no block structure was detected, check if current dir contains a set of data
+        elif is_tdtblock(self.dirname.parent):
+            segment_names.append(str(self.dirname.stem))
+            tankname = None
 
         nb_segment = len(segment_names)
+        if nb_segment == 0:
+            warnings.warn(f'Could not find any data set belonging to {self.dirname}')
 
         # TBK (channel info)
         info_channel_groups = None
         for seg_index, segment_name in enumerate(segment_names):
-            path = os.path.join(self.dirname, segment_name)
 
             # TBK contain channels
-            tbk_filename = os.path.join(path, tankname + '_' + segment_name + '.Tbk')
+            tbk_filename = self._get_filestem(segment_name).with_suffix('.Tbk')
             _info_channel_groups = read_tbk(tbk_filename)
             if info_channel_groups is None:
                 info_channel_groups = _info_channel_groups
@@ -81,9 +107,8 @@ def _parse_header(self):
         # TEV (mixed data)
         self._tev_datas = []
         for seg_index, segment_name in enumerate(segment_names):
-            path = os.path.join(self.dirname, segment_name)
-            tev_filename = os.path.join(path, tankname + '_' + segment_name + '.tev')
-            if os.path.exists(tev_filename):
+            tev_filename = self._get_filestem(segment_name).with_suffix('.tev')
+            if tev_filename.exists():
                 tev_data = np.memmap(tev_filename, mode='r', offset=0, dtype='uint8')
             else:
                 tev_data = None
@@ -94,8 +119,7 @@ def _parse_header(self):
         self._seg_t_starts = []
         self._seg_t_stops = []
         for seg_index, segment_name in enumerate(segment_names):
-            path = os.path.join(self.dirname, segment_name)
-            tsq_filename = os.path.join(path, tankname + '_' + segment_name + '.tsq')
+            tsq_filename = self._get_filestem(segment_name).with_suffix('.tsq')
             tsq = np.fromfile(tsq_filename, dtype=tsq_dtype)
             self._tsq.append(tsq)
             # Start and stop times are only found in the second
@@ -115,9 +139,13 @@ def _parse_header(self):
             #  (generated after offline sorting)
             if self.sortname != '':
                 try:
-                    for file in os.listdir(os.path.join(path, 'sort', sortname)):
+                    if self.tdt_block_mode == 'multi':
+                        path = self.dirname
+                    else:
+                        path = self.dirname.parent
+                    for file in os.listdir(path / 'sort' / self.sortname):
                         if file.endswith(".SortResult"):
-                            sortresult_filename = os.path.join(path, 'sort', sortname, file)
+                            sortresult_filename = path / 'sort' / self.sortname / file
                             # get new sortcode
                             newsortcode = np.fromfile(sortresult_filename, 'int8')[
                                 1024:]  # first 1024 bytes are header
@@ -181,15 +209,24 @@ def _parse_header(self):
                         assert self._sigs_lengths[seg_index][stream_index] == size
 
                     # signal start time, relative to start of segment
-                    t_start = data_index['timestamp'][0]
+                    if len(data_index['timestamp']):
+                        t_start = data_index['timestamp'][0]
+                    else:
+                        # if no signal present use segment t_start as dummy value
+                        t_start = self._seg_t_starts[seg_index]
                     if stream_index not in self._sigs_t_start[seg_index]:
                         self._sigs_t_start[seg_index][stream_index] = t_start
                     else:
                         assert self._sigs_t_start[seg_index][stream_index] == t_start
 
                     # sampling_rate and dtype
-                    _sampling_rate = float(data_index['frequency'][0])
-                    _dtype = data_formats[data_index['dataformat'][0]]
+                    if len(data_index):
+                        _sampling_rate = float(data_index['frequency'][0])
+                        _dtype = data_formats[data_index['dataformat'][0]]
+                    else:
+                        # if no signal present use dummy values
+                        _sampling_rate = 1.
+                        _dtype = int
                     if sampling_rate is None:
                         sampling_rate = _sampling_rate
                         dtype = _dtype
@@ -202,11 +239,23 @@ def _parse_header(self):
                         assert dtype == _dtype, 'sampling is changing!!!'
 
                     # data buffer test if SEV file exists otherwise TEV
-                    path = os.path.join(self.dirname, segment_name)
-                    sev_filename = os.path.join(path, tankname + '_' + segment_name + '_'
-                                                + info['StoreName'].decode('ascii')
-                                                + '_ch' + str(chan_id) + '.sev')
-                    if os.path.exists(sev_filename):
+                    # path = self.dirname / segment_name
+                    if self.tdt_block_mode == 'multi':
+                        # for multi block datasets the names of sev files are fixed
+                        store = info['StoreName'].decode('ascii')
+                        sev_stem = f'{tankname}_{segment_name}_{store}_ch{chan_id}'
+                        sev_filename = (path / sev_stem).with_suffix('.sev')
+                    else:
+                        # for single block datasets the exact name of sev files in not known
+                        sev_regex = f".*_ch{chan_id}.sev"
+                        sev_filename = list(self.dirname.parent.glob(str(sev_regex)))
+
+                        # in case non or multiple sev files are found for current stream + channel
+                        if len(sev_filename) != 1:
+                            warnings.warn(f'Could not identify sev file for channel {chan_id}.')
+                            sev_filename = None
+
+                    if (sev_filename is not None) and sev_filename.exists():
                         data = np.memmap(sev_filename, mode='r', offset=0, dtype='uint8')
                     else:
                         data = self._tev_datas[seg_index]
@@ -526,10 +575,10 @@ def read_tbk(tbk_filename):
 def is_tdtblock(blockpath):
     """Is tha path a  TDT block (=neo.Segment) ?"""
     file_ext = list()
-    if os.path.isdir(blockpath):
+    if blockpath.is_dir():
         # for every file, get extension, convert to lowercase and append
-        for file in os.listdir(blockpath):
-            file_ext.append(os.path.splitext(file)[1].lower())
+        for file in blockpath.iterdir():
+            file_ext.append(file.suffix.lower())
 
     file_ext = set(file_ext)
     tdt_ext = {'.tbk', '.tdx', '.tev', '.tsq'}

neo/test/rawiotest/test_tdtrawio.py

@@ -1,4 +1,6 @@
 import unittest
+from pathlib import Path
+from numpy.testing import assert_array_equal, assert_
 
 from neo.rawio.tdtrawio import TdtRawIO
 from neo.test.rawiotest.common_rawio_test import BaseTestRawIO
@@ -10,9 +12,49 @@ class TestTdtRawIO(BaseTestRawIO, unittest.TestCase, ):
         'tdt'
     ]
     entities_to_test = [
-        'tdt/aep_05'
+        'tdt/aep_05',
+        'tdt/aep_05/Block-1/aep_05_Block-1.Tdx'
     ]
 
+    def test_invalid_dirname(self):
+        invalid_name = 'random_non_existant_tdt_filename'
+        assert not Path(invalid_name).exists()
+
+        with self.assertRaises(ValueError):
+            TdtRawIO(invalid_name)
+
+    def test_compare_load_multi_single_block(self):
+        dirname = self.get_local_path('tdt/aep_05')
+        filename = self.get_local_path('tdt/aep_05/Block-1/aep_05_Block-1.Tdx')
+
+        io_single = TdtRawIO(filename)
+        io_multi = TdtRawIO(dirname)
+
+        io_single.parse_header()
+        io_multi.parse_header()
+
+        self.assertEqual(io_single.tdt_block_mode, 'single')
+        self.assertEqual(io_multi.tdt_block_mode, 'multi')
+
+        self.assertEqual(io_single.block_count(), 1)
+        self.assertEqual(io_multi.block_count(), 1)
+
+        self.assertEqual(io_single.segment_count(0), 1)
+        self.assertEqual(io_multi.segment_count(0), 2)
+
+        # compare header infos
+        assert_array_equal(io_single.header['signal_streams'], io_multi.header['signal_streams'])
+        assert_array_equal(io_single.header['signal_channels'], io_multi.header['signal_channels'])
+        assert_array_equal(io_single.header['event_channels'], io_multi.header['event_channels'])
+
+        # not all spiking channels are present in first tdt block (segment)
+        for spike_channel in io_single.header['spike_channels']:
+            self.assertIn(spike_channel, io_multi.header['spike_channels'])
+
+        # check that extracted signal chunks are identical
+        assert_array_equal(io_single.get_analogsignal_chunk(0, 0, 0, 100, 0),
+                           io_multi.get_analogsignal_chunk(0, 0, 0, 100, 0))
+
 
 if __name__ == "__main__":
     unittest.main()