Split large rec files between multiple SpikegadgetsRawIO iterators

src/spikegadgets_to_nwb/convert_dios.py

@@ -70,24 +70,26 @@ def add_dios(nwbfile: NWBFile, recfile: list[str], metadata: dict) -> None:
             prefix = "ECU_"
             break
 
-    # compile data for all dio channels in all files
-    all_timestamps = [np.array([], dtype=np.float64) for i in channel_name_map]
-    all_state_changes = [np.array([], dtype=np.float64) for i in channel_name_map]
-    for io in neo_io:
-        for i, channel_name in enumerate(channel_name_map):
+    for channel_name in channel_name_map:
+        # merge streams from multiple files
+        all_timestamps = []
+        all_state_changes = []
+        for io in neo_io:
             timestamps, state_changes = io.get_digitalsignal(
                 stream_name, prefix + channel_name
             )
-            all_timestamps[i] = np.concatenate((all_timestamps[i], timestamps))
-            all_state_changes[i] = np.concatenate((all_state_changes[i], state_changes))
-    # Add each channel as a behavioral event time series
-    for i, channel_name in enumerate(channel_name_map):
+            all_timestamps.append(timestamps)
+            all_state_changes.append(state_changes)
+        all_timestamps = np.concatenate(all_timestamps)
+        all_state_changes = np.concatenate(all_state_changes)
+        assert isinstance(all_timestamps[0], np.float64)
+        assert isinstance(all_timestamps, np.ndarray)
         ts = TimeSeries(
             name=channel_name_map[channel_name],
             description=channel_name,
-            data=all_state_changes[i],
+            data=all_state_changes,
             unit="-1",  # TODO change to "N/A",
-            timestamps=all_timestamps[i],  # TODO adjust timestamps
+            timestamps=all_timestamps,  # TODO adjust timestamps
         )
         beh_events.add_timeseries(ts)
 

src/spikegadgets_to_nwb/convert_ephys.py

@@ -10,12 +10,13 @@
 
 from spikegadgets_to_nwb import convert_rec_header
 
-from .spike_gadgets_raw_io import SpikeGadgetsRawIO
+from .spike_gadgets_raw_io import SpikeGadgetsRawIO, SpikeGadgetsRawIOPartial
 
 MICROVOLTS_PER_VOLT = 1e6
 VOLTS_PER_MICROVOLT = 1e-6
 MILLISECONDS_PER_SECOND = 1e3
 NANOSECONDS_PER_SECOND = 1e9
+MAXIMUM_ITERATOR_SIZE = int(30000 * 60 * 30)  # 30 min of data at 30 kHz
 
 
 class RecFileDataChunkIterator(GenericDataChunkIterator):
@@ -83,6 +84,43 @@
         else:
             self.nwb_hw_channel_order = nwb_hw_channel_order
 
+        """split excessively large iterators into smaller ones
+        """
+        iterator_size = [neo_io._raw_memmap.shape[0] for neo_io in self.neo_io]
+        iterator_size.reverse()
+        for i, size in enumerate(
+            iterator_size
+        ):  # iterate backwards so can insert new iterators
+            if size > MAXIMUM_ITERATOR_SIZE:
+                # split into smaller iterators
+                sub_iterators = []
+                j = 0
+                previous_multiplex_state = None
+                iterator_loc = len(iterator_size) - i - 1
+                while j < size:
+                    sub_iterators.append(
+                        SpikeGadgetsRawIOPartial(
+                            self.neo_io[iterator_loc],
+                            start_index=j,
+                            stop_index=j + MAXIMUM_ITERATOR_SIZE,
+                            previous_multiplex_state=previous_multiplex_state,
+                        )
+                    )
+                    if self.n_multiplexed_channel > 0:
+                        partial_size = sub_iterators[-1]._raw_memmap.shape[0]
+                        previous_multiplex_state = sub_iterators[
+                            -1
+                        ].get_analogsignal_multiplexed_partial(
+                            i_start=partial_size - 10,
+                            i_stop=partial_size,
+                            padding=30000,
+                        )[
+                            -1
+                        ]
+                    j += MAXIMUM_ITERATOR_SIZE
+                self.neo_io.pop(iterator_loc)
+                self.neo_io[iterator_loc:iterator_loc] = sub_iterators
+        logger.info(f"# iterators: {len(self.neo_io)}")
         # NOTE: this will read all the timestamps from the rec file, which can be slow
         if timestamps is not None:
             self.timestamps = timestamps

