diff --git a/doc/changes/dev/14237.newfeature.rst b/doc/changes/dev/14237.newfeature.rst new file mode 100644 index 00000000000..4de36460eb5 --- /dev/null +++ b/doc/changes/dev/14237.newfeature.rst @@ -0,0 +1 @@ +Speed up preloading Raw data and reading uniform EDF and BDF data by avoiding unnecessary setup and using pure-NumPy record-stride decoders, by `Bruno Aristimunha`_. diff --git a/mne/io/base.py b/mne/io/base.py index 22e4e1353ed..e9d0063ec2f 100644 --- a/mne/io/base.py +++ b/mne/io/base.py @@ -650,9 +650,10 @@ def _preload_data(self, preload): data_buffer = preload if isinstance(preload, bool | np.bool_) and not preload: data_buffer = None - t = self.times + n_times = self.n_times + last_time = (n_times - 1) / self.info["sfreq"] logger.info( - f"Reading 0 ... {len(t) - 1} = {0.0:9.3f} ... {t[-1]:9.3f} secs..." + f"Reading 0 ... {n_times - 1} = {0.0:9.3f} ... {last_time:9.3f} secs..." ) self._data = self._read_segment(data_buffer=data_buffer) assert len(self._data) == self.info["nchan"] diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 390289ae6bb..443a3517975 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -595,11 +595,37 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): assert dtype is not None # BDF if subtype == "bdf": - ch_data = read_from_file_or_buffer(fid, dtype=dtype, count=samp * dtype_byte) - ch_data = ch_data.reshape(-1, 3).astype(INT32) - ch_data = (ch_data[:, 0]) + (ch_data[:, 1] << 8) + (ch_data[:, 2] << 16) - # 24th bit determines the sign - ch_data[ch_data >= (1 << 23)] -= 1 << 24 + assert dtype_byte == 3 + expected = samp * dtype_byte + try: + raw = read_from_file_or_buffer(fid, dtype=dtype, count=expected) + except ValueError as err: + raise RuntimeError( + f"Could not read {expected} requested BDF bytes" + ) from err + if raw.size != expected: + raise RuntimeError( + f"Only {raw.size} of {expected} requested BDF bytes could be read" + ) + # Read each 3-byte sample as the low bytes of an overlapping 4-byte + # word, mask off the byte borrowed from the next sample, then move the + # sign bit to bit 31 and shift back down to sign-extend it. The last + # sample has no next sample to borrow from, so it is done by hand. + # This is equivalent to, and ~3x faster than, the readable version: + # + # ch_data = raw.reshape(-1, 3).astype(INT32) + # ch_data = ch_data[:, 0] | (ch_data[:, 1] << 8) | (ch_data[:, 2] << 16) + # ch_data <<= 8 # sign-extend bit 23 + # ch_data >>= 8 + ch_data = np.empty(samp, dtype=INT32) + packed = np.ndarray( + (max(samp - 1, 0),), dtype=">= 8 # GDF data and EDF data else: @@ -608,10 +634,11 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): return ch_data +_EDF_CHUNK_BYTES = 10 * 1024 * 1024 # read roughly this much per chunk + + def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, mult): """Read a chunk of raw data.""" - from scipy.interpolate import interp1d - n_samps = raw_extras["n_samps"] buf_len = int(raw_extras["max_samp"]) dtype = raw_extras["dtype_np"] @@ -634,11 +661,20 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # We could read this one EDF block at a time, which would be this: ch_offsets = np.cumsum(np.concatenate([[0], n_samps]), dtype=np.int64) - block_start_idx, r_lims, _ = _blk_read_lims(start, stop, buf_len) + block_start_idx, r_lims, d_lims = _blk_read_lims(start, stop, buf_len) # But to speed it up, we really need to read multiple blocks at once, # Otherwise we can end up with e.g. 18,181 chunks for a 20 MB file! - # Let's do ~10 MB chunks: - n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * dtype_byte), 1) + n_per = max(_EDF_CHUNK_BYTES // (ch_offsets[-1] * dtype_byte), 1) + + # When every picked channel stores buf_len samples per data record there is + # nothing to resample, so the picks form a plain (n_picks, n_times) block we + # can calibrate in one go straight into `data`. Mixed sampling rates, a + # projector, or a stim channel that needs interpolating use the per-channel + # loop below instead. + n_picks = len(idx_arr) + picks = read_sel[:n_picks] # picked signal channels, in output row order + uniform = mult is None and bool((n_samps[picks] == buf_len).all()) + stim_rows = [j for j, i in enumerate(idx_arr) if i in stim_channel_idxs] with _gdf_edf_get_fid(filenames, buffering=0) as fid: # Extract data @@ -647,7 +683,9 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # first read everything into the `ones` array. For channels with # lower sampling frequency, there will be zeros left at the end of the # row. Ignore TAL/annotations channel and only store `orig_sel` - ones = np.zeros((len(orig_sel), data.shape[-1]), dtype=data.dtype) + # `ones` has no rows on the fast path, which writes into `data` itself + n_stage = 0 if uniform else len(orig_sel) + ones = np.zeros((n_stage, data.shape[-1]), dtype=data.dtype) # save how many samples have already been read per channel n_smp_read = [0 for _ in range(len(orig_sel))] @@ -663,6 +701,23 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, r_sidx = r_lims[ai][0] r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] + if uniform: + block = np.empty((n_picks, n_read, buf_len), many_chunk.dtype) + for j, ci in enumerate(picks): + block[j] = many_chunk[:, ch_offsets[ci] : ch_offsets[ci + 1]] + for ci in read_sel[n_picks:]: # annotation channels + tal_data.append( + many_chunk[:, ch_offsets[ci] : ch_offsets[ci + 1]].copy() + ) + out = data[:, d_lims[ai][0] : d_lims[ai + n_read - 1][1]] + flat = block.reshape(n_picks, n_read * buf_len)[:, r_sidx:r_eidx] + np.multiply(flat, cal[idx_arr, np.newaxis], out=out) + out += offsets[idx_arr, np.newaxis] + out *= gains[idx_arr, np.newaxis] + for j in stim_rows: + out[j] = np.bitwise_and(out[j].astype(int), 2**17 - 1) + continue + # loop over selected channels, ci=channel selection for ii, ci in enumerate(read_sel): # This now has size (n_chunks_read, n_samp[ci]) @@ -682,6 +737,8 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, if n_samps[ci] != buf_len: if orig_idx in stim_channel_idxs: + from scipy.interpolate import interp1d + # Stim channel will be interpolated old = np.linspace(0, 1, n_samps[ci] + 1, True) new = np.linspace(0, 1, buf_len, False) @@ -735,6 +792,11 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, _mult_cal_one(data[:, :], ones, idx, cals, mult) + if uniform: + # stands in for the `data_view *= cals` that _mult_cal_one applies; the + # block above is skipped because the fast path leaves n_smp_read zero + data *= cals + if len(tal_data) > 1: tal_data = np.concatenate([tal.ravel() for tal in tal_data]) tal_data = tal_data[np.newaxis, :] diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 1672a131fc0..99d443ee0ea 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -4,6 +4,7 @@ import datetime import gc +import sys from contextlib import nullcontext from functools import partial from io import BytesIO @@ -67,6 +68,52 @@ misc = ["EXG1", "EXG5", "EXG8", "M1", "M2"] +def _multi_record(source, destination, n_records=6): + """Give a one-record EDF/BDF file several records, so windows cross them.""" + blob = bytearray(source.read_bytes()) + if int(blob[236:244]) != 1: + return source # already spans several records + header_nbytes = int(blob[184:192]) + blob[236:244] = f"{n_records:<8}".encode("ascii") + destination.write_bytes(blob[:header_nbytes] + blob[header_nbytes:] * n_records) + return destination + + +@pytest.mark.parametrize("chunk_bytes", (10 * 1024 * 1024, 4096)) # one chunk, several +@pytest.mark.parametrize( + "reader, path", + ( + # EDF+: an annotation channel at its own rate, plus two stim channels + (read_raw_edf, data_dir / "test_stim_channel.edf"), + (read_raw_edf, edf_stim_channel_path), + (read_raw_bdf, bdf_path), # BDF: 24-bit samples and a Status channel + ), + ids=("edf+", "edf", "bdf"), +) +def test_uniform_decode_matches_loop(reader, path, chunk_bytes, monkeypatch, tmp_path): + """Test uniform-rate records decode exactly like the per-channel loop.""" + path = _multi_record(path, tmp_path / f"uniform{path.suffix}") + monkeypatch.setattr(edf.edf, "_EDF_CHUNK_BYTES", chunk_bytes) + raw = reader(path, preload=False, verbose="error") + buf_len = int(raw._raw_extras[0]["max_samp"]) + n_channels = len(raw.ch_names) + picks = np.array([n_channels - 1, 1, 0]) # reordered, and includes any stim channel + kwargs = dict(picks=picks, start=buf_len // 2, stop=4 * buf_len + 13) + # the fast path writes into `data` itself, the per-channel loop stages the + # result and calls _mult_cal_one -- so counting it tells the two apart + staged = [] + stage = edf.edf._mult_cal_one + monkeypatch.setattr( + edf.edf, "_mult_cal_one", lambda *a: (staged.append(1), stage(*a))[1] + ) + want = raw.get_data(**kwargs) + assert not staged + # an identity projector sends the same values down the per-channel loop + raw._projector = np.eye(n_channels) + assert_array_equal(want, raw.get_data(**kwargs)) + assert staged + + def test_orig_units(): """Test exposure of original channel units.""" raw = read_raw_edf(edf_path, preload=True) @@ -319,8 +366,10 @@ def test_edf_data_broken(tmp_path): @pytest.mark.parametrize("method", ("constructor", "load_data")) -def test_edf_preload_memmap_ownership(method, tmp_path): +def test_edf_preload_memmap_ownership(method, tmp_path, monkeypatch): """Test ownership of a populated EDF preload memmap.""" + # uniformly sampled EDF must not need (and hence not import) interpolation + monkeypatch.setitem(sys.modules, "scipy.interpolate", None) memmap_fname = tmp_path / f"edf-{method}-memmap.dat" if method == "constructor": raw = read_raw_edf(edf_stim_channel_path, preload=memmap_fname) @@ -349,6 +398,37 @@ def test_duplicate_channel_labels_edf(): assert raw.ch_names == EXPECTED_CHANNEL_NAMES +@pytest.mark.parametrize( + "values", + ( + np.array([], dtype=np.int32), + np.array([-(1 << 23), -1, 0, 1, (1 << 23) - 1], dtype=np.int32), + np.random.default_rng(42).integers( + -(1 << 23), 1 << 23, size=1000, dtype=np.int32 + ), + ), +) +def test_read_ch_bdf_int24(values): + """Test exact signed 24-bit BDF decoding, including the last sample.""" + unsigned = values.astype(np.int64) & ((1 << 24) - 1) + packed = np.column_stack((unsigned, unsigned >> 8, unsigned >> 16)).astype(np.uint8) + with BytesIO(packed.tobytes()) as fid: + got = _read_ch( + fid, subtype="bdf", samp=len(values), dtype_byte=3, dtype=np.uint8 + ) + assert_array_equal(got, values) + + +def test_read_ch_bdf_short_read(tmp_path): + """Test truncated signed 24-bit BDF data raises instead of mixing bytes.""" + fname = tmp_path / "truncated.bdf" + fname.write_bytes(b"\x00" * 5) + # np.frombuffer raises on a short buffer, np.fromfile just returns less + for fid in (BytesIO(b"\x00" * 5), fname.open("rb")): + with fid, pytest.raises(RuntimeError, match="requested BDF bytes"): + _read_ch(fid, subtype="bdf", samp=2, dtype_byte=3, dtype=np.uint8) + + def test_parse_annotation(tmp_path): """Test parsing the tal channel.""" # test the parser diff --git a/mne/io/eeglab/eeglab.py b/mne/io/eeglab/eeglab.py index a16f1ae9779..ca97e6058cf 100644 --- a/mne/io/eeglab/eeglab.py +++ b/mne/io/eeglab/eeglab.py @@ -16,7 +16,7 @@ from ..._fiff.meas_info import create_info from ..._fiff.pick import _PICK_TYPES_KEYS from ..._fiff.utils import _find_channels, _mult_cal_one, _read_segments_file -from ...annotations import Annotations, read_annotations +from ...annotations import Annotations from ...channels import make_dig_montage from ...defaults import DEFAULTS from ...epochs import BaseEpochs @@ -496,7 +496,7 @@ def __init__( ) # create event_ch from annotations - annot = read_annotations(input_fname, uint16_codec=uint16_codec) + annot = _read_annotations_eeglab(eeg) self.set_annotations(annot) _check_boundary(annot, None) diff --git a/mne/io/eeglab/tests/test_eeglab.py b/mne/io/eeglab/tests/test_eeglab.py index f8c1aa79bfc..c5016951c06 100644 --- a/mne/io/eeglab/tests/test_eeglab.py +++ b/mne/io/eeglab/tests/test_eeglab.py @@ -5,6 +5,7 @@ import os import shutil from copy import deepcopy +from unittest.mock import Mock import numpy as np import pytest @@ -497,7 +498,7 @@ def test_eeglab_annotations(fname): @testing.requires_testing_data -def test_eeglab_read_annotations(): +def test_eeglab_read_annotations(monkeypatch): """Test annotations onsets are timestamps (+ validate some).""" annotations = read_annotations(raw_fname_mat) validation_samples = [0, 1, 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31] @@ -524,9 +525,12 @@ def test_eeglab_read_annotations(): ) # test if event durations are imported correctly + check_load_mat = Mock(wraps=mne.io.eeglab.eeglab._check_load_mat) + monkeypatch.setattr(mne.io.eeglab.eeglab, "_check_load_mat", check_load_mat) raw = read_raw_eeglab(raw_fname_event_duration, preload=True, montage_units="dm") # file contains 3 annotations with 0.5 s (64 samples) duration each assert_allclose(raw.annotations.duration, np.ones(3) * 0.5) + assert check_load_mat.call_count == 1 @testing.requires_testing_data diff --git a/mne/io/tests/test_raw.py b/mne/io/tests/test_raw.py index b6593046df8..bc63f7598c2 100644 --- a/mne/io/tests/test_raw.py +++ b/mne/io/tests/test_raw.py @@ -50,8 +50,14 @@ ) -def _fail_if_times_materialized(*args, **kwargs): - pytest.fail("The full Raw.times vector was materialized") +@pytest.fixture +def fail_if_times_materialized(monkeypatch): + """Fail the test if the full Raw.times vector is ever constructed.""" + + def _fail(*args, **kwargs): + pytest.fail("The full Raw.times vector was materialized") + + monkeypatch.setattr("mne.io.base._arange_div", _fail) def assert_named_constants(info): @@ -105,11 +111,16 @@ def test_orig_units(): BaseRaw(info, last_samps=[1], orig_units=True) -def test_set_annotations_does_not_materialize_times(monkeypatch): +def test_preload_does_not_materialize_times(fail_if_times_materialized): + """Test preloading does not construct the full time vector.""" + raw = read_raw_fif(raw_fname, preload=True, verbose="error") + assert raw.preload + + +def test_set_annotations_does_not_materialize_times(fail_if_times_materialized): """Test annotation bounds use the scalar recording endpoint.""" raw = read_raw_fif(raw_fname, preload=False, verbose="error") annotations = Annotations([0.0], [0.1], ["test"]) - monkeypatch.setattr("mne.io.base._arange_div", _fail_if_times_materialized) raw.set_annotations(annotations) assert len(raw.annotations) == 1 diff --git a/tools/vulture_allowlist.py b/tools/vulture_allowlist.py index 5d529ceb9b1..6625c2494f5 100644 --- a/tools/vulture_allowlist.py +++ b/tools/vulture_allowlist.py @@ -9,6 +9,7 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. +fail_if_times_materialized numba_conditional options_3d invisible_fig