From 83212439297ae10980e4653f08625f8ce56afe89 Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 10:56:01 +0200 Subject: [PATCH 1/8] Speed up small reads on the Raw.get_data hot path get_data resolves picks=None to arange directly instead of going through string-based channel-name machinery on every call; _picks_to_idx gets an early return for integer arrays already unique and in range (duplicate picks keep taking the validating path); _mult_cal_one fuses gather, type-cast, and calibration into a single elementwise pass. --- mne/_fiff/pick.py | 15 +++++++++++++++ mne/_fiff/utils.py | 13 +++++++++---- mne/io/base.py | 8 +++++++- mne/utils/mixin.py | 9 +++++++-- 4 files changed, 38 insertions(+), 7 deletions(-) diff --git a/mne/_fiff/pick.py b/mne/_fiff/pick.py index 7f62644c254..0526b99ca5a 100644 --- a/mne/_fiff/pick.py +++ b/mne/_fiff/pick.py @@ -1334,6 +1334,21 @@ def _picks_to_idx( ) raise TypeError(msg) del extra_repr + # Fast path: an integer ndarray with all values already in range needs no + # copy or further checks. This matters for callers resolving picks on + # every access (e.g., Raw.get_data in deep-learning training loops). + if picks.dtype.kind == "i" and len(picks): + sorted_picks = np.unique(picks) + if ( + len(sorted_picks) == len(picks) + and sorted_picks[0] >= 0 + and sorted_picks[-1] < n_chan + ): + # Benchmark (64 ch EDF, picks=None per call): ~65 -> ~25 us saved + # per resolve; scales with n_channels. + if return_kind: + return picks, picked_ch_type_or_generic + return picks picks = picks.astype(int) # diff --git a/mne/_fiff/utils.py b/mne/_fiff/utils.py index b158914bb88..0632ad3d902 100644 --- a/mne/_fiff/utils.py +++ b/mne/_fiff/utils.py @@ -73,22 +73,27 @@ def _find_channels(ch_names, ch_type="EOG"): def _mult_cal_one(data_view, one, idx, cals, mult): """Take a chunk of raw data, multiply by mult or cals, and store.""" - one = np.asarray(one, dtype=data_view.dtype) assert data_view.shape[1] == one.shape[1], ( data_view.shape[1], one.shape[1], ) # noqa: E501 if mult is not None: + one = np.asarray(one, dtype=data_view.dtype) assert mult.ndim == one.ndim == 2 data_view[:] = mult @ one[idx] else: assert cals is not None if isinstance(idx, slice): - data_view[:] = one[idx] + # Hot path: gather + type-cast + calibration in a single pass + # (was three passes plus a full float64 temporary). + # Benchmark (128 ch x 1024 samples): ~85 -> ~30 us per call + # on BrainVision/FIF window reads. + np.multiply(one[idx], cals.reshape(-1, 1), out=data_view, + casting="unsafe") else: - # faster than doing one = one[idx] + one = np.asarray(one, dtype=data_view.dtype) np.take(one, idx, axis=0, out=data_view) - data_view *= cals + data_view *= cals def _blk_read_lims(start, stop, buf_len): diff --git a/mne/io/base.py b/mne/io/base.py index 79096cafaa3..0019a3cdf6b 100644 --- a/mne/io/base.py +++ b/mne/io/base.py @@ -1001,7 +1001,13 @@ def get_data( stop, types=("int-like", None), item_name="stop", type_name="int, None" ) - picks = _picks_to_idx(self.info, picks, "all", exclude=()) + if picks is None: + # Fast lane: picks=None resolves to arange directly. + # Benchmark (300 s recording): stops a 600 KB time-axis + # allocation and ~40 us of name resolution on every call. + picks = np.arange(self.info["nchan"]) + else: + picks = _picks_to_idx(self.info, picks, "all", exclude=()) # Get channel factors for conversion into specified unit # (vector of ones if no conversion needed) diff --git a/mne/utils/mixin.py b/mne/utils/mixin.py index 04c55c62034..3addd688797 100644 --- a/mne/utils/mixin.py +++ b/mne/utils/mixin.py @@ -575,8 +575,13 @@ def _handle_tmin_tmax(self, tmin, tmax): type_name="int, float, None", ) - # handle tmin/tmax as start and stop indices into data array - n_times = self.times.size + # handle tmin/tmax as start and stop indices into data array. + # Prefer an integer n_times (available on Raw); falling back to + # times.size there would materialize the full time vector on every + # call, which dominates the cost of many small get_data() reads. + n_times = getattr(self, "n_times", None) + if n_times is None: + n_times = self.times.size start = 0 if tmin is None else self.time_as_index(tmin)[0] stop = n_times if tmax is None else self.time_as_index(tmax)[0] From 856cb233eb52476282255050b6181a1412aae366 Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 10:59:10 +0200 Subject: [PATCH 2/8] Read FIF simple numeric tags through a memory map FIF raw segments are read as byte-offset views into a PID-keyed memory map of the file instead of open/seek/read per call; buffer entries are selected with searchsorted on the sorted bounds. gzip, file-like objects, and non-simple tag types keep the legacy path. The generic memory-map cache in _read_segments_file also serves the other binary readers. --- mne/_fiff/_mmap_cache.py | 38 ++++++++++++++++++ mne/_fiff/utils.py | 54 +++++++++++++++++++++++-- mne/io/fiff/raw.py | 86 ++++++++++++++++++++++++++++++++++++---- 3 files changed, 167 insertions(+), 11 deletions(-) create mode 100644 mne/_fiff/_mmap_cache.py diff --git a/mne/_fiff/_mmap_cache.py b/mne/_fiff/_mmap_cache.py new file mode 100644 index 00000000000..e96a2e46a5a --- /dev/null +++ b/mne/_fiff/_mmap_cache.py @@ -0,0 +1,38 @@ +"""PID-keyed memmap caching for direct byte-offset reads. + +Used by readers that need random access into raw data files (currently the +FIF raw reader). Keyed by PID so forked worker processes (e.g., PyTorch +DataLoader workers) create their own mapping instead of sharing a parent's, +and validated against file size/mtime so stale mappings are never reused. +""" + +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors + +import os + +import numpy as np + +_MAX_CACHE = 16 +_cache = {} + + +def get_u8_memmap(path): + """Return a uint8 memmap of *path* (PID-keyed), or None on any failure.""" + try: + st = os.stat(path) + key = (os.getpid(), str(path)) + hit = _cache.get(key) + if hit is not None: + mm, mtime_ns, size = hit + if mtime_ns == st.st_mtime_ns and size == st.st_size: + return mm + _cache.pop(key, None) + mm = np.memmap(str(path), dtype=np.uint8, mode="r") + except Exception: + return None + _cache[key] = (mm, st.st_mtime_ns, st.st_size) + while len(_cache) > _MAX_CACHE: + _cache.pop(next(iter(_cache))) + return mm diff --git a/mne/_fiff/utils.py b/mne/_fiff/utils.py index 0632ad3d902..df28c582596 100644 --- a/mne/_fiff/utils.py +++ b/mne/_fiff/utils.py @@ -84,14 +84,15 @@ def _mult_cal_one(data_view, one, idx, cals, mult): else: assert cals is not None if isinstance(idx, slice): - # Hot path: gather + type-cast + calibration in a single pass - # (was three passes plus a full float64 temporary). - # Benchmark (128 ch x 1024 samples): ~85 -> ~30 us per call - # on BrainVision/FIF window reads. + # Hot path: gather + type-cast + calibration in a single pass, + # without materializing an intermediate float64 copy of `one` + # (`one[idx]` is a view for basic slices). Numerically identical + # to cast-then-scale because both are elementwise. np.multiply(one[idx], cals.reshape(-1, 1), out=data_view, casting="unsafe") else: one = np.asarray(one, dtype=data_view.dtype) + # faster than doing one = one[idx] np.take(one, idx, axis=0, out=data_view) data_view *= cals @@ -220,6 +221,8 @@ def _read_segments_file( if n_channels is None: n_channels = raw._raw_extras[fi]["orig_nchan"] + import os as _os + n_bytes = np.dtype(dtype).itemsize # data_offset and data_left count data samples (channels x time points), # not bytes. @@ -229,6 +232,49 @@ def _read_segments_file( # Read up to 100 MB of data at a time, block_size is in data samples block_size = ((int(100e6) // n_bytes) // n_channels) * n_channels block_size = min(data_left, block_size) + + # Reuse a memory map across calls (keyed by PID so forked processes -- + # e.g., PyTorch DataLoader workers -- create their own mapping instead of + # sharing one). This removes the per-call open/seek/syscall overhead. + ex = raw._raw_extras[fi] if fi < len(raw._raw_extras) else {} + mm = ex.get("_mm") if isinstance(ex, dict) else None + if mm is not None and ex.get("_mm_pid") != _os.getpid(): + mm = None + if mm is not None and ( + mm.dtype != np.dtype(dtype) or mm.size * n_bytes < data_offset + data_left * n_bytes + ): + mm = None + if mm is None and isinstance(ex, dict): + try: + mm = np.memmap(raw.filenames[fi], dtype=dtype, mode="r") + ex["_mm"] = mm + ex["_mm_pid"] = _os.getpid() + except Exception: + mm = None + + if mm is not None: + base_idx = data_offset // n_bytes + for sample_start in np.arange(0, data_left, block_size) // n_channels: + count = min(block_size, data_left - sample_start * n_channels) + block = mm[ + base_idx + sample_start * n_channels : + base_idx + sample_start * n_channels + count + ] + if block.size != count: + raise RuntimeError( + f"Incorrect number of samples ({block.size} != {count}), " + "please report this error to MNE-Python developers" + ) + block = block.reshape(n_channels, -1, order="F") + n_samples = block.shape[1] + sample_stop = sample_start + n_samples + if trigger_ch is not None: + stim_ch = trigger_ch[start:stop][sample_start:sample_stop] + block = np.vstack((block, stim_ch)) + data_view = data[:, sample_start:sample_stop] + _mult_cal_one(data_view, block, idx, cals, mult) + return + with open(raw.filenames[fi], "rb", buffering=0) as fid: fid.seek(data_offset) # extract data in chunks diff --git a/mne/io/fiff/raw.py b/mne/io/fiff/raw.py index 95c6db5dbec..87102f77544 100644 --- a/mne/io/fiff/raw.py +++ b/mne/io/fiff/raw.py @@ -9,10 +9,11 @@ import numpy as np +from ..._fiff._mmap_cache import get_u8_memmap from ..._fiff.constants import FIFF from ..._fiff.meas_info import read_meas_info from ..._fiff.open import _fiff_get_fid, _get_next_fname, fiff_open -from ..._fiff.tag import _call_dict, read_tag +from ..._fiff.tag import _call_dict, _simple_dict, read_tag from ..._fiff.tree import dir_tree_find from ..._fiff.utils import _mult_cal_one from ...annotations import Annotations, _read_annotations_fif @@ -403,13 +404,84 @@ def _dtype(self): def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): """Read a segment of data from a file.""" n_bad = 0 - with _fiff_get_fid(self._raw_extras[fi]["filename"]) as fid: - bounds = self._raw_extras[fi]["bounds"] - ents = self._raw_extras[fi]["ent"] - nchan = self._raw_extras[fi]["orig_nchan"] - use = (stop > bounds[:-1]) & (start < bounds[1:]) + bounds = self._raw_extras[fi]["bounds"] + ents = self._raw_extras[fi]["ent"] + nchan = self._raw_extras[fi]["orig_nchan"] + fname = self._raw_extras[fi]["filename"] + # Entries overlapping [start, stop) via binary search on sorted bounds + # (O(log n) instead of a mask over every entry; matters for long + # recordings with thousands of buffer entries). + eis = range( + max(np.searchsorted(bounds, start, side="right") - 1, 0), + min(np.searchsorted(bounds, stop, side="left"), len(bounds) - 1), + ) + + # Fast path: read tag payloads directly through a PID-keyed memory map, + # skipping per-call open/seek/read syscalls. Only taken for uncompressed + # real files whose touched tags are simple numeric types with the + # expected sizes; everything else falls back to the legacy loop below. + mm = None + if ( + isinstance(fname, Path) + and len(fname.suffixes) > 0 + and fname.suffixes[-1] != ".gz" + ): + mm = get_u8_memmap(fname) + if mm is not None: + for ei in eis: + ent = ents[ei] + if ent is None or ent.type not in _simple_dict: + mm = None + break + nsamp_ei = bounds[ei + 1] - bounds[ei] + itemsize = np.dtype(_simple_dict[ent.type]).itemsize + if getattr(ent, "size", None) != nsamp_ei * nchan * itemsize: + mm = None + break + if mm is not None: + offset = 0 + for ei in eis: + first = bounds[ei] + last = bounds[ei + 1] + nsamp = last - first + ent = ents[ei] + first_pick = max(start - first, 0) + last_pick = min(nsamp, stop - first) + picksamp = last_pick - first_pick + this_start = offset + offset += picksamp + this_stop = offset + if ent is None: + continue # gaps were zero-initialized by the caller + dtype_s = _simple_dict[ent.type] + itemsize = np.dtype(dtype_s).itemsize + nbytes = picksamp * nchan * itemsize + base = ent.pos + 16 + first_pick * nchan * itemsize + one = np.frombuffer( + mm[base : base + nbytes], dtype=dtype_s, count=picksamp * nchan + ) + if one.size != picksamp * nchan: + n_bad += picksamp + continue + one = one.reshape(picksamp, nchan) + _mult_cal_one( + data[:, this_start:this_stop], + one.T, + idx, + cals, + mult, + ) + if n_bad: + warn( + f"FIF raw buffer could not be read, acquisition error " + f"likely: {n_bad} samples set to zero" + ) + assert offset == stop - start + return + + with _fiff_get_fid(fname) as fid: offset = 0 - for ei in np.where(use)[0]: + for ei in eis: first = bounds[ei] last = bounds[ei + 1] nsamp = last - first From 8dc167a27623c42499c151e2aea0693ce299833d Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 10:59:42 +0200 Subject: [PATCH 3/8] Vectorized EDF/BDF window decode fast path Uniform-sampling-rate EDF/BDF windows decode through a strided gather plus an optional-numba kernel writing calibrated samples directly into the caller's output buffer, replacing the per-channel Python loop. Big-endian chunks are byteswapped to native order first; uniform stim channels apply the legacy truncating bitmask on the fast path; projector/compensation reads keep the exact legacy route. EDF/BDF/GDF file handles persist per process (PID-keyed LRU). --- mne/io/edf/_bdf_numba.py | 33 +++++ mne/io/edf/_edf_numba.py | 69 +++++++++++ mne/io/edf/_open.py | 55 ++++++++- mne/io/edf/edf.py | 231 ++++++++++++++++++++++++++++++++++- mne/io/edf/tests/test_edf.py | 13 ++ 5 files changed, 393 insertions(+), 8 deletions(-) create mode 100644 mne/io/edf/_bdf_numba.py create mode 100644 mne/io/edf/_edf_numba.py diff --git a/mne/io/edf/_bdf_numba.py b/mne/io/edf/_bdf_numba.py new file mode 100644 index 00000000000..8315915bace --- /dev/null +++ b/mne/io/edf/_bdf_numba.py @@ -0,0 +1,33 @@ +"""Numba-accelerated BDF (24-bit little-endian) sample decoding. + +Optional acceleration: falls back to the vectorized-numpy path in +``mne.io.edf.edf._read_ch`` when numba is unavailable. +""" + +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors + +import numpy as np + +from ..._numba import jit + + +@jit() +def decode_int24(buf): # pragma: no cover + """Decode packed 24-bit little-endian samples to int32. + + ``buf`` is a ``(n_samples, 3)`` uint8 array whose rows hold the low, + middle, and high bytes of each signed sample. + """ + n = buf.shape[0] + out = np.empty(n, dtype=np.int32) + for i in range(n): + # plain-Python integer arithmetic so the non-numba fallback follows + # the same semantics as the jitted version (values stay within + # [-2**23, 2**23) after the sign fix, so int32 stores never overflow) + v = int(buf[i, 0]) | (int(buf[i, 1]) << 8) | (int(buf[i, 2]) << 16) + if v >= (1 << 23): + v -= 1 << 24 + out[i] = v + return out diff --git a/mne/io/edf/_edf_numba.py b/mne/io/edf/_edf_numba.py new file mode 100644 index 00000000000..7910185f491 --- /dev/null +++ b/mne/io/edf/_edf_numba.py @@ -0,0 +1,69 @@ +"""Numba-accelerated EDF/BDF digital-to-physical window decoding. + +Optional acceleration: falls back to the vectorized-numpy path in +``mne.io.edf.edf._read_segment_file`` when numba is unavailable. +""" + +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors + +from ..._numba import jit + + +@jit(fastmath=False) +def decode_window(digital, cal_v, off_v, gain_v, out): # pragma: no cover + """Decode a block of digital samples to physical units. + + ``digital`` is a ``(k, n_blocks, buf_len)`` (possibly strided) integer + view of raw digital samples; ``cal_v``, ``off_v``, and ``gain_v`` are + length-``k`` float64 vectors; ``out`` is a ``(k, n_blocks * buf_len)`` + float64 array whose rows are filled with blocks concatenated along the + sample axis such that:: + + out[i, b * buf_len + j] = ((digital[i, b, j] * cal[i]) + off[i]) * gain[i] + + replicating exactly the operation order of the vectorized-numpy fallback + (hence ``fastmath=False``: no FMA contraction or reassociation is + allowed, so results are bit-identical to separate multiply/add rounding). + """ + k = digital.shape[0] + n_blk = digital.shape[1] + n_smp = digital.shape[2] + for i in range(k): + cal = cal_v[i] + off = off_v[i] + gain = gain_v[i] + out_i = out[i] + for b in range(n_blk): + base = b * n_smp + for j in range(n_smp): + out_i[base + j] = ((digital[i, b, j] * cal) + off) * gain + + +@jit(fastmath=False) +def decode_window_into( + dst, digital, cal_v, off_v, gain_v, s0, w +): # pragma: no cover + """Decode into a possibly-strided 2-D destination. + + ``dst`` is ``(k, w)`` with arbitrary strides (e.g., a column slice of the + caller's output buffer); ``digital`` is the ``(k, n_blocks, buf_len)`` + strided integer view covering whole data records; ``s0`` is the first + sample to take from that view and ``w`` the number of samples to write, + so edge records at window boundaries are handled without temporaries. + Elementwise op order is identical to :func:`decode_window`. + """ + k = digital.shape[0] + n_smp = digital.shape[2] + for i in range(k): + cal = cal_v[i] + off = off_v[i] + gain = gain_v[i] + dst_i = dst[i] + dig_i = digital[i] + for t in range(w): + g = s0 + t + b = g // n_smp + j = g - b * n_smp + dst_i[t] = ((dig_i[b, j] * cal) + off) * gain diff --git a/mne/io/edf/_open.py b/mne/io/edf/_open.py index 38f8be4113d..1d8f1106dfe 100644 --- a/mne/io/edf/_open.py +++ b/mne/io/edf/_open.py @@ -1,12 +1,55 @@ # Authors: The MNE-Python contributors. # License: BSD-3-Clause -# Copyright the MNE-Python contributors. +# Copyright the MNE-Python contributors +import os from pathlib import Path from ..._fiff.open import _NoCloseRead from ...utils import _file_like, _validate_type, logger +# Persistent read handles for EDF/BDF/GDF files. Readers seek before every +# read, so a shared handle is safe; keying by PID keeps forked worker +# processes (e.g., PyTorch DataLoader workers) from sharing file-offset state +# through an inherited descriptor. +_HANDLE_CACHE = {} +_MAX_HANDLES = 8 + + +class _NoCloseCached(_NoCloseRead): + """A file object whose context manager detaches instead of closing. + + Used for handles shared through the per-process LRU cache: leaving the + reader's ``with`` block must not close a descriptor other reads may still + use. + """ + + def close(self): # noqa: D102 + pass + + def __exit__(self, *args): # noqa: D105 + # detach rather than close; the cache owns the lifetime + return False + + +def _get_cached_fid(fname): + """Return a persistent binary handle for *fname* (per process).""" + key = (os.getpid(), str(fname)) + hit = _HANDLE_CACHE.get(key) + if hit is not None: + hit.seek(0) # match fresh-open semantics + return hit + fid = open(fname, "rb") + cached = _NoCloseCached(fid) + _HANDLE_CACHE[key] = cached + while len(_HANDLE_CACHE) > _MAX_HANDLES: + old_key = next(iter(_HANDLE_CACHE)) + try: + _HANDLE_CACHE.pop(old_key).fid.close() + except Exception: + pass + return cached + def _gdf_edf_get_fid(fname, **kwargs): """Open a EDF/BDF/GDF file with no additional parsing.""" @@ -14,8 +57,8 @@ def _gdf_edf_get_fid(fname, **kwargs): logger.debug("Using file-like I/O") fid = _NoCloseRead(fname) fid.seek(0) - else: - _validate_type(fname, [Path, str], "fname", extra="or file-like") - logger.debug("Using normal I/O") - fid = open(fname, "rb", **kwargs) # Open in binary mode - return fid + return fid + _validate_type(fname, [Path, str], "fname", extra="or file-like") + logger.debug("Using normal I/O") + kwargs.pop("buffering", None) # cached handle manages its own buffering + return _get_cached_fid(Path(fname)) diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 390289ae6bb..5b4bd58a852 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -590,12 +590,36 @@ def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): ) +_decode_int24 = None + + +def _get_int24_decoder(): + """Return the numba int24 decoder if available, else False.""" + global _decode_int24 + if _decode_int24 is None: + dec = False + try: + from mne._numba import has_numba + + # only use the jitted decoder when numba is actually enabled, + # otherwise the decorated function would run as slow pure Python + if has_numba: + from mne.io.edf._bdf_numba import decode_int24 as dec + except Exception: + dec = False + _decode_int24 = dec + return _decode_int24 + + def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): """Read a number of samples for a single channel.""" assert dtype is not None # BDF if subtype == "bdf": ch_data = read_from_file_or_buffer(fid, dtype=dtype, count=samp * dtype_byte) + dec = _get_int24_decoder() + if dec is not False: + return dec(ch_data.reshape(-1, 3)) 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 @@ -610,8 +634,6 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): 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"] @@ -640,6 +662,190 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # Let's do ~10 MB chunks: n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * dtype_byte), 1) + # Fast path: uniform sampling rate among all requested channels, no + # annotations (TAL) channel and no stim-channel special casing. This is + # the common case for large DL corpora; it replaces the per-channel + # Python loop with a handful of vectorized operations. + # The vectorized path removes per-channel Python overhead, which dominates + # small reads (deep-learning window access); for very large outputs the + # legacy loop's cache-friendly working set is slightly faster, so gate on + # decoded size. + _fast_path_output_limit = int(32e6) + uni_mask = [n_samps[ci] == buf_len for ci in read_sel] + slow_ii = [ii for ii, u in enumerate(uni_mask) if not u] + if ( + subtype in ("edf", "bdf") + and len(tal_idx) == 0 + and all(uni_mask) # mixed-sfreq partial decode is future work + and (stop - start) * len(idx_arr) * 8 <= _fast_path_output_limit + ): + with _gdf_edf_get_fid(filenames, buffering=0) as fid: + start_offset = ( + data_offset + block_start_idx * ch_offsets[-1] * dtype_byte + ) + k = len(idx_arr) + uni_list = [ii for ii in range(k) if uni_mask[ii]] + sel_uni = [read_sel[ii] for ii in uni_list] + idx_arr_uni = idx_arr[uni_list] + k_u = len(uni_list) + cal_v = cal[idx_arr][:, np.newaxis, np.newaxis] + off_v = offsets[idx_arr][:, np.newaxis, np.newaxis] + gain_v = gains[idx_arr][:, np.newaxis, np.newaxis] + cal_u = cal[idx_arr_uni][:, np.newaxis, np.newaxis] + off_u = offsets[idx_arr_uni][:, np.newaxis, np.newaxis] + gain_u = gains[idx_arr_uni][:, np.newaxis, np.newaxis] + off_v = offsets[idx_arr][:, np.newaxis, np.newaxis] + gain_v = gains[idx_arr][:, np.newaxis, np.newaxis] + # With no projector/compensation and unit cals (always true for + # EDF/BDF), decoded samples can go straight into the output. + write_direct = ( + mult is None and bool(np.all(cals == 1)) and not slow_ii + ) + ones = None if write_direct else np.zeros( + (len(orig_sel), data.shape[-1]), dtype=data.dtype + ) + dec = _get_window_decoder() + dec_into = None + if dec is not False and k_u: + from ._edf_numba import decode_window_into + + dec_into = decode_window_into + cal_1d = cal_u.ravel() + off_1d = off_u.ravel() + gain_1d = gain_u.ravel() + # Uniform stim channels can stay on the fast path: legacy applies + # a truncating bitmask (float -> int cast, mask 2**17-1) AFTER + # calibration, which we replicate per row below. + stim_rows = [ + ii + for ii, orig in enumerate(idx_arr) + if int(orig) in stim_channel_idxs and uni_mask[ii] + ] + pos = 0 + for ai in range(0, len(r_lims), n_per): + block_offset = ai * ch_offsets[-1] * dtype_byte + n_read = min(len(r_lims) - ai, n_per) + fid.seek(start_offset + block_offset, 0) + many_chunk = _read_ch( + fid, subtype, ch_offsets[-1] * n_read, dtype_byte, dtype + ) + arr3 = many_chunk.reshape(n_read, len(n_samps), buf_len) + r_sidx = r_lims[ai][0] + r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] + # gather this call's channels as a strided view + # (k, n_read, buf_len), each row one requested channel; + # when all channels are requested this is a zero-copy + # transpose + if not slow_ii and k == len(n_samps): + view = arr3.transpose(1, 0, 2) + elif slow_ii: + view = arr3[:, sel_uni, :].transpose(1, 0, 2) + else: + view = arr3[:, read_sel, :].transpose(1, 0, 2) + # digital -> physical, preserving the legacy op order; + # when possible decode straight into the destination slice so + # the block never round-trips through a temporary + width = r_eidx - r_sidx + if dec_into is not None and write_direct: + if not view.dtype.isnative: + # numba only types native byteorder; real-world EDF is + # big-endian, so swap into a native copy per chunk + view = view.astype(view.dtype.newbyteorder("=")) + dst_full = ( + data[:, pos : pos + width] + if write_direct + else ones[idx_arr_uni, pos : pos + width] + ) + dec_into(dst_full, view, cal_1d, off_1d, gain_1d, + r_sidx, width) + block = None # values already in place + elif ones is None: + # no kernel and direct write: decode to a temp then copy + # once into the destination slice + one = np.empty((k, n_read, buf_len), dtype=np.float64) + np.multiply(view, cal_v, out=one) + one += off_v + one *= gain_v + dst_full = data[:, pos : pos + width] + dst_full[...] = one.reshape(k, -1)[:, r_sidx:r_eidx] + block = None # values already in place + else: + one = np.empty((k, n_read, buf_len), dtype=np.float64) + np.multiply(view, cal_v, out=one) + one += off_v + one *= gain_v + block = one.reshape(k, -1)[:, r_sidx:r_eidx] + if stim_rows: + for ii in stim_rows: + src_row = block[ii] if block is not None else dst_full[ii] + row_i = src_row.astype(int) + np.bitwise_and(row_i, 2**17 - 1, out=row_i) + if block is not None: + block[ii] = row_i + else: + dst_full[ii] = row_i + if block is not None: + if write_direct: + data[:, pos : pos + width] = block + elif slow_ii: + ones[idx_arr_uni, pos : pos + width] = block + else: + ones[idx_arr, pos : pos + width] = block + # legacy per-channel treatment for non-uniform rows: their + # columns live inside the same records we just read + for ii in slow_ii: + ci = read_sel[ii] + orig_idx = idx_arr[ii] + ch_data = many_chunk[ + :, ch_offsets[ci] : ch_offsets[ci + 1] + ].copy() + o_i = orig_idx + ch_data = ch_data * cal[o_i] + ch_data += offsets[o_i] + ch_data *= gains[o_i] + if int(orig_idx) in stim_channel_idxs: + from scipy.interpolate import interp1d + + s_n = n_samps[ci] + oldg = np.linspace(0, 1, s_n + 1, True) + newg = np.linspace(0, 1, buf_len, False) + ch_data = np.append( + ch_data, np.zeros((len(ch_data), 1)), -1 + ) + ch_data = interp1d(oldg, ch_data, kind="zero", axis=-1)( + newg + ) + one_i = ch_data.ravel()[r_sidx:r_eidx] + w0 = pos - width # first output column of this chunk + ones[o_i, w0 : w0 + len(one_i)] = one_i + pos += width + if slow_ii: + smp_exp = data.shape[-1] + resampled = False + for ii in slow_ii: + row = int(idx_arr[ii]) + ci = read_sel[ii] + if n_samps[ci] != buf_len and width != smp_exp: + resampled = True + ones[row, :] = resample( + ones[row, :width].astype(np.float64), + smp_exp, + width, + npad=0, + axis=-1, + ) + if resampled and raw_extras["nsamples"] != (stop - start): + warn( + "Loading an EDF with mixed sampling frequencies and " + "preload=False will result in edge artifacts. " + "It is recommended to use preload=True." + "See also " + "https://github.com/mne-tools/mne-python/issues/10635" + ) + if not write_direct: + _mult_cal_one(data[:, :], ones, idx, cals, mult) + return tal_data + with _gdf_edf_get_fid(filenames, buffering=0) as fid: # Extract data start_offset = data_offset + block_start_idx * ch_offsets[-1] * dtype_byte @@ -683,6 +889,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: # Stim channel will be interpolated + from scipy.interpolate import interp1d + old = np.linspace(0, 1, n_samps[ci] + 1, True) new = np.linspace(0, 1, buf_len, False) ch_data = np.append(ch_data, np.zeros((len(ch_data), 1)), -1) @@ -2361,3 +2569,22 @@ def _get_annotations_gdf(edf_info, sfreq): desc = events[2] return onset, duration, desc + + +_decode_window = None + + +def _get_window_decoder(): + """Return the numba fused window decoder if available, else False.""" + global _decode_window + if _decode_window is None: + dec = False + try: + from mne._numba import has_numba + + if has_numba: + from mne.io.edf._edf_numba import decode_window as dec + except Exception: + dec = False + _decode_window = dec + return _decode_window diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 211641dc726..850ad1d0744 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -8,6 +8,7 @@ from io import BytesIO from pathlib import Path +import mne import numpy as np import pytest from numpy.testing import ( @@ -1262,3 +1263,15 @@ def test_edf_read_from_file_like(): ] assert raw.ch_names == channels + + + +def requires_edfio(func): + import pytest + + return pytest.mark.skipif( + __import__("importlib.util", fromlist=["util"]).find_spec("edfio") is None, + reason="Requires edfio", + )(func) + + From 0454a956072df8d411e791d428fcf0439e873d3a Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:41:38 +0000 Subject: [PATCH 4/8] [autofix.ci] apply automated fixes --- mne/_fiff/_mmap_cache.py | 2 +- mne/io/edf/_bdf_numba.py | 2 +- mne/io/edf/_edf_numba.py | 2 +- mne/io/edf/_open.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/mne/_fiff/_mmap_cache.py b/mne/_fiff/_mmap_cache.py index e96a2e46a5a..04bd826985b 100644 --- a/mne/_fiff/_mmap_cache.py +++ b/mne/_fiff/_mmap_cache.py @@ -8,7 +8,7 @@ # Authors: The MNE-Python contributors. # License: BSD-3-Clause -# Copyright the MNE-Python contributors +# Copyright the MNE-Python contributors. import os diff --git a/mne/io/edf/_bdf_numba.py b/mne/io/edf/_bdf_numba.py index 8315915bace..790a6f45cea 100644 --- a/mne/io/edf/_bdf_numba.py +++ b/mne/io/edf/_bdf_numba.py @@ -6,7 +6,7 @@ # Authors: The MNE-Python contributors. # License: BSD-3-Clause -# Copyright the MNE-Python contributors +# Copyright the MNE-Python contributors. import numpy as np diff --git a/mne/io/edf/_edf_numba.py b/mne/io/edf/_edf_numba.py index 7910185f491..fddeafbf04d 100644 --- a/mne/io/edf/_edf_numba.py +++ b/mne/io/edf/_edf_numba.py @@ -6,7 +6,7 @@ # Authors: The MNE-Python contributors. # License: BSD-3-Clause -# Copyright the MNE-Python contributors +# Copyright the MNE-Python contributors. from ..._numba import jit diff --git a/mne/io/edf/_open.py b/mne/io/edf/_open.py index 1d8f1106dfe..5587367ca44 100644 --- a/mne/io/edf/_open.py +++ b/mne/io/edf/_open.py @@ -1,6 +1,6 @@ # Authors: The MNE-Python contributors. # License: BSD-3-Clause -# Copyright the MNE-Python contributors +# Copyright the MNE-Python contributors. import os from pathlib import Path From f983e5a843f51d0f2681475a3520662abbe06983 Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 14:36:27 +0200 Subject: [PATCH 5/8] Use explicit variable names in the EDF fast path cal_v/off_v/gain_v -> channel_calibration/channel_offset/channel_gain, the uniform-row subsets get the uniform_ prefix, one_flat -> decoded_block, arr3 -> record_grid, and the strided gather is named digital_view. No behavior change. --- mne/io/edf/edf.py | 42 +++++++++++++++++++++--------------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 5b4bd58a852..d50cdd7c06f 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -688,14 +688,14 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, sel_uni = [read_sel[ii] for ii in uni_list] idx_arr_uni = idx_arr[uni_list] k_u = len(uni_list) - cal_v = cal[idx_arr][:, np.newaxis, np.newaxis] - off_v = offsets[idx_arr][:, np.newaxis, np.newaxis] - gain_v = gains[idx_arr][:, np.newaxis, np.newaxis] - cal_u = cal[idx_arr_uni][:, np.newaxis, np.newaxis] - off_u = offsets[idx_arr_uni][:, np.newaxis, np.newaxis] - gain_u = gains[idx_arr_uni][:, np.newaxis, np.newaxis] - off_v = offsets[idx_arr][:, np.newaxis, np.newaxis] - gain_v = gains[idx_arr][:, np.newaxis, np.newaxis] + channel_calibration = cal[idx_arr][:, np.newaxis, np.newaxis] + channel_offset = offsets[idx_arr][:, np.newaxis, np.newaxis] + channel_gain = gains[idx_arr][:, np.newaxis, np.newaxis] + uniform_calibration = cal[idx_arr_uni][:, np.newaxis, np.newaxis] + uniform_offset = offsets[idx_arr_uni][:, np.newaxis, np.newaxis] + uniform_gain = gains[idx_arr_uni][:, np.newaxis, np.newaxis] + channel_offset = offsets[idx_arr][:, np.newaxis, np.newaxis] + channel_gain = gains[idx_arr][:, np.newaxis, np.newaxis] # With no projector/compensation and unit cals (always true for # EDF/BDF), decoded samples can go straight into the output. write_direct = ( @@ -710,9 +710,9 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, from ._edf_numba import decode_window_into dec_into = decode_window_into - cal_1d = cal_u.ravel() - off_1d = off_u.ravel() - gain_1d = gain_u.ravel() + cal_1d = uniform_calibration.ravel() + off_1d = uniform_offset.ravel() + gain_1d = uniform_gain.ravel() # Uniform stim channels can stay on the fast path: legacy applies # a truncating bitmask (float -> int cast, mask 2**17-1) AFTER # calibration, which we replicate per row below. @@ -729,7 +729,7 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, many_chunk = _read_ch( fid, subtype, ch_offsets[-1] * n_read, dtype_byte, dtype ) - arr3 = many_chunk.reshape(n_read, len(n_samps), buf_len) + record_grid = many_chunk.reshape(n_read, len(n_samps), buf_len) r_sidx = r_lims[ai][0] r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] # gather this call's channels as a strided view @@ -737,11 +737,11 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # when all channels are requested this is a zero-copy # transpose if not slow_ii and k == len(n_samps): - view = arr3.transpose(1, 0, 2) + view = record_grid.transpose(1, 0, 2) elif slow_ii: - view = arr3[:, sel_uni, :].transpose(1, 0, 2) + view = record_grid[:, sel_uni, :].transpose(1, 0, 2) else: - view = arr3[:, read_sel, :].transpose(1, 0, 2) + view = record_grid[:, read_sel, :].transpose(1, 0, 2) # digital -> physical, preserving the legacy op order; # when possible decode straight into the destination slice so # the block never round-trips through a temporary @@ -763,17 +763,17 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # no kernel and direct write: decode to a temp then copy # once into the destination slice one = np.empty((k, n_read, buf_len), dtype=np.float64) - np.multiply(view, cal_v, out=one) - one += off_v - one *= gain_v + np.multiply(view, channel_calibration, out=one) + one += channel_offset + one *= channel_gain dst_full = data[:, pos : pos + width] dst_full[...] = one.reshape(k, -1)[:, r_sidx:r_eidx] block = None # values already in place else: one = np.empty((k, n_read, buf_len), dtype=np.float64) - np.multiply(view, cal_v, out=one) - one += off_v - one *= gain_v + np.multiply(view, channel_calibration, out=one) + one += channel_offset + one *= channel_gain block = one.reshape(k, -1)[:, r_sidx:r_eidx] if stim_rows: for ii in stim_rows: From d7c3000cb86dc4871190e813b021c03b05dc1e61 Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 20:36:24 +0200 Subject: [PATCH 6/8] Vectorize uniform EDF and BDF window reads --- doc/changes/dev/14214.newfeature.rst | 1 + mne/io/edf/edf.py | 94 ++++++++ mne/io/edf/tests/test_edf.py | 316 +++++++++++++++++++++++++++ 3 files changed, 411 insertions(+) create mode 100644 doc/changes/dev/14214.newfeature.rst diff --git a/doc/changes/dev/14214.newfeature.rst b/doc/changes/dev/14214.newfeature.rst new file mode 100644 index 00000000000..a8b58a3525d --- /dev/null +++ b/doc/changes/dev/14214.newfeature.rst @@ -0,0 +1 @@ +Speed up eligible EDF and BDF window reads with a pure-NumPy record-stride decoder, by `Bruno Aristimunha`_. diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 390289ae6bb..fb13d4bff65 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -608,8 +608,102 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): return ch_data +_EDF_STRIDE_MAX_OUTPUT_BYTES = 32 * 1024**2 +_EDF_STRIDE_MAX_EXTRA_BYTES = 64 * 1024**2 + + +def _read_uniform_segment( + data, idx, start, stop, raw_extras, filenames, cals, mult +) -> bool: + """Read a uniformly sampled EDF or BDF segment.""" + subtype = raw_extras["subtype"] + if subtype not in ("edf", "bdf") or not isinstance(filenames, str | Path): + return False + if len(raw_extras.get("tal_idx", ())) != 0: + return False + + idx_arr = ( + np.arange(idx.start, idx.stop) if isinstance(idx, slice) else np.asarray(idx) + ) + if len(idx_arr) == 0 or len(np.unique(idx_arr)) != len(idx_arr): + return False + + n_samps = raw_extras["n_samps"] + buf_len = int(raw_extras["max_samp"]) + if ( + not np.all(n_samps == buf_len) + or mult is not None + or data.nbytes > _EDF_STRIDE_MAX_OUTPUT_BYTES + ): + return False + + dtype = raw_extras["dtype_np"] + dtype_byte = raw_extras["dtype_byte"] + data_offset = raw_extras["data_offset"] + stim_channel_idxs = raw_extras["stim_channel_idxs"] + orig_sel = raw_extras["sel"] + cal = raw_extras["cal"] + offsets = raw_extras["offsets"] + gains = raw_extras["units"] + read_sel = orig_sel[idx_arr] + + ch_offsets = np.cumsum(np.concatenate([[0], n_samps]), dtype=np.int64) + block_start_idx, r_lims, d_lims = _blk_read_lims(start, stop, buf_len) + n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * dtype_byte), 1) + max_records = min(len(r_lims), n_per) + n_values = len(idx_arr) * max_records * buf_len + estimated_incremental_bytes = n_values * np.dtype(np.float64).itemsize + physical_order = np.arange(len(n_samps)) + if not np.array_equal(read_sel, physical_order): + gather_bytes_per_value = 2 if subtype == "edf" else 4 + estimated_incremental_bytes += n_values * gather_bytes_per_value + if any(orig_idx in stim_channel_idxs for orig_idx in idx_arr): + estimated_incremental_bytes += max_records * buf_len * np.dtype(int).itemsize + if estimated_incremental_bytes > _EDF_STRIDE_MAX_EXTRA_BYTES: + return False + + with _gdf_edf_get_fid(filenames, buffering=0) as fid: + start_offset = data_offset + block_start_idx * ch_offsets[-1] * dtype_byte + ones = np.zeros((len(orig_sel), data.shape[-1]), dtype=data.dtype) + for ai in range(0, len(r_lims), n_per): + block_offset = ai * ch_offsets[-1] * dtype_byte + n_read = min(len(r_lims) - ai, n_per) + fid.seek(start_offset + block_offset, 0) + many_chunk = _read_ch( + fid, subtype, ch_offsets[-1] * n_read, dtype_byte, dtype + ) + record_grid = many_chunk.reshape(n_read, len(n_samps), buf_len) + if np.array_equal(read_sel, physical_order): + view = record_grid.transpose(1, 0, 2) + else: + view = record_grid[:, read_sel, :].transpose(1, 0, 2) + + one = np.empty(view.shape, dtype=np.float64) + np.multiply(view, cal[idx_arr, np.newaxis, np.newaxis], out=one) + one += offsets[idx_arr, np.newaxis, np.newaxis] + one *= gains[idx_arr, np.newaxis, np.newaxis] + r_sidx = r_lims[ai][0] + r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] + block = one.reshape(len(idx_arr), -1)[:, r_sidx:r_eidx] + for row, orig_idx in enumerate(idx_arr): + if orig_idx in stim_channel_idxs: + stim = block[row].astype(int) + np.bitwise_and(stim, 2**17 - 1, out=stim) + block[row] = stim + d_start = d_lims[ai][0] + d_stop = d_lims[ai + n_read - 1][1] + assert d_stop - d_start == block.shape[1] + ones[idx_arr, d_start:d_stop] = block + + _mult_cal_one(data, ones, idx, cals, mult) + return True + + def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, mult): """Read a chunk of raw data.""" + if _read_uniform_segment(data, idx, start, stop, raw_extras, filenames, cals, mult): + return [] + from scipy.interpolate import interp1d n_samps = raw_extras["n_samps"] diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 211641dc726..2e66d4069d4 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -3,6 +3,7 @@ # Copyright the MNE-Python contributors. import datetime +from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from functools import partial from io import BytesIO @@ -20,6 +21,7 @@ from mne import Annotations, pick_types from mne._fiff.pick import channel_indices_by_type, get_channel_type_constants +from mne._fiff.utils import _blk_read_lims from mne.annotations import _ndarray_ch_names, events_from_annotations, read_annotations from mne.datasets import testing from mne.io import edf, read_raw_bdf, read_raw_edf, read_raw_fif, read_raw_gdf @@ -66,6 +68,320 @@ misc = ["EXG1", "EXG5", "EXG8", "M1", "M2"] +def _repeat_edf_records(source, destination, n_records=6): + """Repeat a one-record EDF/BDF payload for boundary tests.""" + blob = bytearray(source.read_bytes()) + header_nbytes = int(blob[184:192]) + assert int(blob[236:244]) == 1 + blob[236:244] = f"{n_records:<8}".encode("ascii") + destination.write_bytes(blob[:header_nbytes] + blob[header_nbytes:] * n_records) + + +def _assert_stride_matches_legacy(monkeypatch, raw, picks, start, stop): + helper = edf.edf._read_uniform_segment + monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) + want = raw.get_data(picks=picks, start=start, stop=stop) + used = [] + + def _record_use(*args, **kwargs): + result = helper(*args, **kwargs) + used.append(result) + return result + + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + got = raw.get_data(picks=picks, start=start, stop=stop) + assert used and all(used) + assert_array_equal(got, want) + return want + + +def _assert_stride_falls_back(monkeypatch, read_data, *, expect_mult=False): + helper = edf.edf._read_uniform_segment + monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) + want = read_data() + used = [] + mults = [] + + def _record_use(*args, **kwargs): + result = helper(*args, **kwargs) + used.append(result) + mults.append(args[-1] if args else kwargs["mult"]) + return result + + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + got = read_data() + assert used and not any(used) + if expect_mult: + assert all(mult is not None for mult in mults) + assert_array_equal(got, want) + + +@pytest.mark.parametrize( + "reader, source, suffix", + [ + (read_raw_edf, edf_stim_channel_path, ".edf"), + (read_raw_bdf, bdf_path, ".bdf"), + ], +) +@pytest.mark.parametrize("pick_kind", ("all", "subset", "reversed", "permuted")) +@pytest.mark.parametrize("window_kind", ("within", "boundary", "multiple")) +def test_uniform_stride_decode( + reader, source, suffix, pick_kind, window_kind, monkeypatch, tmp_path +): + """Test exact pure-NumPy stride decoding against legacy.""" + repeated = tmp_path / f"uniform{suffix}" + _repeat_edf_records(source, repeated) + raw = reader(repeated, preload=False, verbose="error") + n_channels = len(raw.ch_names) + buf_len = int(raw._raw_extras[0]["max_samp"]) + picks = { + "all": np.arange(n_channels), + "subset": np.array([0, n_channels // 2, n_channels - 1]), + "reversed": np.arange(n_channels - 1, -1, -1), + "permuted": np.array([n_channels - 1, 1, n_channels // 2, 0]), + }[pick_kind] + start, stop = { + "within": (7, buf_len - 5), + "boundary": (buf_len - 7, buf_len + 11), + "multiple": (buf_len // 2, 4 * buf_len + 13), + }[window_kind] + _assert_stride_matches_legacy(monkeypatch, raw, picks, start, stop) + + +@pytest.mark.parametrize( + "reader, source, suffix", + ( + (read_raw_edf, edf_stim_channel_path, ".edf"), + (read_raw_bdf, bdf_path, ".bdf"), + ), +) +def test_uniform_stride_multiple_chunks(reader, source, suffix, monkeypatch, tmp_path): + """Test reads that cross the internal record-chunk boundary.""" + source_raw = reader(source, preload=False, verbose="error") + extras = source_raw._raw_extras[0] + record_nbytes = int(extras["n_samps"].sum() * extras["dtype_byte"]) + n_per = max(10 * 1024 * 1024 // record_nbytes, 1) + repeated = tmp_path / f"multiple_chunks{suffix}" + _repeat_edf_records(source, repeated, n_records=n_per + 2) + + raw = reader(repeated, preload=False, verbose="error") + buf_len = int(raw._raw_extras[0]["max_samp"]) + start = buf_len // 2 + stop = (n_per + 1) * buf_len + buf_len // 2 + _, r_lims, _ = _blk_read_lims(start, stop, buf_len) + assert len(r_lims) > n_per + assert stop <= raw.n_times + picks = np.array([0, len(raw.ch_names) // 2]) + output_nbytes = len(picks) * (stop - start) * np.dtype(np.float64).itemsize + assert output_nbytes < 4 * 1024**2 + _assert_stride_matches_legacy(monkeypatch, raw, picks, start, stop) + + +@pytest.mark.parametrize( + "reader, fname", ((read_raw_edf, edf_stim_channel_path), (read_raw_bdf, bdf_path)) +) +def test_uniform_stride_preload(reader, fname, monkeypatch): + """Test that eligible initial eager decoding matches legacy.""" + helper = edf.edf._read_uniform_segment + monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) + want = reader(fname, preload=True, verbose="error").get_data() + used = [] + + def _record_use(*args, **kwargs): + result = helper(*args, **kwargs) + used.append(result) + return result + + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + got = reader(fname, preload=True, verbose="error").get_data() + assert used and all(used) + assert_array_equal(got, want) + + +@pytest.mark.parametrize( + "reader, fname, reader_kwargs", + ( + pytest.param( + read_raw_edf, + edf_stim_channel_path, + {"stim_channel": -1}, + id="edf", + ), + pytest.param(read_raw_bdf, bdf_path, {}, id="bdf"), + ), +) +def test_uniform_stride_stim_only(reader, fname, reader_kwargs, monkeypatch): + """Test exact stim scaling and masking on the stride path.""" + raw = reader(fname, preload=False, verbose="error", **reader_kwargs) + picks = pick_types(raw.info, meg=False, stim=True) + assert len(picks) == 1 + assert raw.get_channel_types(picks=picks) == ["stim"] + assert_array_equal(picks, raw._raw_extras[0]["stim_channel_idxs"]) + if reader is read_raw_edf: + stim_idx = picks[0] + extras = raw._raw_extras[0] + assert extras["cal"][stim_idx] != 1.0 + assert extras["offsets"][stim_idx] != 0.0 + want = _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) + assert_array_equal(want, np.bitwise_and(want.astype(int), 2**17 - 1)) + if reader is read_raw_edf: + assert_array_equal(np.unique(want), [0.0, 100.0]) + + +def test_uniform_stride_cals(monkeypatch): + """Test non-unit Raw calibrations on the stride path.""" + raw = read_raw_bdf(bdf_path, preload=False, verbose="error") + picks = np.array([0, 17, 55]) + raw._cals[picks] *= np.array([0.5, 2.0, 4.0]) + assert np.all(raw._cals[picks] != 1.0) + _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) + + +def test_uniform_stride_excluded_physical_channel(monkeypatch): + """Test logical picks after excluding a physical record row.""" + raw = read_raw_bdf(bdf_path, exclude=["AF7"], preload=False, verbose="error") + physical_sel = raw._raw_extras[0]["sel"] + assert raw.ch_names[1] == "AF3" + assert physical_sel[1] == 2 + assert not np.array_equal(physical_sel, np.arange(len(raw.ch_names))) + picks = np.array([0, 1, len(raw.ch_names) - 1]) + _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) + + +@pytest.mark.parametrize( + "fallback_kind", + ( + "mixed_rate", + "tal", + "file_like", + "projection", + "output_limit", + "extra_limit", + "duplicate_picks", + ), +) +def test_uniform_stride_fallbacks(fallback_kind, monkeypatch): + """Test that unsafe stride layouts and transforms use legacy.""" + cutoff = { + "output_limit": "_EDF_STRIDE_MAX_OUTPUT_BYTES", + "extra_limit": "_EDF_STRIDE_MAX_EXTRA_BYTES", + }.get(fallback_kind) + if cutoff is not None: + monkeypatch.setattr(edf.edf, cutoff, 0) + + if fallback_kind == "file_like": + helper = edf.edf._read_uniform_segment + blob = bdf_path.read_bytes() + monkeypatch.setattr( + edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False + ) + want = read_raw_bdf(BytesIO(blob), preload=True, verbose="error").get_data() + used = [] + + def _record_use(*args, **kwargs): + result = helper(*args, **kwargs) + used.append(result) + return result + + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + got = read_raw_bdf(BytesIO(blob), preload=True, verbose="error").get_data() + assert used and not any(used) + assert_array_equal(got, want) + return + + raw = None + picks = None + start, stop = 100, 900 + expect_mult = False + if fallback_kind == "mixed_rate": + raw = read_raw_edf(edf_uneven_path, preload=False, verbose="error") + extras = raw._raw_extras[0] + n_samps = extras["n_samps"] + assert np.unique(n_samps).size > 1 + high_rate_physical = np.flatnonzero(n_samps == extras["max_samp"])[0] + picks = np.flatnonzero(extras["sel"] == high_rate_physical) + assert len(picks) == 1 + assert n_samps[extras["sel"][picks[0]]] == extras["max_samp"] + elif fallback_kind == "tal": + raw = read_raw_edf(edf_path, preload=False, verbose="error") + extras = raw._raw_extras[0] + picks = np.array([0, len(raw.ch_names) - 1]) + assert len(extras["tal_idx"]) > 0 + assert np.all(extras["n_samps"][extras["sel"][picks]] == extras["max_samp"]) + elif fallback_kind == "projection": + raw = read_raw_bdf(bdf_path, preload=False, verbose="error") + raw.set_eeg_reference(projection=True, verbose="error").apply_proj( + verbose="error" + ) + assert raw._projector is not None + expect_mult = True + elif fallback_kind in ("output_limit", "extra_limit"): + raw = read_raw_bdf(bdf_path, preload=False, verbose="error") + else: + assert fallback_kind == "duplicate_picks" + helper = edf.edf._read_uniform_segment + raw = read_raw_bdf(bdf_path, preload=False, verbose="error") + picks = [2, 0, 2] + monkeypatch.setattr( + edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False + ) + with pytest.raises(ValueError) as want_error: + raw.get_data(picks=picks, start=start, stop=stop) + used = [] + + def _record_use(*args, **kwargs): + result = helper(*args, **kwargs) + used.append(result) + return result + + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + with pytest.raises(ValueError) as got_error: + raw.get_data(picks=picks, start=start, stop=stop) + assert used and not any(used) + assert str(got_error.value) == str(want_error.value) + return + + _assert_stride_falls_back( + monkeypatch, + partial(raw.get_data, picks=picks, start=start, stop=stop), + expect_mult=expect_mult, + ) + + +@pytest.mark.parametrize( + "reader, source", + ((read_raw_edf, edf_stim_channel_path), (read_raw_bdf, bdf_path)), +) +def test_uniform_stride_concurrent(reader, source, monkeypatch): + """Test that stride reads share no seekable file handle.""" + raw = reader(source, preload=False, verbose="error") + helper = edf.edf._read_uniform_segment + monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) + reference = reader(source, preload=True, verbose="error").get_data() + windows = [ + (start, min(start + 64, raw.n_times)) for start in range(0, raw.n_times, 31) + ] + used = [] + + def _record_use(*args, **kwargs): + result = helper(*args, **kwargs) + used.append(result) + return result + + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + with ThreadPoolExecutor(max_workers=8) as pool: + got = list( + pool.map( + lambda limits: raw.get_data(start=limits[0], stop=limits[1]), + windows * 4, + ) + ) + assert len(used) == len(windows) * 4 + assert all(used) + for data, (start, stop) in zip(got, windows * 4): + assert_array_equal(data, reference[:, start:stop]) + + def test_orig_units(): """Test exposure of original channel units.""" raw = read_raw_edf(edf_path, preload=True) From d94dbb6df3709aa94b7cb39834ade8949e2d6560 Mon Sep 17 00:00:00 2001 From: Bru Date: Wed, 26 Aug 2026 10:27:06 +0200 Subject: [PATCH 7/8] Limit uniform stride decoding to EDF --- doc/changes/dev/14214.newfeature.rst | 2 +- mne/io/edf/edf.py | 7 +- mne/io/edf/tests/test_edf.py | 411 ++++++++++++++++++++++----- 3 files changed, 345 insertions(+), 75 deletions(-) diff --git a/doc/changes/dev/14214.newfeature.rst b/doc/changes/dev/14214.newfeature.rst index a8b58a3525d..9af4c4b5b87 100644 --- a/doc/changes/dev/14214.newfeature.rst +++ b/doc/changes/dev/14214.newfeature.rst @@ -1 +1 @@ -Speed up eligible EDF and BDF window reads with a pure-NumPy record-stride decoder, by `Bruno Aristimunha`_. +Speed up eligible EDF window reads with a pure-NumPy record-stride decoder, by `Bruno Aristimunha`_. diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index fb13d4bff65..7e86cf5c44f 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -615,9 +615,9 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): def _read_uniform_segment( data, idx, start, stop, raw_extras, filenames, cals, mult ) -> bool: - """Read a uniformly sampled EDF or BDF segment.""" + """Read a uniformly sampled EDF segment.""" subtype = raw_extras["subtype"] - if subtype not in ("edf", "bdf") or not isinstance(filenames, str | Path): + if subtype != "edf" or not isinstance(filenames, str | Path): return False if len(raw_extras.get("tal_idx", ())) != 0: return False @@ -655,8 +655,7 @@ def _read_uniform_segment( estimated_incremental_bytes = n_values * np.dtype(np.float64).itemsize physical_order = np.arange(len(n_samps)) if not np.array_equal(read_sel, physical_order): - gather_bytes_per_value = 2 if subtype == "edf" else 4 - estimated_incremental_bytes += n_values * gather_bytes_per_value + estimated_incremental_bytes += n_values * dtype_byte if any(orig_idx in stim_channel_idxs for orig_idx in idx_arr): estimated_incremental_bytes += max_records * buf_len * np.dtype(int).itemsize if estimated_incremental_bytes > _EDF_STRIDE_MAX_EXTRA_BYTES: diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 2e66d4069d4..842dc194af9 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -3,6 +3,7 @@ # Copyright the MNE-Python contributors. import datetime +import gc from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from functools import partial @@ -116,22 +117,13 @@ def _record_use(*args, **kwargs): assert_array_equal(got, want) -@pytest.mark.parametrize( - "reader, source, suffix", - [ - (read_raw_edf, edf_stim_channel_path, ".edf"), - (read_raw_bdf, bdf_path, ".bdf"), - ], -) @pytest.mark.parametrize("pick_kind", ("all", "subset", "reversed", "permuted")) @pytest.mark.parametrize("window_kind", ("within", "boundary", "multiple")) -def test_uniform_stride_decode( - reader, source, suffix, pick_kind, window_kind, monkeypatch, tmp_path -): +def test_uniform_stride_decode(pick_kind, window_kind, monkeypatch, tmp_path): """Test exact pure-NumPy stride decoding against legacy.""" - repeated = tmp_path / f"uniform{suffix}" - _repeat_edf_records(source, repeated) - raw = reader(repeated, preload=False, verbose="error") + repeated = tmp_path / "uniform.edf" + _repeat_edf_records(edf_stim_channel_path, repeated) + raw = read_raw_edf(repeated, stim_channel=-1, preload=False, verbose="error") n_channels = len(raw.ch_names) buf_len = int(raw._raw_extras[0]["max_samp"]) picks = { @@ -148,23 +140,18 @@ def test_uniform_stride_decode( _assert_stride_matches_legacy(monkeypatch, raw, picks, start, stop) -@pytest.mark.parametrize( - "reader, source, suffix", - ( - (read_raw_edf, edf_stim_channel_path, ".edf"), - (read_raw_bdf, bdf_path, ".bdf"), - ), -) -def test_uniform_stride_multiple_chunks(reader, source, suffix, monkeypatch, tmp_path): +def test_uniform_stride_multiple_chunks(monkeypatch, tmp_path): """Test reads that cross the internal record-chunk boundary.""" - source_raw = reader(source, preload=False, verbose="error") + source_raw = read_raw_edf( + edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" + ) extras = source_raw._raw_extras[0] record_nbytes = int(extras["n_samps"].sum() * extras["dtype_byte"]) n_per = max(10 * 1024 * 1024 // record_nbytes, 1) - repeated = tmp_path / f"multiple_chunks{suffix}" - _repeat_edf_records(source, repeated, n_records=n_per + 2) + repeated = tmp_path / "multiple_chunks.edf" + _repeat_edf_records(edf_stim_channel_path, repeated, n_records=n_per + 2) - raw = reader(repeated, preload=False, verbose="error") + raw = read_raw_edf(repeated, stim_channel=-1, preload=False, verbose="error") buf_len = int(raw._raw_extras[0]["max_samp"]) start = buf_len // 2 stop = (n_per + 1) * buf_len + buf_len // 2 @@ -177,14 +164,16 @@ def test_uniform_stride_multiple_chunks(reader, source, suffix, monkeypatch, tmp _assert_stride_matches_legacy(monkeypatch, raw, picks, start, stop) -@pytest.mark.parametrize( - "reader, fname", ((read_raw_edf, edf_stim_channel_path), (read_raw_bdf, bdf_path)) -) -def test_uniform_stride_preload(reader, fname, monkeypatch): +def test_uniform_stride_preload(monkeypatch): """Test that eligible initial eager decoding matches legacy.""" helper = edf.edf._read_uniform_segment monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) - want = reader(fname, preload=True, verbose="error").get_data() + want = read_raw_edf( + edf_stim_channel_path, + stim_channel=-1, + preload=True, + verbose="error", + ).get_data() used = [] def _record_use(*args, **kwargs): @@ -193,45 +182,40 @@ def _record_use(*args, **kwargs): return result monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got = reader(fname, preload=True, verbose="error").get_data() + got = read_raw_edf( + edf_stim_channel_path, + stim_channel=-1, + preload=True, + verbose="error", + ).get_data() assert used and all(used) assert_array_equal(got, want) -@pytest.mark.parametrize( - "reader, fname, reader_kwargs", - ( - pytest.param( - read_raw_edf, - edf_stim_channel_path, - {"stim_channel": -1}, - id="edf", - ), - pytest.param(read_raw_bdf, bdf_path, {}, id="bdf"), - ), -) -def test_uniform_stride_stim_only(reader, fname, reader_kwargs, monkeypatch): +def test_uniform_stride_stim_only(monkeypatch): """Test exact stim scaling and masking on the stride path.""" - raw = reader(fname, preload=False, verbose="error", **reader_kwargs) + raw = read_raw_edf( + edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" + ) picks = pick_types(raw.info, meg=False, stim=True) assert len(picks) == 1 assert raw.get_channel_types(picks=picks) == ["stim"] assert_array_equal(picks, raw._raw_extras[0]["stim_channel_idxs"]) - if reader is read_raw_edf: - stim_idx = picks[0] - extras = raw._raw_extras[0] - assert extras["cal"][stim_idx] != 1.0 - assert extras["offsets"][stim_idx] != 0.0 + stim_idx = picks[0] + extras = raw._raw_extras[0] + assert extras["cal"][stim_idx] != 1.0 + assert extras["offsets"][stim_idx] != 0.0 want = _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) assert_array_equal(want, np.bitwise_and(want.astype(int), 2**17 - 1)) - if reader is read_raw_edf: - assert_array_equal(np.unique(want), [0.0, 100.0]) + assert_array_equal(np.unique(want), [0.0, 100.0]) def test_uniform_stride_cals(monkeypatch): """Test non-unit Raw calibrations on the stride path.""" - raw = read_raw_bdf(bdf_path, preload=False, verbose="error") - picks = np.array([0, 17, 55]) + raw = read_raw_edf( + edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" + ) + picks = np.array([0, 12, 23]) raw._cals[picks] *= np.array([0.5, 2.0, 4.0]) assert np.all(raw._cals[picks] != 1.0) _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) @@ -239,15 +223,154 @@ def test_uniform_stride_cals(monkeypatch): def test_uniform_stride_excluded_physical_channel(monkeypatch): """Test logical picks after excluding a physical record row.""" - raw = read_raw_bdf(bdf_path, exclude=["AF7"], preload=False, verbose="error") + raw = read_raw_edf( + edf_stim_channel_path, + exclude=["EEG Fp2"], + stim_channel=-1, + preload=False, + verbose="error", + ) physical_sel = raw._raw_extras[0]["sel"] - assert raw.ch_names[1] == "AF3" + assert raw.ch_names[1] == "EEG F7" assert physical_sel[1] == 2 assert not np.array_equal(physical_sel, np.arange(len(raw.ch_names))) picks = np.array([0, 1, len(raw.ch_names) - 1]) _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) +@pytest.mark.parametrize( + "mode, pick_kind", + ( + pytest.param("lazy", "all", id="lazy-all"), + pytest.param("lazy", "subset", id="lazy-first-four"), + pytest.param("lazy", "reversed", id="lazy-reversed-all"), + pytest.param("lazy", "permuted", id="lazy-permuted"), + pytest.param("preload", "all", id="preload-true"), + pytest.param("path", "all", id="preload-path"), + ), +) +def test_uniform_stride_bdf_always_falls_back(mode, pick_kind, monkeypatch, tmp_path): + """Test that every BDF read mode and channel order uses legacy.""" + helper = edf.edf._read_uniform_segment + used = [] + + def _record_use(*args, **kwargs): + result = helper(*args, **kwargs) + used.append(result) + return result + + if mode == "lazy": + raw = read_raw_bdf(bdf_path, preload=False, verbose="error") + n_channels = len(raw.ch_names) + all_picks = np.arange(n_channels) + picks = { + "all": all_picks, + "subset": np.arange(4), + "reversed": all_picks[::-1], + "permuted": np.array([n_channels - 1, 1, n_channels // 2, 0]), + }[pick_kind] + monkeypatch.setattr( + edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False + ) + want = raw.get_data(picks=picks, start=100, stop=900) + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + got = raw.get_data(picks=picks, start=100, stop=900) + raw.close() + assert_array_equal(got, want) + elif mode == "preload": + monkeypatch.setattr( + edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False + ) + want_raw = read_raw_bdf(bdf_path, preload=True, verbose="error") + want = want_raw.get_data() + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + got_raw = read_raw_bdf(bdf_path, preload=True, verbose="error") + got = got_raw.get_data() + want_raw.close() + got_raw.close() + assert_array_equal(got, want) + else: + assert mode == "path" + destination = tmp_path / "bdf-preload.dat" + forced_destination = tmp_path / "bdf-forced-preload.dat" + monkeypatch.setattr( + edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False + ) + want_raw = read_raw_bdf(bdf_path, preload=forced_destination, verbose="error") + want = want_raw.get_data() + assert forced_destination.is_file() + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) + got_raw = read_raw_bdf(bdf_path, preload=destination, verbose="error") + got = got_raw.get_data() + assert destination.is_file() + assert_array_equal(got, want) + want_raw.close() + got_raw.close() + assert forced_destination.is_file() + assert destination.is_file() + del want, got + del want_raw._data + del got_raw._data + del want_raw, got_raw + gc.collect() + assert forced_destination.is_file() + assert destination.is_file() + forced_destination.unlink() + destination.unlink() + + assert used and not any(used) + + +class _ExplodingIndex: + @property + def start(self): + raise AssertionError("index inspected") + + def __array__(self, *args, **kwargs): + raise AssertionError("index converted") + + +class _SentinelData: + def __init__(self): + self.mutated = False + + @property + def nbytes(self): + raise AssertionError("data inspected") + + def __setitem__(self, key, value): + self.mutated = True + raise AssertionError("data mutated") + + +@pytest.mark.parametrize("subtype", ("bdf", "gdf")) +def test_uniform_stride_rejects_non_edf_before_work(subtype, monkeypatch): + """Test that non-EDF formats are rejected before any work.""" + + def _explode(*args, **kwargs): + raise AssertionError("allocation or I/O attempted") + + sentinel = np.empty(0) + sentinel.flags.writeable = False + data = _SentinelData() + for name in ("arange", "asarray", "unique", "empty", "zeros"): + monkeypatch.setattr(edf.edf.np, name, _explode) + monkeypatch.setattr(edf.edf, "_gdf_edf_get_fid", _explode) + monkeypatch.setattr(edf.edf, "_read_ch", _explode) + assert not edf.edf._read_uniform_segment( + data, + _ExplodingIndex(), + 0, + 1, + {"subtype": subtype}, + Path("never-opened"), + sentinel, + None, + ) + assert not data.mutated + assert sentinel.size == 0 and not sentinel.flags.writeable + + @pytest.mark.parametrize( "fallback_kind", ( @@ -271,11 +394,19 @@ def test_uniform_stride_fallbacks(fallback_kind, monkeypatch): if fallback_kind == "file_like": helper = edf.edf._read_uniform_segment - blob = bdf_path.read_bytes() + blob = edf_stim_channel_path.read_bytes() monkeypatch.setattr( edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False ) - want = read_raw_bdf(BytesIO(blob), preload=True, verbose="error").get_data() + want_raw = read_raw_edf( + BytesIO(blob), stim_channel=-1, preload=True, verbose="error" + ) + extras = want_raw._raw_extras[0] + assert extras["subtype"] == "edf" + assert want_raw.filenames == (None,) + assert len(extras["tal_idx"]) == 0 + assert np.all(extras["n_samps"] == extras["max_samp"]) + want = want_raw.get_data() used = [] def _record_use(*args, **kwargs): @@ -284,9 +415,14 @@ def _record_use(*args, **kwargs): return result monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got = read_raw_bdf(BytesIO(blob), preload=True, verbose="error").get_data() + got_raw = read_raw_edf( + BytesIO(blob), stim_channel=-1, preload=True, verbose="error" + ) + got = got_raw.get_data() assert used and not any(used) assert_array_equal(got, want) + want_raw.close() + got_raw.close() return raw = None @@ -297,6 +433,9 @@ def _record_use(*args, **kwargs): raw = read_raw_edf(edf_uneven_path, preload=False, verbose="error") extras = raw._raw_extras[0] n_samps = extras["n_samps"] + assert extras["subtype"] == "edf" + assert isinstance(raw.filenames[0], Path) + assert len(extras["tal_idx"]) == 0 assert np.unique(n_samps).size > 1 high_rate_physical = np.flatnonzero(n_samps == extras["max_samp"])[0] picks = np.flatnonzero(extras["sel"] == high_rate_physical) @@ -306,21 +445,67 @@ def _record_use(*args, **kwargs): raw = read_raw_edf(edf_path, preload=False, verbose="error") extras = raw._raw_extras[0] picks = np.array([0, len(raw.ch_names) - 1]) + assert extras["subtype"] == "edf" + assert isinstance(raw.filenames[0], Path) assert len(extras["tal_idx"]) > 0 - assert np.all(extras["n_samps"][extras["sel"][picks]] == extras["max_samp"]) + assert np.all(extras["n_samps"] == extras["max_samp"]) elif fallback_kind == "projection": - raw = read_raw_bdf(bdf_path, preload=False, verbose="error") + raw = read_raw_edf( + edf_stim_channel_path, + stim_channel=-1, + preload=False, + verbose="error", + ) + extras = raw._raw_extras[0] + assert extras["subtype"] == "edf" + assert isinstance(raw.filenames[0], Path) + assert len(extras["tal_idx"]) == 0 + assert np.all(extras["n_samps"] == extras["max_samp"]) raw.set_eeg_reference(projection=True, verbose="error").apply_proj( verbose="error" ) assert raw._projector is not None expect_mult = True elif fallback_kind in ("output_limit", "extra_limit"): - raw = read_raw_bdf(bdf_path, preload=False, verbose="error") + raw = read_raw_edf( + edf_stim_channel_path, + stim_channel=-1, + preload=False, + verbose="error", + ) + extras = raw._raw_extras[0] + assert extras["subtype"] == "edf" + assert isinstance(raw.filenames[0], Path) + assert len(extras["tal_idx"]) == 0 + assert np.all(extras["n_samps"] == extras["max_samp"]) + picks = np.array([0, 1]) + output_nbytes = len(picks) * (stop - start) * 8 + if fallback_kind == "output_limit": + assert output_nbytes > edf.edf._EDF_STRIDE_MAX_OUTPUT_BYTES == 0 + else: + buf_len = int(extras["max_samp"]) + _, r_lims, _ = _blk_read_lims(start, stop, buf_len) + n_per = max( + 10 * 1024 * 1024 // (extras["n_samps"].sum() * extras["dtype_byte"]), + 1, + ) + n_values = len(picks) * min(len(r_lims), n_per) * buf_len + estimated_extra = 10 * n_values + assert estimated_extra > edf.edf._EDF_STRIDE_MAX_EXTRA_BYTES == 0 else: assert fallback_kind == "duplicate_picks" helper = edf.edf._read_uniform_segment - raw = read_raw_bdf(bdf_path, preload=False, verbose="error") + raw = read_raw_edf( + edf_stim_channel_path, + stim_channel=-1, + preload=False, + verbose="error", + ) + extras = raw._raw_extras[0] + assert extras["subtype"] == "edf" + assert isinstance(raw.filenames[0], Path) + assert len(extras["tal_idx"]) == 0 + assert np.all(extras["n_samps"] == extras["max_samp"]) picks = [2, 0, 2] monkeypatch.setattr( edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False @@ -348,16 +533,102 @@ def _record_use(*args, **kwargs): ) -@pytest.mark.parametrize( - "reader, source", - ((read_raw_edf, edf_stim_channel_path), (read_raw_bdf, bdf_path)), -) -def test_uniform_stride_concurrent(reader, source, monkeypatch): +def _assert_uniform_limit_result( + monkeypatch, real_helper, *, output_limit, extra_limit, used +): + raw = read_raw_edf( + edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" + ) + picks = np.array([0, 1]) + start, stop = 100, 900 + extras = raw._raw_extras[0] + assert extras["subtype"] == "edf" + assert len(extras["tal_idx"]) == 0 + assert extras["max_samp"] == 1228 + assert np.all(extras["n_samps"] == 1228) + read_sel = extras["sel"][picks] + _, r_lims, _ = _blk_read_lims(start, stop, 1228) + n_per = max( + 10 * 1024 * 1024 // (extras["n_samps"].sum() * extras["dtype_byte"]), + 1, + ) + max_records = min(len(r_lims), n_per) + n_values = len(picks) * max_records * 1228 + estimated_extra = 10 * n_values + assert not np.array_equal(read_sel, np.arange(len(extras["n_samps"]))) + assert not any(idx in extras["stim_channel_idxs"] for idx in picks) + assert max_records == 1 + assert estimated_extra == 24560 + monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) + want = raw.get_data(picks=picks, start=start, stop=stop) + assert want.nbytes == 2 * 800 * 8 == 12800 + + calls = [] + + def _record(*args, **kwargs): + result = real_helper(*args, **kwargs) + calls.append(result) + return result + + monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_OUTPUT_BYTES", output_limit) + monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_EXTRA_BYTES", extra_limit) + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record) + got = raw.get_data(picks=picks, start=start, stop=stop) + assert calls == [used] + assert_array_equal(got, want) + + +def test_uniform_stride_output_limit_boundary(monkeypatch): + """Test that the output limit is inclusive.""" + real_helper = edf.edf._read_uniform_segment + _assert_uniform_limit_result( + monkeypatch, + real_helper, + output_limit=12800, + extra_limit=24560, + used=True, + ) + _assert_uniform_limit_result( + monkeypatch, + real_helper, + output_limit=12799, + extra_limit=24560, + used=False, + ) + + +def test_uniform_stride_extra_limit_boundary(monkeypatch): + """Test that the incremental-allocation limit is inclusive.""" + real_helper = edf.edf._read_uniform_segment + _assert_uniform_limit_result( + monkeypatch, + real_helper, + output_limit=12800, + extra_limit=24560, + used=True, + ) + _assert_uniform_limit_result( + monkeypatch, + real_helper, + output_limit=12800, + extra_limit=24559, + used=False, + ) + + +def test_uniform_stride_concurrent(monkeypatch): """Test that stride reads share no seekable file handle.""" - raw = reader(source, preload=False, verbose="error") + raw = read_raw_edf( + edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" + ) helper = edf.edf._read_uniform_segment monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) - reference = reader(source, preload=True, verbose="error").get_data() + reference = read_raw_edf( + edf_stim_channel_path, + stim_channel=-1, + preload=True, + verbose="error", + ).get_data() windows = [ (start, min(start + 64, raw.n_times)) for start in range(0, raw.n_times, 31) ] From 0182754b44018bbd1624134be024c778c54711d0 Mon Sep 17 00:00:00 2001 From: Bru Date: Wed, 26 Aug 2026 13:03:51 +0200 Subject: [PATCH 8/8] Withdraw EDF stride decoder after benchmark validation --- doc/changes/dev/14214.newfeature.rst | 1 - mne/io/edf/edf.py | 93 ----- mne/io/edf/tests/test_edf.py | 587 --------------------------- 3 files changed, 681 deletions(-) delete mode 100644 doc/changes/dev/14214.newfeature.rst diff --git a/doc/changes/dev/14214.newfeature.rst b/doc/changes/dev/14214.newfeature.rst deleted file mode 100644 index 9af4c4b5b87..00000000000 --- a/doc/changes/dev/14214.newfeature.rst +++ /dev/null @@ -1 +0,0 @@ -Speed up eligible EDF window reads with a pure-NumPy record-stride decoder, by `Bruno Aristimunha`_. diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 7e86cf5c44f..390289ae6bb 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -608,101 +608,8 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): return ch_data -_EDF_STRIDE_MAX_OUTPUT_BYTES = 32 * 1024**2 -_EDF_STRIDE_MAX_EXTRA_BYTES = 64 * 1024**2 - - -def _read_uniform_segment( - data, idx, start, stop, raw_extras, filenames, cals, mult -) -> bool: - """Read a uniformly sampled EDF segment.""" - subtype = raw_extras["subtype"] - if subtype != "edf" or not isinstance(filenames, str | Path): - return False - if len(raw_extras.get("tal_idx", ())) != 0: - return False - - idx_arr = ( - np.arange(idx.start, idx.stop) if isinstance(idx, slice) else np.asarray(idx) - ) - if len(idx_arr) == 0 or len(np.unique(idx_arr)) != len(idx_arr): - return False - - n_samps = raw_extras["n_samps"] - buf_len = int(raw_extras["max_samp"]) - if ( - not np.all(n_samps == buf_len) - or mult is not None - or data.nbytes > _EDF_STRIDE_MAX_OUTPUT_BYTES - ): - return False - - dtype = raw_extras["dtype_np"] - dtype_byte = raw_extras["dtype_byte"] - data_offset = raw_extras["data_offset"] - stim_channel_idxs = raw_extras["stim_channel_idxs"] - orig_sel = raw_extras["sel"] - cal = raw_extras["cal"] - offsets = raw_extras["offsets"] - gains = raw_extras["units"] - read_sel = orig_sel[idx_arr] - - ch_offsets = np.cumsum(np.concatenate([[0], n_samps]), dtype=np.int64) - block_start_idx, r_lims, d_lims = _blk_read_lims(start, stop, buf_len) - n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * dtype_byte), 1) - max_records = min(len(r_lims), n_per) - n_values = len(idx_arr) * max_records * buf_len - estimated_incremental_bytes = n_values * np.dtype(np.float64).itemsize - physical_order = np.arange(len(n_samps)) - if not np.array_equal(read_sel, physical_order): - estimated_incremental_bytes += n_values * dtype_byte - if any(orig_idx in stim_channel_idxs for orig_idx in idx_arr): - estimated_incremental_bytes += max_records * buf_len * np.dtype(int).itemsize - if estimated_incremental_bytes > _EDF_STRIDE_MAX_EXTRA_BYTES: - return False - - with _gdf_edf_get_fid(filenames, buffering=0) as fid: - start_offset = data_offset + block_start_idx * ch_offsets[-1] * dtype_byte - ones = np.zeros((len(orig_sel), data.shape[-1]), dtype=data.dtype) - for ai in range(0, len(r_lims), n_per): - block_offset = ai * ch_offsets[-1] * dtype_byte - n_read = min(len(r_lims) - ai, n_per) - fid.seek(start_offset + block_offset, 0) - many_chunk = _read_ch( - fid, subtype, ch_offsets[-1] * n_read, dtype_byte, dtype - ) - record_grid = many_chunk.reshape(n_read, len(n_samps), buf_len) - if np.array_equal(read_sel, physical_order): - view = record_grid.transpose(1, 0, 2) - else: - view = record_grid[:, read_sel, :].transpose(1, 0, 2) - - one = np.empty(view.shape, dtype=np.float64) - np.multiply(view, cal[idx_arr, np.newaxis, np.newaxis], out=one) - one += offsets[idx_arr, np.newaxis, np.newaxis] - one *= gains[idx_arr, np.newaxis, np.newaxis] - r_sidx = r_lims[ai][0] - r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] - block = one.reshape(len(idx_arr), -1)[:, r_sidx:r_eidx] - for row, orig_idx in enumerate(idx_arr): - if orig_idx in stim_channel_idxs: - stim = block[row].astype(int) - np.bitwise_and(stim, 2**17 - 1, out=stim) - block[row] = stim - d_start = d_lims[ai][0] - d_stop = d_lims[ai + n_read - 1][1] - assert d_stop - d_start == block.shape[1] - ones[idx_arr, d_start:d_stop] = block - - _mult_cal_one(data, ones, idx, cals, mult) - return True - - def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, mult): """Read a chunk of raw data.""" - if _read_uniform_segment(data, idx, start, stop, raw_extras, filenames, cals, mult): - return [] - from scipy.interpolate import interp1d n_samps = raw_extras["n_samps"] diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 842dc194af9..211641dc726 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -3,8 +3,6 @@ # Copyright the MNE-Python contributors. import datetime -import gc -from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from functools import partial from io import BytesIO @@ -22,7 +20,6 @@ from mne import Annotations, pick_types from mne._fiff.pick import channel_indices_by_type, get_channel_type_constants -from mne._fiff.utils import _blk_read_lims from mne.annotations import _ndarray_ch_names, events_from_annotations, read_annotations from mne.datasets import testing from mne.io import edf, read_raw_bdf, read_raw_edf, read_raw_fif, read_raw_gdf @@ -69,590 +66,6 @@ misc = ["EXG1", "EXG5", "EXG8", "M1", "M2"] -def _repeat_edf_records(source, destination, n_records=6): - """Repeat a one-record EDF/BDF payload for boundary tests.""" - blob = bytearray(source.read_bytes()) - header_nbytes = int(blob[184:192]) - assert int(blob[236:244]) == 1 - blob[236:244] = f"{n_records:<8}".encode("ascii") - destination.write_bytes(blob[:header_nbytes] + blob[header_nbytes:] * n_records) - - -def _assert_stride_matches_legacy(monkeypatch, raw, picks, start, stop): - helper = edf.edf._read_uniform_segment - monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) - want = raw.get_data(picks=picks, start=start, stop=stop) - used = [] - - def _record_use(*args, **kwargs): - result = helper(*args, **kwargs) - used.append(result) - return result - - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got = raw.get_data(picks=picks, start=start, stop=stop) - assert used and all(used) - assert_array_equal(got, want) - return want - - -def _assert_stride_falls_back(monkeypatch, read_data, *, expect_mult=False): - helper = edf.edf._read_uniform_segment - monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) - want = read_data() - used = [] - mults = [] - - def _record_use(*args, **kwargs): - result = helper(*args, **kwargs) - used.append(result) - mults.append(args[-1] if args else kwargs["mult"]) - return result - - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got = read_data() - assert used and not any(used) - if expect_mult: - assert all(mult is not None for mult in mults) - assert_array_equal(got, want) - - -@pytest.mark.parametrize("pick_kind", ("all", "subset", "reversed", "permuted")) -@pytest.mark.parametrize("window_kind", ("within", "boundary", "multiple")) -def test_uniform_stride_decode(pick_kind, window_kind, monkeypatch, tmp_path): - """Test exact pure-NumPy stride decoding against legacy.""" - repeated = tmp_path / "uniform.edf" - _repeat_edf_records(edf_stim_channel_path, repeated) - raw = read_raw_edf(repeated, stim_channel=-1, preload=False, verbose="error") - n_channels = len(raw.ch_names) - buf_len = int(raw._raw_extras[0]["max_samp"]) - picks = { - "all": np.arange(n_channels), - "subset": np.array([0, n_channels // 2, n_channels - 1]), - "reversed": np.arange(n_channels - 1, -1, -1), - "permuted": np.array([n_channels - 1, 1, n_channels // 2, 0]), - }[pick_kind] - start, stop = { - "within": (7, buf_len - 5), - "boundary": (buf_len - 7, buf_len + 11), - "multiple": (buf_len // 2, 4 * buf_len + 13), - }[window_kind] - _assert_stride_matches_legacy(monkeypatch, raw, picks, start, stop) - - -def test_uniform_stride_multiple_chunks(monkeypatch, tmp_path): - """Test reads that cross the internal record-chunk boundary.""" - source_raw = read_raw_edf( - edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" - ) - extras = source_raw._raw_extras[0] - record_nbytes = int(extras["n_samps"].sum() * extras["dtype_byte"]) - n_per = max(10 * 1024 * 1024 // record_nbytes, 1) - repeated = tmp_path / "multiple_chunks.edf" - _repeat_edf_records(edf_stim_channel_path, repeated, n_records=n_per + 2) - - raw = read_raw_edf(repeated, stim_channel=-1, preload=False, verbose="error") - buf_len = int(raw._raw_extras[0]["max_samp"]) - start = buf_len // 2 - stop = (n_per + 1) * buf_len + buf_len // 2 - _, r_lims, _ = _blk_read_lims(start, stop, buf_len) - assert len(r_lims) > n_per - assert stop <= raw.n_times - picks = np.array([0, len(raw.ch_names) // 2]) - output_nbytes = len(picks) * (stop - start) * np.dtype(np.float64).itemsize - assert output_nbytes < 4 * 1024**2 - _assert_stride_matches_legacy(monkeypatch, raw, picks, start, stop) - - -def test_uniform_stride_preload(monkeypatch): - """Test that eligible initial eager decoding matches legacy.""" - helper = edf.edf._read_uniform_segment - monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) - want = read_raw_edf( - edf_stim_channel_path, - stim_channel=-1, - preload=True, - verbose="error", - ).get_data() - used = [] - - def _record_use(*args, **kwargs): - result = helper(*args, **kwargs) - used.append(result) - return result - - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got = read_raw_edf( - edf_stim_channel_path, - stim_channel=-1, - preload=True, - verbose="error", - ).get_data() - assert used and all(used) - assert_array_equal(got, want) - - -def test_uniform_stride_stim_only(monkeypatch): - """Test exact stim scaling and masking on the stride path.""" - raw = read_raw_edf( - edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" - ) - picks = pick_types(raw.info, meg=False, stim=True) - assert len(picks) == 1 - assert raw.get_channel_types(picks=picks) == ["stim"] - assert_array_equal(picks, raw._raw_extras[0]["stim_channel_idxs"]) - stim_idx = picks[0] - extras = raw._raw_extras[0] - assert extras["cal"][stim_idx] != 1.0 - assert extras["offsets"][stim_idx] != 0.0 - want = _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) - assert_array_equal(want, np.bitwise_and(want.astype(int), 2**17 - 1)) - assert_array_equal(np.unique(want), [0.0, 100.0]) - - -def test_uniform_stride_cals(monkeypatch): - """Test non-unit Raw calibrations on the stride path.""" - raw = read_raw_edf( - edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" - ) - picks = np.array([0, 12, 23]) - raw._cals[picks] *= np.array([0.5, 2.0, 4.0]) - assert np.all(raw._cals[picks] != 1.0) - _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) - - -def test_uniform_stride_excluded_physical_channel(monkeypatch): - """Test logical picks after excluding a physical record row.""" - raw = read_raw_edf( - edf_stim_channel_path, - exclude=["EEG Fp2"], - stim_channel=-1, - preload=False, - verbose="error", - ) - physical_sel = raw._raw_extras[0]["sel"] - assert raw.ch_names[1] == "EEG F7" - assert physical_sel[1] == 2 - assert not np.array_equal(physical_sel, np.arange(len(raw.ch_names))) - picks = np.array([0, 1, len(raw.ch_names) - 1]) - _assert_stride_matches_legacy(monkeypatch, raw, picks, 100, 900) - - -@pytest.mark.parametrize( - "mode, pick_kind", - ( - pytest.param("lazy", "all", id="lazy-all"), - pytest.param("lazy", "subset", id="lazy-first-four"), - pytest.param("lazy", "reversed", id="lazy-reversed-all"), - pytest.param("lazy", "permuted", id="lazy-permuted"), - pytest.param("preload", "all", id="preload-true"), - pytest.param("path", "all", id="preload-path"), - ), -) -def test_uniform_stride_bdf_always_falls_back(mode, pick_kind, monkeypatch, tmp_path): - """Test that every BDF read mode and channel order uses legacy.""" - helper = edf.edf._read_uniform_segment - used = [] - - def _record_use(*args, **kwargs): - result = helper(*args, **kwargs) - used.append(result) - return result - - if mode == "lazy": - raw = read_raw_bdf(bdf_path, preload=False, verbose="error") - n_channels = len(raw.ch_names) - all_picks = np.arange(n_channels) - picks = { - "all": all_picks, - "subset": np.arange(4), - "reversed": all_picks[::-1], - "permuted": np.array([n_channels - 1, 1, n_channels // 2, 0]), - }[pick_kind] - monkeypatch.setattr( - edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False - ) - want = raw.get_data(picks=picks, start=100, stop=900) - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got = raw.get_data(picks=picks, start=100, stop=900) - raw.close() - assert_array_equal(got, want) - elif mode == "preload": - monkeypatch.setattr( - edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False - ) - want_raw = read_raw_bdf(bdf_path, preload=True, verbose="error") - want = want_raw.get_data() - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got_raw = read_raw_bdf(bdf_path, preload=True, verbose="error") - got = got_raw.get_data() - want_raw.close() - got_raw.close() - assert_array_equal(got, want) - else: - assert mode == "path" - destination = tmp_path / "bdf-preload.dat" - forced_destination = tmp_path / "bdf-forced-preload.dat" - monkeypatch.setattr( - edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False - ) - want_raw = read_raw_bdf(bdf_path, preload=forced_destination, verbose="error") - want = want_raw.get_data() - assert forced_destination.is_file() - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got_raw = read_raw_bdf(bdf_path, preload=destination, verbose="error") - got = got_raw.get_data() - assert destination.is_file() - assert_array_equal(got, want) - want_raw.close() - got_raw.close() - assert forced_destination.is_file() - assert destination.is_file() - del want, got - del want_raw._data - del got_raw._data - del want_raw, got_raw - gc.collect() - assert forced_destination.is_file() - assert destination.is_file() - forced_destination.unlink() - destination.unlink() - - assert used and not any(used) - - -class _ExplodingIndex: - @property - def start(self): - raise AssertionError("index inspected") - - def __array__(self, *args, **kwargs): - raise AssertionError("index converted") - - -class _SentinelData: - def __init__(self): - self.mutated = False - - @property - def nbytes(self): - raise AssertionError("data inspected") - - def __setitem__(self, key, value): - self.mutated = True - raise AssertionError("data mutated") - - -@pytest.mark.parametrize("subtype", ("bdf", "gdf")) -def test_uniform_stride_rejects_non_edf_before_work(subtype, monkeypatch): - """Test that non-EDF formats are rejected before any work.""" - - def _explode(*args, **kwargs): - raise AssertionError("allocation or I/O attempted") - - sentinel = np.empty(0) - sentinel.flags.writeable = False - data = _SentinelData() - for name in ("arange", "asarray", "unique", "empty", "zeros"): - monkeypatch.setattr(edf.edf.np, name, _explode) - monkeypatch.setattr(edf.edf, "_gdf_edf_get_fid", _explode) - monkeypatch.setattr(edf.edf, "_read_ch", _explode) - assert not edf.edf._read_uniform_segment( - data, - _ExplodingIndex(), - 0, - 1, - {"subtype": subtype}, - Path("never-opened"), - sentinel, - None, - ) - assert not data.mutated - assert sentinel.size == 0 and not sentinel.flags.writeable - - -@pytest.mark.parametrize( - "fallback_kind", - ( - "mixed_rate", - "tal", - "file_like", - "projection", - "output_limit", - "extra_limit", - "duplicate_picks", - ), -) -def test_uniform_stride_fallbacks(fallback_kind, monkeypatch): - """Test that unsafe stride layouts and transforms use legacy.""" - cutoff = { - "output_limit": "_EDF_STRIDE_MAX_OUTPUT_BYTES", - "extra_limit": "_EDF_STRIDE_MAX_EXTRA_BYTES", - }.get(fallback_kind) - if cutoff is not None: - monkeypatch.setattr(edf.edf, cutoff, 0) - - if fallback_kind == "file_like": - helper = edf.edf._read_uniform_segment - blob = edf_stim_channel_path.read_bytes() - monkeypatch.setattr( - edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False - ) - want_raw = read_raw_edf( - BytesIO(blob), stim_channel=-1, preload=True, verbose="error" - ) - extras = want_raw._raw_extras[0] - assert extras["subtype"] == "edf" - assert want_raw.filenames == (None,) - assert len(extras["tal_idx"]) == 0 - assert np.all(extras["n_samps"] == extras["max_samp"]) - want = want_raw.get_data() - used = [] - - def _record_use(*args, **kwargs): - result = helper(*args, **kwargs) - used.append(result) - return result - - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - got_raw = read_raw_edf( - BytesIO(blob), stim_channel=-1, preload=True, verbose="error" - ) - got = got_raw.get_data() - assert used and not any(used) - assert_array_equal(got, want) - want_raw.close() - got_raw.close() - return - - raw = None - picks = None - start, stop = 100, 900 - expect_mult = False - if fallback_kind == "mixed_rate": - raw = read_raw_edf(edf_uneven_path, preload=False, verbose="error") - extras = raw._raw_extras[0] - n_samps = extras["n_samps"] - assert extras["subtype"] == "edf" - assert isinstance(raw.filenames[0], Path) - assert len(extras["tal_idx"]) == 0 - assert np.unique(n_samps).size > 1 - high_rate_physical = np.flatnonzero(n_samps == extras["max_samp"])[0] - picks = np.flatnonzero(extras["sel"] == high_rate_physical) - assert len(picks) == 1 - assert n_samps[extras["sel"][picks[0]]] == extras["max_samp"] - elif fallback_kind == "tal": - raw = read_raw_edf(edf_path, preload=False, verbose="error") - extras = raw._raw_extras[0] - picks = np.array([0, len(raw.ch_names) - 1]) - assert extras["subtype"] == "edf" - assert isinstance(raw.filenames[0], Path) - assert len(extras["tal_idx"]) > 0 - assert np.all(extras["n_samps"] == extras["max_samp"]) - elif fallback_kind == "projection": - raw = read_raw_edf( - edf_stim_channel_path, - stim_channel=-1, - preload=False, - verbose="error", - ) - extras = raw._raw_extras[0] - assert extras["subtype"] == "edf" - assert isinstance(raw.filenames[0], Path) - assert len(extras["tal_idx"]) == 0 - assert np.all(extras["n_samps"] == extras["max_samp"]) - raw.set_eeg_reference(projection=True, verbose="error").apply_proj( - verbose="error" - ) - assert raw._projector is not None - expect_mult = True - elif fallback_kind in ("output_limit", "extra_limit"): - raw = read_raw_edf( - edf_stim_channel_path, - stim_channel=-1, - preload=False, - verbose="error", - ) - extras = raw._raw_extras[0] - assert extras["subtype"] == "edf" - assert isinstance(raw.filenames[0], Path) - assert len(extras["tal_idx"]) == 0 - assert np.all(extras["n_samps"] == extras["max_samp"]) - picks = np.array([0, 1]) - output_nbytes = len(picks) * (stop - start) * 8 - if fallback_kind == "output_limit": - assert output_nbytes > edf.edf._EDF_STRIDE_MAX_OUTPUT_BYTES == 0 - else: - buf_len = int(extras["max_samp"]) - _, r_lims, _ = _blk_read_lims(start, stop, buf_len) - n_per = max( - 10 * 1024 * 1024 // (extras["n_samps"].sum() * extras["dtype_byte"]), - 1, - ) - n_values = len(picks) * min(len(r_lims), n_per) * buf_len - estimated_extra = 10 * n_values - assert estimated_extra > edf.edf._EDF_STRIDE_MAX_EXTRA_BYTES == 0 - else: - assert fallback_kind == "duplicate_picks" - helper = edf.edf._read_uniform_segment - raw = read_raw_edf( - edf_stim_channel_path, - stim_channel=-1, - preload=False, - verbose="error", - ) - extras = raw._raw_extras[0] - assert extras["subtype"] == "edf" - assert isinstance(raw.filenames[0], Path) - assert len(extras["tal_idx"]) == 0 - assert np.all(extras["n_samps"] == extras["max_samp"]) - picks = [2, 0, 2] - monkeypatch.setattr( - edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False - ) - with pytest.raises(ValueError) as want_error: - raw.get_data(picks=picks, start=start, stop=stop) - used = [] - - def _record_use(*args, **kwargs): - result = helper(*args, **kwargs) - used.append(result) - return result - - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - with pytest.raises(ValueError) as got_error: - raw.get_data(picks=picks, start=start, stop=stop) - assert used and not any(used) - assert str(got_error.value) == str(want_error.value) - return - - _assert_stride_falls_back( - monkeypatch, - partial(raw.get_data, picks=picks, start=start, stop=stop), - expect_mult=expect_mult, - ) - - -def _assert_uniform_limit_result( - monkeypatch, real_helper, *, output_limit, extra_limit, used -): - raw = read_raw_edf( - edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" - ) - picks = np.array([0, 1]) - start, stop = 100, 900 - extras = raw._raw_extras[0] - assert extras["subtype"] == "edf" - assert len(extras["tal_idx"]) == 0 - assert extras["max_samp"] == 1228 - assert np.all(extras["n_samps"] == 1228) - read_sel = extras["sel"][picks] - _, r_lims, _ = _blk_read_lims(start, stop, 1228) - n_per = max( - 10 * 1024 * 1024 // (extras["n_samps"].sum() * extras["dtype_byte"]), - 1, - ) - max_records = min(len(r_lims), n_per) - n_values = len(picks) * max_records * 1228 - estimated_extra = 10 * n_values - assert not np.array_equal(read_sel, np.arange(len(extras["n_samps"]))) - assert not any(idx in extras["stim_channel_idxs"] for idx in picks) - assert max_records == 1 - assert estimated_extra == 24560 - monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) - want = raw.get_data(picks=picks, start=start, stop=stop) - assert want.nbytes == 2 * 800 * 8 == 12800 - - calls = [] - - def _record(*args, **kwargs): - result = real_helper(*args, **kwargs) - calls.append(result) - return result - - monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_OUTPUT_BYTES", output_limit) - monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_EXTRA_BYTES", extra_limit) - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record) - got = raw.get_data(picks=picks, start=start, stop=stop) - assert calls == [used] - assert_array_equal(got, want) - - -def test_uniform_stride_output_limit_boundary(monkeypatch): - """Test that the output limit is inclusive.""" - real_helper = edf.edf._read_uniform_segment - _assert_uniform_limit_result( - monkeypatch, - real_helper, - output_limit=12800, - extra_limit=24560, - used=True, - ) - _assert_uniform_limit_result( - monkeypatch, - real_helper, - output_limit=12799, - extra_limit=24560, - used=False, - ) - - -def test_uniform_stride_extra_limit_boundary(monkeypatch): - """Test that the incremental-allocation limit is inclusive.""" - real_helper = edf.edf._read_uniform_segment - _assert_uniform_limit_result( - monkeypatch, - real_helper, - output_limit=12800, - extra_limit=24560, - used=True, - ) - _assert_uniform_limit_result( - monkeypatch, - real_helper, - output_limit=12800, - extra_limit=24559, - used=False, - ) - - -def test_uniform_stride_concurrent(monkeypatch): - """Test that stride reads share no seekable file handle.""" - raw = read_raw_edf( - edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" - ) - helper = edf.edf._read_uniform_segment - monkeypatch.setattr(edf.edf, "_read_uniform_segment", lambda *args, **kwargs: False) - reference = read_raw_edf( - edf_stim_channel_path, - stim_channel=-1, - preload=True, - verbose="error", - ).get_data() - windows = [ - (start, min(start + 64, raw.n_times)) for start in range(0, raw.n_times, 31) - ] - used = [] - - def _record_use(*args, **kwargs): - result = helper(*args, **kwargs) - used.append(result) - return result - - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record_use) - with ThreadPoolExecutor(max_workers=8) as pool: - got = list( - pool.map( - lambda limits: raw.get_data(start=limits[0], stop=limits[1]), - windows * 4, - ) - ) - assert len(used) == len(windows) * 4 - assert all(used) - for data, (start, stop) in zip(got, windows * 4): - assert_array_equal(data, reference[:, start:stop]) - - def test_orig_units(): """Test exposure of original channel units.""" raw = read_raw_edf(edf_path, preload=True)