add mearec option to load only recordings or only sorting

Author: h-mayorquin

neo/rawio/mearecrawio.py

@@ -36,53 +36,63 @@ class MEArecRawIO(BaseRawIO):
     extensions = ['h5']
     rawmode = 'one-file'
 
-    def __init__(self, filename=''):
+    def __init__(self, filename='', load_spiketrains=True, load_recordings=True):
         BaseRawIO.__init__(self)
         self.filename = filename
-
+        self.load_spiketrains = load_spiketrains
+        self.load_recordings = load_recordings
+        
     def _source_name(self):
         return self.filename
 
     def _parse_header(self):
         import MEArec as mr
+        load = ['channel_positions'] 
+        if self.load_recordings:
+            load.append("recordings")
+        if self.load_spiketrains:
+            load.append("spiketrains")
+        
         self._recgen = mr.load_recordings(recordings=self.filename, return_h5_objects=True,
                                           check_suffix=False,
-                                          load=['recordings', 'spiketrains', 'channel_positions'],
+                                          load=load,
                                           load_waveforms=False)
         self._sampling_rate = self._recgen.info['recordings']['fs']
-        self._recordings = self._recgen.recordings
-        self._num_frames, self._num_channels = self._recordings.shape
 
         signal_streams = np.array([('Signals', '0')], dtype=_signal_stream_dtype)
 
         sig_channels = []
-        for c in range(self._num_channels):
-            ch_name = 'ch{}'.format(c)
-            chan_id = str(c + 1)
-            sr = self._sampling_rate  # Hz
-            dtype = self._recordings.dtype
-            units = 'uV'
-            gain = 1.
-            offset = 0.
-            stream_id = '0'
-            sig_channels.append((ch_name, chan_id, sr, dtype, units, gain, offset, stream_id))
+        if self.load_recordings:
+            self._recordings = self._recgen.recordings
+            self._num_frames, self._num_channels = self._recordings.shape
+            for c in range(self._num_channels):
+                ch_name = 'ch{}'.format(c)
+                chan_id = str(c + 1)
+                sr = self._sampling_rate  # Hz
+                dtype = self._recordings.dtype
+                units = 'uV'
+                gain = 1.
+                offset = 0.
+                stream_id = '0'
+                sig_channels.append((ch_name, chan_id, sr, dtype, units, gain, offset, stream_id))
         sig_channels = np.array(sig_channels, dtype=_signal_channel_dtype)
 
         # creating units channels
         spike_channels = []
-        self._spiketrains = self._recgen.spiketrains
-        for c in range(len(self._spiketrains)):
-            unit_name = 'unit{}'.format(c)
-            unit_id = '#{}'.format(c)
-            # if spiketrains[c].waveforms is not None:
-            wf_units = ''
-            wf_gain = 1.
-            wf_offset = 0.
-            wf_left_sweep = 0
-            wf_sampling_rate = self._sampling_rate
-            spike_channels.append((unit_name, unit_id, wf_units, wf_gain,
-                                  wf_offset, wf_left_sweep, wf_sampling_rate))
-        spike_channels = np.array(spike_channels, dtype=_spike_channel_dtype)
+        if self.load_spiketrains:
+            self._spiketrains = self._recgen.spiketrains
+            for c in range(len(self._spiketrains)):
+                unit_name = 'unit{}'.format(c)
+                unit_id = '#{}'.format(c)
+                # if spiketrains[c].waveforms is not None:
+                wf_units = ''
+                wf_gain = 1.
+                wf_offset = 0.
+                wf_left_sweep = 0
+                wf_sampling_rate = self._sampling_rate
+                spike_channels.append((unit_name, unit_id, wf_units, wf_gain,
+                                    wf_offset, wf_left_sweep, wf_sampling_rate))
+            spike_channels = np.array(spike_channels, dtype=_spike_channel_dtype)
 
         event_channels = []
         event_channels = np.array(event_channels, dtype=_event_channel_dtype)
@@ -119,6 +129,10 @@ def _get_signal_t_start(self, block_index, seg_index, stream_index):
 
     def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop,
                                 stream_index, channel_indexes):
+        
+        if not self.load_recordings:
+            raise AttributeError("Recordings not loaded. Set load_recordings=True in MEArecRawIO constructor")
+        
         if i_start is None:
             i_start = 0
         if i_stop is None:
@@ -127,23 +141,31 @@ def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop,
         if channel_indexes is None:
             channel_indexes = slice(self._num_channels)
         if isinstance(channel_indexes, slice):
-            raw_signals = self._recgen.recordings[i_start:i_stop, channel_indexes]
+            raw_signals = self._recordings[i_start:i_stop, channel_indexes]
         else:
             # sort channels because h5py neeeds sorted indexes
             if np.any(np.diff(channel_indexes) < 0):
                 sorted_channel_indexes = np.sort(channel_indexes)
                 sorted_idx = np.array([list(sorted_channel_indexes).index(ch)
                                        for ch in channel_indexes])
-                raw_signals = self._recgen.recordings[i_start:i_stop, sorted_channel_indexes]
+                raw_signals = self._recordings[i_start:i_stop, sorted_channel_indexes]
                 raw_signals = raw_signals[:, sorted_idx]
             else:
-                raw_signals = self._recgen.recordings[i_start:i_stop, channel_indexes]
+                raw_signals = self._recordings[i_start:i_stop, channel_indexes]
         return raw_signals
 
     def _spike_count(self, block_index, seg_index, unit_index):
+        
+        if not self.load_spiketrains:
+            raise AttributeError("Spiketrains not loaded. Set load_spiketrains=True in MEArecRawIO constructor")
+        
         return len(self._spiketrains[unit_index])
 
     def _get_spike_timestamps(self, block_index, seg_index, unit_index, t_start, t_stop):
+        
+        if not self.load_spiketrains:
+            raise AttributeError("Spiketrains not loaded. Set load_spiketrains=True in MEArecRawIO constructor")
+        
         spike_timestamps = self._spiketrains[unit_index].times.magnitude
         if t_start is None:
             t_start = self._segment_t_start(block_index, seg_index)