Implement get_analogsignal_chunk() generically when a rawio class has_buffer_description_api=True

This should also solve the memmap and memmory leak problem.

Author: samuelgarcia

neo/rawio/baserawio.py

@@ -77,6 +77,8 @@
 
 from neo import logging_handler
 
+from .utils import get_memmap_chunk_from_opened_file
+
 
 possible_raw_modes = [
     "one-file",
@@ -656,7 +658,12 @@ def get_signal_size(self, block_index: int, seg_index: int, stream_index: int |
 
         """
         stream_index = self._get_stream_index_from_arg(stream_index)
-        return self._get_signal_size(block_index, seg_index, stream_index)
+
+        if not self.has_buffer_description_api:
+            return self._get_signal_size(block_index, seg_index, stream_index)
+        else:
+            # use the buffer description
+            return self._get_signal_size_generic(block_index, seg_index, stream_index)
 
     def get_signal_t_start(self, block_index: int, seg_index: int, stream_index: int | None = None):
         """
@@ -805,7 +812,11 @@ def get_analogsignal_chunk(
             if np.all(np.diff(channel_indexes) == 1):
                 channel_indexes = slice(channel_indexes[0], channel_indexes[-1] + 1)
 
-        raw_chunk = self._get_analogsignal_chunk(block_index, seg_index, i_start, i_stop, stream_index, channel_indexes)
+        if not self.has_buffer_description_api:
+            raw_chunk = self._get_analogsignal_chunk(block_index, seg_index, i_start, i_stop, stream_index, channel_indexes)
+        else:
+            # use the buffer description
+            raw_chunk = self._get_analogsignal_chunk_generic(block_index, seg_index, i_start, i_stop, stream_index, channel_indexes)
 
         return raw_chunk
 
@@ -1277,6 +1288,7 @@ def _get_signal_size(self, block_index: int, seg_index: int, stream_index: int):
 
         All channels indexed must have the same size and t_start.
         """
+        # must NOT be implemented if has_buffer_description_api=True
         raise (NotImplementedError)
 
     def _get_signal_t_start(self, block_index: int, seg_index: int, stream_index: int):
@@ -1304,7 +1316,7 @@ def _get_analogsignal_chunk(
         -------
             array of samples, with each requested channel in a column
         """
-
+        # must NOT be implemented if has_buffer_description_api=True
         raise (NotImplementedError)
 
     ###
@@ -1355,6 +1367,59 @@ def get_analogsignal_buffer_description(self, block_index: int = 0, seg_index: i
     def _get_analogsignal_buffer_description(self, block_index, seg_index, stream_index):
         raise (NotImplementedError)
 
+    def _get_signal_size_generic(self, block_index, seg_index, stream_index):
+        # When has_buffer_description_api=True this used to avoid to write _get_analogsignal_chunk())
+
+        buffer_desc = self._get_signal_size(block_index, seg_index, stream_index)
+        return buffer_desc['shape'][0]
+
+    def _get_analogsignal_chunk_generic(
+        self,
+        block_index: int,
+        seg_index: int,
+        i_start: int | None,
+        i_stop: int | None,
+        stream_index: int,
+        channel_indexes: list[int] | None,
+    ):
+        # When has_buffer_description_api=True this used to avoid to write _get_analogsignal_chunk())
+        buffer_desc = self._get_analogsignal_buffer_description( block_index, seg_index, stream_index)
+
+        if buffer_desc['type'] == 'binary':
+
+            # open files on demand and keep reference to opened file 
+            if not hasattr(self, '_memmap_analogsignal_streams'):
+                self._memmap_analogsignal_streams = {}
+            if block_index not in self._memmap_analogsignal_streams:
+                self._memmap_analogsignal_streams[block_index] = {}
+            if seg_index not in self._memmap_analogsignal_streams[block_index]:
+                self._memmap_analogsignal_streams[block_index][seg_index] = {}
+            if stream_index not in self._memmap_analogsignal_streams[block_index][seg_index]:
+                fid = open(buffer_desc['file_path'], mode='rb')
+                self._memmap_analogsignal_streams[block_index][seg_index] = fid
+            else:
+                fid = self._memmap_analogsignal_streams[block_index][seg_index]
+
+            
+            i_start = i_start or 0
+            i_stop = i_stop or  buffer_desc['shape'][0]
+
+            num_channels = buffer_desc['shape'][1]
+            
+            raw_sigs = get_memmap_chunk_from_opened_file(fid, num_channels,  i_start, i_stop, np.dtype(buffer_desc['dtype']), file_offset=buffer_desc['file_offset'])
+
+            # this is a pre slicing when the stream do not contain all channels (for instance spikeglx when load_sync_channel=False)
+            channel_slice = buffer_desc.get('channel_slice', None)
+            if channel_slice is not None:
+                raw_sigs = raw_sigs[:, channel_slice]
+
+            # channel slice requested
+            if channel_indexes is not None:
+                raw_sigs = raw_sigs[:, channel_indexes]
+        else:
+            raise NotImplementedError()
+
+        return raw_sigs
 
 
 

neo/rawio/spikeglxrawio.py

@@ -63,7 +63,7 @@
     _spike_channel_dtype,
     _event_channel_dtype,
 )
-from .utils import get_memmap_shape, get_memmap_chunk_from_open_file
+from .utils import get_memmap_shape
 
 
 class SpikeGLXRawIO(BaseRawIO):
@@ -127,7 +127,7 @@ def _parse_header(self):
         nb_segment = np.unique([info["seg_index"] for info in self.signals_info_list]).size
 
 
-        self._memmaps = {}
+        # self._memmaps = {}
         self.signals_info_dict = {}
         # on block
         self._buffer_descriptions = {0 :{}}
@@ -139,11 +139,11 @@ def _parse_header(self):
             self.signals_info_dict[key] = info
 
             # create memmap
-            data = np.memmap(info["bin_file"], dtype="int16", mode="r", offset=0, order="C")
+            # data = np.memmap(info["bin_file"], dtype="int16", mode="r", offset=0, order="C")
             # this should be (info['sample_length'], info['num_chan'])
             # be some file are shorten
-            data = data.reshape(-1, info["num_chan"])
-            self._memmaps[key] = data
+            # data = data.reshape(-1, info["num_chan"])
+            # self._memmaps[key] = data
 
             stream_index = stream_names.index(info["stream_name"])
             if seg_index not in self._buffer_descriptions[0]:
@@ -273,42 +273,42 @@ def _segment_t_start(self, block_index, seg_index):
     def _segment_t_stop(self, block_index, seg_index):
         return self._t_stops[seg_index]
 
-    def _get_signal_size(self, block_index, seg_index, stream_index):
-        stream_id = self.header["signal_streams"][stream_index]["id"]
-        memmap = self._memmaps[seg_index, stream_id]
-        return int(memmap.shape[0])
+    # def _get_signal_size(self, block_index, seg_index, stream_index):
+    #     stream_id = self.header["signal_streams"][stream_index]["id"]
+    #     memmap = self._memmaps[seg_index, stream_id]
+    #     return int(memmap.shape[0])
 
     def _get_signal_t_start(self, block_index, seg_index, stream_index):
         return 0.0
 
-    def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes):
-        stream_id = self.header["signal_streams"][stream_index]["id"]
-        memmap = self._memmaps[seg_index, stream_id]
-        stream_name = self.header["signal_streams"]["name"][stream_index]
-
-        # take care of sync channel
-        info = self.signals_info_dict[0, stream_name]
-        if not self.load_sync_channel and info["has_sync_trace"]:
-            memmap = memmap[:, :-1]
-
-        # since we cut the memmap, we can simplify the channel selection
-        if channel_indexes is None:
-            channel_selection = slice(None)
-        elif isinstance(channel_indexes, slice):
-            channel_selection = channel_indexes
-        elif not isinstance(channel_indexes, slice):
-            if np.all(np.diff(channel_indexes) == 1):
-                # consecutive channel then slice this avoid a copy (because of ndarray.take(...)
-                # and so keep the underlying memmap
-                channel_selection = slice(channel_indexes[0], channel_indexes[0] + len(channel_indexes))
-            else:
-                channel_selection = channel_indexes
-        else:
-            raise ValueError("get_analogsignal_chunk : channel_indexes" "must be slice or list or array of int")
-
-        raw_signals = memmap[slice(i_start, i_stop), channel_selection]
-
-        return raw_signals
+    # def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes):
+    #     stream_id = self.header["signal_streams"][stream_index]["id"]
+    #     memmap = self._memmaps[seg_index, stream_id]
+    #     stream_name = self.header["signal_streams"]["name"][stream_index]
+
+    #     # take care of sync channel
+    #     info = self.signals_info_dict[0, stream_name]
+    #     if not self.load_sync_channel and info["has_sync_trace"]:
+    #         memmap = memmap[:, :-1]
+
+    #     # since we cut the memmap, we can simplify the channel selection
+    #     if channel_indexes is None:
+    #         channel_selection = slice(None)
+    #     elif isinstance(channel_indexes, slice):
+    #         channel_selection = channel_indexes
+    #     elif not isinstance(channel_indexes, slice):
+    #         if np.all(np.diff(channel_indexes) == 1):
+    #             # consecutive channel then slice this avoid a copy (because of ndarray.take(...)
+    #             # and so keep the underlying memmap
+    #             channel_selection = slice(channel_indexes[0], channel_indexes[0] + len(channel_indexes))
+    #         else:
+    #             channel_selection = channel_indexes
+    #     else:
+    #         raise ValueError("get_analogsignal_chunk : channel_indexes" "must be slice or list or array of int")
+
+    #     raw_signals = memmap[slice(i_start, i_stop), channel_selection]
+
+    #     return raw_signals
 
     def _event_count(self, event_channel_idx, block_index=None, seg_index=None):
         timestamps, _, _ = self._get_event_timestamps(block_index, seg_index, event_channel_idx, None, None)

neo/rawio/utils.py

@@ -17,7 +17,7 @@ def get_memmap_shape(filename, dtype, num_channels=None, offset=0):
         return shape
 
 
-def get_memmap_chunk_from_open_file(fid, num_channels,  start, stop, dtype, file_offset=0):
+def get_memmap_chunk_from_opened_file(fid, num_channels,  start, stop, dtype, file_offset=0):
     """
     Utility fonction to get a chunk as a memmap array directly from an opened file.
     Using this instead memmap can avoid memmory consumption when multiprocessing.