src/spikegadgets_to_nwb/spike_gadgets_raw_io.py

@@ -14,6 +14,8 @@
 )
 from scipy.stats import linregress
 
+INT_16_CONVERSION = 256
+
 
 class SpikeGadgetsRawIO(BaseRawIO):
     extensions = ["rec"]
@@ -526,7 +528,7 @@
         # read the data into int16
         data = (
             self._raw_memmap[:, data_offsets[:, 0]]
-            + self._raw_memmap[:, data_offsets[:, 0] + 1] * 256
+            + self._raw_memmap[:, data_offsets[:, 0] + 1] * INT_16_CONVERSION
         )
         # initialize the first row
         analog_multiplexed_data[0] = data[0]
@@ -537,6 +539,94 @@
             )
         return analog_multiplexed_data
 
+    def get_analogsignal_multiplexed_partial(
+        self,
+        i_start: int,
+        i_stop: int,
+        channel_names: list = None,
+        padding: int = 30000,
+    ) -> np.ndarray:
+        """Alternative method to access part of the multiplexed data.
+        Not memory efficient for many calls because it reads a buffer chunk before the requested data.
+        Better than get_analogsignal_multiplexed when need one call to specific time region
+
+        Parameters
+        ----------
+        i_start : int
+            index start
+        i_stop : int
+            index stop
+        channel_names : list[str], optional
+            channels to get, by default None will get all multiplex channels
+        padding : int, optional
+            how many packets before the desired series to load to ensure every channel receives update before requested,
+            by default 30000
+
+        Returns
+        -------
+        np.ndarray
+            multiplex data
+
+        Raises
+        ------
+        ValueError
+            _description_
+        """
+        print("compute multiplex cache", self.filename)
+        if channel_names is None:
+            # read all multiplexed channels
+            channel_names = list(self.multiplexed_channel_xml.keys())
+        else:
+            for ch_name in channel_names:
+                if ch_name not in self.multiplexed_channel_xml:
+                    raise ValueError(f"Channel name '{ch_name}' not found in file.")
+        # determine which packets to get from data
+        padding = min(padding, i_start)
+        i_start = i_start - padding
+        if i_stop is None:
+            i_stop = self._raw_memmap.shape[0]
+
+        # Make object to hold data
+        num_packet = i_stop - i_start
+        analog_multiplexed_data = np.empty(
+            (num_packet, len(channel_names)), dtype=np.int16
+        )
+
+        # precompute the static data offsets
+        data_offsets = np.empty((len(channel_names), 3), dtype=int)
+        for j, ch_name in enumerate(channel_names):
+            ch_xml = self.multiplexed_channel_xml[ch_name]
+            data_offsets[j, 0] = int(
+                self._multiplexed_byte_start + int(ch_xml.attrib["startByte"])
+            )
+            data_offsets[j, 1] = int(ch_xml.attrib["interleavedDataIDByte"])
+            data_offsets[j, 2] = int(ch_xml.attrib["interleavedDataIDBit"])
+        interleaved_data_id_byte_values = self._raw_memmap[
+            i_start:i_stop, data_offsets[:, 1]
+        ]
+        interleaved_data_id_bit_values = (
+            interleaved_data_id_byte_values >> data_offsets[:, 2]
+        ) & 1
+        # calculate which packets encode for which channel
+        initialize_stream_mask = np.logical_or(
+            (np.arange(num_packet) == 0)[:, None], interleaved_data_id_bit_values == 1
+        )
+        # read the data into int16
+        data = (
+            self._raw_memmap[i_start:i_stop, data_offsets[:, 0]]
+            + self._raw_memmap[i_start:i_stop, data_offsets[:, 0] + 1]
+            * INT_16_CONVERSION
+        )
+        # initialize the first row
+        analog_multiplexed_data[0] = data[0]
+        # for packets that do not have an update for a channel, use the previous value
+        # this method assumes that every channel has an update within the buffer
+        for i in range(1, num_packet):
+            analog_multiplexed_data[i] = np.where(
+                initialize_stream_mask[i], data[i], analog_multiplexed_data[i - 1]
+            )
+        return analog_multiplexed_data[padding:]
+
     def get_digitalsignal(self, stream_id, channel_id):
         # stream_id = self.header["signal_streams"][stream_index]["id"]
 
