Speed up EDF and BDF reading (#14237)

Co-authored-by: Eric Larson <larson.eric.d@gmail.com>

Author: bruAristimunha

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`_.

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"]

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="<u4", buffer=raw, strides=(dtype_byte,)
+        )
+        np.bitwise_and(packed, (1 << 24) - 1, out=ch_data[:-1])
+        if samp:
+            ch_data[-1] = int(raw[-3]) | int(raw[-2]) << 8 | int(raw[-1]) << 16
+        ch_data <<= 8
+        ch_data >>= 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, :]

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

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
@@ -501,7 +501,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)
 

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

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
 

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