diff --git a/neo/rawio/biocamrawio.py b/neo/rawio/biocamrawio.py index d62cd6581..488fa4525 100644 --- a/neo/rawio/biocamrawio.py +++ b/neo/rawio/biocamrawio.py @@ -61,6 +61,14 @@ class BiocamRawIO(BaseRawIO): extensions = ["h5", "brw"] rawmode = "one-file" + # Ceiling on the temporary all-channel buffer used when only a subset of channels is requested. + # The dense layouts read a contiguous slice, so their cost scales with bytes and a small buffer + # stays in cache. The event-based sparse decoder re-reads both tables of contents and re-walks + # the event blob on every call, so its cost is per call rather than per byte and it wants a + # buffer large enough that it almost never splits. + max_read_bytes_dense = 4 * 1024 * 1024 + max_read_bytes_sparse = 512 * 1024 * 1024 + def __init__(self, filename="", fill_gaps_strategy="zeros"): BaseRawIO.__init__(self) self.filename = filename @@ -136,12 +144,8 @@ def _get_signal_t_start(self, block_index, seg_index, stream_index): raise ValueError("`stream_index must be 0") return self._segment_t_start(block_index, seg_index) - def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes): - if i_start is None: - i_start = 0 - if i_stop is None: - i_stop = self._num_frames - + def _read_block(self, i_start, i_stop): + """Read frames [i_start, i_stop) for all channels, as (n_frames, n_channels).""" # read functions are different based on the version of biocam if self._read_function is readHDF5t_brw4_sparse: if self._fill_gaps_strategy is None: @@ -163,34 +167,49 @@ def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, strea else: data = self._read_function(self._filehandle, i_start, i_stop, self._num_channels) - # older style data returns array of (n_samples, n_channels), should be a view - # but if memory issues come up we should doublecheck out how the file is being stored - if data.ndim > 1: - if channel_indexes is None: - channel_indexes = slice(None) - sig_chunk = data[:, channel_indexes] - - # newer style data returns an initial flat array (n_samples * n_channels) - # we iterate through channels rather than slicing - # Due to the fact that Neo and SpikeInterface tend to prefer slices we need to add - # some careful checks around slicing of None in the case we need to iterate through - # channels. First check if None. Then check if slice and only if slice check that it is slice(None) + # Newer style data is read as a flat (n_samples * n_channels) array. The h5py read is + # C-contiguous, so this reshape is a free view: it copies nothing and allocates nothing. + # Older style data is already (n_samples, n_channels). + if data.ndim == 1: + data = data.reshape(i_stop - i_start, self._num_channels) + return data + + def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes): + if i_start is None: + i_start = 0 + if i_stop is None: + i_stop = self._num_frames + n_frames = i_stop - i_start + if channel_indexes is None: + channel_indexes = slice(None) + # Number of selected channels, for every form `channel_indexes` can take: a slice, an + # integer array or a boolean mask. + n_ch = np.arange(self._num_channels)[channel_indexes].size + dtype = self.header["signal_channels"]["dtype"][0] + itemsize = np.dtype(dtype).itemsize + read_bytes = n_frames * self._num_channels * itemsize + out_bytes = n_frames * n_ch * itemsize + max_read_bytes = ( + self.max_read_bytes_sparse if self._read_function is readHDF5t_brw4_sparse else self.max_read_bytes_dense + ) + + # Reading a subset of channels from a very large recording. + # The BioCam hdf5 layout is such that all channels are chunked together, so they cannot be + # read independently. In the event that only a subset of channels of a very large recording + # is needed, we cannot slice a single channel. Instead, we must loop and read all channels, + # slicing as we go. We will do this chunking if the read is large and the output to be + # returned is smaller. + read_by_chunks = read_bytes > max_read_bytes and out_bytes < read_bytes // 2 + if read_by_chunks: + # Read in chunks, copying desired channels into a preallocated array. + out = np.empty((n_frames, n_ch), dtype=dtype) + max_frames = max(1, max_read_bytes // (itemsize * self._num_channels)) + for start in range(i_start, i_stop, max_frames): + stop = min(start + max_frames, i_stop) + out[start - i_start : stop - i_start] = self._read_block(start, stop)[:, channel_indexes] else: - if channel_indexes is None: - channel_indexes = [ch for ch in range(self._num_channels)] - elif isinstance(channel_indexes, slice): - start = channel_indexes.start or 0 - stop = channel_indexes.stop or self._num_channels - step = channel_indexes.step or 1 - channel_indexes = [ch for ch in range(start, stop, step)] - - sig_chunk = np.zeros((i_stop - i_start, len(channel_indexes)), dtype=data.dtype) - # iterate through channels to prevent loading all channels into memory which can cause - # memory exhaustion. See https://github.com/SpikeInterface/spikeinterface/issues/3303 - for index, channel_index in enumerate(channel_indexes): - sig_chunk[:, index] = data[channel_index :: self._num_channels] - - return sig_chunk + out = self._read_block(i_start, i_stop)[:, channel_indexes] + return out def open_biocam_file_header(filename) -> dict: diff --git a/neo/test/rawiotest/test_biocamrawio.py b/neo/test/rawiotest/test_biocamrawio.py index 12ebf39ad..7a59d4821 100644 --- a/neo/test/rawiotest/test_biocamrawio.py +++ b/neo/test/rawiotest/test_biocamrawio.py @@ -24,8 +24,76 @@ class TestBiocamRawIO( entities_to_test = [ "biocam/biocam_hw3.0_fw1.6.brw", "biocam/biocam_hw3.0_fw1.7.0.12_raw.brw", + "biocam/raw_shape_variants/version_100.brw", + "biocam/raw_shape_variants/version_101.brw", ] + # The two files in `raw_shape_variants` hold the same recording, 64 channels over 1000 samples, + # written in the two shapes that the BRW v3.x specification defines for `3BData/Raw`: a T x W + # matrix for version 100 and a flat interleaved array for version 101. The value stored at a + # given position is `sample * 64 + channel`, so no two positions share a value and a misread + # returns a number that says where it actually came from. + shape_variants = [ + "biocam/raw_shape_variants/version_100.brw", + "biocam/raw_shape_variants/version_101.brw", + ] + variant_num_channels = 64 + variant_num_frames = 1000 + + # Some of the ways that neo and SpikeInterface are known to ask for channels. + channel_queries = [None, slice(None), slice(2, 6), [0, 1, 2], [3, 0, 2], [5], np.array([1, 4])] + # None means the whole segment. The rest include a non-zero start, a single frame, and the + # empty range at the end of the recording. + frame_ranges = [(None, None), (0, 5), (7, 13), (999, 1000), (1000, 1000)] + + def test_get_analogsignal_chunk_values(self): + """Values, shape and dtype are correct for every way of selecting channels.""" + expected_all = np.arange(self.variant_num_frames * self.variant_num_channels, dtype="uint16") + expected_all = expected_all.reshape(self.variant_num_frames, self.variant_num_channels) + + for entity in self.shape_variants: + reader = BiocamRawIO(filename=self.get_local_path(entity)) + reader.parse_header() + for channel_indexes in self.channel_queries: + for i_start, i_stop in self.frame_ranges: + with self.subTest(entity=entity, channel_indexes=channel_indexes, i_start=i_start, i_stop=i_stop): + sig = reader.get_analogsignal_chunk( + block_index=0, + seg_index=0, + i_start=i_start, + i_stop=i_stop, + channel_indexes=channel_indexes, + ) + start = 0 if i_start is None else i_start + stop = self.variant_num_frames if i_stop is None else i_stop + columns = slice(None) if channel_indexes is None else channel_indexes + expected = expected_all[start:stop][:, columns] + assert sig.dtype == expected.dtype + np.testing.assert_array_equal(sig, expected) + + def test_chunked_read_matches_direct_read(self): + """Shrinking the read budget must not change the data that comes back. + + A request for a few channels over a long stretch of a large recording is served by + reading all the channels in blocks and copying out the requested ones, the BioCAM + layout interleaving the channels so that they cannot be read independently. + """ + for entity in self.entities_to_test: + reader = BiocamRawIO(filename=self.get_local_path(entity)) + reader.parse_header() + for channel_indexes in self.channel_queries: + with self.subTest(entity=entity, channel_indexes=channel_indexes): + kwargs = dict(block_index=0, seg_index=0, i_start=0, i_stop=100, channel_indexes=channel_indexes) + direct = reader.get_analogsignal_chunk(**kwargs) + # A budget of ten frames splits the hundred frames requested above into ten + # blocks that the reader has to stitch back together. + reader.max_read_bytes_dense = 10 * reader._num_channels * np.dtype("uint16").itemsize + chunked = reader.get_analogsignal_chunk(**kwargs) + del reader.max_read_bytes_dense + + assert chunked.dtype == direct.dtype + np.testing.assert_array_equal(chunked, direct) + def test_biocamrawio_gain(tmp_path): """Test that BiocamRawIO correctly reads the gain from a Biocam HDF5 file.