@@ -713,3 +803,143 @@
         # if list of indeces
         else:
             return self.mapped_index[index]
+
+
+class SpikeGadgetsRawIOPartial(SpikeGadgetsRawIO):
+    extensions = ["rec"]
+    rawmode = "one-file"
+
+    def __init__(
+        self,
+        full_io: SpikeGadgetsRawIO,
+        start_index: int,
+        stop_index: int,
+        previous_multiplex_state: np.ndarray = None,
+    ):
+        """Initialize a partial SpikeGadgetsRawIO object.
+
+        Parameters
+        ----------
+        full_io : SpikeGadgetsRawIO
+            The SpikeGadgetsRawIO for the complete rec file
+        start_index : int
+            Where this partial file starts in the complete file
+        stop_index : int
+            Where this partial file stops in the complete file
+        previous_multiplex_state : np.ndarray, optional
+            The last multiplex state in the previous partial file.
+            If None, will default to behavior of SpikeGadgetsRawIO.
+            Use None if first partial iterator for the rec file or if not accessing multiplex data.
+            By default None
+        """
+        # initialization from the base class
+        BaseRawIO.__init__(self)
+        self.filename = full_io.filename
+        self.selected_streams = full_io.selected_streams
+        self.interpolate_dropped_packets = full_io.interpolate_dropped_packets
+
+        # define some key information
+        self.interpolate_index = None
+        self.previous_multiplex_state = previous_multiplex_state
+
+        # copy conserved information from parsed_header from full_io
+        self.header = full_io.header
+        self.system_time_at_creation = full_io.system_time_at_creation
+        self.timestamp_at_creation = full_io.timestamp_at_creation
+        self._sampling_rate = full_io._sampling_rate
+        self.sysClock_byte = full_io.sysClock_byte
+        self._timestamp_byte = full_io._timestamp_byte
+        self._mask_channels_ids = full_io._mask_channels_ids
+        self._mask_channels_bytes = full_io._mask_channels_bytes
+        self._mask_channels_bits = full_io._mask_channels_bits
+        self.multiplexed_channel_xml = full_io.multiplexed_channel_xml
+        self._multiplexed_byte_start = full_io._multiplexed_byte_start
+        self._mask_streams = full_io._mask_streams
+        self.selected_streams = full_io.selected_streams
+        self._generate_minimal_annotations()
+
+        # crop key information to range of interest
+        header_size = None
+        with open(self.filename, mode="rb") as f:
+            while True:
+                line = f.readline()
+                if b"</Configuration>" in line:
+                    header_size = f.tell()
+                    break
+
+            if header_size is None:
+                ValueError(
+                    "SpikeGadgets: the xml header does not contain '</Configuration>'"
+                )
+
+        raw_memmap = np.memmap(self.filename, mode="r", offset=header_size, dtype="<u1")
+        packet_size = full_io._raw_memmap.shape[1]
+        num_packet = raw_memmap.size // packet_size
+        raw_memmap = raw_memmap[: num_packet * packet_size]
+        self._raw_memmap = raw_memmap.reshape(-1, packet_size)
+        self._raw_memmap = self._raw_memmap[start_index:stop_index]
+
+    @functools.lru_cache(maxsize=2)
+    def get_analogsignal_multiplexed(self, channel_names=None) -> np.ndarray:
+        """
+        Overide of the superclass to use the last state of the previous file segment
+        to define the first state of the current file segment.
+        """
+        print("compute multiplex cache", self.filename)
+        if channel_names is None:
+            # read all multiplexed channels
+            channel_names = list(self.multiplexed_channel_xml.keys())
+        else:
+            for ch_name in channel_names:
+                if ch_name not in self.multiplexed_channel_xml:
+                    raise ValueError(f"Channel name '{ch_name}' not found in file.")
+
+        # because of the encoding scheme, it is easiest to read all the data in sequence
+        # one packet at a time
+        num_packet = self._raw_memmap.shape[0]
+        analog_multiplexed_data = np.empty(
+            (num_packet, len(channel_names)), dtype=np.int16
+        )
+
+        # precompute the static data offsets
+        data_offsets = np.empty((len(channel_names), 3), dtype=int)
+        for j, ch_name in enumerate(channel_names):
+            ch_xml = self.multiplexed_channel_xml[ch_name]
+            data_offsets[j, 0] = int(
+                self._multiplexed_byte_start + int(ch_xml.attrib["startByte"])
+            )
+            data_offsets[j, 1] = int(ch_xml.attrib["interleavedDataIDByte"])
+            data_offsets[j, 2] = int(ch_xml.attrib["interleavedDataIDBit"])
+        interleaved_data_id_byte_values = self._raw_memmap[:, data_offsets[:, 1]]
+        interleaved_data_id_bit_values = (
+            interleaved_data_id_byte_values >> data_offsets[:, 2]
+        ) & 1
+        # calculate which packets encode for which channel
+        initialize_stream_mask = np.logical_or(
+            (np.arange(num_packet) == 0)[:, None], interleaved_data_id_bit_values == 1
+        )
+        # read the data into int16
+        data = (
+            self._raw_memmap[:, data_offsets[:, 0]]
+            + self._raw_memmap[:, data_offsets[:, 0] + 1] * INT_16_CONVERSION
+        )
+        # initialize the first row
+        # if no previous state, assume first segment. Default to superclass behavior
+        analog_multiplexed_data[0] = data[0]
+        if not self.previous_multiplex_state is None:
+            # if previous state, use it to initialize elements of first row not updated in that packet
+            ind = np.where(initialize_stream_mask[0])[0]
+            analog_multiplexed_data[0][ind] = self.previous_multiplex_state[ind]
+        # for packets that do not have an update for a channel, use the previous value
+        for i in range(1, num_packet):
+            analog_multiplexed_data[i] = np.where(
+                initialize_stream_mask[i], data[i], analog_multiplexed_data[i - 1]
+            )
+        return analog_multiplexed_data
+
+    def get_digitalsignal(self, stream_id, channel_id):
+        dio_change_times, change_dir_trim = super().get_digitalsignal(
+            stream_id, channel_id
+        )
+        # clip the setting of the first state
+        return dio_change_times[1:], change_dir_trim[1:]

