From 1744ac8a38fc6ec899c3452b3563a64c7bcf9bf3 Mon Sep 17 00:00:00 2001 From: Bru Date: Wed, 26 Aug 2026 17:53:22 +0200 Subject: [PATCH 01/12] Avoid materializing Raw time vectors during setup --- mne/io/base.py | 14 ++++++------- mne/io/tests/test_raw.py | 45 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 7 deletions(-) diff --git a/mne/io/base.py b/mne/io/base.py index 47b7b3af8fe..5c8c765b98c 100644 --- a/mne/io/base.py +++ b/mne/io/base.py @@ -633,9 +633,10 @@ def _preload_data(self, preload): data_buffer = preload if isinstance(preload, bool | np.bool_) and not preload: data_buffer = None - t = self.times + n_times = self.n_times + last_time = (n_times - 1) / self.info["sfreq"] logger.info( - f"Reading 0 ... {len(t) - 1} = {0.0:9.3f} ... {t[-1]:9.3f} secs..." + f"Reading 0 ... {n_times - 1} = {0.0:9.3f} ... {last_time:9.3f} secs..." ) self._data = self._read_segment(data_buffer=data_buffer) assert len(self._data) == self.info["nchan"] @@ -793,17 +794,16 @@ def set_annotations( "of the raw object." ) - delta = 1.0 / self.info["sfreq"] + sfreq = self.info["sfreq"] + annotation_end = (self.n_times - 1) / sfreq + 1.0 / sfreq new_annotations = annotations.copy() new_annotations._prune_ch_names(self.info, on_missing) if annotations.orig_time is None: - new_annotations.crop( - 0, self.times[-1] + delta, emit_warning=emit_warning - ) + new_annotations.crop(0, annotation_end, emit_warning=emit_warning) new_annotations.onset += self._first_time else: tmin = meas_date + timedelta(0, self._first_time) - tmax = tmin + timedelta(seconds=self.times[-1] + delta) + tmax = tmin + timedelta(seconds=annotation_end) new_annotations.crop(tmin=tmin, tmax=tmax, emit_warning=emit_warning) new_annotations.onset -= ( meas_date - new_annotations.orig_time diff --git a/mne/io/tests/test_raw.py b/mne/io/tests/test_raw.py index 0f4b47d6e80..1a564433f9a 100644 --- a/mne/io/tests/test_raw.py +++ b/mne/io/tests/test_raw.py @@ -48,6 +48,20 @@ ) +def _fail_if_times_materialized(*args, **kwargs): + pytest.fail("The full Raw.times vector was materialized") + + +class _CropRecorder: + def __init__(self): + self.args = None + self.kwargs = None + + def crop(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + def assert_named_constants(info): """Assert that info['chs'] has named constants.""" # for now we just check one @@ -99,6 +113,37 @@ def test_orig_units(): BaseRaw(info, last_samps=[1], orig_units=True) +def test_preload_does_not_materialize_times(monkeypatch): + """Test preloading does not construct the full time vector.""" + monkeypatch.setattr("mne.io.base._arange_div", _fail_if_times_materialized) + raw = read_raw_fif(raw_fname, preload=True, verbose="error") + assert raw.preload + + +def test_set_annotations_does_not_materialize_times(monkeypatch): + """Test annotation bounds use the scalar recording endpoint.""" + raw = read_raw_fif(raw_fname, preload=False, verbose="error") + annotations = Annotations([0.0], [0.1], ["test"]) + monkeypatch.setattr("mne.io.base._arange_div", _fail_if_times_materialized) + raw.set_annotations(annotations) + assert len(raw.annotations) == 1 + + +def test_set_annotations_preserves_endpoint_arithmetic(monkeypatch): + """Test annotation bounds preserve the prior floating-point operations.""" + raw = RawArray(np.zeros((1, 6)), create_info(1, 100.0), verbose="error") + annotations = Annotations([0.0], [0.0], ["test"]) + recorder = _CropRecorder() + monkeypatch.setattr(Annotations, "crop", recorder.crop) + + raw.set_annotations(annotations) + + endpoint = (raw.n_times - 1) / raw.info["sfreq"] + 1.0 / raw.info["sfreq"] + assert endpoint != raw.duration + assert recorder.args == (0, endpoint) + assert recorder.kwargs == {"emit_warning": True} + + def _test_raw_reader( reader, test_preloading=True, From c082b9dfb9a6fc8f008aee98ca076900a0e95333 Mon Sep 17 00:00:00 2001 From: Bru Date: Wed, 26 Aug 2026 17:53:27 +0200 Subject: [PATCH 02/12] Speed up uniform EDF and BDF decoding --- mne/io/edf/edf.py | 202 +++++++++++++++++++++++- mne/io/edf/tests/test_edf.py | 288 +++++++++++++++++++++++++++++++++++ 2 files changed, 484 insertions(+), 6 deletions(-) diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 390289ae6bb..8e41d99c60d 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -595,11 +595,28 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): assert dtype is not None # BDF if subtype == "bdf": - ch_data = read_from_file_or_buffer(fid, dtype=dtype, count=samp * dtype_byte) - ch_data = ch_data.reshape(-1, 3).astype(INT32) - ch_data = (ch_data[:, 0]) + (ch_data[:, 1] << 8) + (ch_data[:, 2] << 16) - # 24th bit determines the sign - ch_data[ch_data >= (1 << 23)] -= 1 << 24 + assert dtype_byte == 3 + expected = samp * dtype_byte + try: + raw = read_from_file_or_buffer(fid, dtype=dtype, count=expected) + except ValueError as err: + raise RuntimeError( + f"Could not read {expected} requested BDF bytes" + ) from err + if raw.size != expected: + raise RuntimeError( + f"Only {raw.size} of {expected} requested BDF bytes could be read" + ) + ch_data = np.empty(samp, dtype=INT32) + packed = np.ndarray( + (max(samp - 1, 0),), dtype=">= 8 # GDF data and EDF data else: @@ -608,9 +625,163 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): return ch_data +_EDF_STRIDE_MAX_EXTRA_BYTES = 64 * 1024**2 + + +def _calibrate_uniform_edf(data, source, cal, offsets, gains, cals): + """Calibrate uniform EDF/BDF samples directly into their output buffer.""" + broadcast_shape = (len(data),) + (1,) * (data.ndim - 1) + np.multiply(source, cal.reshape(broadcast_shape), out=data, casting="unsafe") + data += offsets.reshape(broadcast_shape) + data *= gains.reshape(broadcast_shape) + if cals is not None: + data *= cals.reshape(broadcast_shape) + + +def _read_uniform_segment( + data, idx, start, stop, raw_extras, filenames, cals, mult +) -> bool: + """Read a uniformly sampled EDF/BDF segment using NumPy strides.""" + 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_is_slice = isinstance(idx, slice) + if idx_is_slice and idx.step not in (None, 1): + return False + idx_arr = np.arange(idx.start, idx.stop) if idx_is_slice else np.asarray(idx) + if len(idx_arr) == 0 or ( + not idx_is_slice and len(np.unique(idx_arr)) != len(idx_arr) + ): + return False + + stride_layout = raw_extras.get("stride_layout") + if stride_layout is None or mult is not None: + return False + n_samps = raw_extras["n_samps"] + buf_len = int(raw_extras["max_samp"]) + ch_offsets, n_per, sel_in_physical_order = stride_layout + + 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] + cal = cal[idx_arr, np.newaxis, np.newaxis] + offsets = offsets[idx_arr, np.newaxis, np.newaxis] + gains = gains[idx_arr, np.newaxis, np.newaxis] + stim_rows = ( + np.flatnonzero(np.isin(idx_arr, stim_channel_idxs)) + if len(stim_channel_idxs) + else () + ) + + block_start_idx, r_lims, d_lims = _blk_read_lims(start, stop, buf_len) + max_records = min(len(r_lims), n_per) + n_values = len(idx_arr) * max_records * buf_len + in_physical_order = ( + sel_in_physical_order + and idx_is_slice + and idx.start == 0 + and idx.stop == len(n_samps) + ) + direct_output = ( + in_physical_order + and not len(stim_rows) + and len(r_lims) > 1 + and n_per > 1 + and r_lims[0][0] == 0 + and r_lims[-1][1] == buf_len + ) + direct_output = direct_output and bool(np.all(cals == 1.0)) + estimated_incremental_bytes = 0 + if not direct_output: + estimated_incremental_bytes += n_values * np.dtype(np.float64).itemsize + if not in_physical_order: + decoded_itemsize = np.dtype(INT32).itemsize if subtype == "bdf" else dtype_byte + estimated_incremental_bytes += n_values * decoded_itemsize + if len(stim_rows): + 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 + 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 in_physical_order: + view = record_grid.transpose(1, 0, 2) + else: + view = record_grid[:, read_sel, :].transpose(1, 0, 2) + + r_sidx = r_lims[ai][0] + r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] + d_start = d_lims[ai][0] + d_stop = d_lims[ai + n_read - 1][1] + assert d_stop - d_start == r_eidx - r_sidx + if direct_output and n_read > 1: + assert r_sidx == 0 and r_eidx == n_read * buf_len + output = data[:, d_start:d_stop].reshape(view.shape) + assert np.shares_memory(output, data) + _calibrate_uniform_edf( + output, + view, + cal, + offsets, + gains, + None, + ) + continue + if ( + n_read == 1 + and not len(stim_rows) + and (r_sidx != 0 or r_eidx != buf_len) + ): + _calibrate_uniform_edf( + data[:, d_start:d_stop], + view[:, 0, r_sidx:r_eidx], + cal[:, 0, 0], + offsets[:, 0, 0], + gains[:, 0, 0], + cals, + ) + continue + + one = np.empty(view.shape, dtype=np.float64) + np.multiply(view, cal, out=one) + one += offsets + one *= gains + block = one.reshape(len(idx_arr), -1)[:, r_sidx:r_eidx] + for row in stim_rows: + stim = block[row].astype(int) + np.bitwise_and(stim, 2**17 - 1, out=stim) + block[row] = stim + assert d_stop - d_start == block.shape[1] + np.multiply( + block, + cals, + out=data[:, d_start:d_stop], + casting="unsafe", + ) + return True + + 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 + if _read_uniform_segment(data, idx, start, stop, raw_extras, filenames, cals, mult): + return [] n_samps = raw_extras["n_samps"] buf_len = int(raw_extras["max_samp"]) @@ -682,6 +853,8 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, if n_samps[ci] != buf_len: if orig_idx in stim_channel_idxs: + from scipy.interpolate import interp1d + # Stim channel will be interpolated old = np.linspace(0, 1, n_samps[ci] + 1, True) new = np.linspace(0, 1, buf_len, False) @@ -910,6 +1083,23 @@ def _get_info( edf_info["max_samp"] = max_samp = n_samps[picks].max() else: edf_info["max_samp"] = max_samp = n_samps.max() + all_n_samps = edf_info["n_samps"] + edf_info["stride_layout"] = None + if ( + edf_info["subtype"] in ("edf", "bdf") + and len(edf_info.get("tal_idx", ())) == 0 + and np.all(all_n_samps == max_samp) + ): + ch_offsets = np.cumsum(np.concatenate([[0], all_n_samps]), dtype=np.int64) + n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * edf_info["dtype_byte"]), 1) + sel_in_physical_order = np.array_equal( + edf_info["sel"], np.arange(len(all_n_samps)) + ) + edf_info["stride_layout"] = ( + ch_offsets, + n_per, + sel_in_physical_order, + ) # Info structure # ------------------------------------------------------------------------- diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 1672a131fc0..becca9a09d1 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -4,6 +4,8 @@ import datetime import gc +import sys +from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from functools import partial from io import BytesIO @@ -67,6 +69,234 @@ misc = ["EXG1", "EXG5", "EXG8", "M1", "M2"] +def _repeat_edf_records(source, destination, n_records=6): + """Repeat a one-record EDF 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 _disable_uniform_stride(*args, **kwargs): + return False + + +def _get_data_window(raw, limits): + return raw.get_data(start=limits[0], stop=limits[1]) + + +class _StrideRecorder: + def __init__(self, helper): + self.helper = helper + self.results = [] + + def __call__(self, *args, **kwargs): + result = self.helper(*args, **kwargs) + self.results.append(result) + return result + + +def _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop): + helper = edf.edf._read_uniform_segment + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) + want = raw.get_data(picks=picks, start=start, stop=stop) + recorder = _StrideRecorder(helper) + monkeypatch.setattr(edf.edf, "_read_uniform_segment", recorder) + got = raw.get_data(picks=picks, start=start, stop=stop) + assert recorder.results == [True] + assert_array_equal(got, want) + return 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 uniform EDF stride decoding across picks and boundaries.""" + 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) + buffer_length = 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, buffer_length - 5), + "boundary": (buffer_length - 7, buffer_length + 11), + "multiple": (buffer_length // 2, 4 * buffer_length + 13), + }[window_kind] + _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop) + + +def test_uniform_stride_calibration_and_stim(monkeypatch): + """Test exact Raw calibration and stim masking on the stride path.""" + raw = read_raw_edf( + edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" + ) + picks = np.array([0, 12, len(raw.ch_names) - 1]) + raw._cals[picks] *= np.array([0.5, 2.0, 4.0]) + want = _assert_uniform_stride_matches(monkeypatch, raw, picks, 100, 900) + stim = want[-1] / raw._cals[picks[-1]] + assert_array_equal(stim, np.bitwise_and(stim.astype(int), 2**17 - 1)) + + +@pytest.mark.parametrize("pick_kind", ("all", "subset", "reversed")) +def test_uniform_stride_direct_calibration(pick_kind, monkeypatch, tmp_path): + """Test exact direct calibration for a single-record window.""" + repeated = tmp_path / "uniform.edf" + _repeat_edf_records(edf_stim_channel_path, repeated) + raw = read_raw_edf(repeated, stim_channel=None, preload=False, verbose="error") + n_channels = len(raw.ch_names) + picks = { + "all": np.arange(n_channels), + "subset": np.array([0, n_channels // 2, n_channels - 1]), + "reversed": np.arange(n_channels - 1, -1, -1), + }[pick_kind] + raw._cals[picks] *= np.linspace(0.5, 2.0, len(picks)) + buffer_length = int(raw._raw_extras[0]["max_samp"]) + _assert_uniform_stride_matches(monkeypatch, raw, picks, 7, buffer_length - 5) + + +@pytest.mark.parametrize("calibration", (1.0, 2.0)) +def test_uniform_stride_full_single_record_uses_buffer(monkeypatch, calibration): + """Test full single records retain buffered calibration.""" + raw = read_raw_edf( + edf_stim_channel_path, stim_channel=None, preload=False, verbose="error" + ) + raw._cals *= calibration + helper = edf.edf._read_uniform_segment + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) + want = raw.get_data() + calibration = _StrideRecorder(edf.edf._calibrate_uniform_edf) + monkeypatch.setattr(edf.edf, "_read_uniform_segment", helper) + monkeypatch.setattr(edf.edf, "_calibrate_uniform_edf", calibration) + got = raw.get_data() + assert calibration.results == [] + assert_array_equal(got.view(np.uint64), want.view(np.uint64)) + + +@pytest.mark.parametrize( + "reader, source, extension", + ((read_raw_edf, edf_stim_channel_path, "edf"), (read_raw_bdf, bdf_path, "bdf")), + ids=("edf", "bdf"), +) +def test_uniform_stride_aligned_memmap_direct_output( + monkeypatch, tmp_path, reader, source, extension +): + """Test an aligned full-record preload needs no float work buffer.""" + repeated = tmp_path / f"uniform.{extension}" + _repeat_edf_records(source, repeated) + raw = reader(repeated, stim_channel=None, preload=False, verbose="error") + helper = edf.edf._read_uniform_segment + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) + want = raw.get_data() + + recorder = _StrideRecorder(helper) + monkeypatch.setattr(edf.edf, "_read_uniform_segment", recorder) + monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_EXTRA_BYTES", 0) + mmap_path = tmp_path / "uniform.dat" + got = reader( + repeated, + stim_channel=None, + preload=mmap_path, + verbose="error", + ) + assert recorder.results == [True] + assert isinstance(got._data, np.memmap) + assert Path(got._data.filename) == mmap_path + assert_array_equal(got._data.view(np.uint64), want.view(np.uint64)) + got._data._mmap.close() + got._data = None + + +def test_uniform_calibration_identity_cals_bit_exact(): + """Test omitting identity output calibrations preserves every bit.""" + source = np.array( + [[np.iinfo(np.int16).min, -1, 0], [1, 17, np.iinfo(np.int16).max]], + dtype=np.int16, + ) + cal = np.array([0.125, 0.3]) + offsets = np.array([-0.0, 1.1]) + gains = np.array([1e-6, -2.5]) + want = np.empty(source.shape, dtype=np.float64) + np.multiply(source, cal[:, np.newaxis], out=want) + want += offsets[:, np.newaxis] + want *= gains[:, np.newaxis] + want *= np.ones((len(source), 1)) + got = np.empty_like(want) + edf.edf._calibrate_uniform_edf(got, source, cal, offsets, gains, None) + assert_array_equal(got.view(np.uint64), want.view(np.uint64)) + + +@pytest.mark.parametrize("pick_kind", ("all", "subset", "reversed", "permuted")) +@pytest.mark.parametrize("window_kind", ("within", "boundary", "multiple")) +def test_uniform_stride_bdf(pick_kind, window_kind, monkeypatch, tmp_path): + """Test exact BDF stride decoding across picks and record boundaries.""" + repeated = tmp_path / "uniform.bdf" + _repeat_edf_records(bdf_path, repeated) + raw = read_raw_bdf(repeated, preload=False, verbose="error") + n_channels = len(raw.ch_names) + buffer_length = 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, buffer_length - 5), + "boundary": (buffer_length - 7, buffer_length + 11), + "multiple": (buffer_length // 2, 4 * buffer_length + 13), + }[window_kind] + _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop) + + +@pytest.mark.parametrize("fallback_kind", ("mixed_rate", "tal", "projection", "memory")) +def test_uniform_stride_edf_falls_back(fallback_kind, monkeypatch): + """Test EDF layouts requiring special handling retain legacy behavior.""" + path = { + "mixed_rate": edf_uneven_path, + "tal": edf_path, + "projection": edf_stim_channel_path, + "memory": edf_stim_channel_path, + }[fallback_kind] + raw = read_raw_edf(path, preload=False, verbose="error") + if fallback_kind == "projection": + raw.set_eeg_reference(projection=True, verbose="error").apply_proj( + verbose="error" + ) + elif fallback_kind == "memory": + monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_EXTRA_BYTES", 0) + picks = np.array([0]) + helper = edf.edf._read_uniform_segment + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) + want = raw.get_data(picks=picks, start=100, stop=900) + recorder = _StrideRecorder(helper) + monkeypatch.setattr(edf.edf, "_read_uniform_segment", recorder) + got = raw.get_data(picks=picks, start=100, stop=900) + assert recorder.results == [False] + assert_array_equal(got, want) + + +def test_uniform_stride_concurrent(monkeypatch, tmp_path): + """Test concurrent stride reads do not share seekable file state.""" + repeated = tmp_path / "concurrent.edf" + _repeat_edf_records(edf_stim_channel_path, repeated) + raw = read_raw_edf(repeated, stim_channel=-1, preload=False, verbose="error") + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) + reference = raw.get_data() + monkeypatch.undo() + windows = [(start, start + 64) for start in range(0, raw.n_times - 64, 31)] + with ThreadPoolExecutor(max_workers=8) as pool: + got = list(pool.map(partial(_get_data_window, raw), windows * 4)) + 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) @@ -91,6 +321,15 @@ def test_orig_units(): assert set(raw_back._orig_units) == set(raw.ch_names) +def test_uniform_edf_does_not_import_interpolation(monkeypatch): + """Test uniform EDF decoding does not load interpolation support.""" + monkeypatch.setitem(sys.modules, "scipy.interpolate", None) + raw = read_raw_edf( + edf_stim_channel_path, stim_channel=-1, preload=True, verbose="error" + ) + assert raw.preload + + def test_units_params(): """Test enforcing original channel units.""" with pytest.raises( @@ -349,6 +588,55 @@ def test_duplicate_channel_labels_edf(): assert raw.ch_names == EXPECTED_CHANNEL_NAMES +@pytest.mark.parametrize( + "values", + ( + np.array([], dtype=np.int32), + np.array([-(1 << 23), -1, 0, 1, (1 << 23) - 1], dtype=np.int32), + np.array([-(1 << 23)], dtype=np.int32), + np.array([-1, 1], dtype=np.int32), + np.random.default_rng(42).integers( + -(1 << 23), 1 << 23, size=1000, dtype=np.int32 + ), + ), +) +@pytest.mark.parametrize("file_kind", ("buffer", "disk")) +def test_read_ch_bdf_int24(tmp_path, values, file_kind): + """Test exact signed 24-bit BDF decoding, including buffer boundaries.""" + unsigned = values.astype(np.int64) & ((1 << 24) - 1) + packed = np.column_stack((unsigned, unsigned >> 8, unsigned >> 16)).astype(np.uint8) + if file_kind == "buffer": + fid = BytesIO(packed.tobytes()) + else: + fname = tmp_path / "samples.bdf" + fname.write_bytes(packed.tobytes()) + fid = fname.open("rb") + with fid: + got = _read_ch( + fid, + subtype="bdf", + samp=len(values), + dtype_byte=3, + dtype=np.uint8, + ) + assert_array_equal(got, values) + + +@pytest.mark.parametrize("missing", (1, 2)) +@pytest.mark.parametrize("file_kind", ("buffer", "disk")) +def test_read_ch_bdf_short_read(tmp_path, missing, file_kind): + """Test truncated signed 24-bit BDF data raises instead of mixing bytes.""" + data = b"\x00" * (6 - missing) + if file_kind == "buffer": + fid = BytesIO(data) + else: + fname = tmp_path / "truncated.bdf" + fname.write_bytes(data) + fid = fname.open("rb") + with fid, pytest.raises(RuntimeError, match="requested BDF bytes"): + _read_ch(fid, subtype="bdf", samp=2, dtype_byte=3, dtype=np.uint8) + + def test_parse_annotation(tmp_path): """Test parsing the tal channel.""" # test the parser From fcb8b9977cee7abae8d88012d84bad41411ec49d Mon Sep 17 00:00:00 2001 From: Bru Date: Wed, 26 Aug 2026 17:53:33 +0200 Subject: [PATCH 03/12] Use cache-sized blocks for BrainVision reads --- mne/_fiff/tests/test_utils.py | 47 ++++++++++++++++++++++++++++++- mne/_fiff/utils.py | 8 ++++-- mne/io/brainvision/brainvision.py | 1 + 3 files changed, 53 insertions(+), 3 deletions(-) diff --git a/mne/_fiff/tests/test_utils.py b/mne/_fiff/tests/test_utils.py index ba6748826a3..8627be8efd6 100644 --- a/mne/_fiff/tests/test_utils.py +++ b/mne/_fiff/tests/test_utils.py @@ -4,7 +4,13 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. -from mne._fiff.utils import _check_orig_units +from types import SimpleNamespace + +import numpy as np +import pytest +from numpy.testing import assert_allclose, assert_array_equal + +from mne._fiff.utils import _check_orig_units, _read_segments_file def test_check_orig_units(): @@ -16,3 +22,42 @@ def test_check_orig_units(): assert orig_units["Pz"] == "µV" assert orig_units["greekMu"] == "µV" assert orig_units["microSign"] == "µV" + + +@pytest.mark.parametrize("use_mult", (False, True)) +def test_read_segments_file_max_block_bytes(tmp_path, use_mult): + """Test reading in configurable complete channel frames.""" + source = np.arange(20, dtype=" Date: Wed, 26 Aug 2026 17:53:38 +0200 Subject: [PATCH 04/12] Reuse parsed EEGLAB metadata for annotations --- mne/io/eeglab/eeglab.py | 4 +-- mne/io/eeglab/tests/test_eeglab.py | 53 ++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/mne/io/eeglab/eeglab.py b/mne/io/eeglab/eeglab.py index a16f1ae9779..ca97e6058cf 100644 --- a/mne/io/eeglab/eeglab.py +++ b/mne/io/eeglab/eeglab.py @@ -16,7 +16,7 @@ from ..._fiff.meas_info import create_info from ..._fiff.pick import _PICK_TYPES_KEYS from ..._fiff.utils import _find_channels, _mult_cal_one, _read_segments_file -from ...annotations import Annotations, read_annotations +from ...annotations import Annotations from ...channels import make_dig_montage from ...defaults import DEFAULTS from ...epochs import BaseEpochs @@ -496,7 +496,7 @@ def __init__( ) # create event_ch from annotations - annot = read_annotations(input_fname, uint16_codec=uint16_codec) + annot = _read_annotations_eeglab(eeg) self.set_annotations(annot) _check_boundary(annot, None) diff --git a/mne/io/eeglab/tests/test_eeglab.py b/mne/io/eeglab/tests/test_eeglab.py index f8c1aa79bfc..2dfa1f30b89 100644 --- a/mne/io/eeglab/tests/test_eeglab.py +++ b/mne/io/eeglab/tests/test_eeglab.py @@ -5,6 +5,7 @@ import os import shutil from copy import deepcopy +from unittest.mock import Mock import numpy as np import pytest @@ -529,6 +530,58 @@ def test_eeglab_read_annotations(): assert_allclose(raw.annotations.duration, np.ones(3) * 0.5) +@pytest.mark.parametrize( + "data_kind, preload_kind", + ( + ("embedded", "ram"), + ("external", "lazy"), + ("external", "ram"), + ("external", "mmap"), + ), +) +def test_raw_eeglab_reuses_metadata_for_annotations( + tmp_path, monkeypatch, data_kind, preload_kind +): + """Test that raw annotations reuse the already-loaded EEG structure.""" + fname = tmp_path / "single-parse.set" + events = np.array( + [("first", 2.0, 2.0), ("second", 4.0, 4.0)], + dtype=[("type", "O"), ("latency", "f8"), ("duration", "f8")], + ) + data = np.arange(10.0, dtype=" Date: Wed, 26 Aug 2026 23:39:11 +0200 Subject: [PATCH 05/12] Pipeline large EDF, BDF, and BrainVision preloads --- mne/_fiff/tests/test_utils.py | 110 ++++++++++++++++++++++++++++++ mne/_fiff/utils.py | 98 +++++++++++++++++++++----- mne/io/brainvision/brainvision.py | 7 +- mne/io/edf/edf.py | 48 ++++++++++--- mne/io/edf/tests/test_edf.py | 11 +++ 5 files changed, 248 insertions(+), 26 deletions(-) diff --git a/mne/_fiff/tests/test_utils.py b/mne/_fiff/tests/test_utils.py index 8627be8efd6..de061dfce3a 100644 --- a/mne/_fiff/tests/test_utils.py +++ b/mne/_fiff/tests/test_utils.py @@ -4,12 +4,14 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. +import threading from types import SimpleNamespace import numpy as np import pytest from numpy.testing import assert_allclose, assert_array_equal +from mne._fiff import utils as fiff_utils from mne._fiff.utils import _check_orig_units, _read_segments_file @@ -61,3 +63,111 @@ def test_read_segments_file_max_block_bytes(tmp_path, use_mult): assert_allclose(data, want, rtol=1e-15) else: assert_array_equal(data, want) + + +def test_read_segments_file_mmap_threaded(tmp_path, monkeypatch): + """Test mapped blocks are calibrated on worker threads.""" + source = np.arange(40, dtype="= max_block_bytes: + try: + mapped = mmap.mmap(fid.fileno(), 0, access=mmap.ACCESS_READ) + except (OSError, ValueError): + pass fid.seek(data_offset) # extract data in chunks - for sample_start in np.arange(0, data_left, block_size) // n_channels: - count = min(block_size, data_left - sample_start * n_channels) - block = np.fromfile(fid, dtype, 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] # = count // n_channels - 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) + executor = None + pending = [] + worker_error = None + if ( + mapped is not None + and n_jobs > 1 + and data.nbytes >= _READ_SEGMENTS_FILE_THREAD_MIN_BYTES + ): + executor = ThreadPoolExecutor(max_workers=n_jobs) + try: + for sample_start in np.arange(0, data_left, block_size) // n_channels: + count = min(block_size, data_left - sample_start * n_channels) + raw_block = None + block = None + try: + if mapped is None: + raw_block = np.fromfile(fid, dtype, count) + else: + byte_offset = data_offset + sample_start * n_channels * n_bytes + if len(mapped) - byte_offset < count * n_bytes: + raw_block = np.empty(0, dtype=dtype) + else: + raw_block = np.frombuffer( + mapped, dtype, count=count, offset=byte_offset + ) + if raw_block.size != count: + raise RuntimeError( + f"Incorrect number of samples ({raw_block.size} != " + f"{count}), please report this error to MNE-Python " + "developers" + ) + block = raw_block.reshape(n_channels, -1, order="F") + n_samples = block.shape[1] # = count // n_channels + 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] + if executor is None: + _mult_cal_one(data_view, block, idx, cals, mult) + else: + future = executor.submit( + _mult_cal_one, data_view, block, idx, cals, mult + ) + pending.append((future, block, raw_block)) + finally: + del block, raw_block + for pending_item in pending: + error = pending_item[0].exception() + if worker_error is None and error is not None: + worker_error = error + if pending: + del pending_item + finally: + if executor is not None: + executor.shutdown(wait=True) + # Worker tracebacks retain their mapped NumPy inputs. Clear frame locals + # before closing the mapping, but keep the traceback locations intact. + for pending_item in pending: + error = pending_item[0].exception() + if error is not None and error.__traceback__ is not None: + clear_frames(error.__traceback__) + if pending: + del pending_item + pending.clear() + if mapped is not None: + mapped.close() + if worker_error is not None: + raise worker_error def read_str(fid, count=1): diff --git a/mne/io/brainvision/brainvision.py b/mne/io/brainvision/brainvision.py index 9ca420b680c..ac4f2ca583a 100644 --- a/mne/io/brainvision/brainvision.py +++ b/mne/io/brainvision/brainvision.py @@ -35,6 +35,9 @@ ) from ..base import BaseRaw +_BRAINVISION_BLOCK_BYTES = 8 * 1024**2 +_BRAINVISION_READ_WORKERS = 4 + @fill_doc class RawBrainVision(BaseRaw): @@ -191,7 +194,9 @@ def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): mult, dtype=dtype, n_channels=n_data_ch, - max_block_bytes=8 * 1024**2, + max_block_bytes=_BRAINVISION_BLOCK_BYTES, + use_mmap=True, + n_jobs=_BRAINVISION_READ_WORKERS, ) else: offsets = self._raw_extras[fi]["offsets"] diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 8e41d99c60d..fe2b7c518bf 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -6,6 +6,8 @@ import os import re +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext from datetime import UTC, date, datetime, timedelta from enum import Enum from pathlib import Path @@ -626,6 +628,8 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): _EDF_STRIDE_MAX_EXTRA_BYTES = 64 * 1024**2 +_EDF_FROMFILE_THREAD_MIN_BYTES = 64 * 1024**2 +_EDF_FROMFILE_WORKERS = 2 def _calibrate_uniform_edf(data, source, cal, offsets, gains, cals): @@ -711,7 +715,18 @@ def _read_uniform_segment( if estimated_incremental_bytes > _EDF_STRIDE_MAX_EXTRA_BYTES: return False - with _gdf_edf_get_fid(filenames, buffering=0) as fid: + threaded = ( + subtype in ("edf", "bdf") + and direct_output + and data.nbytes >= _EDF_FROMFILE_THREAD_MIN_BYTES + ) + worker_context = ( + ThreadPoolExecutor(max_workers=_EDF_FROMFILE_WORKERS) + if threaded + else nullcontext() + ) + with worker_context as executor, _gdf_edf_get_fid(filenames, buffering=0) as fid: + pending = [] start_offset = data_offset + block_start_idx * ch_offsets[-1] * dtype_byte for ai in range(0, len(r_lims), n_per): block_offset = ai * ch_offsets[-1] * dtype_byte @@ -735,14 +750,27 @@ def _read_uniform_segment( assert r_sidx == 0 and r_eidx == n_read * buf_len output = data[:, d_start:d_stop].reshape(view.shape) assert np.shares_memory(output, data) - _calibrate_uniform_edf( - output, - view, - cal, - offsets, - gains, - None, - ) + if executor is None: + _calibrate_uniform_edf( + output, + view, + cal, + offsets, + gains, + None, + ) + else: + pending.append( + executor.submit( + _calibrate_uniform_edf, + output, + view, + cal, + offsets, + gains, + None, + ) + ) continue if ( n_read == 1 @@ -775,6 +803,8 @@ def _read_uniform_segment( out=data[:, d_start:d_stop], casting="unsafe", ) + for future in pending: + future.result() return True diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index becca9a09d1..78206e1730b 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -5,6 +5,7 @@ import datetime import gc import sys +import threading from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from functools import partial @@ -198,6 +199,15 @@ def test_uniform_stride_aligned_memmap_direct_output( recorder = _StrideRecorder(helper) monkeypatch.setattr(edf.edf, "_read_uniform_segment", recorder) monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_EXTRA_BYTES", 0) + thread_ids = set() + calibrate = edf.edf._calibrate_uniform_edf + + def _record_thread(*args): + thread_ids.add(threading.get_ident()) + return calibrate(*args) + + monkeypatch.setattr(edf.edf, "_EDF_FROMFILE_THREAD_MIN_BYTES", 0) + monkeypatch.setattr(edf.edf, "_calibrate_uniform_edf", _record_thread) mmap_path = tmp_path / "uniform.dat" got = reader( repeated, @@ -208,6 +218,7 @@ def test_uniform_stride_aligned_memmap_direct_output( assert recorder.results == [True] assert isinstance(got._data, np.memmap) assert Path(got._data.filename) == mmap_path + assert threading.get_ident() not in thread_ids assert_array_equal(got._data.view(np.uint64), want.view(np.uint64)) got._data._mmap.close() got._data = None From e9770d654f5f1dc9b85c77de80b0e5a22aa11d6c Mon Sep 17 00:00:00 2001 From: Bru Date: Fri, 28 Aug 2026 13:33:00 +0200 Subject: [PATCH 06/12] Add changelog entry for the cold-path IO speedups --- doc/changes/dev/14237.newfeature.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 doc/changes/dev/14237.newfeature.rst diff --git a/doc/changes/dev/14237.newfeature.rst b/doc/changes/dev/14237.newfeature.rst new file mode 100644 index 00000000000..4de36460eb5 --- /dev/null +++ b/doc/changes/dev/14237.newfeature.rst @@ -0,0 +1 @@ +Speed up preloading Raw data and reading uniform EDF and BDF data by avoiding unnecessary setup and using pure-NumPy record-stride decoders, by `Bruno Aristimunha`_. From ea3b9ae54b0877d314260b9d3ecad26279aa9002 Mon Sep 17 00:00:00 2001 From: Bru Date: Fri, 28 Aug 2026 14:05:05 +0200 Subject: [PATCH 07/12] Let the stride path handle EDF+ annotation channels The stride decoder reshaped a whole data record into an (n_channels, max_samp) matrix, which requires every channel to share a sampling rate. An EDF+ annotation channel usually does not: for a 64-channel 256 Hz file with a 60-sample TAL channel a record holds 16444 values while the reshape needs 16640, so 'stride_layout' was None and every EDF+ read fell back to the per-channel loop. Gate on the *selected* channels instead, and when a record is not rectangular gather each picked channel by its own offset. Files whose records are rectangular keep the existing reshape unchanged. Fast-path coverage across the EDF/BDF corpus goes from 6/22 files to 18/22. EDF+ windowed reads 1.76x, EDF+ preload 1.28x, BDF+ windowed reads 1.42x; files without an annotation channel are unaffected. Output stays bit-identical to main: 73/73 array snapshots and 52/52 annotation snapshots across 26 files. --- mne/io/edf/edf.py | 35 ++++++++++++++++++++++------------- mne/io/edf/tests/test_edf.py | 13 +++++++++++-- 2 files changed, 33 insertions(+), 15 deletions(-) diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index fe2b7c518bf..4d35fc2b00f 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -649,9 +649,6 @@ def _read_uniform_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_is_slice = isinstance(idx, slice) if idx_is_slice and idx.step not in (None, 1): return False @@ -666,7 +663,7 @@ def _read_uniform_segment( return False n_samps = raw_extras["n_samps"] buf_len = int(raw_extras["max_samp"]) - ch_offsets, n_per, sel_in_physical_order = stride_layout + ch_offsets, n_per, sel_in_physical_order, rectangular = stride_layout dtype = raw_extras["dtype_np"] dtype_byte = raw_extras["dtype_byte"] @@ -735,11 +732,19 @@ def _read_uniform_segment( 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 in_physical_order: - view = record_grid.transpose(1, 0, 2) + if rectangular: + record_grid = many_chunk.reshape(n_read, len(n_samps), buf_len) + if in_physical_order: + view = record_grid.transpose(1, 0, 2) + else: + view = record_grid[:, read_sel, :].transpose(1, 0, 2) else: - view = record_grid[:, read_sel, :].transpose(1, 0, 2) + # a record is not a matrix (e.g. an EDF+ annotation channel is + # stored at its own rate), so slice out each picked channel + records = many_chunk.reshape(n_read, -1) + view = np.empty((len(read_sel), n_read, buf_len), many_chunk.dtype) + for j, ci in enumerate(read_sel): + view[j] = records[:, ch_offsets[ci] : ch_offsets[ci + 1]] r_sidx = r_lims[ai][0] r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] @@ -1115,20 +1120,24 @@ def _get_info( edf_info["max_samp"] = max_samp = n_samps.max() all_n_samps = edf_info["n_samps"] edf_info["stride_layout"] = None - if ( - edf_info["subtype"] in ("edf", "bdf") - and len(edf_info.get("tal_idx", ())) == 0 - and np.all(all_n_samps == max_samp) + # Only the *selected* channels have to share a sampling rate. An EDF+ + # annotation channel usually does not, which is why `rectangular` is checked + # separately: it says whether a data record is a plain (n_channels, max_samp) + # matrix that can simply be reshaped, or has to be gathered channel by channel. + if edf_info["subtype"] in ("edf", "bdf") and np.all( + all_n_samps[edf_info["sel"]] == max_samp ): ch_offsets = np.cumsum(np.concatenate([[0], all_n_samps]), dtype=np.int64) n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * edf_info["dtype_byte"]), 1) - sel_in_physical_order = np.array_equal( + rectangular = bool(np.all(all_n_samps == max_samp)) + sel_in_physical_order = rectangular and np.array_equal( edf_info["sel"], np.arange(len(all_n_samps)) ) edf_info["stride_layout"] = ( ch_offsets, n_per, sel_in_physical_order, + rectangular, ) # Info structure diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 78206e1730b..500cfcc575f 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -266,12 +266,11 @@ def test_uniform_stride_bdf(pick_kind, window_kind, monkeypatch, tmp_path): _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop) -@pytest.mark.parametrize("fallback_kind", ("mixed_rate", "tal", "projection", "memory")) +@pytest.mark.parametrize("fallback_kind", ("mixed_rate", "projection", "memory")) def test_uniform_stride_edf_falls_back(fallback_kind, monkeypatch): """Test EDF layouts requiring special handling retain legacy behavior.""" path = { "mixed_rate": edf_uneven_path, - "tal": edf_path, "projection": edf_stim_channel_path, "memory": edf_stim_channel_path, }[fallback_kind] @@ -293,6 +292,16 @@ def test_uniform_stride_edf_falls_back(fallback_kind, monkeypatch): assert_array_equal(got, want) +def test_uniform_stride_edf_annotations(monkeypatch): + """Test an EDF+ annotation channel does not disable the stride path.""" + raw = read_raw_edf(edf_path, preload=False, verbose="error") + buffer_length = int(raw._raw_extras[0]["max_samp"]) + picks = np.array([0, 2]) + _assert_uniform_stride_matches( + monkeypatch, raw, picks, buffer_length // 3, 3 * buffer_length + 7 + ) + + def test_uniform_stride_concurrent(monkeypatch, tmp_path): """Test concurrent stride reads do not share seekable file state.""" repeated = tmp_path / "concurrent.edf" From 049ff5b58793a490582bef0addc132b3631f8387 Mon Sep 17 00:00:00 2001 From: Bru Date: Fri, 28 Aug 2026 14:52:32 +0200 Subject: [PATCH 08/12] Trim the PR: drop the byte budget, threading and test bloat Four review passes (reuse, simplification, efficiency, altitude) against the shape of #14216. Measured rather than assumed: - byte-budget heuristic: never fires. n_per already caps a chunk near 10 MiB of source bytes, so temporaries stay under 50 MB of the 64 MB cap even for an adversarial 2-channel/6-hour file with reversed picks. - n_read == 1 branch: 0/200 hits on the windowed benchmark, and the general branch produces the same values. - threading: real (13-22% on preload) but reachable only via direct_output, and it is one of only two ThreadPoolExecutor sites in MNE. mne.parallel already has parallel_func(prefer='threads'). Deferred to its own PR with a benchmark. - direct_output: kept, it measures 27-28% on full preload. _read_segments_file mmap/threading and the BrainVision block sizing are reverted to main and move to a follow-up: no MNE fixture is large enough to reach the 64 MB threading threshold (largest .eeg is 3.7 MB), so it was untested at any realistic scale. Tests 549 -> 156 lines: four near-identical stride tests merged into one parametrized test, redundant file_kind axes dropped, and a test asserting x * 1.0 == x bit-exactly removed. +884/-29 -> +330/-18. Output stays bit-identical to main: 77/77 array snapshots and 56/56 annotation snapshots across 28 files. --- mne/_fiff/tests/test_utils.py | 157 +-------------- mne/_fiff/utils.py | 106 ++-------- mne/io/brainvision/brainvision.py | 6 - mne/io/edf/edf.py | 133 +++---------- mne/io/edf/tests/test_edf.py | 301 +++++++---------------------- mne/io/eeglab/tests/test_eeglab.py | 55 +----- mne/io/tests/test_raw.py | 25 --- 7 files changed, 120 insertions(+), 663 deletions(-) diff --git a/mne/_fiff/tests/test_utils.py b/mne/_fiff/tests/test_utils.py index de061dfce3a..ba6748826a3 100644 --- a/mne/_fiff/tests/test_utils.py +++ b/mne/_fiff/tests/test_utils.py @@ -4,15 +4,7 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. -import threading -from types import SimpleNamespace - -import numpy as np -import pytest -from numpy.testing import assert_allclose, assert_array_equal - -from mne._fiff import utils as fiff_utils -from mne._fiff.utils import _check_orig_units, _read_segments_file +from mne._fiff.utils import _check_orig_units def test_check_orig_units(): @@ -24,150 +16,3 @@ def test_check_orig_units(): assert orig_units["Pz"] == "µV" assert orig_units["greekMu"] == "µV" assert orig_units["microSign"] == "µV" - - -@pytest.mark.parametrize("use_mult", (False, True)) -def test_read_segments_file_max_block_bytes(tmp_path, use_mult): - """Test reading in configurable complete channel frames.""" - source = np.arange(20, dtype="= max_block_bytes: - try: - mapped = mmap.mmap(fid.fileno(), 0, access=mmap.ACCESS_READ) - except (OSError, ValueError): - pass fid.seek(data_offset) # extract data in chunks - executor = None - pending = [] - worker_error = None - if ( - mapped is not None - and n_jobs > 1 - and data.nbytes >= _READ_SEGMENTS_FILE_THREAD_MIN_BYTES - ): - executor = ThreadPoolExecutor(max_workers=n_jobs) - try: - for sample_start in np.arange(0, data_left, block_size) // n_channels: - count = min(block_size, data_left - sample_start * n_channels) - raw_block = None - block = None - try: - if mapped is None: - raw_block = np.fromfile(fid, dtype, count) - else: - byte_offset = data_offset + sample_start * n_channels * n_bytes - if len(mapped) - byte_offset < count * n_bytes: - raw_block = np.empty(0, dtype=dtype) - else: - raw_block = np.frombuffer( - mapped, dtype, count=count, offset=byte_offset - ) - if raw_block.size != count: - raise RuntimeError( - f"Incorrect number of samples ({raw_block.size} != " - f"{count}), please report this error to MNE-Python " - "developers" - ) - block = raw_block.reshape(n_channels, -1, order="F") - n_samples = block.shape[1] # = count // n_channels - 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] - if executor is None: - _mult_cal_one(data_view, block, idx, cals, mult) - else: - future = executor.submit( - _mult_cal_one, data_view, block, idx, cals, mult - ) - pending.append((future, block, raw_block)) - finally: - del block, raw_block - for pending_item in pending: - error = pending_item[0].exception() - if worker_error is None and error is not None: - worker_error = error - if pending: - del pending_item - finally: - if executor is not None: - executor.shutdown(wait=True) - # Worker tracebacks retain their mapped NumPy inputs. Clear frame locals - # before closing the mapping, but keep the traceback locations intact. - for pending_item in pending: - error = pending_item[0].exception() - if error is not None and error.__traceback__ is not None: - clear_frames(error.__traceback__) - if pending: - del pending_item - pending.clear() - if mapped is not None: - mapped.close() - if worker_error is not None: - raise worker_error + for sample_start in np.arange(0, data_left, block_size) // n_channels: + count = min(block_size, data_left - sample_start * n_channels) + block = np.fromfile(fid, dtype, 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] # = count // n_channels + 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) def read_str(fid, count=1): diff --git a/mne/io/brainvision/brainvision.py b/mne/io/brainvision/brainvision.py index ac4f2ca583a..9c9d978a861 100644 --- a/mne/io/brainvision/brainvision.py +++ b/mne/io/brainvision/brainvision.py @@ -35,9 +35,6 @@ ) from ..base import BaseRaw -_BRAINVISION_BLOCK_BYTES = 8 * 1024**2 -_BRAINVISION_READ_WORKERS = 4 - @fill_doc class RawBrainVision(BaseRaw): @@ -194,9 +191,6 @@ def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): mult, dtype=dtype, n_channels=n_data_ch, - max_block_bytes=_BRAINVISION_BLOCK_BYTES, - use_mmap=True, - n_jobs=_BRAINVISION_READ_WORKERS, ) else: offsets = self._raw_extras[fi]["offsets"] diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 4d35fc2b00f..02914b30df0 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -6,8 +6,6 @@ import os import re -from concurrent.futures import ThreadPoolExecutor -from contextlib import nullcontext from datetime import UTC, date, datetime, timedelta from enum import Enum from pathlib import Path @@ -627,19 +625,12 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): return ch_data -_EDF_STRIDE_MAX_EXTRA_BYTES = 64 * 1024**2 -_EDF_FROMFILE_THREAD_MIN_BYTES = 64 * 1024**2 -_EDF_FROMFILE_WORKERS = 2 - - def _calibrate_uniform_edf(data, source, cal, offsets, gains, cals): """Calibrate uniform EDF/BDF samples directly into their output buffer.""" - broadcast_shape = (len(data),) + (1,) * (data.ndim - 1) - np.multiply(source, cal.reshape(broadcast_shape), out=data, casting="unsafe") - data += offsets.reshape(broadcast_shape) - data *= gains.reshape(broadcast_shape) - if cals is not None: - data *= cals.reshape(broadcast_shape) + np.multiply(source, cal, out=data, casting="unsafe") + data += offsets + data *= gains + data *= cals def _read_uniform_segment( @@ -663,7 +654,7 @@ def _read_uniform_segment( return False n_samps = raw_extras["n_samps"] buf_len = int(raw_extras["max_samp"]) - ch_offsets, n_per, sel_in_physical_order, rectangular = stride_layout + ch_offsets, n_per, sel_in_physical_order = stride_layout dtype = raw_extras["dtype_np"] dtype_byte = raw_extras["dtype_byte"] @@ -684,8 +675,6 @@ def _read_uniform_segment( ) block_start_idx, r_lims, d_lims = _blk_read_lims(start, stop, buf_len) - max_records = min(len(r_lims), n_per) - n_values = len(idx_arr) * max_records * buf_len in_physical_order = ( sel_in_physical_order and idx_is_slice @@ -700,30 +689,7 @@ def _read_uniform_segment( and r_lims[0][0] == 0 and r_lims[-1][1] == buf_len ) - direct_output = direct_output and bool(np.all(cals == 1.0)) - estimated_incremental_bytes = 0 - if not direct_output: - estimated_incremental_bytes += n_values * np.dtype(np.float64).itemsize - if not in_physical_order: - decoded_itemsize = np.dtype(INT32).itemsize if subtype == "bdf" else dtype_byte - estimated_incremental_bytes += n_values * decoded_itemsize - if len(stim_rows): - estimated_incremental_bytes += max_records * buf_len * np.dtype(int).itemsize - if estimated_incremental_bytes > _EDF_STRIDE_MAX_EXTRA_BYTES: - return False - - threaded = ( - subtype in ("edf", "bdf") - and direct_output - and data.nbytes >= _EDF_FROMFILE_THREAD_MIN_BYTES - ) - worker_context = ( - ThreadPoolExecutor(max_workers=_EDF_FROMFILE_WORKERS) - if threaded - else nullcontext() - ) - with worker_context as executor, _gdf_edf_get_fid(filenames, buffering=0) as fid: - pending = [] + with _gdf_edf_get_fid(filenames, buffering=0) as fid: start_offset = data_offset + block_start_idx * ch_offsets[-1] * dtype_byte for ai in range(0, len(r_lims), n_per): block_offset = ai * ch_offsets[-1] * dtype_byte @@ -732,15 +698,13 @@ def _read_uniform_segment( many_chunk = _read_ch( fid, subtype, ch_offsets[-1] * n_read, dtype_byte, dtype ) - if rectangular: - record_grid = many_chunk.reshape(n_read, len(n_samps), buf_len) - if in_physical_order: - view = record_grid.transpose(1, 0, 2) - else: - view = record_grid[:, read_sel, :].transpose(1, 0, 2) + if in_physical_order: + view = many_chunk.reshape(n_read, len(n_samps), buf_len).transpose( + 1, 0, 2 + ) else: - # a record is not a matrix (e.g. an EDF+ annotation channel is - # stored at its own rate), so slice out each picked channel + # a record need not be a matrix (e.g. an EDF+ annotation channel + # is stored at its own rate), so slice out each picked channel records = many_chunk.reshape(n_read, -1) view = np.empty((len(read_sel), n_read, buf_len), many_chunk.dtype) for j, ci in enumerate(read_sel): @@ -755,40 +719,8 @@ def _read_uniform_segment( assert r_sidx == 0 and r_eidx == n_read * buf_len output = data[:, d_start:d_stop].reshape(view.shape) assert np.shares_memory(output, data) - if executor is None: - _calibrate_uniform_edf( - output, - view, - cal, - offsets, - gains, - None, - ) - else: - pending.append( - executor.submit( - _calibrate_uniform_edf, - output, - view, - cal, - offsets, - gains, - None, - ) - ) - continue - if ( - n_read == 1 - and not len(stim_rows) - and (r_sidx != 0 or r_eidx != buf_len) - ): _calibrate_uniform_edf( - data[:, d_start:d_stop], - view[:, 0, r_sidx:r_eidx], - cal[:, 0, 0], - offsets[:, 0, 0], - gains[:, 0, 0], - cals, + output, view, cal, offsets, gains, cals[:, :, np.newaxis] ) continue @@ -808,8 +740,6 @@ def _read_uniform_segment( out=data[:, d_start:d_stop], casting="unsafe", ) - for future in pending: - future.result() return True @@ -838,13 +768,16 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # actually one of the requested channels idx_arr = np.arange(idx.start, idx.stop) if isinstance(idx, slice) else idx - # We could read this one EDF block at a time, which would be this: - ch_offsets = np.cumsum(np.concatenate([[0], n_samps]), dtype=np.int64) + # We could read this one EDF block at a time, but to speed it up we really + # need to read multiple blocks at once, otherwise we can end up with e.g. + # 18,181 chunks for a 20 MB file! Let's do ~10 MB chunks: + stride_layout = raw_extras.get("stride_layout") + if stride_layout is None: + ch_offsets = np.cumsum(np.concatenate([[0], n_samps]), dtype=np.int64) + n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * dtype_byte), 1) + else: + ch_offsets, n_per = stride_layout[:2] block_start_idx, r_lims, _ = _blk_read_lims(start, stop, buf_len) - # But to speed it up, we really need to read multiple blocks at once, - # Otherwise we can end up with e.g. 18,181 chunks for a 20 MB file! - # Let's do ~10 MB chunks: - n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * dtype_byte), 1) with _gdf_edf_get_fid(filenames, buffering=0) as fid: # Extract data @@ -1120,25 +1053,19 @@ def _get_info( edf_info["max_samp"] = max_samp = n_samps.max() all_n_samps = edf_info["n_samps"] edf_info["stride_layout"] = None - # Only the *selected* channels have to share a sampling rate. An EDF+ - # annotation channel usually does not, which is why `rectangular` is checked - # separately: it says whether a data record is a plain (n_channels, max_samp) - # matrix that can simply be reshaped, or has to be gathered channel by channel. + # Only the *selected* channels have to share a sampling rate for striding. if edf_info["subtype"] in ("edf", "bdf") and np.all( all_n_samps[edf_info["sel"]] == max_samp ): ch_offsets = np.cumsum(np.concatenate([[0], all_n_samps]), dtype=np.int64) n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * edf_info["dtype_byte"]), 1) - rectangular = bool(np.all(all_n_samps == max_samp)) - sel_in_physical_order = rectangular and np.array_equal( - edf_info["sel"], np.arange(len(all_n_samps)) - ) - edf_info["stride_layout"] = ( - ch_offsets, - n_per, - sel_in_physical_order, - rectangular, - ) + # A data record is a plain (n_channels, max_samp) matrix -- and hence + # reshapeable -- only when *every* channel shares the rate; an EDF+ + # annotation channel usually does not. + sel_in_physical_order = bool( + np.all(all_n_samps == max_samp) + ) and np.array_equal(edf_info["sel"], np.arange(len(all_n_samps))) + edf_info["stride_layout"] = (ch_offsets, n_per, sel_in_physical_order) # Info structure # ------------------------------------------------------------------------- diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 500cfcc575f..c323f55f0cf 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -5,8 +5,6 @@ import datetime import gc import sys -import threading -from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from functools import partial from io import BytesIO @@ -71,7 +69,7 @@ def _repeat_edf_records(source, destination, n_records=6): - """Repeat a one-record EDF payload for boundary tests.""" + """Repeat a one-record EDF/BDF payload so windows can cross boundaries.""" blob = bytearray(source.read_bytes()) header_nbytes = int(blob[184:192]) assert int(blob[236:244]) == 1 @@ -79,244 +77,106 @@ def _repeat_edf_records(source, destination, n_records=6): destination.write_bytes(blob[:header_nbytes] + blob[header_nbytes:] * n_records) -def _disable_uniform_stride(*args, **kwargs): +def _disable_uniform_stride(*args): return False -def _get_data_window(raw, limits): - return raw.get_data(start=limits[0], stop=limits[1]) - - -class _StrideRecorder: - def __init__(self, helper): - self.helper = helper - self.results = [] - - def __call__(self, *args, **kwargs): - result = self.helper(*args, **kwargs) - self.results.append(result) - return result - - -def _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop): +def _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop, used=True): + """Compare a read against the same read with the stride path disabled.""" helper = edf.edf._read_uniform_segment monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) want = raw.get_data(picks=picks, start=start, stop=stop) - recorder = _StrideRecorder(helper) - monkeypatch.setattr(edf.edf, "_read_uniform_segment", recorder) + calls = [] + + def _record(*args): + calls.append(helper(*args)) + return calls[-1] + + monkeypatch.setattr(edf.edf, "_read_uniform_segment", _record) got = raw.get_data(picks=picks, start=start, stop=stop) - assert recorder.results == [True] + assert calls == [used] assert_array_equal(got, want) return 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 uniform EDF stride decoding across picks and boundaries.""" - 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") +@pytest.mark.parametrize( + "reader, path", + ((read_raw_edf, edf_stim_channel_path), (read_raw_bdf, bdf_path)), + ids=("edf", "bdf"), +) +@pytest.mark.parametrize("pick_kind", ("all", "permuted")) +@pytest.mark.parametrize("window_kind", ("boundary", "multiple")) +def test_uniform_stride_decode( + reader, path, pick_kind, window_kind, monkeypatch, tmp_path +): + """Test exact uniform EDF/BDF striding, incl. calibration and stim masking.""" + repeated = tmp_path / f"uniform{path.suffix}" + _repeat_edf_records(path, repeated) + raw = reader(repeated, stim_channel=-1, preload=False, verbose="error") n_channels = len(raw.ch_names) - buffer_length = 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] + # powers of two, so dividing the calibration back out below stays exact + raw._cals[picks] *= 2.0 ** (np.arange(len(picks)) % 3 - 1) + buf_len = int(raw._raw_extras[0]["max_samp"]) start, stop = { - "within": (7, buffer_length - 5), - "boundary": (buffer_length - 7, buffer_length + 11), - "multiple": (buffer_length // 2, 4 * buffer_length + 13), + "boundary": (buf_len - 7, buf_len + 11), + "multiple": (buf_len // 2, 4 * buf_len + 13), }[window_kind] - _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop) - - -def test_uniform_stride_calibration_and_stim(monkeypatch): - """Test exact Raw calibration and stim masking on the stride path.""" - raw = read_raw_edf( - edf_stim_channel_path, stim_channel=-1, preload=False, verbose="error" - ) - picks = np.array([0, 12, len(raw.ch_names) - 1]) - raw._cals[picks] *= np.array([0.5, 2.0, 4.0]) - want = _assert_uniform_stride_matches(monkeypatch, raw, picks, 100, 900) - stim = want[-1] / raw._cals[picks[-1]] + want = _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop) + (row,) = np.flatnonzero(np.isin(picks, raw._raw_extras[0]["stim_channel_idxs"])) + stim = want[row] / raw._cals[picks[row]] assert_array_equal(stim, np.bitwise_and(stim.astype(int), 2**17 - 1)) -@pytest.mark.parametrize("pick_kind", ("all", "subset", "reversed")) -def test_uniform_stride_direct_calibration(pick_kind, monkeypatch, tmp_path): - """Test exact direct calibration for a single-record window.""" +def test_uniform_stride_aligned_memmap_direct_output(monkeypatch, tmp_path): + """Test an aligned full-record preload calibrates straight into the memmap.""" repeated = tmp_path / "uniform.edf" _repeat_edf_records(edf_stim_channel_path, repeated) - raw = read_raw_edf(repeated, stim_channel=None, preload=False, verbose="error") - n_channels = len(raw.ch_names) - picks = { - "all": np.arange(n_channels), - "subset": np.array([0, n_channels // 2, n_channels - 1]), - "reversed": np.arange(n_channels - 1, -1, -1), - }[pick_kind] - raw._cals[picks] *= np.linspace(0.5, 2.0, len(picks)) - buffer_length = int(raw._raw_extras[0]["max_samp"]) - _assert_uniform_stride_matches(monkeypatch, raw, picks, 7, buffer_length - 5) - - -@pytest.mark.parametrize("calibration", (1.0, 2.0)) -def test_uniform_stride_full_single_record_uses_buffer(monkeypatch, calibration): - """Test full single records retain buffered calibration.""" - raw = read_raw_edf( - edf_stim_channel_path, stim_channel=None, preload=False, verbose="error" - ) - raw._cals *= calibration - helper = edf.edf._read_uniform_segment - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) - want = raw.get_data() - calibration = _StrideRecorder(edf.edf._calibrate_uniform_edf) - monkeypatch.setattr(edf.edf, "_read_uniform_segment", helper) - monkeypatch.setattr(edf.edf, "_calibrate_uniform_edf", calibration) - got = raw.get_data() - assert calibration.results == [] - assert_array_equal(got.view(np.uint64), want.view(np.uint64)) - - -@pytest.mark.parametrize( - "reader, source, extension", - ((read_raw_edf, edf_stim_channel_path, "edf"), (read_raw_bdf, bdf_path, "bdf")), - ids=("edf", "bdf"), -) -def test_uniform_stride_aligned_memmap_direct_output( - monkeypatch, tmp_path, reader, source, extension -): - """Test an aligned full-record preload needs no float work buffer.""" - repeated = tmp_path / f"uniform.{extension}" - _repeat_edf_records(source, repeated) - raw = reader(repeated, stim_channel=None, preload=False, verbose="error") - helper = edf.edf._read_uniform_segment + kwargs = dict(stim_channel=None, verbose="error") monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) - want = raw.get_data() - - recorder = _StrideRecorder(helper) - monkeypatch.setattr(edf.edf, "_read_uniform_segment", recorder) - monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_EXTRA_BYTES", 0) - thread_ids = set() + want = read_raw_edf(repeated, preload=True, **kwargs).get_data() + monkeypatch.undo() calibrate = edf.edf._calibrate_uniform_edf + direct = [] - def _record_thread(*args): - thread_ids.add(threading.get_ident()) + def _record(*args): + direct.append(args[0].shape) return calibrate(*args) - monkeypatch.setattr(edf.edf, "_EDF_FROMFILE_THREAD_MIN_BYTES", 0) - monkeypatch.setattr(edf.edf, "_calibrate_uniform_edf", _record_thread) - mmap_path = tmp_path / "uniform.dat" - got = reader( - repeated, - stim_channel=None, - preload=mmap_path, - verbose="error", - ) - assert recorder.results == [True] - assert isinstance(got._data, np.memmap) - assert Path(got._data.filename) == mmap_path - assert threading.get_ident() not in thread_ids - assert_array_equal(got._data.view(np.uint64), want.view(np.uint64)) - got._data._mmap.close() - got._data = None - - -def test_uniform_calibration_identity_cals_bit_exact(): - """Test omitting identity output calibrations preserves every bit.""" - source = np.array( - [[np.iinfo(np.int16).min, -1, 0], [1, 17, np.iinfo(np.int16).max]], - dtype=np.int16, - ) - cal = np.array([0.125, 0.3]) - offsets = np.array([-0.0, 1.1]) - gains = np.array([1e-6, -2.5]) - want = np.empty(source.shape, dtype=np.float64) - np.multiply(source, cal[:, np.newaxis], out=want) - want += offsets[:, np.newaxis] - want *= gains[:, np.newaxis] - want *= np.ones((len(source), 1)) - got = np.empty_like(want) - edf.edf._calibrate_uniform_edf(got, source, cal, offsets, gains, None) - assert_array_equal(got.view(np.uint64), want.view(np.uint64)) - - -@pytest.mark.parametrize("pick_kind", ("all", "subset", "reversed", "permuted")) -@pytest.mark.parametrize("window_kind", ("within", "boundary", "multiple")) -def test_uniform_stride_bdf(pick_kind, window_kind, monkeypatch, tmp_path): - """Test exact BDF stride decoding across picks and record boundaries.""" - repeated = tmp_path / "uniform.bdf" - _repeat_edf_records(bdf_path, repeated) - raw = read_raw_bdf(repeated, preload=False, verbose="error") - n_channels = len(raw.ch_names) - buffer_length = 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, buffer_length - 5), - "boundary": (buffer_length - 7, buffer_length + 11), - "multiple": (buffer_length // 2, 4 * buffer_length + 13), - }[window_kind] - _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop) + monkeypatch.setattr(edf.edf, "_calibrate_uniform_edf", _record) + raw = read_raw_edf(repeated, preload=tmp_path / "uniform.dat", **kwargs) + assert len(direct) == 1 # calibrated in place, no float64 work buffer + assert isinstance(raw._data, np.memmap) + assert_array_equal(raw._data.view(np.uint64), want.view(np.uint64)) + raw._data._mmap.close() + raw._data = None -@pytest.mark.parametrize("fallback_kind", ("mixed_rate", "projection", "memory")) +@pytest.mark.parametrize("fallback_kind", ("mixed_rate", "projection")) def test_uniform_stride_edf_falls_back(fallback_kind, monkeypatch): """Test EDF layouts requiring special handling retain legacy behavior.""" - path = { - "mixed_rate": edf_uneven_path, - "projection": edf_stim_channel_path, - "memory": edf_stim_channel_path, - }[fallback_kind] + path = edf_uneven_path if fallback_kind == "mixed_rate" else edf_stim_channel_path raw = read_raw_edf(path, preload=False, verbose="error") if fallback_kind == "projection": raw.set_eeg_reference(projection=True, verbose="error").apply_proj( verbose="error" ) - elif fallback_kind == "memory": - monkeypatch.setattr(edf.edf, "_EDF_STRIDE_MAX_EXTRA_BYTES", 0) picks = np.array([0]) - helper = edf.edf._read_uniform_segment - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) - want = raw.get_data(picks=picks, start=100, stop=900) - recorder = _StrideRecorder(helper) - monkeypatch.setattr(edf.edf, "_read_uniform_segment", recorder) - got = raw.get_data(picks=picks, start=100, stop=900) - assert recorder.results == [False] - assert_array_equal(got, want) + _assert_uniform_stride_matches(monkeypatch, raw, picks, 100, 900, used=False) def test_uniform_stride_edf_annotations(monkeypatch): """Test an EDF+ annotation channel does not disable the stride path.""" raw = read_raw_edf(edf_path, preload=False, verbose="error") - buffer_length = int(raw._raw_extras[0]["max_samp"]) - picks = np.array([0, 2]) + buf_len = int(raw._raw_extras[0]["max_samp"]) _assert_uniform_stride_matches( - monkeypatch, raw, picks, buffer_length // 3, 3 * buffer_length + 7 + monkeypatch, raw, np.array([0, 2]), buf_len // 3, 3 * buf_len + 7 ) -def test_uniform_stride_concurrent(monkeypatch, tmp_path): - """Test concurrent stride reads do not share seekable file state.""" - repeated = tmp_path / "concurrent.edf" - _repeat_edf_records(edf_stim_channel_path, repeated) - raw = read_raw_edf(repeated, stim_channel=-1, preload=False, verbose="error") - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) - reference = raw.get_data() - monkeypatch.undo() - windows = [(start, start + 64) for start in range(0, raw.n_times - 64, 31)] - with ThreadPoolExecutor(max_workers=8) as pool: - got = list(pool.map(partial(_get_data_window, raw), windows * 4)) - 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) @@ -341,15 +201,6 @@ def test_orig_units(): assert set(raw_back._orig_units) == set(raw.ch_names) -def test_uniform_edf_does_not_import_interpolation(monkeypatch): - """Test uniform EDF decoding does not load interpolation support.""" - monkeypatch.setitem(sys.modules, "scipy.interpolate", None) - raw = read_raw_edf( - edf_stim_channel_path, stim_channel=-1, preload=True, verbose="error" - ) - assert raw.preload - - def test_units_params(): """Test enforcing original channel units.""" with pytest.raises( @@ -578,8 +429,10 @@ def test_edf_data_broken(tmp_path): @pytest.mark.parametrize("method", ("constructor", "load_data")) -def test_edf_preload_memmap_ownership(method, tmp_path): +def test_edf_preload_memmap_ownership(method, tmp_path, monkeypatch): """Test ownership of a populated EDF preload memmap.""" + # uniformly sampled EDF must not need (and hence not import) interpolation + monkeypatch.setitem(sys.modules, "scipy.interpolate", None) memmap_fname = tmp_path / f"edf-{method}-memmap.dat" if method == "constructor": raw = read_raw_edf(edf_stim_channel_path, preload=memmap_fname) @@ -613,48 +466,30 @@ def test_duplicate_channel_labels_edf(): ( np.array([], dtype=np.int32), np.array([-(1 << 23), -1, 0, 1, (1 << 23) - 1], dtype=np.int32), - np.array([-(1 << 23)], dtype=np.int32), - np.array([-1, 1], dtype=np.int32), np.random.default_rng(42).integers( -(1 << 23), 1 << 23, size=1000, dtype=np.int32 ), ), ) -@pytest.mark.parametrize("file_kind", ("buffer", "disk")) -def test_read_ch_bdf_int24(tmp_path, values, file_kind): - """Test exact signed 24-bit BDF decoding, including buffer boundaries.""" +def test_read_ch_bdf_int24(values): + """Test exact signed 24-bit BDF decoding, including the last sample.""" unsigned = values.astype(np.int64) & ((1 << 24) - 1) packed = np.column_stack((unsigned, unsigned >> 8, unsigned >> 16)).astype(np.uint8) - if file_kind == "buffer": - fid = BytesIO(packed.tobytes()) - else: - fname = tmp_path / "samples.bdf" - fname.write_bytes(packed.tobytes()) - fid = fname.open("rb") - with fid: + with BytesIO(packed.tobytes()) as fid: got = _read_ch( - fid, - subtype="bdf", - samp=len(values), - dtype_byte=3, - dtype=np.uint8, + fid, subtype="bdf", samp=len(values), dtype_byte=3, dtype=np.uint8 ) assert_array_equal(got, values) -@pytest.mark.parametrize("missing", (1, 2)) -@pytest.mark.parametrize("file_kind", ("buffer", "disk")) -def test_read_ch_bdf_short_read(tmp_path, missing, file_kind): +def test_read_ch_bdf_short_read(tmp_path): """Test truncated signed 24-bit BDF data raises instead of mixing bytes.""" - data = b"\x00" * (6 - missing) - if file_kind == "buffer": - fid = BytesIO(data) - else: - fname = tmp_path / "truncated.bdf" - fname.write_bytes(data) - fid = fname.open("rb") - with fid, pytest.raises(RuntimeError, match="requested BDF bytes"): - _read_ch(fid, subtype="bdf", samp=2, dtype_byte=3, dtype=np.uint8) + fname = tmp_path / "truncated.bdf" + fname.write_bytes(b"\x00" * 5) + # np.frombuffer raises on a short buffer, np.fromfile just returns less + for fid in (BytesIO(b"\x00" * 5), fname.open("rb")): + with fid, pytest.raises(RuntimeError, match="requested BDF bytes"): + _read_ch(fid, subtype="bdf", samp=2, dtype_byte=3, dtype=np.uint8) def test_parse_annotation(tmp_path): diff --git a/mne/io/eeglab/tests/test_eeglab.py b/mne/io/eeglab/tests/test_eeglab.py index 2dfa1f30b89..c5016951c06 100644 --- a/mne/io/eeglab/tests/test_eeglab.py +++ b/mne/io/eeglab/tests/test_eeglab.py @@ -498,7 +498,7 @@ def test_eeglab_annotations(fname): @testing.requires_testing_data -def test_eeglab_read_annotations(): +def test_eeglab_read_annotations(monkeypatch): """Test annotations onsets are timestamps (+ validate some).""" annotations = read_annotations(raw_fname_mat) validation_samples = [0, 1, 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31] @@ -525,61 +525,12 @@ def test_eeglab_read_annotations(): ) # test if event durations are imported correctly + check_load_mat = Mock(wraps=mne.io.eeglab.eeglab._check_load_mat) + monkeypatch.setattr(mne.io.eeglab.eeglab, "_check_load_mat", check_load_mat) raw = read_raw_eeglab(raw_fname_event_duration, preload=True, montage_units="dm") # file contains 3 annotations with 0.5 s (64 samples) duration each assert_allclose(raw.annotations.duration, np.ones(3) * 0.5) - - -@pytest.mark.parametrize( - "data_kind, preload_kind", - ( - ("embedded", "ram"), - ("external", "lazy"), - ("external", "ram"), - ("external", "mmap"), - ), -) -def test_raw_eeglab_reuses_metadata_for_annotations( - tmp_path, monkeypatch, data_kind, preload_kind -): - """Test that raw annotations reuse the already-loaded EEG structure.""" - fname = tmp_path / "single-parse.set" - events = np.array( - [("first", 2.0, 2.0), ("second", 4.0, 4.0)], - dtype=[("type", "O"), ("latency", "f8"), ("duration", "f8")], - ) - data = np.arange(10.0, dtype=" Date: Fri, 28 Aug 2026 15:10:49 +0200 Subject: [PATCH 09/12] Simplify the EDF stride decode to one branch in the existing loop The decode lived in a 118-line _read_uniform_segment that duplicated the chunk loop, a _calibrate_uniform_edf helper, and a stride_layout tuple cached in the header and unpacked far away, gated by a rectangular -> sel_in_physical_order -> in_physical_order -> direct_output chain. A reviewer had to hold two copies of the same loop in their head. It is now one branch inside the loop that was already there: gather the picked channels by their record offsets, calibrate the block straight into 'data', and mask any stim rows. No new helper functions, no precomputed tuple, no flag chain. edf.py goes from 235 to 64 added lines and is measurably faster -- EDF preload 1.26x -> 1.40x and EDF+ preload from parity to 1.47x, because every case now writes into 'data' directly, not just the aligned one that direct_output required. The tests use the two hooks the reader already has instead of recorder classes: an identity projector routes the same values through the per-channel loop, and counting _mult_cal_one tells the two paths apart. 308 -> 81 added lines, and mutation testing catches 15/15 injected defects where the longer version caught 11. Output stays bit-identical to main: 77/77 array snapshots and 56/56 annotation snapshots across 28 files. --- mne/io/edf/edf.py | 187 ++++++++--------------------------- mne/io/edf/tests/test_edf.py | 129 +++++++----------------- 2 files changed, 75 insertions(+), 241 deletions(-) diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 02914b30df0..d789a194036 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -625,129 +625,11 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): return ch_data -def _calibrate_uniform_edf(data, source, cal, offsets, gains, cals): - """Calibrate uniform EDF/BDF samples directly into their output buffer.""" - np.multiply(source, cal, out=data, casting="unsafe") - data += offsets - data *= gains - data *= cals - - -def _read_uniform_segment( - data, idx, start, stop, raw_extras, filenames, cals, mult -) -> bool: - """Read a uniformly sampled EDF/BDF segment using NumPy strides.""" - subtype = raw_extras["subtype"] - if subtype not in ("edf", "bdf") or not isinstance(filenames, str | Path): - return False - idx_is_slice = isinstance(idx, slice) - if idx_is_slice and idx.step not in (None, 1): - return False - idx_arr = np.arange(idx.start, idx.stop) if idx_is_slice else np.asarray(idx) - if len(idx_arr) == 0 or ( - not idx_is_slice and len(np.unique(idx_arr)) != len(idx_arr) - ): - return False - - stride_layout = raw_extras.get("stride_layout") - if stride_layout is None or mult is not None: - return False - n_samps = raw_extras["n_samps"] - buf_len = int(raw_extras["max_samp"]) - ch_offsets, n_per, sel_in_physical_order = stride_layout - - 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] - cal = cal[idx_arr, np.newaxis, np.newaxis] - offsets = offsets[idx_arr, np.newaxis, np.newaxis] - gains = gains[idx_arr, np.newaxis, np.newaxis] - stim_rows = ( - np.flatnonzero(np.isin(idx_arr, stim_channel_idxs)) - if len(stim_channel_idxs) - else () - ) - - block_start_idx, r_lims, d_lims = _blk_read_lims(start, stop, buf_len) - in_physical_order = ( - sel_in_physical_order - and idx_is_slice - and idx.start == 0 - and idx.stop == len(n_samps) - ) - direct_output = ( - in_physical_order - and not len(stim_rows) - and len(r_lims) > 1 - and n_per > 1 - and r_lims[0][0] == 0 - and r_lims[-1][1] == buf_len - ) - with _gdf_edf_get_fid(filenames, buffering=0) as fid: - start_offset = data_offset + block_start_idx * ch_offsets[-1] * dtype_byte - 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 - ) - if in_physical_order: - view = many_chunk.reshape(n_read, len(n_samps), buf_len).transpose( - 1, 0, 2 - ) - else: - # a record need not be a matrix (e.g. an EDF+ annotation channel - # is stored at its own rate), so slice out each picked channel - records = many_chunk.reshape(n_read, -1) - view = np.empty((len(read_sel), n_read, buf_len), many_chunk.dtype) - for j, ci in enumerate(read_sel): - view[j] = records[:, ch_offsets[ci] : ch_offsets[ci + 1]] - - r_sidx = r_lims[ai][0] - r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] - d_start = d_lims[ai][0] - d_stop = d_lims[ai + n_read - 1][1] - assert d_stop - d_start == r_eidx - r_sidx - if direct_output and n_read > 1: - assert r_sidx == 0 and r_eidx == n_read * buf_len - output = data[:, d_start:d_stop].reshape(view.shape) - assert np.shares_memory(output, data) - _calibrate_uniform_edf( - output, view, cal, offsets, gains, cals[:, :, np.newaxis] - ) - continue - - one = np.empty(view.shape, dtype=np.float64) - np.multiply(view, cal, out=one) - one += offsets - one *= gains - block = one.reshape(len(idx_arr), -1)[:, r_sidx:r_eidx] - for row in stim_rows: - stim = block[row].astype(int) - np.bitwise_and(stim, 2**17 - 1, out=stim) - block[row] = stim - assert d_stop - d_start == block.shape[1] - np.multiply( - block, - cals, - out=data[:, d_start:d_stop], - casting="unsafe", - ) - return True +_EDF_CHUNK_BYTES = 10 * 1024 * 1024 # read roughly this much per chunk def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, mult): """Read a chunk of raw data.""" - if _read_uniform_segment(data, idx, start, stop, raw_extras, filenames, cals, mult): - return [] - n_samps = raw_extras["n_samps"] buf_len = int(raw_extras["max_samp"]) dtype = raw_extras["dtype_np"] @@ -768,16 +650,22 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # actually one of the requested channels idx_arr = np.arange(idx.start, idx.stop) if isinstance(idx, slice) else idx - # We could read this one EDF block at a time, but to speed it up we really - # need to read multiple blocks at once, otherwise we can end up with e.g. - # 18,181 chunks for a 20 MB file! Let's do ~10 MB chunks: - stride_layout = raw_extras.get("stride_layout") - if stride_layout is None: - ch_offsets = np.cumsum(np.concatenate([[0], n_samps]), dtype=np.int64) - n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * dtype_byte), 1) - else: - ch_offsets, n_per = stride_layout[:2] - block_start_idx, r_lims, _ = _blk_read_lims(start, stop, buf_len) + # We could read this one EDF block at a time, which would be this: + ch_offsets = np.cumsum(np.concatenate([[0], n_samps]), dtype=np.int64) + block_start_idx, r_lims, d_lims = _blk_read_lims(start, stop, buf_len) + # But to speed it up, we really need to read multiple blocks at once, + # Otherwise we can end up with e.g. 18,181 chunks for a 20 MB file! + n_per = max(_EDF_CHUNK_BYTES // (ch_offsets[-1] * dtype_byte), 1) + + # When every picked channel stores buf_len samples per data record there is + # nothing to resample, so the picks form a plain (n_picks, n_times) block we + # can calibrate in one go straight into `data`. Mixed sampling rates, a + # projector, or a stim channel that needs interpolating use the per-channel + # loop below instead. + n_picks = len(idx_arr) + picks = read_sel[:n_picks] # picked signal channels, in output row order + uniform = mult is None and bool((n_samps[picks] == buf_len).all()) + stim_rows = [j for j, i in enumerate(idx_arr) if i in stim_channel_idxs] with _gdf_edf_get_fid(filenames, buffering=0) as fid: # Extract data @@ -786,7 +674,9 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # first read everything into the `ones` array. For channels with # lower sampling frequency, there will be zeros left at the end of the # row. Ignore TAL/annotations channel and only store `orig_sel` - ones = np.zeros((len(orig_sel), data.shape[-1]), dtype=data.dtype) + # `ones` has no rows on the fast path, which writes into `data` itself + n_stage = 0 if uniform else len(orig_sel) + ones = np.zeros((n_stage, data.shape[-1]), dtype=data.dtype) # save how many samples have already been read per channel n_smp_read = [0 for _ in range(len(orig_sel))] @@ -802,6 +692,23 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, r_sidx = r_lims[ai][0] r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] + if uniform: + block = np.empty((n_picks, n_read, buf_len), many_chunk.dtype) + for j, ci in enumerate(picks): + block[j] = many_chunk[:, ch_offsets[ci] : ch_offsets[ci + 1]] + for ci in read_sel[n_picks:]: # annotation channels + tal_data.append( + many_chunk[:, ch_offsets[ci] : ch_offsets[ci + 1]].copy() + ) + out = data[:, d_lims[ai][0] : d_lims[ai + n_read - 1][1]] + flat = block.reshape(n_picks, n_read * buf_len)[:, r_sidx:r_eidx] + np.multiply(flat, cal[idx_arr, np.newaxis], out=out) + out += offsets[idx_arr, np.newaxis] + out *= gains[idx_arr, np.newaxis] + for j in stim_rows: + out[j] = np.bitwise_and(out[j].astype(int), 2**17 - 1) + continue + # loop over selected channels, ci=channel selection for ii, ci in enumerate(read_sel): # This now has size (n_chunks_read, n_samp[ci]) @@ -876,6 +783,11 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, _mult_cal_one(data[:, :], ones, idx, cals, mult) + if uniform: + # stands in for the `data_view *= cals` that _mult_cal_one applies; the + # block above is skipped because the fast path leaves n_smp_read zero + data *= cals + if len(tal_data) > 1: tal_data = np.concatenate([tal.ravel() for tal in tal_data]) tal_data = tal_data[np.newaxis, :] @@ -1051,21 +963,6 @@ def _get_info( edf_info["max_samp"] = max_samp = n_samps[picks].max() else: edf_info["max_samp"] = max_samp = n_samps.max() - all_n_samps = edf_info["n_samps"] - edf_info["stride_layout"] = None - # Only the *selected* channels have to share a sampling rate for striding. - if edf_info["subtype"] in ("edf", "bdf") and np.all( - all_n_samps[edf_info["sel"]] == max_samp - ): - ch_offsets = np.cumsum(np.concatenate([[0], all_n_samps]), dtype=np.int64) - n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * edf_info["dtype_byte"]), 1) - # A data record is a plain (n_channels, max_samp) matrix -- and hence - # reshapeable -- only when *every* channel shares the rate; an EDF+ - # annotation channel usually does not. - sel_in_physical_order = bool( - np.all(all_n_samps == max_samp) - ) and np.array_equal(edf_info["sel"], np.arange(len(all_n_samps))) - edf_info["stride_layout"] = (ch_offsets, n_per, sel_in_physical_order) # Info structure # ------------------------------------------------------------------------- diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index c323f55f0cf..99d443ee0ea 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -68,113 +68,50 @@ misc = ["EXG1", "EXG5", "EXG8", "M1", "M2"] -def _repeat_edf_records(source, destination, n_records=6): - """Repeat a one-record EDF/BDF payload so windows can cross boundaries.""" +def _multi_record(source, destination, n_records=6): + """Give a one-record EDF/BDF file several records, so windows cross them.""" blob = bytearray(source.read_bytes()) + if int(blob[236:244]) != 1: + return source # already spans several records header_nbytes = int(blob[184:192]) - 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) + return destination -def _disable_uniform_stride(*args): - return False - - -def _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop, used=True): - """Compare a read against the same read with the stride path disabled.""" - helper = edf.edf._read_uniform_segment - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) - want = raw.get_data(picks=picks, start=start, stop=stop) - calls = [] - - def _record(*args): - calls.append(helper(*args)) - return calls[-1] - - 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) - return want - - +@pytest.mark.parametrize("chunk_bytes", (10 * 1024 * 1024, 4096)) # one chunk, several @pytest.mark.parametrize( "reader, path", - ((read_raw_edf, edf_stim_channel_path), (read_raw_bdf, bdf_path)), - ids=("edf", "bdf"), + ( + # EDF+: an annotation channel at its own rate, plus two stim channels + (read_raw_edf, data_dir / "test_stim_channel.edf"), + (read_raw_edf, edf_stim_channel_path), + (read_raw_bdf, bdf_path), # BDF: 24-bit samples and a Status channel + ), + ids=("edf+", "edf", "bdf"), ) -@pytest.mark.parametrize("pick_kind", ("all", "permuted")) -@pytest.mark.parametrize("window_kind", ("boundary", "multiple")) -def test_uniform_stride_decode( - reader, path, pick_kind, window_kind, monkeypatch, tmp_path -): - """Test exact uniform EDF/BDF striding, incl. calibration and stim masking.""" - repeated = tmp_path / f"uniform{path.suffix}" - _repeat_edf_records(path, repeated) - raw = reader(repeated, stim_channel=-1, preload=False, verbose="error") - n_channels = len(raw.ch_names) - picks = { - "all": np.arange(n_channels), - "permuted": np.array([n_channels - 1, 1, n_channels // 2, 0]), - }[pick_kind] - # powers of two, so dividing the calibration back out below stays exact - raw._cals[picks] *= 2.0 ** (np.arange(len(picks)) % 3 - 1) +def test_uniform_decode_matches_loop(reader, path, chunk_bytes, monkeypatch, tmp_path): + """Test uniform-rate records decode exactly like the per-channel loop.""" + path = _multi_record(path, tmp_path / f"uniform{path.suffix}") + monkeypatch.setattr(edf.edf, "_EDF_CHUNK_BYTES", chunk_bytes) + raw = reader(path, preload=False, verbose="error") buf_len = int(raw._raw_extras[0]["max_samp"]) - start, stop = { - "boundary": (buf_len - 7, buf_len + 11), - "multiple": (buf_len // 2, 4 * buf_len + 13), - }[window_kind] - want = _assert_uniform_stride_matches(monkeypatch, raw, picks, start, stop) - (row,) = np.flatnonzero(np.isin(picks, raw._raw_extras[0]["stim_channel_idxs"])) - stim = want[row] / raw._cals[picks[row]] - assert_array_equal(stim, np.bitwise_and(stim.astype(int), 2**17 - 1)) - - -def test_uniform_stride_aligned_memmap_direct_output(monkeypatch, tmp_path): - """Test an aligned full-record preload calibrates straight into the memmap.""" - repeated = tmp_path / "uniform.edf" - _repeat_edf_records(edf_stim_channel_path, repeated) - kwargs = dict(stim_channel=None, verbose="error") - monkeypatch.setattr(edf.edf, "_read_uniform_segment", _disable_uniform_stride) - want = read_raw_edf(repeated, preload=True, **kwargs).get_data() - monkeypatch.undo() - calibrate = edf.edf._calibrate_uniform_edf - direct = [] - - def _record(*args): - direct.append(args[0].shape) - return calibrate(*args) - - monkeypatch.setattr(edf.edf, "_calibrate_uniform_edf", _record) - raw = read_raw_edf(repeated, preload=tmp_path / "uniform.dat", **kwargs) - assert len(direct) == 1 # calibrated in place, no float64 work buffer - assert isinstance(raw._data, np.memmap) - assert_array_equal(raw._data.view(np.uint64), want.view(np.uint64)) - raw._data._mmap.close() - raw._data = None - - -@pytest.mark.parametrize("fallback_kind", ("mixed_rate", "projection")) -def test_uniform_stride_edf_falls_back(fallback_kind, monkeypatch): - """Test EDF layouts requiring special handling retain legacy behavior.""" - path = edf_uneven_path if fallback_kind == "mixed_rate" else edf_stim_channel_path - raw = read_raw_edf(path, preload=False, verbose="error") - if fallback_kind == "projection": - raw.set_eeg_reference(projection=True, verbose="error").apply_proj( - verbose="error" - ) - picks = np.array([0]) - _assert_uniform_stride_matches(monkeypatch, raw, picks, 100, 900, used=False) - - -def test_uniform_stride_edf_annotations(monkeypatch): - """Test an EDF+ annotation channel does not disable the stride path.""" - raw = read_raw_edf(edf_path, preload=False, verbose="error") - buf_len = int(raw._raw_extras[0]["max_samp"]) - _assert_uniform_stride_matches( - monkeypatch, raw, np.array([0, 2]), buf_len // 3, 3 * buf_len + 7 + n_channels = len(raw.ch_names) + picks = np.array([n_channels - 1, 1, 0]) # reordered, and includes any stim channel + kwargs = dict(picks=picks, start=buf_len // 2, stop=4 * buf_len + 13) + # the fast path writes into `data` itself, the per-channel loop stages the + # result and calls _mult_cal_one -- so counting it tells the two apart + staged = [] + stage = edf.edf._mult_cal_one + monkeypatch.setattr( + edf.edf, "_mult_cal_one", lambda *a: (staged.append(1), stage(*a))[1] ) + want = raw.get_data(**kwargs) + assert not staged + # an identity projector sends the same values down the per-channel loop + raw._projector = np.eye(n_channels) + assert_array_equal(want, raw.get_data(**kwargs)) + assert staged def test_orig_units(): From a124497808719f0f0c50a847fdc7558639ab517e Mon Sep 17 00:00:00 2001 From: Bru Date: Fri, 28 Aug 2026 16:37:20 +0200 Subject: [PATCH 10/12] Address review: document the readable decode, add a times fixture Two maintainability points from @larsoner: - The BDF int24 unpack now carries the equivalent three-liner in a comment, so the next person has a conceptual model and something to swap in if they suspect the optimized version. Verified equivalent across sizes 1-100000 and at the range boundaries; the readable form is ~3x slower. - _fail_if_times_materialized becomes a fail_if_times_materialized fixture that patches and yields, used by both tests. The annotations test now has the patch active during the read too, so it is slightly stricter than before. --- mne/io/edf/edf.py | 11 ++++++++++- mne/io/tests/test_raw.py | 16 ++++++++++------ 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index d789a194036..443a3517975 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -607,6 +607,16 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): raise RuntimeError( f"Only {raw.size} of {expected} requested BDF bytes could be read" ) + # Read each 3-byte sample as the low bytes of an overlapping 4-byte + # word, mask off the byte borrowed from the next sample, then move the + # sign bit to bit 31 and shift back down to sign-extend it. The last + # sample has no next sample to borrow from, so it is done by hand. + # This is equivalent to, and ~3x faster than, the readable version: + # + # ch_data = raw.reshape(-1, 3).astype(INT32) + # ch_data = ch_data[:, 0] | (ch_data[:, 1] << 8) | (ch_data[:, 2] << 16) + # ch_data <<= 8 # sign-extend bit 23 + # ch_data >>= 8 ch_data = np.empty(samp, dtype=INT32) packed = np.ndarray( (max(samp - 1, 0),), dtype=">= 8 diff --git a/mne/io/tests/test_raw.py b/mne/io/tests/test_raw.py index 256db6e319a..bc63f7598c2 100644 --- a/mne/io/tests/test_raw.py +++ b/mne/io/tests/test_raw.py @@ -50,8 +50,14 @@ ) -def _fail_if_times_materialized(*args, **kwargs): - pytest.fail("The full Raw.times vector was materialized") +@pytest.fixture +def fail_if_times_materialized(monkeypatch): + """Fail the test if the full Raw.times vector is ever constructed.""" + + def _fail(*args, **kwargs): + pytest.fail("The full Raw.times vector was materialized") + + monkeypatch.setattr("mne.io.base._arange_div", _fail) def assert_named_constants(info): @@ -105,18 +111,16 @@ def test_orig_units(): BaseRaw(info, last_samps=[1], orig_units=True) -def test_preload_does_not_materialize_times(monkeypatch): +def test_preload_does_not_materialize_times(fail_if_times_materialized): """Test preloading does not construct the full time vector.""" - monkeypatch.setattr("mne.io.base._arange_div", _fail_if_times_materialized) raw = read_raw_fif(raw_fname, preload=True, verbose="error") assert raw.preload -def test_set_annotations_does_not_materialize_times(monkeypatch): +def test_set_annotations_does_not_materialize_times(fail_if_times_materialized): """Test annotation bounds use the scalar recording endpoint.""" raw = read_raw_fif(raw_fname, preload=False, verbose="error") annotations = Annotations([0.0], [0.1], ["test"]) - monkeypatch.setattr("mne.io.base._arange_div", _fail_if_times_materialized) raw.set_annotations(annotations) assert len(raw.annotations) == 1 From fe4c5f11b8eb02fddfdf0c1e55ac9d0c2aea0d79 Mon Sep 17 00:00:00 2001 From: Bru Date: Fri, 28 Aug 2026 20:09:46 +0200 Subject: [PATCH 11/12] Allowlist the fail_if_times_materialized fixture for vulture vulture does not model pytest fixtures: the fixture function looks uncalled and the two test parameters that request it look like unused variables. Same treatment as the other fixture names already in the allowlist. --- tools/vulture_allowlist.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/vulture_allowlist.py b/tools/vulture_allowlist.py index 5d529ceb9b1..6625c2494f5 100644 --- a/tools/vulture_allowlist.py +++ b/tools/vulture_allowlist.py @@ -9,6 +9,7 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. +fail_if_times_materialized numba_conditional options_3d invisible_fig From 354e33643a095e6a467ad37269fb855da61af184 Mon Sep 17 00:00:00 2001 From: Bru Date: Fri, 28 Aug 2026 20:19:08 +0200 Subject: [PATCH 12/12] Re-run CI (nilearn intersphinx inventory was unreachable)