src/spikegadgets_to_nwb/tests/test_convert_intervals.py

@@ -16,17 +16,14 @@ def test_add_epochs():
     metadata, _ = load_metadata(metadata_path, [])
     nwbfile = initialize_nwb(metadata, default_test_xml_tree())
     file_info = get_file_info(data_path)
-    rec_to_nwb_file = data_path / "minirec20230622_.nwb"  # comparison file
-    # get all streams for all files
-    neo_io = [
-        SpikeGadgetsRawIO(filename=file)
-        for file in file_info[file_info.file_extension == ".rec"].full_path
-    ]
-    [neo_io.parse_header() for neo_io in neo_io]
-
     file_info = file_info[file_info.animal == "sample"]
     file_info = file_info[file_info.date == 20230622]
-    add_epochs(nwbfile, file_info, neo_io)
+    rec_to_nwb_file = data_path / "minirec20230622_.nwb"  # comparison file
+    # get all streams for all files
+    rec_dci = RecFileDataChunkIterator(
+        file_info[file_info.file_extension == ".rec"].full_path.to_list()
+    )
+    add_epochs(nwbfile, file_info, rec_dci.neo_io)
     epochs_df = nwbfile.epochs.to_dataframe()
     # load old nwb version
     io = NWBHDF5IO(rec_to_nwb_file, "r")