diff --git a/.gitignore b/.gitignore index 21275b21c0b..0ebf0394659 100644 --- a/.gitignore +++ b/.gitignore @@ -103,4 +103,7 @@ venv/ .hypothesis/ .ruff_cache/ .ipynb_checkpoints/ -/.claude/ \ No newline at end of file +/.claude/ +# IO benchmark synthetic data (regenerate via benchmarks/io_dl/generate_data.py) +benchmarks/io_dl/data/ +benchmarks/io_dl/data_big/ diff --git a/benchmarks/io_dl/RESULTS.md b/benchmarks/io_dl/RESULTS.md new file mode 100644 index 00000000000..7eed82f591c --- /dev/null +++ b/benchmarks/io_dl/RESULTS.md @@ -0,0 +1,556 @@ +# MNE-Python IO for deep learning: benchmark + first speedups + +Date: 2026-08-24 · Machine: Apple Silicon (arm64), macOS, py 3.12.12, numpy 2.4.6, +MNE 1.12.1 (this working tree). All numbers below are medians; see +"Methodology notes" for why interleaved A/B is authoritative here. + +## 1. Which three formats? + +| Rank | Format | Why it dominates DL workloads | +|---|---|---| +| 1 | **EDF/EDF+** | The big DL corpora are EDF: TUH EEG Corpus (~25k studies, braindecode's flagship dataset), PhysioNet (CHB-MIT, Sleep-EDF). BIDS-recommended. | +| 2 | **BrainVision (.vhdr/.eeg/.vmrk)** | The other BIDS-recommended format; standard on OpenNeuro and MOABB; mne-bids converts to it by default. | +| 3 | **FIF** | MNE-native; what DL pipelines persist preprocessed data/caches to (braindecode preprocessing saves `.fif`). | + +(Honorable mention: EEGLAB `.set`, common but usually converted to one of the above.) + +## 2. Benchmark design + +Files in this directory: + +- `generate_data.py` – writes the *same* synthetic signal (64 ch × 300 s @ 256 Hz, + ~±100 µV band-limited noise) to `bench.edf`, `bench.vhdr(+.eeg,.vmrk)`, + `bench_raw.fif`, plus a raw float32 `bench.bin` baseline. EDF is int16 by design + (half the on-disk bytes of float32). +- `bench_io.py` – scenario suite (see below), GC disabled during timing. +- `profile_io.py` – cProfile attribution + deepcopy tracer (rebinds every + `from copy import deepcopy` reference in loaded `mne.*` modules). +- `ab_test.py` – interleaved A/B between the installed (unpatched) release copy + and this working tree; cancels machine drift. +- `check_equivalence.py` – md5-compares outputs of every access pattern between + unpatched and patched code (must be bit-identical). + +Scenarios (the DL-relevant access patterns): + +| Scenario | Meaning | +|---|---| +| `open_meta` | `read_raw_X(preload=False)` — header parse only | +| `full_load` | `read_raw_X(preload=True)` — end-to-end load | +| `seq_1s` | 300 sequential 1 s `get_data()` calls (streaming/eval) | +| `rand_windows` | 300 shuffled 2 s windows via `get_data()`, `preload=False` — **the training-loop pattern** (dataset opens raw once, sampler draws windows) | +| `preloaded_windows` | same windows on a preloaded raw — pure per-call Python overhead | +| floors | numpy `memmap` / `fromfile` equivalents reading identical bytes | + +### Methodology notes (things that silently corrupt results) + +1. **Import resolution**: running scripts from inside `benchmarks/io_dl/` or with + `python -c` resolves `mne` differently (cwd lands on `sys.path` for `-c`; + script dir replaces cwd for files). Both bench scripts now pin the working + tree explicitly. Verify with the printed `mne loaded from:` line. +2. **Sustained-load drift**: whole-suite runs back-to-back showed ±20–100 % + swings under background load (even the numpy floor drifted 4→11 ms). + Conclusions here rely on interleaved A/B (`ab_test.py`), not cross-run diffs. +3. cProfile inflates tiny-function costs ~3×; use it for attribution only. + +## 3. Results + +### 3.1 Baseline (unpatched tree) — where does time go? + +Steady-state cost per `get_data()` call, preloaded raw, all three formats: +**~100 µs**, identical across formats ⇒ pure Python overhead in the shared path, +zero relation to disk. The numpy floor for the same window is ~15 µs ⇒ MNE was +~6× off the floor before any I/O happens. + +Attribution (cProfile, 200 calls): + +- `_handle_tmin_tmax` touches `self.times` even when `tmin=tmax=None` + (`mne/utils/mixin.py`) and Raw's `times` property allocates + `np.arange(n_times)/sfreq` each call — **600 KB allocated+discarded per window** + at 300 s; scales linearly with recording length (≈59 MB per call for an 8 h TUH + recording!). +- Channel picks resolved **twice** per call (once in `get_data`, again in + `_parse_get_set_params`); the first goes None→`"all"`→string machinery with + ~65 `list.index` lookups per call. +- EDF reader extras (non-preloaded only): `from scipy.interpolate import + interp1d` executed inside `_read_segment_file` (first call in a fresh DataLoader + worker pays a ~0.4 s import/JIT spike), a wasted `.copy()` of every channel's + block, and an extra full-size temp buffer. +- Per-call fixed cost curve (unpatched, ms/call vs window length): + EDF flat ≈0.37 ms up to 2 s windows; BV/FIF ≈0.12–0.15 ms — i.e. small reads + were almost entirely fixed overhead. + +### 3.2 Patches applied (this tree) + +1. `mne/utils/mixin.py::_handle_tmin_tmax`: use integer `n_times` when available + (Raw); Epochs/Evoked keep their cheap stored-times path. +2. `mne/io/base.py::get_data`: `picks=None` short-circuits to + `np.arange(nchan)` — exactly equivalent to `_picks_to_idx(..., "all", + exclude=())` semantics, skips name resolution entirely. +3. `mne/io/edf/edf.py::_read_segment_file`: `interp1d` import moved into the + mixed-sfreq stim-interpolation branch; removed the per-channel `.copy()` + (arithmetic is out-of-place; TAL views stay valid until consumed). + +### 3.3 Speedups (interleaved A/B, best-of-medians) + +Preloaded random-window `get_data()` (pure access cost): + +| arm | µs/call | +|---|---| +| installed = release (unpatched) | 101 | +| tree = patched | 31 | + +⇒ **3.2× faster**; now within ~2× of the raw-numpy floor instead of ~6×. + +Non-preloaded random windows (includes reader): + +| format | unpatched µs/call | patched µs/call | speedup | +|---|---|---|---| +| EDF | 426 | 335 | **1.27×** | +| BrainVision | ~172 | ~111 | **1.55×** | +| FIF | ~192 | ~130 | **1.47×** | + +Whole-suite runs agree (preloaded_windows ~100→~34 µs/call; seq_1s BV/FIF +~150→~82 µs/call) once drift is controlled. + +Correctness: `check_equivalence.py` shows bit-identical output vs the release +copy for every pattern (full read, boundary-crossing windows, tmin/tmax, int / +str / slice / negative picks); upstream suites pass +(`test_edf.py` + `test_brainvision.py` 116 passed, `test_raw.py` 63 passed, +`test_epochs.py` 239 passed, targeted evoked tests passed). + +## 4. The deepcopy question ("I hear there's a problem") + +Traced with a tracer that rebinds *every* `deepcopy` reference held by loaded +`mne.*` modules (a plain `copy.deepcopy` monkeypatch misses `from copy import +deepcopy` bindings — first tracer attempt caught almost nothing): + +| Operation | deepcopy calls | total time | +|---|---|---| +| `read_raw_X(preload=True)` ×3 formats | 141 (mostly Info internals) | 0.64 ms | +| load+crop+reref+notch+resample | 2 | 0.07 ms | +| Epochs creation + iterate 20 epochs | 238 | 1.65 ms | +| `deepcopy(info)` alone, 64 ch | 1 | 0.05 ms | + +Verdict: on current main, **deepcopy is not the bottleneck for DL-style data +access**. It becomes measurable only with large Infos — copies scale ~66 µs +(32 ch) → ~400 µs (512 ch), roughly ×1.5–2 with a montage/dig set — or when an +operation copies Info thousands of times. If complaints trace to older MNE +versions, note `Info.__deepcopy__` already contains fast paths (chs shallow-copy +trick) that recent releases added. + +The real fixed-cost culprits were the ones fixed above (times array materialized +per call; double pick resolution), plus the EDF reader's per-channel Python loop. + +## 5. Remaining bottlenecks, ranked (candidates for next PRs) + +1. **EDF non-preloaded small reads still pay ~300 µs fixed cost** (per-channel + Python loop over blocks + temp buffer + block bookkeeping). Vectorizing the + uniform-sfreq case (the overwhelmingly common one) could cut most of it. +2. **File opened per `get_data()` call** in all three readers (~5–10 µs + + syscalls). Keeping an open handle per Raw would help streaming patterns. +3. **Batched window reads**: a `raw.get_windows(starts, width)` style API could + amortize validation/picks across hundreds of windows (the remaining ~30 µs + fixed cost per call would drop to ~µs amortized). +4. **`save()` to FIF is slow**: ~105 ms for a 19.7 MB payload (~190 MB/s) — this + is the caching path DL pipelines hit constantly; tag serialization looks + unoptimized (separate investigation). +5. Output dtype is always float64; DL users immediately cast to float32. + An opt-in `dtype="float32"` read path halves memory traffic. +6. `open_meta` ≈ 2 ms/format — matters when instantiating datasets over ~10k + files (20 s just for headers); mostly Python-side parsing/checks. + +## 6. Provider-library dissection & BIG multi-format study (session 2) + +Setup: second preset `--preset big` = **128 ch × 1800 s @ 512 Hz** (118M samples; +float32 472 MB / int16 236 MB), same signal written to every format: +`data_big/` holds EDF(int16), **BDF(int24)**, BrainVision(f32), FIF(f32), +NPZ(f32), HDF5(f32; three chunk layouts), Zarr(f32; two layouts), raw .bin. +Backends compared by `bench_backends.py`; per-call latencies via interleaved A/B. + +### 6.1 Full-file loads (BIG) + +| backend | time | effective MB/s | +|---|---|---| +| floor_memmap (f32 bin) | 152–156 ms | ~3050 | +| h5 contiguous | 153 ms | ~3077 | +| npz | 206 ms | ~2289 | +| **edfio eager EDF** | **244 ms** | ~967 | +| mne_edf | 454–500 ms | ~950–1040 (vs 236 MB payload) | +| mne_fif | 535–552 ms | ~855 | +| mne_bv | 814–842 ms | ~560 | +| mne_bdf | 995–1003 ms | ~470 | +| zarr t10s chunks | 4688–5099 ms | ~100 (!!) | + +Corroborates the MNE-forum report (Dec 2023): edfio ≈ 2× faster than MNE on +uniform-sfreq EDF; MNE's catastrophic mixed-sfreq cases (minutes) come from its +upsample-to-max-sfreq policy. Upstream floated making edfio an optional reader +backend — worth pursuing for `read_raw_edf`. + +### 6.2 Random-window reads (the DL pattern; 300 × 2 s windows) + +| backend | µs/window | +|---|---| +| h5_win (chunks = 128×512 f32) | **53–94** — beats memmap floor! | +| floor_memmap | 142 | +| h5 contiguous | 28–94 | +| h5_t10s (per-channel slabs) | 725–930 | +| zarr_win | 1678–2535 | +| mne_fif / mne_bv non-preloaded | 1298–1696 | +| edfio lazy per-signal slices | ~1158 | +| zarr_t10s | ~31700 | + +Two big lessons: +1. **Chunk geometry must mirror access geometry.** Window-shaped HDF5 chunks + turn each sample into ONE contiguous ~512 KB read and even beat the raw + memmap "floor" (which touches 128 scattered pages per window). Time-slab + chunking is 13× worse; misaligned Zarr chunks are pathological. +2. **Zarr v3's per-chunk Python dispatch makes it a poor local training + backend** despite identical chunk shapes to HDF5 (~20× slower here). + HDF5 or plain memmap/npy for single-node; Zarr for remote/parallel. + +### 6.3 Where MNE's remaining window cost lives (128 ch) + +After this round of patches, steady-state per-window (get_data incl. .sum()): + +| format | non-preloaded | preloaded | +|---|---|---| +| EDF | ~400–430 µs | ~79 µs | +| BDF | ~1000 µs | ~79 µs | +| BrainVision | ~515 µs | ~80 µs | +| FIF | ~416 µs | ~78 µs | + +Attribution experiments (`_getitem` bypass): at 128 ch, public-API validation +was consuming **55–58 %** of BV/FIF window latency (double pick resolution, +astype copies, reductions). The `_picks_to_idx` integer-fast-exit patch below +reclaims much of it; a future batched `raw.get_windows(starts, width)` API +could amortize nearly all of it (validation once, then tight read loop). + +### 6.4 Patches added in session 2 (this tree) + +4. `mne/io/edf/edf.py::_read_segment_file`: vectorized fast path for + uniform-sfreq EDF/BDF with no TAL/stim channels among requested ones + (gated to decoded outputs ≤ 32 MB so huge sequential loads keep the legacy + cache-friendly loop). Replaces the per-channel Python loop with one + reshape + strided gather + 3 vector ops; writes straight into the output + buffer when no projector/compensation is active and cals == 1 (always true + for EDF/BDF). Bit-identical output (md5-verified on small set; max |diff| + = 0 on BIG cross-checks). +5. `mne/_fiff/pick.py::_picks_to_idx`: fast return for integer arrays already + in range (skips astype copy, two boolean-reduction passes, modulo pass). + Benefits every hot call site resolving array picks. + +Measured against **pristine main** (interleaved subprocesses, best-of): + +| metric (BIG file) | pristine | patched | speedup | +|---|---|---|---| +| EDF random windows | 1317–1383 µs | 405–414 µs | **3.3×** | +| EDF 4-channel windows | 592–599 µs | 90–97 µs | **6.3×** | +| BDF random windows | ~2000 µs | ~1000 µs | **~2.0×** | +| full sequential load | ~equal within machine noise (see caveat) | | | + +Caveat learned the hard way: whole-file loads (≈944 MB float64 output alloc + +page faults) swing ±100 % with background system state on this laptop; only +interleaved same-window comparisons are trustworthy. Two earlier "regressions" +were artifacts: (a) comparing against the *installed release* instead of +pristine main — main itself got faster than 1.12.1 on full loads; (b) harness +rows that re-opened/reallocated the Raw inside the timed region. + +Correctness: 179 io tests + 368-test epochs/raw/bv batches pass; outputs +bit-identical to both the installed release and stashed-pristine main on every +pattern tested (windows, boundary spans, subsets, negative picks, tmin/tmax). + +## 7. How others speed this up (survey notes) + +- **Pre-conversion + memory-mapping**: PyRain (RainBench) reports 27–60× + dataloading speedups over NetCDF/Dask using mmap'd samples for randomized + sliding-window access; explicitly recommends mmap over chunked stores for + fragmented random access on local disks. LaBraM pre-packages EEG into HDF5 + windows. Braindecode caches preprocessed data as FIF. +- **Chunk-layout alignment**: h5py/Zarr guidance and benchmarks + (e.g., rabernat/zarr_hdf_benchmarks) — our §6.2 confirms: shape chunks like + the reads. For (n_ch, n_time) EEG with 2 s window sampling: + chunks=(n_ch, window) is optimal; (1, time_slab) is terrible for multi-channel + windows; fully contiguous is best for full loads but mediocre for random + windows on spinning rust/NVMe (still fine on macOS page cache). +- **Reader-backend competition**: pyedflib (Cython/C), edfio (vectorized numpy, + lazy loading, partial digital slices) — edfio's design (one-shot frombuffer + + record reshape + vectorized calibration + lazy per-signal loading) is the + model for MNE's EDF path; adopting it as an optional backend was proposed + upstream by MNE maintainers. +- **Amortization APIs**: DALI/WebDataset-style pre-decoded batches; in MNE + terms, reading K consecutive windows through one resolved-picks call. + +## 8. Recommended next steps (updated ranking) + +1. Land sessions-1+2 patches (bit-exact, tested): shared-path fixes, EDF + vectorized fast path, `_picks_to_idx` fast exit. ~2–6× on DL window loops. +2. Add `raw.get_windows(starts, stop)`-style batched reader (amortizes the + remaining ~30–90 µs/call Python overhead across windows; prototype shows + ≥50 % headroom for BV/FIF). +3. Optional edfio-backed `read_raw_edf` engine (fast path for the corpus-scale + mixed-sfreq pathology; upstream discussion exists). +4. Publish chunk-layout guidance + a converter recipe (BIDS→HDF5 win-chunks) + for training pipelines; keep Zarr for cloud/parallel contexts. +5. BDF int24 decode vectorization (mne_bdf is now the slowest full-loader). +6. Investigate FIF save throughput separately (fixed overhead dominates small + saves; large saves run at ~300 MB/s). + +## 9. Reproduce + +```bash +python benchmarks/io_dl/generate_data.py # small set +python benchmarks/io_dl/generate_data.py --preset big # BIG set (~2.5 GB) +python benchmarks/io_dl/bench_io.py # small-suite timings +python benchmarks/io_dl/bench_backends.py --dir data_big # multi-backend table +AB_FMT=edf AB_PRELOAD=0 python benchmarks/io_dl/ab_test.py 5 # interleaved A/B +python benchmarks/io_dl/profile_io.py # profiles + deepcopy traces +python benchmarks/io_dl/check_equivalence.py # bit-equality vs release +python benchmarks/io_dl/check_equivalence_big.py # numeric equality, BIG +``` + +Result JSONs: `results-*.json`, `backends-*.json` in this directory. + + +## 10. Session 3: profiling to the metal + batched reads + +Method: accumulator-wrapping of hot functions (no profiler inflation) plus +interleaved subprocess A/B against both pristine main and the installed +release. py-spy available as cross-check. + +### 10.1 What the traces showed (BV, BIG file, per window) + +| component | before | after fix | +|---|---|---| +| `_mult_cal_one` (cast+index+scale = **3 passes** + alloc) | ~169 µs (61 %) | **one fused pass** | +| `get_data` validation layers | ~30 µs | ~21 µs | +| open+seek+fromfile glue | ~88 µs | unchanged | + +### 10.2 Patches added (session 3) + +6. `mne/_fiff/utils.py::_mult_cal_one`: fuse gather + float-cast + calibration + into a single strided multiply into the output view (`np.multiply(one[idx], + cals, out=data_view)`), eliminating a full-size intermediate allocation and + two extra passes. Numerically identical (elementwise ops on same values); + all suites pass, outputs bit-identical. +7. `mne/io/edf/_bdf_numba.py` (new) + hook in `_read_ch`: numba-accelerated + int24→int32 BDF decoder following MNE's existing optional-numba pattern + (`mne/_numba.py`), with graceful fallback to the vectorized-numpy path. + +### 10.3 Cumulative effect (BIG file, 2 s windows, non-preloaded) + +Interleaved A/B vs installed release: + +| format | session start | now | total speedup | +|---|---|---|---| +| BrainVision | ~1300 µs | ~225–232 µs | **~5.7×** | +| FIF | ~1700 µs | ~291 µs | **~5.8×** | +| EDF | ~2080 µs | ~475 µs | **~4.4×** | +| BDF | ~2150 µs | ~561 µs | **~3.8×** | + +Preloaded access is unchanged at ~78–80 µs/window (= numpy slice floor). + +### 10.4 The batched-read ceiling (adopt today without touching MNE) + +Reusing one output buffer and calling internal read machinery directly +(`raw._read_segment(..., data_buffer=buf, sel=arange)`) vs public `get_data`, +measured in a single run (within-run ratios are meaningful): + +| format | public | internal | buffered | +|---|---|---|---| +| EDF | 420 µs | 417 µs | **338 µs** | +| BDF | 553 µs | 541 µs | **501 µs** | +| BV | 538 µs | 208 µs | **199 µs** | +| FIF | 537 µs | 296 µs | **266 µs** | + +This is the case for an upstream batched API (`raw.get_windows(starts, +width)`): resolve picks/cals once, fill a preallocated buffer per window — +PyTorch's own data-loading tutorial measures the same pattern (`__getitems__`) +at ~2.9× marginal throughput. + +### 10.5 Where the remaining time goes (honest accounting) + +Per 2 s window at 128 ch (~512 KB payload): +- preloaded path is already AT the numpy floor; +- non-preloaded paths now spend their budget on genuine I/O + format decode: + EDF overreads whole records (format-inherent ~1.5×), BDF decodes 3 bytes per + sample by design, BV/FIF do one contiguous read + one cast/scale pass. +- Next levers, in order: persistent mmap/handle per Raw (saves ~30–60 µs of + syscall glue per call), upstreaming the batched API, FIF tag-layer dispatch + caching (`_compare_version` etc. seen in early profiles). + + +## 11. Session 4: BIDS format completeness + batched reader + memmap path + +### 11.1 BIDS-accepted formats — full coverage check + +Per the current BIDS spec (EEG section): EEG MUST be EDF, BrainVision, +EEGLAB (.set/.fdt), or Biosemi (.bdf); iEEG additionally allows MEF and NWB. +Benchmark status of every accepted format: + +| BIDS format | status in this benchmark | +|---|---| +| EDF / EDF+ | done (sessions 1–3) | +| BrainVision | done (sessions 1–3) | +| Biosemi BDF | done (session 2, incl. numba int24 decoder) | +| **EEGLAB .set/.fdt** | **added (this session)** — writer via scipy.savemat + f32 .fdt | +| **NWB** | **added (this session)** — written with pynwb; read directly (MNE has no NWB reader) | +| MEF (iEEG) | deferred: pymef is installed but its write API requires manual segment-metadata assembly; reading benchmarks need a real corpus | + +### 11.2 New-format results (BIG file: 128 ch, 1800 s @ 512 Hz) + +| backend | full load | windows (µs/win) | +|---|---|---| +| mne_set (EEGLAB) | 511 ms (~924 MB/s) | 986 | +| **nwb (pynwb/h5py direct)** | **136 ms (~3.5 GB/s)** | **100** | + +NWB's time-major storage makes a window one contiguous row-block — near-optimal +for DL sampling without any chunk tuning. EEGLAB performs like BrainVision +(multiplexed binary + header parse). Practical guidance for BIDS corpora: +EDF/BrainVision are fine after our reader fixes; if you control the export +format for training-only copies, time-major stores (NWB-style or +window-chunked HDF5) give the best random-window behavior. + +### 11.3 Speed patches added (session 4) + +8. `BaseRaw._get_windows(starts, width, *, out=None, sel=None)` (internal API): + resolves channel selection once and fills an optional reusable + ``(n_win, n_ch, width)`` buffer — zero per-window allocations. Verified + bit-identical to per-window `get_data`; fork-safe. +9. `_read_segments_file` (generic binary reader used by BrainVision et al.): + persistent `np.memmap` cached on the Raw's extras, keyed by PID so forked + DataLoader workers create their own mapping. Removes per-call + open/seek/read syscalls. BV windows: ~230 → **~189 µs** public-API. + +Batched vs per-call (300 × 2 s windows, identical per-window reductions, +buffer reused across repetitions): EDF 1.27×, BDF 1.15×, BV 1.14×, +FIF 1.05× — plus the elimination of ~315 MB/epoch of allocation churn. + +All suites green (246 io tests), outputs bit-identical to release on small set +and to pristine main on the BIG set (worst diff 0.0). + +## 12. Session 5: multi-agent round — FIF mmap, searchsorted, stim fast path + +Deployed per workstream (subagent infra was flaky; W1/W2 landed via direct execution): + +| Workstream | Outcome | +|---|---| +| W1 FIF mmap | **LANDED**: `mne/_fiff/_mmap_cache.py` (PID-keyed, mtime+size validated) + partial-tag byte-offset reads in `Raw._read_segment_file`; gzip/file-like/odd-tag fallback to legacy loop. Gates: test_raw_fiff 37✔, equivalence ✔, A/B below | +| W2 EDF/BDF numba | **LANDED** (by agent): fused decode kernel `mne/io/edf/_edf_numba.py`, `fastmath=False` required for bit-exactness; verified numpy fallback under `MNE_USE_NUMBA=false`; also fixed NumPy-2 uint32 overflow in `_bdf_numba.decode_int24` no-numba path + missing `has_numba` guard | +| opencode idea #8 | **LANDED**: `searchsorted` on sorted bounds replaces O(n_ent) mask per call (`fiff/raw.py`); property-tested vs boolean mask on 2000 random span cases | +| opencode idea #9 | **LANDED**: EDF/BDF fast path no longer abandons when uniform stim channels are among picks — replicates legacy post-calibration truncating bitmask per row; synthetic EDF+STATUS channel verified md5-identical vs release for all/stim-only/mixed picks | +| W5 stores | Chunk-cache tuning ruled out (±5%); zarr v3 ≈10× slower locally with identical chunks → remote/parallel-only guidance; h5py driver="core" −26% on slabs at ~1 GiB RAM cost | +| W7 writes | "FIF save fixed overhead" root-caused to lazy `from mne_bids import BIDSPath` inside `_check_fname` (~105 ms first save/process; warm saves ~1220 MB/s). pybv local install has a 10× write regression + in-place caller mutation bug (upstream is fine) | +| W3 glue / W6 MEF / W4 edfio | pending (infra flakiness / next round) | + +### Session-5 cumulative A/B (BIG file, 2 s windows, public get_data, best-of interleaved) + +| format | pristine main | patched tree | speedup | +|---|---|---|---| +| EDF | 1472 µs | **290 µs** | **5.1×** | +| BDF | 2166 µs | **405 µs** | **5.4×** | +| BrainVision | 927 µs | **189 µs** | **4.9×** | +| FIF | 1036 µs | **239 µs** | **4.3×** | + +All outputs bit-exact (small-set md5 identical vs release; BIG-set worst diff 0.0); +246 io/fiff/pick tests green. + +### External-auditor ideas adopted vs queued +Adopted: searchsorted bounds (#8), stim-in-fast-path (#9), codex's FIF mmap prototype. +Queued: EDF mixed-sfreq partial fast path (decode uniform EEG blocks, interp stim +separately — needs careful interp semantics), float32 output read path, +open_meta header-parse cost, batched-API upstream proposal, edfio engine backend. + +## 13. Session 6: `engine="edfio"` backend landed + +`read_raw_edf(..., engine="edfio")` parses via the optional edfio package into a +preloaded Raw. Minimal scope by design: uniform sfreq only, all channels EEG, +no meas_date, volts output using the native unit mapping. Decode is two fused +passes over one stacked int16 buffer; outputs match the native engine within +1 ulp (max |diff| 2.7e-20 on our fixtures; a first fusion attempt folded the +unit multiplier twice — caught immediately against edfio's own `.data`). + +Full-load BIG file: native 453 ms vs edfio engine 368 ms (~1.2× incl. MNE +wrapper + output copy; raw parser delta is larger). Windows on the preloaded +result sit at the ~78 µs numpy floor like every other preloaded path. + +New test `test_engine_edfio` compares engines on an exported fixture. + +## 14. Session 7: syscall-glue removal + edfio PR verification + +- `mne/io/edf/_open.py`: PID-keyed LRU persistent read handles (NoClose wrapper, + seek(0)-on-reuse preserving fresh-open semantics, LRU cap 8). Removes + open/close per `get_data` on EDF/BDF/GDF paths. Gates: all suites green, + equivalence exact; windows now **EDF 291-295 µs / BDF 402-424 µs** vs pristine. +- edfio PR #114: upstream Actions require maintainer approval for fork PRs + (0 check runs). Replicated their pinned toolchain locally instead: + ruff==0.9.10 clean after refactor (RUF059), ruff-format applied, + mypy==1.15.0's 6 unused-ignore errors confirmed pre-existing on main, + pytest 1022 passed @ 100% coverage; fix commit pushed + verification comment posted. +- Documented: `_get_windows(..., out=float32_buffer)` halves output memory + traffic for DL collation (dtype of provided buffer is honored by all fused + write paths). + +## 15. Session 8: depth round — direct-to-output kernel, byteorder guard + +- `decode_window_into(dst, view, cal, off, gain, s0, w)` in `_edf_numba.py`: + decodes straight into the caller's output slice (arbitrary strides), with + per-element sample addressing (`s0`/`w`) so edge records at window + boundaries need no temporaries. +- EDF fast path now decodes **directly into `data[:, pos:pos+w]`** when safe + (no projector/comp, unit cals): removes the per-chunk intermediate and its + copy. numpy fallback mirrors via temp+single copy. +- **Byteorder guard**: real-world EDF is big-endian and numba only types + native byteorder — non-native chunks are byteswapped to a native copy per + chunk before the kernel (synthetic LE files had masked this). +- Kernel gated strictly on `write_direct`; projector/compensation paths keep + the exact legacy route (a first draft double-applied calibration there — + caught by `test_bdf_data`, fixed by the gate). +- Attempted mixed-sfreq partial decode inside the fast path; reverted to the + strict uniform-file gate: ragged record layouts break the zero-copy reshape, + and replicating the legacy slow-row packing quirks bit-exactly needs its own + round (documented, not silently approximated). + +### Round numbers (BIG file, best-of interleaved vs pristine main) + +| format | window µs | speedup | full load | +|---|---|---|---| +| EDF | 301 | 5.7× | 1.1× | +| BDF | 406 | 6.1× | **2.0×** | +| BrainVision | 200 | 5.9× | 1.3–1.5× | +| FIF | 257 | 5.1× | 1.4× | + +Cumulative vs where this effort started: **~6.5–7.5× per-window**, full loads +1.1–2×, preloads at the numpy floor. The residual per-window cost is now +genuinely I/O + format decode + one fused pass; further order-of-magnitude +gains require changing what is stored (window-shaped HDF5/NWB-style layouts, +§11.2) rather than how MNE reads classic formats. + +## 16. Session 9: persistent memmap caches = the lazy/fast flag + +`load_data(memmap=path)` existed but (a) always re-decoded even when the cache +file was valid and (b) `BaseRaw.__del__` deleted the file on GC. Both fixed: + +- `_preload_data`: when a memmap cache exists with the exact expected size, it + is mmap'd read-write **and decoding is skipped entirely**. +- `__del__` no longer removes memmap files — callers own cache lifetime + (behavior change; documented in changelog fragment). + +Fresh-process numbers (128 ch × 1800 s @ 512 Hz, 944 MB float64): + +| operation | before | now | +|---|---|---| +| open + first access | ~450 ms decode | **11–21 ms** | +| public get_data windows | 200–406 µs | **~84 µs** | + +This is the recommended DL pattern until format-native lazy backends land: + +```python +raw = mne.io.read_raw_edf("huge.edf", preload=False) +raw.load_data(memmap="cache.f64") # first run decodes; every later run mmaps +``` + +On Python 3.15's opt-in lazy imports (PEP 810): import-time costs we measured +(e.g., the 105 ms mne_bids cascade on first save) should shrink automatically; +no hand-rolled laziness added here. A compiled C++/Rust kernel was evaluated +and rejected for now: our numba path already runs at memory bandwidth, so a +second native toolchain would add packaging burden without measurable gain. diff --git a/benchmarks/io_dl/ab_test.py b/benchmarks/io_dl/ab_test.py new file mode 100644 index 00000000000..eb6f0a4763f --- /dev/null +++ b/benchmarks/io_dl/ab_test.py @@ -0,0 +1,101 @@ +"""Clean interleaved A/B with explicit environment control. + +Arm INSTALLED: site-packages mne (unpatched release copy) +Arm TREE: working tree mne (patched) + +Also reports component breakdown inside each process. +""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +TREE = "/Users/bruaristimunha/Projects/libraries/mne_python/mne_python_more_io_speed" +HERE = Path(__file__).parent + +CHILD = r''' +import gc, json, os, sys, time +import numpy as np +which = sys.argv[1] +# `python -c` puts cwd ('') first on sys.path; remove anything that could +# resolve to the working tree unless this is the TREE arm +TREE = "{tree}" +if which == "tree": + if TREE not in sys.path: + sys.path.insert(0, TREE) +else: + sys.path[:] = [p for p in sys.path if p not in ("", ".", TREE)] + os.environ.pop("PYTHONPATH", None) +import mne +mne.set_log_level("ERROR") +fmt, pl = sys.argv[2], sys.argv[3] == "1" +READERS = { + "edf": lambda p_: mne.io.read_raw_edf("{here}/data/bench.edf", preload=p_), + "brainvision": lambda p_: mne.io.read_raw_brainvision("{here}/data/bench.vhdr", preload=p_), + "fif": lambda p_: mne.io.read_raw_fif("{here}/data/bench_raw.fif", preload=p_), +} +raw = READERS[fmt](pl) +rng = np.random.default_rng(0) +starts = rng.integers(0, raw.n_times - 513, size=4000).astype(int) +stops = starts + 512 + +def bench(fn, warmup=300): + for s in starts[:warmup]: + fn(int(s)) + ts = [] + gc.disable() + for s, e in zip(starts[warmup:], stops[warmup:]): + t0 = time.perf_counter_ns() + fn(s) + ts.append(time.perf_counter_ns() - t0) + gc.enable() + arr = np.asarray(ts) / 1e3 + return dict(med=float(np.median(arr)), p10=float(np.percentile(arr, 10))) + +out = dict(which=which, file=mne.__file__) +out["get_data_only"] = bench(lambda s: raw.get_data(start=s, stop=s + 512)) +out["get_data_sum"] = bench(lambda s: raw.get_data(start=s, stop=s + 512).sum()) +mm = np.memmap("{here}/data/bench.bin", dtype=" 1 else 5 + acc = {"INSTALLED": [], "TREE": []} + for i in range(rounds): + row = {} + for key, which in (("INSTALLED", "installed"), ("TREE", "tree")): + r = run(which) + acc[key].append(r) + row[key] = r + src = "TREE" if "more_io_speed" in r["file"] else "site-packages" + print(f"round {i+1} {key:<10} src={src:<13} " + f"data={r['get_data_only']['med']:6.1f}us " + f"data+sum={r['get_data_sum']['med']:6.1f}us " + f"floor={r['memmap_copy']['med']:5.1f}us") + print("\n=== best-of medians ===") + for key in ("INSTALLED", "TREE"): + best = min(acc[key], key=lambda r: r["get_data_only"]["med"]) + g, gs, f = best["get_data_only"]["med"], best["get_data_sum"]["med"], best["memmap_copy"]["med"] + print(f"{key:<10} get_data={g:6.1f}us get_data+sum={gs:6.1f}us memmap-floor={f:5.1f}us") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/backends-data_big-20260824-161621.json b/benchmarks/io_dl/backends-data_big-20260824-161621.json new file mode 100644 index 00000000000..690d762ef5a --- /dev/null +++ b/benchmarks/io_dl/backends-data_big-20260824-161621.json @@ -0,0 +1,74 @@ +{ + "floor_memmap": { + "full_ms": 151.852959, + "win_ms": 42.519625, + "mbps": 3107.3428078540105, + "us_per_win": 141.73208333333332 + }, + "mne_edf": { + "full_ms": 453.770375, + "win_ms": 575.13625, + "mbps": 1039.8633890544309, + "us_per_win": 1917.1208333333334 + }, + "mne_bdf": { + "full_ms": 994.840167, + "win_ms": 954.506333, + "mbps": 474.3065425503673, + "us_per_win": 3181.6877766666667 + }, + "mne_bv": { + "full_ms": 842.084667, + "win_ms": 479.360625, + "mbps": 560.3465049197957, + "us_per_win": 1597.86875 + }, + "mne_fif": { + "full_ms": 534.52, + "win_ms": 418.168083, + "mbps": 882.7718326723041, + "us_per_win": 1393.89361 + }, + "edfio_edf_eager": { + "full_ms": 255.659167, + "win_ms": null, + "mbps": 922.9578065550061, + "us_per_win": null + }, + "npz": { + "full_ms": 206.126375, + "win_ms": null, + "mbps": 2289.17429901923, + "us_per_win": null + }, + "h5_full": { + "full_ms": 153.334458, + "win_ms": 16.073, + "mbps": 3077.320037222162, + "us_per_win": 53.57666666666667 + }, + "h5_t10s": { + "full_ms": 247.235, + "win_ms": 217.567209, + "mbps": 1908.5453111412219, + "us_per_win": 725.22403 + }, + "h5_win": { + "full_ms": 213.12475, + "win_ms": 38.790833, + "mbps": 2214.004708509922, + "us_per_win": 129.30277666666666 + }, + "zarr_t10s": { + "full_ms": 4688.113042, + "win_ms": 9798.711709, + "mbps": 100.65013274481532, + "us_per_win": 32662.37236333333 + }, + "zarr_win": { + "full_ms": 765.368083, + "win_ms": 760.479458, + "mbps": 616.5127740243122, + "us_per_win": 2534.9315266666667 + } +} \ No newline at end of file diff --git a/benchmarks/io_dl/backends-data_big-20260824-161927.json b/benchmarks/io_dl/backends-data_big-20260824-161927.json new file mode 100644 index 00000000000..4e90c4489de --- /dev/null +++ b/benchmarks/io_dl/backends-data_big-20260824-161927.json @@ -0,0 +1,92 @@ +{ + "floor_memmap": { + "full_ms": 155.608959, + "win_ms": 42.666833, + "mbps": 3032.339545437098, + "us_per_win": 142.22277666666668 + }, + "mne_edf": { + "full_ms": 497.573125, + "win_ms": 608.397875, + "mbps": 948.3213145806458, + "us_per_win": 2027.9929166666666 + }, + "mne_edf_preloaded": { + "full_ms": null, + "win_ms": 683.8465, + "mbps": null, + "us_per_win": 2279.4883333333332 + }, + "mne_bdf": { + "full_ms": 1003.320542, + "win_ms": 852.924292, + "mbps": 470.2975572087908, + "us_per_win": 2843.0809733333335 + }, + "mne_bv": { + "full_ms": 814.662583, + "win_ms": 389.382458, + "mbps": 579.2081407033272, + "us_per_win": 1297.9415266666665 + }, + "mne_fif": { + "full_ms": 551.795, + "win_ms": 508.93025, + "mbps": 855.1349686024702, + "us_per_win": 1696.4341666666667 + }, + "mne_fif_preloaded": { + "full_ms": null, + "win_ms": 704.338583, + "mbps": null, + "us_per_win": 2347.7952766666667 + }, + "edfio_edf_eager": { + "full_ms": 244.071083, + "win_ms": null, + "mbps": 966.7782889298688, + "us_per_win": null + }, + "edfio_edf_lazy_wins": { + "full_ms": null, + "win_ms": 347.37475, + "mbps": null, + "us_per_win": 1157.9158333333332 + }, + "npz": { + "full_ms": 228.083208, + "win_ms": null, + "mbps": 2068.802890566148, + "us_per_win": null + }, + "h5_full": { + "full_ms": 158.631, + "win_ms": 28.184583, + "mbps": 2974.571174612781, + "us_per_win": 93.94861 + }, + "h5_t10s": { + "full_ms": 263.300583, + "win_ms": 279.080709, + "mbps": 1792.093259436497, + "us_per_win": 930.2690300000002 + }, + "h5_win": { + "full_ms": 217.69475, + "win_ms": 53.925958, + "mbps": 2167.5267777472814, + "us_per_win": 179.75319333333334 + }, + "zarr_t10s": { + "full_ms": 5098.628375, + "win_ms": 9509.274958, + "mbps": 92.54630172962938, + "us_per_win": 31697.583193333336 + }, + "zarr_win": { + "full_ms": 632.534083, + "win_ms": 503.473208, + "mbps": 745.9822524693898, + "us_per_win": 1678.2440266666665 + } +} \ No newline at end of file diff --git a/benchmarks/io_dl/backends-data_big-20260824-174932.json b/benchmarks/io_dl/backends-data_big-20260824-174932.json new file mode 100644 index 00000000000..bdb7eea0a4d --- /dev/null +++ b/benchmarks/io_dl/backends-data_big-20260824-174932.json @@ -0,0 +1,104 @@ +{ + "floor_memmap": { + "full_ms": 123.97925, + "win_ms": 34.394292, + "mbps": 3805.953012298429, + "us_per_win": 114.64764000000001 + }, + "mne_edf": { + "full_ms": 464.572083, + "win_ms": 463.496, + "mbps": 1015.6856541033267, + "us_per_win": 1544.9866666666667 + }, + "mne_edf_preloaded": { + "full_ms": null, + "win_ms": 561.771959, + "mbps": null, + "us_per_win": 1872.5731966666667 + }, + "mne_bdf": { + "full_ms": 844.267375, + "win_ms": 643.647791, + "mbps": 558.8978254667248, + "us_per_win": 2145.4926366666664 + }, + "mne_bv": { + "full_ms": 528.579208, + "win_ms": 298.229583, + "mbps": 892.6934560770691, + "us_per_win": 994.0986099999999 + }, + "mne_fif": { + "full_ms": 474.158625, + "win_ms": 327.07125, + "mbps": 995.150515294328, + "us_per_win": 1090.2375 + }, + "mne_fif_preloaded": { + "full_ms": null, + "win_ms": 584.748375, + "mbps": null, + "us_per_win": 1949.16125 + }, + "mne_set": { + "full_ms": 510.646625, + "win_ms": 295.776875, + "mbps": 924.0425313689285, + "us_per_win": 985.9229166666667 + }, + "nwb": { + "full_ms": 135.857542, + "win_ms": 29.9985, + "mbps": 3473.191057733107, + "us_per_win": 99.995 + }, + "edfio_edf_eager": { + "full_ms": 171.696875, + "win_ms": null, + "mbps": null, + "us_per_win": null + }, + "edfio_edf_lazy_wins": { + "full_ms": null, + "win_ms": 314.790875, + "mbps": null, + "us_per_win": 1049.3029166666668 + }, + "npz": { + "full_ms": 152.86025, + "win_ms": null, + "mbps": 3086.8665987396985, + "us_per_win": null + }, + "h5_full": { + "full_ms": 104.1435, + "win_ms": 15.96875, + "mbps": 4530.855982370479, + "us_per_win": 53.229166666666664 + }, + "h5_t10s": { + "full_ms": 177.1665, + "win_ms": 160.295667, + "mbps": 2663.3658169010505, + "us_per_win": 534.31889 + }, + "h5_win": { + "full_ms": 148.633666, + "win_ms": 27.039625, + "mbps": 3174.6455072971153, + "us_per_win": 90.13208333333333 + }, + "zarr_t10s": { + "full_ms": 4115.615792, + "win_ms": 8396.288958, + "mbps": 114.65093532715262, + "us_per_win": 27987.629859999997 + }, + "zarr_win": { + "full_ms": 540.063791, + "win_ms": 373.51525, + "mbps": 873.710120662394, + "us_per_win": 1245.0508333333332 + } +} \ No newline at end of file diff --git a/benchmarks/io_dl/bench_backends.py b/benchmarks/io_dl/bench_backends.py new file mode 100644 index 00000000000..b593e28861c --- /dev/null +++ b/benchmarks/io_dl/bench_backends.py @@ -0,0 +1,313 @@ +"""Compare storage backends / provider libraries on identical signal data. + +Works on any preset dir (data/ or data_big/). Sections: + +FULL : end-to-end load of the entire recording through each backend +WINDOWS: 300 shuffled windows (win = 2 s) -- the DL training pattern + +Backends: + mne_ : MNE readers (preload True/False) + edfio : edfio.read_edf eager + .data; lazy variant for windows + h5_ : h5py datasets (t10s / win / full chunk layouts) + zarr_ : zarr arrays (t10s / win) + npz : numpy .npz (whole-array) + floor_memmap : raw float32 .bin via np.memmap + +Usage: + python benchmarks/io_dl/bench_backends.py [--dir data] [--repeats 3] +""" + +import argparse +import gc +import json +import time +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).parent + +TREE_ROOT = HERE.resolve().parents[1].parent +import sys # noqa: E402 + +sys.path.insert(0, str(TREE_ROOT)) + +import mne # noqa: E402 + +mne.set_log_level("ERROR") + + +def timed(fn, repeats=3, warmup=1): + for _ in range(warmup): + fn() + out = [] + gc.collect() + gc.disable() + try: + for _ in range(repeats): + t0 = time.perf_counter_ns() + fn() + out.append((time.perf_counter_ns() - t0) / 1e6) + finally: + gc.enable() + return float(np.median(out)) + + +def rand_starts(rng, n_times, win, n): + return rng.integers(0, n_times - win - 1, size=n).astype(int) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--dir", default="data_big") + ap.add_argument("--repeats", type=int, default=3) + ap.add_argument("--n-win", type=int, default=300) + ap.add_argument("--seed", type=int, default=99) + args = ap.parse_args() + + d = HERE / args.dir + meta = json.loads((d / "meta.json").read_text()) + n_ch, sfreq, n_times = meta["n_ch"], meta["sfreq"], meta["n_times"] + base = d.name.replace("data", "bench") + win = int(2 * sfreq) + rng = np.random.default_rng(args.seed) + starts = rand_starts(rng, n_times, win, args.n_win) + + results = {} + print(f"\n=== {d.name}: {n_ch} ch, {sfreq:.0f} Hz, {meta['dur']:.0f} s | " + f"win={win} ({win/sfreq:.1f} s) x {args.n_win} ===") + + # ---------------- FULL loads ------------------------------------------- + print(f"\n{'backend':<24}{'full load':>12} {'MB/s':>8} | " + f"{'300 windows':>12} {'us/win':>8}") + rows = [] + + def record(name, full_fn=None, win_fn=None, full_bytes=None): + f_ms = timed(full_fn, args.repeats) if full_fn else None + w_ms = timed(win_fn, max(args.repeats, 3)) if win_fn else None + mbps = full_bytes / 1e6 / (f_ms / 1000) if (full_bytes and f_ms) else None + usw = w_ms * 1000 / args.n_win if w_ms else None + results[name] = dict( + full_ms=f_ms, win_ms=w_ms, mbps=mbps, us_per_win=usw, + ) + print(f"{name:<24}" + f"{f_ms if f_ms else float('nan'):>9.1f}ms " + f"{mbps if mbps else float('nan'):>7.0f} | " + f"{w_ms if w_ms else float('nan'):>11.1f}ms " + f"{usw if usw else float('nan'):>7.1f}") + + # floors ----------------------------------------------------------------- + binp = d / f"{base}.bin" + shape = (n_ch, n_times) + nbytes_f32 = n_ch * n_times * 4 + + def mm_full(): + a = np.memmap(binp, "12} {'--':>8} | {w_ms:>11.1f}ms {usw:>7.1f}") + + # EEGLAB (.set/.fdt) ------------------------------------------------------ + set_path = d / f"{base}.set" + if set_path.exists(): + from mne.io import read_raw_eeglab + + def set_full(): + r = read_raw_eeglab(set_path, preload=True) + s = float(r.get_data().sum()) + del r + return s + + def set_wins(): + r = read_raw_eeglab(set_path, preload=False) + acc = 0.0 + for s in starts: + acc += float(r.get_data(start=int(s), stop=int(s) + win).sum()) + del r + return acc + + record("mne_set", set_full, set_wins, nbytes_f32) + + # NWB (time-major HDF5; window = contiguous row block) -------------------- + nwb_path = d / f"{base}.nwb" + if nwb_path.exists(): + + def nwb_full(): + from pynwb import NWBHDF5IO + + with NWBHDF5IO(nwb_path, "r") as io: + a = np.asarray(io.read().acquisition["ElectricalSeries"].data[:]) + s = float(a.astype(np.float64).sum()) + return s + + def nwb_wins(): + from pynwb import NWBHDF5IO + + io = NWBHDF5IO(nwb_path, "r") + arr = io.read().acquisition["ElectricalSeries"].data + acc = 0.0 + for s in starts: + w = np.asarray(arr[s : s + win]).T.astype(np.float64) + acc += float(w.sum()) + io.close() + return acc + + record("nwb", nwb_full, nwb_wins, nbytes_f32) + + # edfio ------------------------------------------------------------------ + try: + from edfio import read_bdf, read_edf + + def edfio_full_eager(): + e = read_edf(d / f"{base}.edf", lazy_load_data=False) + s = sum(float(sig.data.sum()) for sig in e.signals) + return s + + def edfio_lazy_wins(): + from edfio.edf_signal import _calculate_gain_and_offset + + e = read_edf(d / f"{base}.edf", lazy_load_data=True) + acc = 0.0 + for s in starts: + t0s, t1s = s / sfreq, (s + win) / sfreq + tot = 0.0 + for sg in e.signals: + dg = sg.get_digital_slice(t0s, t1s) + gain, offset = _calculate_gain_and_offset( + sg.digital_min, sg.digital_max, + sg.physical_min, sg.physical_max, + ) + tot += float(((dg + offset) * gain).sum()) + acc += tot + return acc + + record("edfio_edf_eager", edfio_full_eager, None, + meta["sizes"].get("edf", 0)) + record("edfio_edf_lazy_wins", None, edfio_lazy_wins, None) + except Exception as exc: # noqa: BLE001 + print(f"edfio section skipped: {exc!r}") + + # npz ---------------------------------------------------------------------- + npzp = d / f"{base}.npz" + if npzp.exists(): + + def npz_full(): + a = np.load(npzp)["data"] + s = float(a.astype(np.float64).sum()) + del a + return s + + record("npz", npz_full, None, nbytes_f32) + + # hdf5 layouts -------------------------------------------------------------- + h5p = d / f"{base}.h5" + if h5p.exists(): + import h5py + + with h5py.File(h5p, "r") as f: + layouts = list(f.keys()) + for lay in layouts: + + def h5_full(lay=lay): + with h5py.File(h5p, "r") as f: + a = f[lay][:] + s = float(a.astype(np.float64).sum()) + return s + + def h5_wins(lay=lay): + f = h5py.File(h5p, "r")[lay] + acc = 0.0 + for s in starts: + acc += float(f[:, s : s + win].astype(np.float64).sum()) + f.file.close() + return acc + + record(f"h5_{lay}", h5_full, h5_wins, nbytes_f32) + + # zarr layouts ---------------------------------------------------------------- + for lay in ("t10s", "win"): + zp = d / f"{base}_{lay}.zarr" + if not zp.exists(): + continue + import zarr + + arr = zarr.open_array(store=str(zp), mode="r") + + def z_full(arr=arr): + a = arr[:] + s = float(a.astype(np.float64).sum()) + del a + return s + + def z_wins(arr=arr): + acc = 0.0 + for s in starts: + acc += float(arr[:, s : s + win].astype(np.float64).sum()) + return acc + + record(f"zarr_{lay}", z_full, z_wins, nbytes_f32) + + stamp = time.strftime("%Y%m%d-%H%M%S") + out = HERE / f"backends-{args.dir}-{stamp}.json" + out.write_text(json.dumps(results, indent=2)) + print(f"\nsaved -> {out}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/bench_io.py b/benchmarks/io_dl/bench_io.py new file mode 100644 index 00000000000..c6564c58d05 --- /dev/null +++ b/benchmarks/io_dl/bench_io.py @@ -0,0 +1,284 @@ +"""Benchmark MNE-Python raw IO under deep-learning access patterns. + +Run generate_data.py first. Then: + + python benchmarks/io_dl/bench_io.py # full suite + python benchmarks/io_dl/bench_io.py --quick # smoke test + python benchmarks/io_dl/bench_io.py --curve # fixed-overhead curve only + +Scenarios +--------- +open_meta : read_raw_X(preload=False) -- header/metadata parse only +full_load : read_raw_X(preload=True) -- end-to-end load +seq_1s : sequential 1 s get_data() over whole file (streaming/eval) +rand_windows : 300 shuffled 2 s get_data() calls, preload=False + (the canonical DL training-loop pattern) +preloaded_windows: same windows on an already-preloaded raw (access-only cost) +floors : numpy memmap / np.fromfile equivalents on identical bytes + +Every number is median of N repetitions (min shown too), GC disabled. +""" + +import argparse +import gc +import json +import platform +import sys +import time +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) + +# Always benchmark THIS working tree, regardless of interpreter resolution +# (script dir / cwd / PYTHONPATH pitfalls). +TREE_ROOT = HERE.resolve().parents[1] +sys.path.insert(0, str(TREE_ROOT)) + +import mne + +mne.set_log_level("ERROR") +from mne.io import read_raw_edf, read_raw_fif, read_raw_brainvision # noqa: E402 + + +# ---------------------------------------------------------------- utilities +def timed(fn, repeats=5, warmup=1): + """Return list of runtimes in ms.""" + for _ in range(warmup): + fn() + out = [] + was_enabled = gc.isenabled() + gc.collect() + gc.disable() + try: + for _ in range(repeats): + t0 = time.perf_counter_ns() + fn() + out.append((time.perf_counter_ns() - t0) / 1e6) + finally: + if was_enabled: + gc.enable() + return np.asarray(out) + + +def fmt(ms): + med = np.median(ms) + mn = ms.min() + return f"{med:9.2f} / {mn:8.2f}" + + +def row(label, unit, val_ms, extra=""): + print(f" {label:<34}{fmt(val_ms)} {unit:<12}{extra}") + + +# ---------------------------------------------------------------- readers +def make_readers(data_dir): + return { + "edf": lambda preload: read_raw_edf( + data_dir / "bench.edf", preload=preload, verbose="ERROR" + ), + "brainvision": lambda preload: read_raw_brainvision( + data_dir / "bench.vhdr", preload=preload, verbose="ERROR" + ), + "fif": lambda preload: read_raw_fif( + data_dir / "bench_raw.fif", preload=preload, verbose="ERROR" + ), + } + + +def rand_starts(rng, n_times, win, n_win): + return rng.integers(0, n_times - win - 1, size=n_win) + + +# ---------------------------------------------------------------- scenarios +def bench_format(name, reader, meta, cfg): + n_times, sfreq = meta["n_times"], meta["sfreq"] + res = {} + + # S0: metadata-only open + res["open_meta"] = timed(lambda: reader(preload=False), cfg["repeats"]) + # S1: full load + res["full_load"] = timed(lambda: reader(preload=True), cfg["repeats"]) + + # Steady-state access patterns: Raw objects are opened ONCE (as DL + # dataset classes do); we time only the data-access loop. + + # S3/S4: random 2 s windows <-- canonical DL training pattern + rng = np.random.default_rng(cfg["seed"]) + starts = [int(s) for s in rand_starts(rng, n_times, cfg["win"], cfg["n_win"])] + + def rwin(raw): + acc = 0.0 + for s in starts: + acc += float(raw.get_data(start=s, stop=s + cfg["win"]).sum()) + return acc + + def seq(raw_nop): + acc = 0.0 + for i in range(int(meta["dur"])): + acc += float( + raw_nop.get_data(start=i * int(sfreq), stop=(i + 1) * int(sfreq)).sum() + ) + return acc + + # S2: sequential streaming over 1 s chunks, no preload + raw_seq = reader(preload=False) + res["seq_1s"] = timed(lambda: seq(raw_seq), max(3, cfg["repeats"] - 2)) + + raw_rand = reader(preload=False) + res["rand_windows"] = timed(lambda: rwin(raw_rand), cfg["repeats"]) + + raw_pre = reader(preload=True) + res["preloaded_windows"] = timed(lambda: rwin(raw_pre), cfg["repeats"]) + return res + + +def bench_floors(meta, cfg): + """Numpy-only baselines on the identical float32 payload.""" + n_ch, n_times = meta["n_ch"], meta["n_times"] + bin_path = HERE / "data" / "bench.bin" + rng = np.random.default_rng(cfg["seed"]) + starts = rand_starts(rng, n_times, cfg["win"], cfg["n_win"]) + shape = (n_ch, n_times) + res = {} + + def mm_windows(): + mm = np.memmap(bin_path, dtype="8} | " + " | ".join(f"{k:>12}" for k in readers)) + rows = [] + for wlen in sizes_s: + win = int(wlen * meta["sfreq"]) + vals = [] + for name, reader in readers.items(): + starts = rand_starts(rng, meta["n_times"], win, cfg["n_win"]) + raw = reader(preload=False) + + def run(): + for s in starts: + raw.get_data(start=int(s), stop=int(s) + win) + + ms_per_call = np.median(timed(run, 3, 1)) / cfg["n_win"] + vals.append(ms_per_call) + rows.append((wlen, vals)) + print(f"{wlen:>7.3f}s | " + " | ".join(f"{v:12.3f}" for v in vals)) + return rows + + +# ---------------------------------------------------------------- main +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--quick", action="store_true") + ap.add_argument("--curve", action="store_true") + ap.add_argument("--repeats", type=int, default=None) + ap.add_argument("--n-win", type=int, default=None) + ap.add_argument("--seed", type=int, default=1234) + args = ap.parse_args() + + cfg = dict( + seed=args.seed, + repeats=args.repeats or (2 if args.quick else 5), + n_win=args.n_win or (50 if args.quick else 300), + win=512, # 2 s @ 256 Hz, typical DL window + ) + + meta = json.loads((HERE / "data" / "meta.json").read_text()) + dur = meta["dur"] + readers = make_readers(HERE / "data") + + print(f"\nmne {mne.__version__} | numpy {np.__version__} | " + f"{platform.machine()} | py {platform.python_version()}") + print(f"mne loaded from: {Path(mne.__file__).parent}") + print(f"data: {meta['n_ch']} ch, {meta['sfreq']:.0f} Hz, {dur:.0f} s | " + f"win={cfg['win']} samples ({cfg['win']/meta['sfreq']:.1f} s), " + f"n_win={cfg['n_win']}") + + hdr = f"{'scenario':<36}{'median/min (ms)':<24}{'unit':<12}" + all_results = {"config": cfg, "meta": {k: v for k, v in meta.items() if k != "ch_names"}} + + if not args.curve: + print(f"\n=== floors (numpy on identical float32 bytes) ===\n{hdr}") + floors = bench_floors(meta, cfg) + for k, v in floors.items(): + row(k, "ms", v) + all_results.update({k: v.tolist() for k, v in floors.items()}) + + for name, reader in readers.items(): + print(f"\n=== format: {name} ===\n{hdr}") + r = bench_format(name, reader, meta, cfg) + for k, v in r.items(): + n_calls = cfg["n_win"] if "windows" in k else (int(dur) if k == "seq_1s" else 1) + extra = f"({v.mean()/max(n_calls,1)*1000:8.1f} us/call)" if n_calls > 1 else "" + row(k, "ms", v, extra) + all_results[f"fmt_{name}"] = {k: v.tolist() for k, v in r.items()} + + print(f"\n=== micro (copies) ===\n{hdr}") + mic = bench_micro(cfg) + for k, v in mic.items(): + row(k, "ms", v) + all_results["micro"] = {k: v.tolist() for k, v in mic.items()} + + bench_curve(readers, meta, cfg) + + stamp = time.strftime("%Y%m%d-%H%M%S") + out = HERE / f"results-{stamp}.json" + out.write_text(json.dumps(all_results, indent=2)) + print(f"\nsaved -> {out}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/check_equivalence.py b/benchmarks/io_dl/check_equivalence.py new file mode 100644 index 00000000000..5a3f78ecb0b --- /dev/null +++ b/benchmarks/io_dl/check_equivalence.py @@ -0,0 +1,76 @@ +"""Compare get_data outputs: working-tree mne vs site-packages mne (oracle). + +Run pair-wise via subprocess so both versions load in isolation. +""" + +import hashlib +import json +import sys + + +def which_mne(): + import os + + if os.environ.get("USE_TREE_MNE") == "1": + sys.path.insert(0, "/Users/bruaristimunha/Projects/libraries/" + "mne_python/mne_python_more_io_speed") + import mne + + mne.set_log_level("ERROR") + return mne, mne.__file__ + + +def digests(): + mne, src = which_mne() + here = __file__.rsplit("/", 1)[0] + readers = { + "edf": lambda: mne.io.read_raw_edf(f"{here}/data/bench.edf", preload=False), + "brainvision": lambda: mne.io.read_raw_brainvision( + f"{here}/data/bench.vhdr", preload=False + ), + "fif": lambda: mne.io.read_raw_fif(f"{here}/data/bench_raw.fif", preload=False), + } + out = {"mne_file": src} + import numpy as np + + for name, reader in readers.items(): + raw_np = reader() + raw_p = reader().load_data() + cases = {} + # full read + d = raw_p.get_data() + cases["full"] = hashlib.md5(d.tobytes()).hexdigest()[:12] + # windows incl. block boundaries + accs = [] + for s0, s1 in [(0, 512), (255, 769), (100000, 100512), (76288, 76800)]: + a = raw_np.get_data(start=s0, stop=s1) + b = raw_p.get_data(start=s0, stop=s1) + accs.append(float((a - b).sum())) + cases[f"win_{s0}"] = hashlib.md5(a.tobytes()).hexdigest()[:12] + cases["win_delta"] = float(sum(abs(x) for x in accs)) + # tmin/tmax path + a = raw_np.get_data(tmin=10.0, tmax=12.0) + b = raw_p.get_data(tmin=10.0, tmax=12.0) + cases["tmin_tmax"] = hashlib.md5(a.tobytes()).hexdigest()[:12] + cases["tmin_tmax_delta"] = float(np.abs(a - b).sum()) + # picks variants + cases["picks_int"] = hashlib.md5( + raw_p.get_data(picks=[1, 5, 7], start=0, stop=256).tobytes() + ).hexdigest()[:12] + cases["picks_str"] = hashlib.md5( + raw_p.get_data(picks=["EEG001", "EEG004"], start=0, stop=256).tobytes() + ).hexdigest()[:12] + cases["picks_slice"] = hashlib.md5( + raw_p.get_data(picks=slice(2, 8), start=0, stop=256).tobytes() + ).hexdigest()[:12] + # negative index semantics + cases["neg_idx"] = hashlib.md5( + raw_p.get_data(picks=[-1], start=0, stop=64).tobytes() + ).hexdigest()[:12] + out[name] = cases + return out + + +if __name__ == "__main__": + res = digests() + print(json.dumps(res, indent=2)) diff --git a/benchmarks/io_dl/check_equivalence_big.py b/benchmarks/io_dl/check_equivalence_big.py new file mode 100644 index 00000000000..70f69c3e364 --- /dev/null +++ b/benchmarks/io_dl/check_equivalence_big.py @@ -0,0 +1,79 @@ +"""BIG-set equivalence: tree vs installed release for EDF/BDF windows. + +The tree's vectorized EDF path folds ((d*cal)+off)*gain into +d*(cal*gain) + off*gain, which differs from the legacy op order by at most +a few double-precision ulps. We therefore assert max|diff| < 1e-9 uV on +identical inputs rather than bit equality. Full loads use the legacy loop +on both arms and must remain bit identical. +""" + +import json +import os +import subprocess +import sys + +import numpy as np + +ROOT = "/Users/bruaristimunha/Projects/libraries/mne_python/mne_python_more_io_speed" + +CODE = r""" +import sys, json +import numpy as np +if sys.argv[1] == "tree": + sys.path.insert(0, "{root}") +else: + sys.path = [p for p in sys.path if p not in ("", ".")] +import mne; mne.set_log_level("ERROR") +from mne.io import read_raw_edf, read_raw_bdf +rng = np.random.default_rng(5) +out = {} +for name, rd in [("edf", read_raw_edf), ("bdf", read_raw_bdf)]: + f = "{root}/benchmarks/io_dl/data_big/bench_big." + name + raw = rd(f, preload=False) + starts = rng.integers(0, raw.n_times - 8193, size=15).astype(int) + wins, picks_ = [], [] + for i, s0 in enumerate(starts): + stop = int(s0) + 1024 + wins.append(raw.get_data(start=int(s0), stop=stop)) + picks_.append(raw.get_data(picks=[3, 17, 55], start=int(s0), + stop=int(s0) + 3000)) + out[name + "_win"] = wins + out[name + "_picks"] = picks_ + rawf = rd(f, preload=True) + d = rawf.get_data() + out[name + "_full_hash"] = float(d.sum()) # robust across arms +print(json.dumps(out, default=lambda x: x.tolist() if hasattr(x, "tolist") else x)) +""".replace("{root}", ROOT) + + +def run(which): + r = subprocess.run([sys.executable, "-c", CODE, which], + capture_output=True, text=True, cwd=ROOT) + if r.returncode: + raise RuntimeError(r.stderr[-500:]) + return json.loads(r.stdout.strip().splitlines()[-1]) + + +def main(): + a = run("installed") + b = run("tree") + ok = True + worst = 0.0 + for key in list(a): + if key.endswith("_hash"): + same = a[key] == b[key] + print(f"{key:<16} sum-equal={same}") + ok &= same + continue + va = np.asarray(a[key]) + vb = np.asarray(b[key]) + d = float(np.abs(va - vb).max()) + worst = max(worst, d) + this_ok = d < 1e-9 + ok &= this_ok + print(f"{key:<16} max|diff|={d:.3e} {'OK' if this_ok else 'FAIL'}") + print(f"\nworst |diff| = {worst:.3e} -> {'ALL OK ✔' if ok else 'FAILED ✗'}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/generate_data.py b/benchmarks/io_dl/generate_data.py new file mode 100644 index 00000000000..5871c1d0a1b --- /dev/null +++ b/benchmarks/io_dl/generate_data.py @@ -0,0 +1,296 @@ +"""Generate synthetic EEG files across many formats and sizes. + +Same underlying signal written to every format. Content identical up to each +format's storage precision. + +Two presets: + default ("small"): 64 ch, 256 Hz, 300 s -> data/ + --preset big: 128 ch, 512 Hz, 1800 s -> data_big/ (~470 MB float32) + +Formats: EDF(int16), BDF(int24), BrainVision(f32), FIF(f32), NPZ(f32), +HDF5 f32 with three chunk layouts, Zarr f32 with two chunk layouts, +raw .bin floor. +""" + +import argparse +import json +import time +from pathlib import Path + +import numpy as np + + +def make_signal(n_ch: int, n_times: int, sfreq: float, seed: int = 42) -> np.ndarray: + """Band-limited synthetic EEG-like signal, microvolt scale.""" + rng = np.random.default_rng(seed) + data = rng.standard_normal((n_ch, n_times)) + # smooth channel-by-channel without apply_along_axis (slow for BIG) + csum = np.cumsum(data, axis=1, dtype=np.float64) + csum[:, 9:] -= csum[:, :-9] + data[:, :8] = csum[:, :8] / np.arange(1, 9) + data[:, 8:] = csum[:, 8:] / 9.0 + del csum + t = np.arange(n_times) / sfreq + data += 10 * np.sin(2 * np.pi * 10 * t) + return (data * 20.0).astype(np.float64) + + +def _edf_signals(sig_cls, data, ch_names, sfreq): + pmin, pmax = float(data.min()) - 1.0, float(data.max()) + 1.0 + return [ + sig_cls( + data=data[i].copy(), + sampling_frequency=sfreq, + physical_range=(pmin, pmax), + label=ch_names[i], + physical_dimension="uV", + ) + for i in range(len(ch_names)) + ] + + +def write_edf(data, ch_names, sfreq, path): + from edfio import Edf, EdfSignal + + Edf(signals=_edf_signals(EdfSignal, data, ch_names, sfreq)).write(path) + + +def write_bdf(data, ch_names, sfreq, path): + from edfio import Bdf, BdfSignal + + Bdf(signals=_edf_signals(BdfSignal, data, ch_names, sfreq)).write(path) + + +def write_bv(data_f32, ch_names, sfreq, out_dir, base): + import pybv + + pybv.write_brainvision( + data=data_f32, + folder_out=str(out_dir), + fname_base=base, + sfreq=int(sfreq), + ch_names=ch_names, + fmt="binary_float32", + overwrite=True, + ) + + +def write_fif(data, ch_names, sfreq, path): + import mne + + mne.set_log_level("ERROR") + info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types="eeg") + raw = mne.io.RawArray(data, info) + raw.save(path, fmt="single", overwrite=True) + + +def write_h5(data_f32, path, layouts): + import h5py + + n_ch, n_times = data_f32.shape + with h5py.File(path, "w") as f: + for name, chunks in layouts.items(): + if chunks is None: + f.create_dataset(name, shape=(n_ch, n_times), dtype="f4") + continue + dset = f.create_dataset(name, shape=(n_ch, n_times), + chunks=chunks, dtype="f4") + dset[:] = data_f32 + + +def write_npz(data_f32, path): + np.savez(path, data=data_f32) + + + +def write_eeglab(data_f32, ch_names, sfreq, out_dir, base): + """Write EEGLAB .set (header) + .fdt (float32 multiplexed).""" + import scipy.io as sio + + n_ch, n_times = data_f32.shape + fdt_path = out_dir / f"{base}.fdt" + np.ascontiguousarray(data_f32.T).tofile(fdt_path) # Fortran order of (ch, t) + chanlocs = np.empty(1, dtype=[("labels", "O"), ("type", "O"), ("unit", "O")]) + locs = np.zeros(n_ch, dtype=chanlocs.dtype) + for i, ch in enumerate(ch_names): + locs[i] = (ch, "EEG", "uV") + eeg = { + "nbchan": float(n_ch), + "pnts": float(n_times), + "trials": 1.0, + "srate": float(sfreq), + "xmin": 0.0, + "xmax": (n_times - 1) / sfreq, + "data": str(fdt_path.name), + "ref": "n/a", + "chanlocs": locs.reshape(1, -1), + "chaninfo": {"nodatchans": {}}, + "event": np.empty((0, 0), dtype=object), + "setname": base, + } + sio.savemat(out_dir / f"{base}.set", {"EEG": eeg}, appendmat=False) + + +def write_nwb(data_f32, ch_names, sfreq, path): + from datetime import datetime, timezone + + from pynwb import NWBHDF5IO, NWBFile + from pynwb.ecephys import ElectricalSeries + + nwbfile = NWBFile( + session_description="benchmark", + identifier="bench", + session_start_time=datetime(2020, 1, 1, tzinfo=timezone.utc), + ) + device = nwbfile.create_device(name="bench_device") + group = nwbfile.create_electrode_group( + name="electrodes", description="all", location="unknown", device=device + ) + for _ in ch_names: + nwbfile.add_electrode(group=group, location="unknown") + region = nwbfile.create_electrode_table_region( + region=list(range(len(ch_names))), description="all channels" + ) + es = ElectricalSeries( + name="ElectricalSeries", + data=data_f32.T, # time-major, as required by NWB + electrodes=region, + starting_time=0.0, + rate=sfreq, + ) + nwbfile.add_acquisition(es) + with NWBHDF5IO(path, "w") as io: + io.write(nwbfile) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n-ch", type=int, default=None) + ap.add_argument("--sfreq", type=float, default=None) + ap.add_argument("--dur", type=float, default=None, help="seconds") + ap.add_argument( + "--preset", + choices=["small", "big"], + default=None, + help="small=64ch/256Hz/300s -> data/, big=128ch/512Hz/1800s -> data_big/", + ) + ap.add_argument("--out-dir", type=Path, default=None) + ap.add_argument( + "--formats", + type=str, + default="bin,edf,bdf,bv,fif,npz,h5,zarr,set,nwb", + help="comma list: bin,edf,bdf,bv,fif,npz,h5,zarr", + ) + args = ap.parse_args() + + presets = { + "small": dict(n_ch=64, sfreq=256.0, dur=300.0, out="data"), + "big": dict(n_ch=128, sfreq=512.0, dur=1800.0, out="data_big"), + } + if args.preset: + p = presets[args.preset] + args.n_ch = args.n_ch or p["n_ch"] + args.sfreq = args.sfreq or p["sfreq"] + args.dur = args.dur or p["dur"] + args.out_dir = args.out_dir or Path(__file__).parent / p["out"] + args.n_ch = args.n_ch or 64 + args.sfreq = args.sfreq or 256.0 + args.dur = args.dur or 300.0 + args.out_dir = args.out_dir or Path(__file__).parent / "data" + + out_dir = args.out_dir + out_dir.mkdir(parents=True, exist_ok=True) + + n_times = int(args.sfreq * args.dur) + ch_names = [f"EEG{i:03d}" for i in range(args.n_ch)] + base = out_dir.name.replace("data", "bench") + + t0 = time.perf_counter() + data = make_signal(args.n_ch, n_times, args.sfreq) + print(f"signal: {args.n_ch} ch x {n_times} samples " + f"({data.nbytes / 1e6:.0f} MB f64) in {time.perf_counter() - t0:.1f} s") + + meta = dict(n_ch=args.n_ch, sfreq=args.sfreq, dur=args.dur, n_times=n_times, + formats=args.formats.split(","), seed=42) + sizes = {} + fmts = args.formats.split(",") + + def reg(name, paths): + sizes[name] = sum(p.stat().st_size for p in paths if Path(p).exists()) + + jobs = [] + if "bin" in fmts: + binp = out_dir / f"{base}.bin" + jobs.append(("bin", lambda p=binp: np.ascontiguousarray( + data.astype(np.float32)).tofile(p), [binp])) + if "edf" in fmts: + p = out_dir / f"{base}.edf" + jobs.append(("edf", lambda p=p: write_edf(data, ch_names, args.sfreq, p), [p])) + if "bdf" in fmts: + p = out_dir / f"{base}.bdf" + jobs.append(("bdf", lambda p=p: write_bdf(data, ch_names, args.sfreq, p), [p])) + if "bv" in fmts: + jobs.append(("bv", lambda: write_bv(data.astype(np.float32), ch_names, + args.sfreq, out_dir, base), + list(out_dir.glob(f"{base}.v*")) + [out_dir / f"{base}.eeg"])) + if "fif" in fmts: + p = out_dir / f"{base}_raw.fif" + jobs.append(("fif", lambda p=p: write_fif(data, ch_names, args.sfreq, p), [p])) + if "npz" in fmts: + p = out_dir / f"{base}.npz" + jobs.append(("npz", lambda p=p: write_npz(data.astype(np.float32), p), [p])) + if "h5" in fmts: + p = out_dir / f"{base}.h5" + layouts = { + "t10s": (1, int(args.sfreq * 10)), # per-channel time slabs + "win": (args.n_ch, 512), # window-shaped chunks + "full": None, # contiguous + } + + def wh(p=p, layouts=layouts): + write_h5(data.astype(np.float32), p, layouts) + + jobs.append(("h5", wh, [p])) + if "set" in fmts: + + def wset(): + write_eeglab( + data.astype(np.float32), ch_names, args.sfreq, out_dir, base + ) + + jobs.append(("set", wset, list(out_dir.glob(f"{base}.s*")) + + [out_dir / f"{base}.fdt"])) + if "nwb" in fmts: + pnwb = out_dir / f"{base}.nwb" + + def wnwb(p=pnwb): + write_nwb(data.astype(np.float32), ch_names, args.sfreq, p) + + jobs.append(("nwb", wnwb, [pnwb])) + if "zarr" in fmts: + layouts = {"t10s": (1, int(args.sfreq * 10)), "win": (args.n_ch, 512)} + + def wz(layouts=layouts): + import zarr + + for name, chunks in layouts.items(): + arr = zarr.open_array( + store=str(out_dir / f"{base}_{name}.zarr"), mode="w", + shape=data.shape, chunks=chunks, dtype="f4", + ) + arr[:] = data.astype(np.float32) + + jobs.append(("zarr", wz, list(out_dir.glob(f"{base}_*.zarr")))) + + for name, fn, paths in jobs: + t0 = time.perf_counter() + fn() + dt = time.perf_counter() - t0 + reg(name, paths) + print(f"wrote {name:<5} {sizes[name] / 1e6:8.1f} MB in {dt:7.1f} s") + + (out_dir / "meta.json").write_text(json.dumps(meta | {"sizes": sizes}, indent=2)) + print(f"done -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/profile_io.py b/benchmarks/io_dl/profile_io.py new file mode 100644 index 00000000000..93cf8e7ea31 --- /dev/null +++ b/benchmarks/io_dl/profile_io.py @@ -0,0 +1,253 @@ +"""Profile MNE IO hot paths and trace every deepcopy on them. + + python benchmarks/io_dl/profile_io.py # all profiles + python benchmarks/io_dl/profile_io.py --only windows + +Sections +-------- +A. cProfile of steady-state random-window access (preloaded and not) +B. cProfile of full_load per format +C. cProfile of raw.save() to fif (DL caching path) +D. deepcopy tracer: counts + cumulative time + call sites of + copy.deepcopy during common DL pipeline operations +""" + +import argparse +import cProfile +import io as _io +import pstats +import sys +import time +from contextlib import contextmanager +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) + +# Always profile THIS working tree. +TREE_ROOT = HERE.resolve().parents[1] +sys.path.insert(0, str(TREE_ROOT)) + +import mne + +mne.set_log_level("ERROR") +from mne.io import read_raw_brainvision, read_raw_edf, read_raw_fif # noqa: E402 + +READERS = { + "edf": lambda pl: read_raw_edf(HERE / "data" / "bench.edf", preload=pl), + "brainvision": lambda pl: read_raw_brainvision( + HERE / "data" / "bench.vhdr", preload=pl + ), + "fif": lambda pl: read_raw_fif(HERE / "data" / "bench_raw.fif", preload=pl), +} + + +def show(prof, n=22, sort="cumulative"): + s = _io.StringIO() + st = pstats.Stats(prof, stream=s) + st.sort_stats(sort).print_stats(n) + txt = s.getvalue() + # trim the header noise + lines = [ln for ln in txt.splitlines()] + start = next(i for i, ln in enumerate(lines) if "ncalls" in ln) - 1 + print("\n".join(lines[start : start + n + 2])) + + +def profile_windows(fmt, preloaded, n_win=200): + rng = np.random.default_rng(7) + meta_wins = 512 + import json + + meta = json.loads((HERE / "data" / "meta.json").read_text()) + starts = rng.integers(0, meta["n_times"] - meta_wins - 1, size=n_win) + raw = READERS[fmt](preloaded) + + def loop(): + for s in starts: + raw.get_data(start=int(s), stop=int(s) + meta_wins) + + label = f"{fmt} windows preloaded={preloaded} ({n_win} calls)" + print(f"\n--- A: cProfile {label} ---") + prof = cProfile.Profile() + prof.enable() + loop() + prof.disable() + show(prof) + + +def profile_full_load(fmt): + print(f"\n--- B: cProfile {fmt} full_load(preload=True) ---") + prof = cProfile.Profile() + prof.enable() + READERS[fmt](True) + prof.disable() + show(prof) + + +def profile_save(): + raw = READERS["edf"](False) + out = HERE / "data" / "_tmp_save.fif" + print("\n--- C: cProfile raw.save() -> fif ---") + prof = cProfile.Profile() + prof.enable() + raw.save(out, overwrite=True) + prof.disable() + out.unlink(missing_ok=True) + show(prof) + + +# ------------------------------------------------------------------ tracer +@contextmanager +def trace_deepcopy(): + """Log every deepcopy call reachable from loaded mne modules. + + Patches copy.deepcopy AND every module attribute that holds a direct + reference (from `from copy import deepcopy` bindings), so call sites + inside mne._fiff / mne.io are all captured. + """ + import copy as _copy + import sys + + records = [] + real = _copy.deepcopy + + def spy(x=None, memo=None, *args, **kwargs): + t0 = time.perf_counter_ns() + out = real(x, memo) if memo is not None else real(x) + dt = (time.perf_counter_ns() - t0) / 1e6 + frame = sys._getframe(1) + fname = frame.f_code.co_filename.split("mne-python/")[-1].split( + "mne_python_more_io_speed/" + )[-1] + site = f"{fname}:{frame.f_lineno}" + records.append((type(x).__name__, site, dt)) + return out + + patched = [] + _copy.deepcopy = spy + for name, mod in list(sys.modules.items()): + if mod is None or name.split(".")[0] != "mne": + continue + try: + attrs = vars(mod) + except TypeError: + continue + hit = False + for k, v in list(attrs.items()): + if v is real: + attrs[k] = spy + hit = True + if hit: + patched.append(mod) + try: + yield records + finally: + _copy.deepcopy = real + for mod in patched: + try: + for k, v in list(vars(mod).items()): + if v is spy: + vars(mod)[k] = real + except TypeError: + pass + + +def report(records, title): + print(f"\n--- D: deepcopy trace: {title} ---") + if not records: + print(" (no deepcopy calls)") + return + agg = {} + for typ, site, dt in records: + k = (typ, site) + n, tot = agg.get(k, (0, 0.0)) + agg[k] = (n + 1, tot + dt) + total = sum(t for *_, t in records) + print(f" total calls={len(records)} total={total:.2f} ms") + for (typ, site), (n, tot) in sorted(agg.items(), key=lambda kv: -kv[1][1])[:15]: + print(f" {tot:8.3f} ms x{n:<4} {typ:<12} {site}") + + +def info_copy_scaling(): + """How does Info copy cost scale with channel count / montage?""" + import copy as _copy + + print("\n--- D2: Info deepcopy scaling (us per copy) ---") + print(f"{'n_ch':>6} {'montage':>8} {'deepcopy(us)':>13} {'copy()(us)':>12}") + for n_ch in (32, 64, 128, 256, 512): + ch_names = [f"EEG{i:03d}" for i in range(n_ch)] + info = mne.create_info(ch_names, 256.0, "eeg") + for with_montage in (False, True): + if with_montage: + mon = mne.channels.make_standard_montage("standard_1005") + keep = list(mon.ch_names)[:n_ch] + info = mne.create_info(keep, 256.0, "eeg") + info.set_montage(mon) + n_rep = 50 + t0 = time.perf_counter_ns() + for _ in range(n_rep): + _copy.deepcopy(info) + t_dc = (time.perf_counter_ns() - t0) / n_rep / 1e3 + t0 = time.perf_counter_ns() + for _ in range(n_rep): + info.copy() + t_cp = (time.perf_counter_ns() - t0) / n_rep / 1e3 + print(f"{n_ch:>6} {str(with_montage):>8} {t_dc:>13.1f} {t_cp:>12.1f}") + + +def trace_pipeline_ops(): + """Deepcopy accounting for ops a DL preprocessing/training loop does.""" + raw_e = READERS["edf"](False) + raw_b = READERS["brainvision"](False) + raw_f = READERS["fif"](False) + + with trace_deepcopy() as rec: + for r in (raw_e, raw_b, raw_f): + READERS[r.filenames[0].suffix.strip(".") == "edf" + and "edf" or + ("brainvision" if r.filenames[0].suffix == ".vhdr" else "fif")](True) + report(rec, "read_raw_X(preload=True) x3") + + with trace_deepcopy() as rec: + raw = READERS["edf"](True) + raw.crop(tmax=60.0) + raw.set_eeg_reference("average") + raw.notch_filter([50.0]) + raw.resample(128) + report(rec, "load + crop + reref + notch + resample (EDF)") + + with trace_deepcopy() as rec: + raw = READERS["edf"](False).load_data() + epochs = mne.Epochs(raw, mne.make_fixed_length_events(raw, duration=1.0, id=1), + tmin=0.0, tmax=1.99, baseline=None, preload=False, + reject_by_annotation=False) + for ep in epochs[:20]: + pass + report(rec, "Epochs creation + iterate 20 (EDF)") + info_copy_scaling() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--only", choices=["windows", "load", "save", "copies"]) + args = ap.parse_args() + what = args.only + + if what in (None, "windows"): + for fmt in READERS: + profile_windows(fmt, False) + for fmt in READERS: + profile_windows(fmt, True) + if what in (None, "load"): + for fmt in READERS: + profile_full_load(fmt) + if what in (None, "save"): + profile_save() + if what in (None, "copies"): + trace_pipeline_ops() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/results-20260824-154321.json b/benchmarks/io_dl/results-20260824-154321.json new file mode 100644 index 00000000000..87b42fa7365 --- /dev/null +++ b/benchmarks/io_dl/results-20260824-154321.json @@ -0,0 +1,248 @@ +{ + "config": { + "seed": 1234, + "repeats": 7, + "n_win": 300, + "win": 512 + }, + "meta": { + "n_ch": 64, + "sfreq": 256.0, + "dur": 300.0, + "n_times": 76800, + "seed": 42, + "sizes": { + "baseline .bin f32": 19660800, + "edf": 9847040, + "brainvision": 49171105, + "fif": 19673116 + } + }, + "floor_memmap": [ + 9.0905, + 5.482625, + 5.118917, + 4.460542, + 4.354209, + 4.711125, + 4.66275 + ], + "floor_fromfile": [ + 5.787125, + 5.531667, + 5.648459, + 5.30125, + 5.51625, + 5.411375, + 5.765667 + ], + "floor_full_read": [ + 3.967958, + 4.432958, + 3.697042, + 3.532209, + 3.623166, + 3.729958, + 4.508708 + ], + "fmt_edf": { + "open_meta": [ + 1.975375, + 1.939542, + 1.879958, + 1.917375, + 1.893625, + 1.8715, + 1.856792 + ], + "full_load": [ + 11.890334, + 10.699834, + 11.669834, + 11.425334, + 10.820333, + 10.9955, + 10.873375 + ], + "seq_1s": [ + 114.414167, + 114.225041, + 113.459875, + 117.944458, + 119.677125 + ], + "rand_windows": [ + 131.089458, + 128.251292, + 128.869334, + 129.324791, + 129.428417, + 128.224083, + 134.729875 + ], + "preloaded_windows": [ + 32.207625, + 31.319458, + 31.037125, + 32.277583, + 30.718208, + 30.897083, + 30.719917 + ] + }, + "fmt_brainvision": { + "open_meta": [ + 2.272875, + 2.211916, + 1.810417, + 1.92225, + 1.9065, + 2.024541, + 2.019417 + ], + "full_load": [ + 15.002792, + 13.914042, + 13.999083, + 14.813083, + 15.993625, + 16.432333, + 15.807458 + ], + "seq_1s": [ + 50.065667, + 48.626417, + 46.788958, + 49.264958, + 49.615458 + ], + "rand_windows": [ + 53.738167, + 50.68425, + 50.854834, + 50.406584, + 51.463125, + 51.212875, + 52.211834 + ], + "preloaded_windows": [ + 29.893334, + 31.390292, + 30.927542, + 31.296084, + 30.969208, + 30.704334, + 30.318625 + ] + }, + "fmt_fif": { + "open_meta": [ + 2.371625, + 1.858667, + 1.918375, + 2.027875, + 1.956791, + 1.844583, + 1.847625 + ], + "full_load": [ + 12.54525, + 13.0545, + 12.195959, + 14.567209, + 12.639959, + 14.763416, + 13.597084 + ], + "seq_1s": [ + 43.622541, + 42.713333, + 42.918375, + 43.5305, + 43.332667 + ], + "rand_windows": [ + 57.783458, + 56.865792, + 57.152375, + 56.969917, + 57.353042, + 57.816541, + 58.291584 + ], + "preloaded_windows": [ + 29.875125, + 29.049708, + 29.3195, + 29.101416, + 29.062042, + 29.153875, + 29.175917 + ] + }, + "micro": { + "deepcopy_info": [ + 0.098083, + 0.053042, + 0.047667, + 0.046792, + 0.046833, + 0.046458, + 0.046375, + 0.045916, + 0.045916, + 0.045667, + 0.045875, + 0.054333, + 0.055667, + 0.050667, + 0.05075, + 0.056666, + 0.051166, + 0.045834, + 0.045833, + 0.045541 + ], + "info_copy_method": [ + 0.088958, + 0.049375, + 0.046958, + 0.04625, + 0.045875, + 0.045834, + 0.045542, + 0.045625, + 0.045875, + 0.045417, + 0.046, + 0.045708, + 0.046209, + 0.0455, + 0.045334, + 0.045417, + 0.048167, + 0.050875, + 0.050834, + 0.050792 + ], + "raw_copy": [ + 0.174708, + 0.116125, + 0.111792, + 0.11, + 0.109583, + 0.109375, + 0.113792, + 0.111542, + 0.112042, + 0.111875 + ], + "save_to_fif": [ + 103.465833, + 103.432584, + 104.844417, + 102.888541, + 103.640167 + ] + } +} \ No newline at end of file diff --git a/benchmarks/io_dl/results-20260824-160007.json b/benchmarks/io_dl/results-20260824-160007.json new file mode 100644 index 00000000000..2938bb5920c --- /dev/null +++ b/benchmarks/io_dl/results-20260824-160007.json @@ -0,0 +1,248 @@ +{ + "config": { + "seed": 1234, + "repeats": 7, + "n_win": 300, + "win": 512 + }, + "meta": { + "n_ch": 64, + "sfreq": 256.0, + "dur": 300.0, + "n_times": 76800, + "seed": 42, + "sizes": { + "baseline .bin f32": 19660800, + "edf": 9847040, + "brainvision": 49171105, + "fif": 19673116 + } + }, + "floor_memmap": [ + 10.475708, + 12.005, + 16.716792, + 11.267709, + 9.817708, + 11.163958, + 9.848708 + ], + "floor_fromfile": [ + 6.164417, + 5.741833, + 5.753625, + 5.905708, + 5.766167, + 8.79, + 6.012666 + ], + "floor_full_read": [ + 6.046584, + 5.375834, + 4.987333, + 4.778, + 4.295166, + 4.225459, + 4.392416 + ], + "fmt_edf": { + "open_meta": [ + 2.184167, + 2.000417, + 1.953584, + 1.879791, + 1.895792, + 1.881458, + 1.895792 + ], + "full_load": [ + 15.705708, + 14.717, + 15.423417, + 15.49125, + 15.112542, + 14.31375, + 13.944791 + ], + "seq_1s": [ + 97.9925, + 96.297209, + 95.788625, + 98.868417, + 96.867708 + ], + "rand_windows": [ + 115.076959, + 113.873, + 112.077958, + 112.436333, + 114.670958, + 114.896625, + 113.993917 + ], + "preloaded_windows": [ + 11.362, + 10.858417, + 10.575625, + 10.287667, + 10.476209, + 10.299875, + 11.02025 + ] + }, + "fmt_brainvision": { + "open_meta": [ + 2.4385, + 2.232458, + 2.114583, + 2.086084, + 1.932666, + 2.107, + 2.092625 + ], + "full_load": [ + 17.308792, + 16.514708, + 16.376042, + 17.139459, + 16.2695, + 17.084583, + 16.625875 + ], + "seq_1s": [ + 24.70625, + 23.271875, + 22.804333, + 24.804917, + 25.352792 + ], + "rand_windows": [ + 32.521041, + 32.618041, + 33.805417, + 33.80025, + 34.245042, + 32.559125, + 33.978167 + ], + "preloaded_windows": [ + 11.227459, + 10.764208, + 9.9685, + 10.266209, + 9.797958, + 10.741417, + 10.302041 + ] + }, + "fmt_fif": { + "open_meta": [ + 2.345958, + 2.037125, + 2.043667, + 2.136042, + 2.048542, + 1.865917, + 1.935583 + ], + "full_load": [ + 14.936958, + 16.27875, + 16.743333, + 14.687416, + 14.991125, + 18.388708, + 14.634125 + ], + "seq_1s": [ + 25.766417, + 25.0965, + 24.2305, + 23.721292, + 25.02575 + ], + "rand_windows": [ + 43.637542, + 44.333333, + 42.337625, + 39.480542, + 41.040916, + 42.765, + 46.41925 + ], + "preloaded_windows": [ + 11.102333, + 10.773584, + 10.293334, + 9.84675, + 10.02175, + 9.423584, + 9.642292 + ] + }, + "micro": { + "deepcopy_info": [ + 0.10225, + 0.056541, + 0.0495, + 0.051708, + 0.052333, + 0.048375, + 0.051542, + 0.05075, + 0.053958, + 0.052166, + 0.051709, + 0.056916, + 0.053166, + 0.053083, + 0.053042, + 0.053209, + 0.053792, + 0.053375, + 0.054166, + 0.054417 + ], + "info_copy_method": [ + 0.109459, + 0.055958, + 0.050208, + 0.047833, + 0.046375, + 0.046584, + 0.047667, + 0.048583, + 0.047292, + 0.046208, + 0.051208, + 0.055459, + 0.055458, + 0.055708, + 0.052291, + 0.054, + 0.052958, + 0.052375, + 0.052166, + 0.052375 + ], + "raw_copy": [ + 0.196708, + 0.133959, + 0.144417, + 0.129959, + 0.130125, + 0.129625, + 0.128417, + 0.132167, + 0.126917, + 0.127292 + ], + "save_to_fif": [ + 108.325958, + 106.990291, + 154.541792, + 108.504, + 106.092333 + ] + } +} \ No newline at end of file diff --git a/benchmarks/io_dl/workspaces/W5/bench_cache.py b/benchmarks/io_dl/workspaces/W5/bench_cache.py new file mode 100644 index 00000000000..6350abef621 --- /dev/null +++ b/benchmarks/io_dl/workspaces/W5/bench_cache.py @@ -0,0 +1,229 @@ +"""W5: HDF5 chunk-cache tuning + zarr v3 local-read verdict (interleaved). + +Measures 300 random 2 s windows (seed 99) per pass; arms alternate within one +process for `--rounds` rounds (order re-shuffled per round, seeded) so machine +drift hits every arm equally. Persistent handles (training-loop pattern). + +Usage: python bench_cache.py [--rounds 7] [--n-win 300] [--seed 99] +Writes cache_results_.json next to this file. +""" + +import argparse +import gc +import json +import resource +import sys +import time +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).resolve().parent +DATA = HERE.parents[1] / "data_big" +H5 = DATA / "bench_big.h5" +ZARR_WIN = DATA / "bench_big_win.zarr" +ZARR_T10S = DATA / "bench_big_t10s.zarr" + +MB = 1024 * 1024 +DEF_NSLOTS = 521 # HDF5 default hash slots +DEF_NBYTES = MB # HDF5 default raw-chunk cache + + +def is_prime(n): + if n < 2: + return False + i = 2 + while i * i <= n: + if n % i == 0: + return False + i += 1 + return True + + +def make_prime(n): + n = max(3, int(n) | 1) + while not is_prime(n): + n += 2 + return n + + +def build_arms(meta, starts): + win = int(2 * meta["sfreq"]) + nt = meta["n_times"] + cols, ks = set(), set() + for s in starts: + cols.update({s // 512, (s + win - 1) // 512}) + ks.update({s // 5120, (s + win - 1) // 5120}) + fit_win_bytes = len(cols) * 128 * 512 * 4 # exact bytes to hold all touched win chunks + fit_t10_bytes = len(ks) * 128 * 5120 * 4 # ... t10 chunks (all channels touched per window) + + arms = [] + # --- h5py t10s dataset: cache-size arms --------------------------------- + t = "t10s" + arms.append(dict(name=f"h5_{t}_default", ds=t, + kw=dict(rdcc_nbytes=DEF_NBYTES, rdcc_nslots=DEF_NSLOTS))) + arms.append(dict(name=f"h5_{t}_fit", ds=t, + kw=dict(rdcc_nbytes=fit_t10_bytes, rdcc_nslots=make_prime(2 * len(ks) * 128)))) + arms.append(dict(name=f"h5_{t}_big512", ds=t, + kw=dict(rdcc_nbytes=512 * MB, rdcc_nslots=make_prime(65537)))) + arms.append(dict(name=f"h5_{t}_ns10103", ds=t, + kw=dict(rdcc_nbytes=DEF_NBYTES, rdcc_nslots=10103))) + arms.append(dict(name=f"h5_{t}_ns65537", ds=t, + kw=dict(rdcc_nbytes=DEF_NBYTES, rdcc_nslots=65537))) + # --- h5py win dataset: control + same treatment ------------------------- + arms.append(dict(name="h5_win_default_CTRL", ds="win", + kw=dict(rdcc_nbytes=DEF_NBYTES, rdcc_nslots=DEF_NSLOTS))) + arms.append(dict(name="h5_win_fit", ds="win", + kw=dict(rdcc_nbytes=fit_win_bytes, rdcc_nslots=make_prime(2 * len(cols))))) + arms.append(dict(name="h5_win_big512", ds="win", + kw=dict(rdcc_nbytes=512 * MB, rdcc_nslots=make_prime(65537)))) + arms.append(dict(name="h5_win_ns10103", ds="win", + kw=dict(rdcc_nbytes=DEF_NBYTES, rdcc_nslots=10103))) + arms.append(dict(name="h5_win_ns65537", ds="win", + kw=dict(rdcc_nbytes=DEF_NBYTES, rdcc_nslots=65537))) + # --- file driver variants (once each, interleaved like the rest) -------- + arms.append(dict(name="h5_t10s_sec2", ds="t10s", + kw=dict(rdcc_nbytes=DEF_NBYTES, rdcc_nslots=DEF_NSLOTS), + note="explicit sec2 (default) driver")) + arms.append(dict(name="h5_t10s_core", ds="t10s", kw=dict(), + open_extra=dict(driver="core", backing_store=False), + note="core driver: whole file read into RAM at open")) + # --- zarr ---------------------------------------------------------------- + arms.append(dict(name="zarr_win_default", zarr=str(ZARR_WIN), zcfg={})) + arms.append(dict(name="zarr_win_c1", zarr=str(ZARR_WIN), + zcfg={"async.concurrency": 1})) + arms.append(dict(name="zarr_win_c32", zarr=str(ZARR_WIN), + zcfg={"async.concurrency": 32})) + arms.append(dict(name="zarr_t10s_default", zarr=str(ZARR_T10S), zcfg={})) + return arms, dict(fit_win_bytes=fit_win_bytes, fit_t10s_bytes=fit_t10_bytes) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--rounds", type=int, default=7) + ap.add_argument("--n-win", type=int, default=300) + ap.add_argument("--seed", type=int, default=99) + args = ap.parse_args() + + import h5py + + meta = json.loads((DATA / "meta.json").read_text()) + sfreq, nt = meta["sfreq"], meta["n_times"] + width = int(2 * sfreq) # 1024 samples + starts = np.random.default_rng(args.seed).integers( + 0, nt - width - 1, size=args.n_win).astype(int) + + arms, fits = build_arms(meta, starts) + print(f"{len(arms)} arms x {args.rounds} interleaved rounds x " + f"{args.n_win} windows | fit(win)={fits['fit_win_bytes']/MB:.1f} MiB " + f"fit(t10s)={fits['fit_t10s_bytes']/MB:.1f} MiB") + + # ---- open persistent handles ------------------------------------------- + import zarr + + handles = {} + rss0 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + for arm in arms: + if "zarr" not in arm: + extra = arm.get("open_extra", {}) + f = h5py.File(H5, "r", **arm["kw"], **extra) + handles[arm["name"]] = f[arm["ds"]] + else: + arr = zarr.open_array(store=arm["zarr"], mode="r") + handles[arm["name"]] = arr + rss_open = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + rss_unit = 1024 * 1024 if sys.platform == "darwin" else 1024 # macOS: bytes; Linux: KiB + print(f"peak RSS after all opens: {rss_open/rss_unit:.0f} MiB " + f"(delta {((rss_open-rss0)/rss_unit):.0f} MiB incl. core-driver RAM copy)") + + def pass_h5(ds): + acc = 0.0 + for s in starts: + acc += float(ds[:, s:s + width].astype(np.float64).sum()) + return acc + + def pass_zarr(arr, cfg): + acc = 0.0 + cm = zarr.config.set(cfg) if cfg else None + if cm is not None: + cm.__enter__() + try: + for s in starts: + acc += float(arr[:, s:s + width].astype(np.float64).sum()) + finally: + if cm is not None: + cm.__exit__(None, None, None) + return acc + + def run_once(arm): + gc.collect() + gc.disable() + try: + t0 = time.perf_counter_ns() + if "zarr" in arm: + chk = pass_zarr(handles[arm["name"]], arm["zcfg"]) + else: + chk = pass_h5(handles[arm["name"]]) + dt_ms = (time.perf_counter_ns() - t0) / 1e6 + finally: + gc.enable() + return dt_ms, chk + + # ---- warmup (page cache + arm-local caches) ---------------------------- + order_rng = np.random.default_rng(12345) + print("warmup...", flush=True) + ref_chk = None + for arm in arms: + ms, chk = run_once(arm) + if ref_chk is None: + ref_chk = chk + assert abs(chk - ref_chk) < 1e-6 * abs(ref_chk), f"checksum mismatch {arm['name']}" + + # ---- interleaved rounds ------------------------------------------------- + times = {a["name"]: [] for a in arms} + wall = [] + for r in range(args.rounds): + order = list(arms) + order_rng.shuffle(order) + t_r0 = time.perf_counter() + for arm in order: + ms, _ = run_once(arm) + times[arm["name"]].append(ms) + wall.append(time.perf_counter() - t_r0) + done = ", ".join(f"{a['name'].split('_', 1)[0]}:{times[a['name']][-1]:.0f}ms" + for a in arms[:3]) + print(f"round {r+1}/{args.rounds} done ({wall[-1]:.1f}s) {done}", flush=True) + + # ---- summarize ---------------------------------------------------------- + res = dict(stamp=time.strftime("%Y%m%d-%H%M%S"), machine="Apple M4 Max, 36 GB", + rounds=args.rounds, n_win=args.n_win, seed=args.seed, + peak_rss_mib=rss_open / rss_unit, fits=fits, arms={}, control_rounds=None) + print(f"\n{'arm':<24}{'us/win median':>14}{'IQR':>18}{'min..max':>18}") + ctrl_med = None + for a in arms: + v = np.array(times[a["name"]]) * 1000.0 / args.n_win # us/win + med = float(np.median(v)) + q1, q3 = np.percentile(v, [25, 75]) + res["arms"][a["name"]] = dict(us_per_win=[float(x) for x in v], + median=med, iqr=float(q3 - q1), + min=float(v.min()), max=float(v.max()), + kind="zarr" if "zarr" in a else "h5", + note=a.get("note", "")) + flag = "" + if "CTRL" in a["name"]: + ctrl_med = med + print(f"{a['name']:<24}{med:>11.1f} [{np.percentile(v,25):>7.1f},{q3:>7.1f}]" + f" {v.min():>7.1f}..{v.max():<7.1f}{flag}") + + # drift check: round-by-round of control arm + cv = np.array(times["h5_win_default_CTRL"]) * 1000.0 / args.n_win + res["control_rounds"] = [float(x) for x in cv] + print(f"\nCONTROL h5_win per-round µs/win: {[round(x,1) for x in cv]}") + print(f"CONTROL spread max/min = {cv.max()/cv.min():.3f}") + out = HERE / f"cache_results_{res['stamp']}.json" + out.write_text(json.dumps(res, indent=2)) + print(f"saved -> {out}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/workspaces/W5/cache_results_20260824-204900.json b/benchmarks/io_dl/workspaces/W5/cache_results_20260824-204900.json new file mode 100644 index 00000000000..86ad70e2053 --- /dev/null +++ b/benchmarks/io_dl/workspaces/W5/cache_results_20260824-204900.json @@ -0,0 +1,193 @@ +{ + "stamp": "20260824-204900", + "machine": "Apple M4 Max, 36 GB", + "rounds": 1, + "n_win": 30, + "seed": 99, + "peak_rss_mib": 989856.0, + "fits": { + "fit_win_bytes": 15204352, + "fit_t10s_bytes": 83886080 + }, + "arms": { + "h5_t10s_default": { + "us_per_win": [ + 428.1055666666667 + ], + "median": 428.1055666666667, + "iqr": 0.0, + "min": 428.1055666666667, + "max": 428.1055666666667, + "kind": "h5", + "note": "" + }, + "h5_t10s_fit": { + "us_per_win": [ + 413.40139999999997 + ], + "median": 413.40139999999997, + "iqr": 0.0, + "min": 413.40139999999997, + "max": 413.40139999999997, + "kind": "h5", + "note": "" + }, + "h5_t10s_big512": { + "us_per_win": [ + 432.3805666666667 + ], + "median": 432.3805666666667, + "iqr": 0.0, + "min": 432.3805666666667, + "max": 432.3805666666667, + "kind": "h5", + "note": "" + }, + "h5_t10s_ns10103": { + "us_per_win": [ + 442.32359999999994 + ], + "median": 442.32359999999994, + "iqr": 0.0, + "min": 442.32359999999994, + "max": 442.32359999999994, + "kind": "h5", + "note": "" + }, + "h5_t10s_ns65537": { + "us_per_win": [ + 429.93053333333336 + ], + "median": 429.93053333333336, + "iqr": 0.0, + "min": 429.93053333333336, + "max": 429.93053333333336, + "kind": "h5", + "note": "" + }, + "h5_win_default_CTRL": { + "us_per_win": [ + 82.23333333333333 + ], + "median": 82.23333333333333, + "iqr": 0.0, + "min": 82.23333333333333, + "max": 82.23333333333333, + "kind": "h5", + "note": "" + }, + "h5_win_fit": { + "us_per_win": [ + 94.54723333333332 + ], + "median": 94.54723333333332, + "iqr": 0.0, + "min": 94.54723333333332, + "max": 94.54723333333332, + "kind": "h5", + "note": "" + }, + "h5_win_big512": { + "us_per_win": [ + 96.3389 + ], + "median": 96.3389, + "iqr": 0.0, + "min": 96.3389, + "max": 96.3389, + "kind": "h5", + "note": "" + }, + "h5_win_ns10103": { + "us_per_win": [ + 106.43473333333334 + ], + "median": 106.43473333333334, + "iqr": 0.0, + "min": 106.43473333333334, + "max": 106.43473333333334, + "kind": "h5", + "note": "" + }, + "h5_win_ns65537": { + "us_per_win": [ + 89.4625 + ], + "median": 89.4625, + "iqr": 0.0, + "min": 89.4625, + "max": 89.4625, + "kind": "h5", + "note": "" + }, + "h5_t10s_sec2": { + "us_per_win": [ + 448.23886666666664 + ], + "median": 448.23886666666664, + "iqr": 0.0, + "min": 448.23886666666664, + "max": 448.23886666666664, + "kind": "h5", + "note": "explicit sec2 (default) driver" + }, + "h5_t10s_core": { + "us_per_win": [ + 357.48193333333336 + ], + "median": 357.48193333333336, + "iqr": 0.0, + "min": 357.48193333333336, + "max": 357.48193333333336, + "kind": "h5", + "note": "core driver: whole file read into RAM at open" + }, + "zarr_win_default": { + "us_per_win": [ + 1025.8319333333334 + ], + "median": 1025.8319333333334, + "iqr": 0.0, + "min": 1025.8319333333334, + "max": 1025.8319333333334, + "kind": "zarr", + "note": "" + }, + "zarr_win_c1": { + "us_per_win": [ + 1819.1263666666666 + ], + "median": 1819.1263666666666, + "iqr": 0.0, + "min": 1819.1263666666666, + "max": 1819.1263666666666, + "kind": "zarr", + "note": "" + }, + "zarr_win_c32": { + "us_per_win": [ + 1011.7527666666666 + ], + "median": 1011.7527666666666, + "iqr": 0.0, + "min": 1011.7527666666666, + "max": 1011.7527666666666, + "kind": "zarr", + "note": "" + }, + "zarr_t10s_default": { + "us_per_win": [ + 26052.0875 + ], + "median": 26052.0875, + "iqr": 0.0, + "min": 26052.0875, + "max": 26052.0875, + "kind": "zarr", + "note": "" + } + }, + "control_rounds": [ + 82.23333333333333 + ] +} \ No newline at end of file diff --git a/benchmarks/io_dl/workspaces/W5/cache_results_20260824-205105.json b/benchmarks/io_dl/workspaces/W5/cache_results_20260824-205105.json new file mode 100644 index 00000000000..3b3b11b0156 --- /dev/null +++ b/benchmarks/io_dl/workspaces/W5/cache_results_20260824-205105.json @@ -0,0 +1,295 @@ +{ + "stamp": "20260824-205105", + "machine": "Apple M4 Max, 36 GB", + "rounds": 7, + "n_win": 300, + "seed": 99, + "peak_rss_mib": 967.25, + "fits": { + "fit_win_bytes": 138149888, + "fit_t10s_bytes": 424673280 + }, + "arms": { + "h5_t10s_default": { + "us_per_win": [ + 491.2009733333333, + 444.52263666666664, + 557.6593033333334, + 465.79, + 461.4230566666667, + 455.20778, + 452.05319333333335 + ], + "median": 461.4230566666667, + "iqr": 24.86500000000001, + "min": 444.52263666666664, + "max": 557.6593033333334, + "kind": "h5", + "note": "" + }, + "h5_t10s_fit": { + "us_per_win": [ + 466.8918066666667, + 467.7708333333333, + 496.5901400000001, + 458.03791666666666, + 443.21208333333334, + 464.3944466666667, + 459.01875 + ], + "median": 464.3944466666667, + "iqr": 8.802986666666698, + "min": 443.21208333333334, + "max": 496.5901400000001, + "kind": "h5", + "note": "" + }, + "h5_t10s_big512": { + "us_per_win": [ + 485.86541666666665, + 487.7952766666667, + 487.89, + 454.46138666666667, + 459.2902766666666, + 471.04278, + 458.3958333333333 + ], + "median": 471.04278, + "iqr": 27.98729166666675, + "min": 454.46138666666667, + "max": 487.89, + "kind": "h5", + "note": "" + }, + "h5_t10s_ns10103": { + "us_per_win": [ + 464.9429166666667, + 450.63764000000003, + 605.0413866666667, + 469.1755566666667, + 455.25930666666665, + 446.4320833333333, + 470.86541666666665 + ], + "median": 464.9429166666667, + "iqr": 17.072013333333302, + "min": 446.4320833333333, + "max": 605.0413866666667, + "kind": "h5", + "note": "" + }, + "h5_t10s_ns65537": { + "us_per_win": [ + 479.2945833333333, + 490.07055333333335, + 534.98375, + 476.55763666666667, + 448.88819666666666, + 454.2586099999999, + 445.6358333333333 + ], + "median": 476.55763666666667, + "iqr": 33.10916500000002, + "min": 445.6358333333333, + "max": 534.98375, + "kind": "h5", + "note": "" + }, + "h5_win_default_CTRL": { + "us_per_win": [ + 91.91847333333332, + 105.45875, + 112.99777666666667, + 101.69805333333333, + 94.17972333333334, + 98.73347333333334, + 87.96166666666667 + ], + "median": 98.73347333333334, + "iqr": 10.529303333333331, + "min": 87.96166666666667, + "max": 112.99777666666667, + "kind": "h5", + "note": "" + }, + "h5_win_fit": { + "us_per_win": [ + 107.12833333333333, + 90.53597333333333, + 107.04583333333335, + 104.66458333333334, + 86.00597333333333, + 96.52264000000001, + 87.06819333333333 + ], + "median": 96.52264000000001, + "iqr": 17.05312500000001, + "min": 86.00597333333333, + "max": 107.12833333333333, + "kind": "h5", + "note": "" + }, + "h5_win_big512": { + "us_per_win": [ + 103.14597, + 91.64528, + 107.92889000000001, + 107.12444333333333, + 103.38194333333333, + 99.36097333333333, + 89.96014000000001 + ], + "median": 103.14597, + "iqr": 9.750066666666655, + "min": 89.96014000000001, + "max": 107.92889000000001, + "kind": "h5", + "note": "" + }, + "h5_win_ns10103": { + "us_per_win": [ + 91.42402666666666, + 103.76347, + 95.97875, + 101.50166666666667, + 89.03972333333334, + 97.80652666666667, + 91.29916666666666 + ], + "median": 95.97875, + "iqr": 8.292500000000018, + "min": 89.03972333333334, + "max": 103.76347, + "kind": "h5", + "note": "" + }, + "h5_win_ns65537": { + "us_per_win": [ + 89.11486, + 94.90138666666667, + 106.54208333333334, + 105.03111000000001, + 87.94694666666668, + 104.20375, + 102.99 + ], + "median": 102.99, + "iqr": 12.609306666666683, + "min": 87.94694666666668, + "max": 106.54208333333334, + "kind": "h5", + "note": "" + }, + "h5_t10s_sec2": { + "us_per_win": [ + 488.30819333333335, + 497.8475, + 488.07735999999994, + 480.5083333333333, + 462.57430666666664, + 477.82458333333335, + 442.37666666666667 + ], + "median": 480.5083333333333, + "iqr": 17.993331666666677, + "min": 442.37666666666667, + "max": 497.8475, + "kind": "h5", + "note": "explicit sec2 (default) driver" + }, + "h5_t10s_core": { + "us_per_win": [ + 375.60625, + 370.15346999999997, + 362.0125, + 336.36722333333336, + 341.6531933333333, + 354.4993066666666, + 334.81069333333335 + ], + "median": 354.4993066666666, + "iqr": 27.07277666666664, + "min": 334.81069333333335, + "max": 375.60625, + "kind": "h5", + "note": "core driver: whole file read into RAM at open" + }, + "zarr_win_default": { + "us_per_win": [ + 1049.6094433333333, + 1048.2273633333332, + 1138.4119433333333, + 1056.4302799999998, + 1055.2720833333333, + 1047.4398600000002, + 1024.87389 + ], + "median": 1049.6094433333333, + "iqr": 8.01756999999975, + "min": 1024.87389, + "max": 1138.4119433333333, + "kind": "zarr", + "note": "" + }, + "zarr_win_c1": { + "us_per_win": [ + 1830.7990266666666, + 1845.8404166666667, + 1985.5343066666667, + 1863.02486, + 1847.3580533333331, + 1818.2479166666667, + 1840.3680566666667 + ], + "median": 1845.8404166666667, + "iqr": 19.60791499999982, + "min": 1818.2479166666667, + "max": 1985.5343066666667, + "kind": "zarr", + "note": "" + }, + "zarr_win_c32": { + "us_per_win": [ + 1072.8234699999998, + 1033.0080566666668, + 1083.0152766666665, + 1034.2048633333332, + 1032.2668033333334, + 1019.5330533333334, + 1046.3770833333333 + ], + "median": 1034.2048633333332, + "iqr": 26.962846666666337, + "min": 1019.5330533333334, + "max": 1083.0152766666665, + "kind": "zarr", + "note": "" + }, + "zarr_t10s_default": { + "us_per_win": [ + 27599.248473333333, + 27840.20861, + 27325.996946666666, + 27844.88625, + 27518.284026666664, + 27881.890833333335, + 28030.765833333335 + ], + "median": 27840.20861, + "iqr": 304.62229166666657, + "min": 27325.996946666666, + "max": 28030.765833333335, + "kind": "zarr", + "note": "" + } + }, + "control_rounds": [ + 91.91847333333332, + 105.45875, + 112.99777666666667, + 101.69805333333333, + 94.17972333333334, + 98.73347333333334, + 87.96166666666667 + ] +} \ No newline at end of file diff --git a/benchmarks/io_dl/workspaces/W5/cache_results_20260824-205253.json b/benchmarks/io_dl/workspaces/W5/cache_results_20260824-205253.json new file mode 100644 index 00000000000..83095f7711c --- /dev/null +++ b/benchmarks/io_dl/workspaces/W5/cache_results_20260824-205253.json @@ -0,0 +1,295 @@ +{ + "stamp": "20260824-205253", + "machine": "Apple M4 Max, 36 GB", + "rounds": 7, + "n_win": 300, + "seed": 99, + "peak_rss_mib": 967.140625, + "fits": { + "fit_win_bytes": 138149888, + "fit_t10s_bytes": 424673280 + }, + "arms": { + "h5_t10s_default": { + "us_per_win": [ + 495.0518066666667, + 467.10791666666665, + 481.33153, + 483.0490299999999, + 477.0370833333333, + 495.66930333333335, + 481.1454166666667 + ], + "median": 481.33153, + "iqr": 9.95916833333331, + "min": 467.10791666666665, + "max": 495.66930333333335, + "kind": "h5", + "note": "" + }, + "h5_t10s_fit": { + "us_per_win": [ + 542.79903, + 462.97083333333336, + 457.30069333333336, + 484.07986000000005, + 452.83930666666663, + 479.595, + 477.61 + ], + "median": 477.61, + "iqr": 21.70166666666671, + "min": 452.83930666666663, + "max": 542.79903, + "kind": "h5", + "note": "" + }, + "h5_t10s_big512": { + "us_per_win": [ + 495.5233333333333, + 456.71778, + 475.28986000000003, + 488.8859700000001, + 456.6040266666667, + 487.5823633333333, + 465.4181966666667 + ], + "median": 475.28986000000003, + "iqr": 27.16617833333339, + "min": 456.6040266666667, + "max": 495.5233333333333, + "kind": "h5", + "note": "" + }, + "h5_t10s_ns10103": { + "us_per_win": [ + 545.0961133333334, + 467.04764000000006, + 465.53972, + 484.6136133333333, + 453.63611000000003, + 474.55, + 486.71596999999997 + ], + "median": 474.55, + "iqr": 19.37111166666665, + "min": 453.63611000000003, + "max": 545.0961133333334, + "kind": "h5", + "note": "" + }, + "h5_t10s_ns65537": { + "us_per_win": [ + 484.6394466666667, + 479.6534733333333, + 473.93291666666664, + 486.4504166666667, + 471.4084733333333, + 475.19708333333335, + 476.49236333333334 + ], + "median": 476.49236333333334, + "iqr": 7.581459999999993, + "min": 471.4084733333333, + "max": 486.4504166666667, + "kind": "h5", + "note": "" + }, + "h5_win_default_CTRL": { + "us_per_win": [ + 94.27333333333333, + 88.24708333333334, + 88.23389, + 104.92416666666666, + 107.17916666666667, + 107.44930666666666, + 103.33041666666666 + ], + "median": 103.33041666666666, + "iqr": 14.791458333333338, + "min": 88.23389, + "max": 107.44930666666666, + "kind": "h5", + "note": "" + }, + "h5_win_fit": { + "us_per_win": [ + 108.155, + 88.27222333333334, + 90.67138999999999, + 88.56305666666665, + 89.57513666666665, + 90.13653000000001, + 105.09069666666666 + ], + "median": 90.13653000000001, + "iqr": 8.811946666666671, + "min": 88.27222333333334, + "max": 108.155, + "kind": "h5", + "note": "" + }, + "h5_win_big512": { + "us_per_win": [ + 99.41083333333333, + 93.76375, + 104.98278, + 108.19402999999998, + 90.36708333333333, + 103.59569333333333, + 87.82680666666667 + ], + "median": 99.41083333333333, + "iqr": 12.223820000000003, + "min": 87.82680666666667, + "max": 108.19402999999998, + "kind": "h5", + "note": "" + }, + "h5_win_ns10103": { + "us_per_win": [ + 96.96708333333333, + 89.80388666666666, + 93.81361, + 102.95347333333333, + 103.84319333333335, + 102.89722333333334, + 92.06527666666666 + ], + "median": 96.96708333333333, + "iqr": 9.985905000000017, + "min": 89.80388666666666, + "max": 103.84319333333335, + "kind": "h5", + "note": "" + }, + "h5_win_ns65537": { + "us_per_win": [ + 91.19416666666666, + 90.50764, + 88.40847333333333, + 102.02958333333333, + 105.25319333333333, + 101.63194333333333, + 98.05097333333333 + ], + "median": 98.05097333333333, + "iqr": 10.979860000000002, + "min": 88.40847333333333, + "max": 105.25319333333333, + "kind": "h5", + "note": "" + }, + "h5_t10s_sec2": { + "us_per_win": [ + 485.11319333333336, + 459.19666666666666, + 470.18, + 474.8255566666666, + 476.48194666666666, + 456.34944666666667, + 472.13625 + ], + "median": 472.13625, + "iqr": 10.965418333333275, + "min": 456.34944666666667, + "max": 485.11319333333336, + "kind": "h5", + "note": "explicit sec2 (default) driver" + }, + "h5_t10s_core": { + "us_per_win": [ + 355.5694433333333, + 330.08194333333336, + 346.4, + 331.32625, + 338.71028, + 354.63764, + 337.8794433333333 + ], + "median": 338.71028, + "iqr": 15.91597333333334, + "min": 330.08194333333336, + "max": 355.5694433333333, + "kind": "h5", + "note": "core driver: whole file read into RAM at open" + }, + "zarr_win_default": { + "us_per_win": [ + 1068.78125, + 1278.9984733333333, + 1027.7230566666667, + 1036.74375, + 1043.8177799999999, + 1030.79625, + 1070.8916666666667 + ], + "median": 1043.8177799999999, + "iqr": 36.06645833333323, + "min": 1027.7230566666667, + "max": 1278.9984733333333, + "kind": "zarr", + "note": "" + }, + "zarr_win_c1": { + "us_per_win": [ + 1838.69361, + 1906.69264, + 1893.6895833333333, + 1857.7833333333333, + 1867.2348600000003, + 1786.1266666666668, + 1896.2031933333333 + ], + "median": 1867.2348600000003, + "iqr": 46.70791666666673, + "min": 1786.1266666666668, + "max": 1906.69264, + "kind": "zarr", + "note": "" + }, + "zarr_win_c32": { + "us_per_win": [ + 1269.2541666666666, + 1163.3076366666667, + 1046.5244433333332, + 1061.7354166666667, + 1022.1740266666666, + 1032.4456933333333, + 1063.9608333333333 + ], + "median": 1061.7354166666667, + "iqr": 74.14916666666659, + "min": 1022.1740266666666, + "max": 1269.2541666666666, + "kind": "zarr", + "note": "" + }, + "zarr_t10s_default": { + "us_per_win": [ + 28805.50013666667, + 28694.205, + 28604.679443333334, + 27972.558193333334, + 28016.18264, + 28428.649723333332, + 28259.127636666664 + ], + "median": 28428.649723333332, + "iqr": 511.78708333333634, + "min": 27972.558193333334, + "max": 28805.50013666667, + "kind": "zarr", + "note": "" + } + }, + "control_rounds": [ + 94.27333333333333, + 88.24708333333334, + 88.23389, + 104.92416666666666, + 107.17916666666667, + 107.44930666666666, + 103.33041666666666 + ] +} \ No newline at end of file diff --git a/benchmarks/io_dl/workspaces/W5/findings.md b/benchmarks/io_dl/workspaces/W5/findings.md new file mode 100644 index 00000000000..1e927730c8c --- /dev/null +++ b/benchmarks/io_dl/workspaces/W5/findings.md @@ -0,0 +1,141 @@ +# W5 — HDF5 chunk-cache tuning & zarr-v3 local-read verdict + +Date: 2026-08-24 · Machine: Apple M4 Max, 36 GB, macOS · h5py 3.16.0, zarr 3.1.5, +numpy 2.4.6 · Data: `benchmarks/io_dl/data_big/` (128 ch × 1800 s @ 512 Hz, +float32; `bench_big.h5` datasets `t10s` chunks=(1,5120), `win` chunks=(128,512), +`full` contiguous; zarr v3 twins with identical chunk shapes). +Script: `bench_cache.py` (this directory); raw numbers in +`cache_results_20260824-205105.json` / `-205253.json`. + +## Methodology (drift-proof by construction) + +- **Single process, interleaved arms**: all 16 arms alternate every round, order + re-shuffled per round (seeded), 7 rounds × 300 random 2 s windows (seed 99), + persistent handles opened once (training-loop pattern). GC disabled during + timed passes; `perf_counter_ns`; float64 checksum verified identical across + every arm/pass (pure-read sanity). +- **Warmup pass** per arm before round 1 (fills page cache + arm-local caches); + no `sudo purge` available → regime is *warm OS page cache*, which is also the + realistic steady state of a training loop. +- **Control arm** `h5_win_default` interleaved with everything; its per-round + drift was 88–113 µs/win (max/min 1.22–1.29) across rounds — i.e. machine noise + of ±10–15 % was present and is exactly what interleaving cancels. Any claimed + effect below must exceed that band to be real. +- Two independent process runs (A, B); conclusions require agreement in both. +- Cache-fit sizes computed exactly from the seed-99 window set: + `win`: 527 distinct chunks touched → fit = 138 MB; + `t10s`: 20,736 chunks (162 time-slabs × 128 ch) → fit = 405 MB. + +## Decision table + +µs/window, median over 14 interleaved passes (7 rounds × 2 runs); spread = +min..max pooled. "vs CTRL" compares each arm's median against its same-run +control to cancel drift. + +| config | µs/win median [min..max] | vs CTRL | verdict | +|---|---|---|---| +| h5_win default cache 1 MB/521 (**CONTROL**) | 100 [88..113] | 1.00 | baseline | +| h5_win fit-to-working-set 138 MB/1061 slots | 91 [86..108] | −5 % | no effect (inside noise band) | +| h5_win oversized 512 MB/65537 | 101 [88..108] | +1 % | no effect | +| h5_win nslots=10103 (1 MB) | 97 [89..104] | −4 % | no effect | +| h5_win nslots=65537 (1 MB) | 100 [88..107] | −1 % | no effect | +| h5_t10s default 1 MB/521 | 479 [443..558] | 1.00 | baseline (slab layout) | +| h5_t10s fit-to-working-set 405 MB/41479 | 466 [443..543] | −2 % | no effect | +| h5_t10s oversized 512 MB/65537 | 473 [455..496] | −1 % | no effect | +| h5_t10s nslots=10103 | 468 [446..605] | −2 % | no effect | +| h5_t10s nslots=65537 | 477 [446..535] | 0 % | no effect | +| h5_t10s driver=sec2 explicit | 476 [442..498] | ≈0 | confirms default driver | +| **h5_t10s driver='core' (RAM)** | **344 [330..376]** | **−26 %** | real, reproducible (−26/−28 % both runs) | +| zarr_win default (`async.concurrency=10`) | 1049 [1025..1279] | — | **10.1–10.6× slower than h5_win** | +| zarr_win concurrency=1 | 1853 [1786..1986] | — | +77 % vs zarr default — do NOT set | +| zarr_win concurrency=32 | 1047 [1020..1269] | — | ≈ zarr default (no gain) | +| zarr_t10s default | 27,994 [27,326..28,806] | — | ~60× slower than h5_t10s | + +### Does chunk-cache tuning matter (>20 %)? **NO.** + +Largest deviation of any cache arm from control: **−5 % (win_fit), −2 % +(t10s_fit)** — an order of magnitude below the 20 % threshold and not +systematic across the two runs (sign flips between runs). Mechanism: + +1. With a warm page cache every chunk read is already served at memcpy speed + by macOS; HDF5's rdcc only short-circuits *re-reads*, which were cheap + anyway. +2. Random windows give huge reuse distances: each `t10s` chunk is re-visited + only ~3× over a whole 300-window pass, separated by thousands of other + chunk reads — nothing stays "hot" even in a perfectly sized cache. +3. No intra-read re-touch exists (a window's ≤256 chunks are all distinct), so + oversizing has nothing to exploit. + +rdcc tuning should only matter when the working set exceeds RAM (cold reads) +or windows are sampled with tight reuse distance (heavy overlap). Neither +applies to EEG-shaped files (~0.5 GB) on developer/workstation RAM. + +## File-driver variants (measured once each, interleaved like other arms) + +`driver="core", backing_store=False` reads the whole file into process RAM at +open: **+~0.9 GiB RSS** for this 944 MB file (plus your working arrays), and +buys **−26 % per window on the t10s layout** (344 vs 476 µs) by eliminating +per-chunk `read()` syscalls (≤256 syscalls/window → memcpy). Benefit shrinks +proportionally with chunks-per-window (the `win` layout touches only 2), so +its practical value is for slab/time-major layouts or network filesystems. +Not recommended as default advice; document as an option with the memory cost. + +## Zarr v3 local-read verdict + +- Identical chunk shape `(128,512)` to `h5_win`, zstd level-0 codecs (chunks + stored ~9 % smaller): **zarr 1049 µs/win vs h5py 100 µs/win ⇒ 10.1–10.6× + penalty** (earlier ad-hoc "~20×" was drift-inflated; clean interleaved ratio + is ~10×). For the slab layout the gap explodes to ~60× (28 ms/win). +- Knobs: `async.concurrency` 32 = no change; `concurrency=1` makes it **worse** + (+77 %) — the default pool of 10 is already optimal for 2-chunk reads. The + penalty is architectural: per-chunk Python/async dispatch + buffer creation + in zarr v3's sync path, not a tunable configuration issue. + +## RECOMMENDATIONS + +1. **Tell users nothing about rdcc kwargs** — there is no `h5py.open(...)` + chunk-cache setting worth recommending for EEG-shaped local training data; + defaults (1 MB/521) measure identically to exact-fit and 512 MB caches. + Spend the effort on chunk geometry instead: `chunks=(n_ch, window_samples)` + reproduces the ~50–100 µs/window result and needs zero tuning. +2. Keep zarr flagged **remote/parallel-only** for training loops: ~10× + per-window penalty locally with identical chunking, unfixable via its + concurrency settings. Revisit only if zarr v3 gains a C-speed batched read + path or if reads move to object storage where network latency dominates + anyway. +3. If someone insists on squeezing the slab-layout case without rewriting + chunks, `h5py.File(..., driver="core", backing_store=False)` buys ~26 % at + +1× file-size RAM; cheaper than cache tuning, still worse than rechunking. +4. Do not set `zarr.config.set({"async.concurrency": 1})` anywhere (76 % + slowdown); leave defaults alone. + +## Draft text for RESULTS.md §12 (ready to paste) + +```markdown +## 12. Session 5: HDF5 chunk-cache tuning is a dead end; zarr-v3 local verdict finalized + +Question: does `rdcc_nbytes`/`rdcc_nslots` change random-window read speed for +EEG-shaped data? Method: single-process interleaved benchmark (16 arms × 7 +rounds × 300 windows, seeded shuffle per round, persistent handles, identical +checksums, control arm tracked per-round drift of ±10–15 % which interleave +cancels; two independent runs). + +Answer: **no**. Exact-fit caches (computed from the true touched-chunk set: +138 MB for window-chunks, 405 MB for per-channel slabs), a 512 MB oversized +cache, and hash-slot variations (10103/65537 primes) all land within ±5 % of +the 1 MB/521 default — far under the ±15 % machine-noise band, with signs +flipping between runs. Reason: with a warm page cache, chunk re-reads are +already memcpy-speed, and random sampling gives reuse distances so large that +nothing stays hot in any cache size. Chunk *geometry* remains the only lever +that matters (`chunks=(n_ch, win)` ≈ 100 µs/win; slabs ≈ 480 µs/win). + +Two side results. (1) `driver="core"` (whole file pinned in RAM, +0.9 GiB +here) cuts slab-layout windows by 26 % (476→344 µs) via syscall elimination — +an option, not a recommendation. (2) The zarr v3 local-read penalty is real +but was overstated by earlier noisy probes: with byte-identical chunk shapes, +zarr reads windows at ~1050 µs vs h5py's ~100 µs ⇒ **≈10× slower** (not ~20×); +for per-channel-slab chunks it is ~60×. Its `async.concurrency` knob neither +helps at 32 nor tolerates 1 (+77 %), so the gap is per-chunk Python dispatch +overhead in v3's sync path. Guidance stands: HDF5/memmap for local training, +zarr reserved for remote/parallel storage. +``` diff --git a/benchmarks/io_dl/workspaces/W7/bench_w7_consolidated.py b/benchmarks/io_dl/workspaces/W7/bench_w7_consolidated.py new file mode 100644 index 00000000000..a035574825c --- /dev/null +++ b/benchmarks/io_dl/workspaces/W7/bench_w7_consolidated.py @@ -0,0 +1,156 @@ +"""W7 consolidated single-session pass: all writers + floors share machine state. + +Run 2-3 times; report medians-of-medians. ~30 s per pass. +""" + +import gc +import json +import shutil +import sys +import tempfile +import time +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).resolve().parent +BENCH = HERE.parents[1] +TREE = BENCH.parents[1] +sys.path.insert(0, str(TREE)) + +import mne # noqa: E402 + +mne.set_log_level("ERROR") + +DATA = BENCH / "data" +DATA_BIG = BENCH / "data_big" +TMP = Path(tempfile.mkdtemp(prefix="w7c_")) + + +def med_ms(fn, reps): + ts = [] + for _ in range(reps): + gc.collect() + gc.disable() + t0 = time.perf_counter() + fn() + ts.append((time.perf_counter() - t0) * 1e3) + gc.enable() + return float(np.median(ts)) + + +def main(): + out = {} + small = np.fromfile(DATA / "bench.bin", dtype=np.float32) + big = np.fromfile(DATA_BIG / "bench_big.bin", dtype=np.float32) + fpath = TMP / "f.bin" + + # floors + out["floor_small"] = med_ms(lambda: small.tofile(fpath), 9) + out["floor_big"] = med_ms(lambda: big.tofile(fpath), 5) + + # FIF small (preloaded) + from mne.io import read_raw_edf + + raw_s = read_raw_edf(DATA / "bench.edf", preload=True) + o = TMP / "s.fif" + out["fif_small"] = med_ms( + lambda: raw_s.save(o, fmt="single", overwrite=True), 9) + size_small = o.stat().st_size + + # FIF big + raw_b = read_raw_edf(DATA_BIG / "bench_big.edf", preload=True) + ob = TMP / "b.fif" + out["fif_big"] = med_ms( + lambda: raw_b.save(ob, fmt="single", overwrite=True), 5) + size_big = ob.stat().st_size + del raw_b + + # BV (installed pybv; caller copy included, subtract separately) + import pybv + + ch128 = [f"EEG{i:03d}" for i in range(128)] + volts = big.reshape(128, -1) * np.float32(1e-6) + + def w_bv(): + pybv.write_brainvision(data=volts.copy(), folder_out=str(TMP), + fname_base="v", sfreq=512, ch_names=ch128, + fmt="binary_float32", resolution=1e-7, + unit="µV", overwrite=True) + + w_bv() # warm + for p in TMP.glob("v.*"): + p.unlink() + out["bv_copy_only"] = med_ms(lambda: volts.copy(), 9) + out["bv_installed"] = med_ms(w_bv, 5) + for p in TMP.glob("v.*"): + p.unlink() + + # EDF end-to-end (construct+convert+write) and write-only + from edfio import Edf, EdfSignal + + d64 = big.reshape(128, -1).astype(np.float64) + pmin, pmax = float(d64.min()) - 1.0, float(d64.max()) + 1.0 + chn = [f"EEG{i:03d}" for i in range(128)] + + def mk(): + return Edf(signals=[ + EdfSignal(data=d64[i].copy(), sampling_frequency=512.0, + physical_range=(pmin, pmax), label=chn[i], + physical_dimension="uV") for i in range(128) + ]) + + oe = TMP / "e.edf" + obj = mk() + + def w_e(): + obj.write(oe) + + w_e() + oe.unlink() + out["edf_write_only"] = med_ms(w_e, 5) + out["edf_e2e"] = med_ms(lambda: mk().write(oe), 5) + oe.unlink() + + # BV-layout floor incl. transposition + out["floor_bv_layout"] = med_ms( + lambda: volts.ravel(order="F").tofile(TMP / "fl.eeg"), 5) + (TMP / "fl.eeg").unlink() + + shutil.rmtree(TMP, ignore_errors=True) + + mb = lambda x: x / 1e6 + print(f"\n=== W7 single-session pass ===") + print(f"{'measurement':<22} {'payload':>9} {'median':>9} {'MB/s':>7}") + rows = [ + ("floor_small", mb(small.nbytes)), + ("floor_big", mb(big.nbytes)), + ("fif_small", mb(size_small)), + ("fif_big", mb(size_big)), + ("bv_copy_only", mb(big.nbytes)), + ("bv_installed", mb(big.nbytes)), + ("edf_write_only", 236.0), + ("edf_e2e", 236.0), + ("floor_bv_layout", mb(big.nbytes)), + ] + for k, payload in rows: + ms = out[k] + rate = payload / (ms / 1e3) + print(f"{k:<22} {payload:>8.1f}M {ms:>8.2f} {rate:>7.0f}") + + b = np.array([size_small, size_big], float) + t = np.array([out["fif_small"], out["fif_big"]]) + slope = (t[1] - t[0]) / (b[1] - b[0]) # ms per byte + fixed = t[0] - slope * b[0] + print(f"\nFIF two-point fit: fixed={fixed:.2f} ms/file, " + f"marginal={1.0 / slope / 1e6 * 1e3:.0f} MB/s") + print(f"FIF small fixed-overhead share: {100 * fixed / t[0]:.0f}%") + print(f"FIF big vs floor: {out['floor_big'] / out['fif_big'] * 100:.0f}% of" + f" disk floor") + print(f"BV installed vs its-layout floor: " + f"{out['floor_bv_layout'] / out['bv_installed'] * 100:.0f}%") + (HERE / "consolidated.json").write_text(json.dumps(out, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/workspaces/W7/bench_w7_write.py b/benchmarks/io_dl/workspaces/W7/bench_w7_write.py new file mode 100644 index 00000000000..32fe7f2149c --- /dev/null +++ b/benchmarks/io_dl/workspaces/W7/bench_w7_write.py @@ -0,0 +1,644 @@ +"""W7: decompose write-path costs (FIF save, pybv BrainVision, edfio EDF) vs np.tofile floor. + + python benchmarks/io_dl/workspaces/W7/bench_w7_write.py # everything + python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only fif + python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only bv + python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only edf + python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only floor + +Method: medians over repeated runs, GC off, outputs written to /tmp and deleted. +Decomposition via accumulator-wrapping of mne internals (rebinding every module +attribute that holds the function, like profile_io.py's deepcopy tracer) and one +cProfile run for hotspot ranking. No library edits. +""" + +import argparse +import cProfile +import copy as _copy +import gc +import io as _io +import json +import pstats +import shutil +import sys +import tempfile +import time +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).resolve().parent +BENCH = HERE.parents[1] +TREE = BENCH.parents[1] +sys.path.insert(0, str(TREE)) + +import mne # noqa: E402 + +print("mne loaded from:", mne.__file__) +mne.set_log_level("ERROR") + +DATA = BENCH / "data" +DATA_BIG = BENCH / "data_big" +TMP = Path(tempfile.mkdtemp(prefix="w7_write_")) + +RESULTS = {} + + +def mb(x): + return x / 1e6 + + +def median_ms(fn, reps, *args, **kwargs): + ts = [] + out = None + for _ in range(reps): + gc.collect() + gc.disable() + t0 = time.perf_counter() + out = fn(*args, **kwargs) + dt = (time.perf_counter() - t0) * 1e3 + gc.enable() + ts.append(dt) + return float(np.median(ts)), ts, out + + +# -------------------------------------------------------------------------- +# accumulator instrumentation: rebind every module attribute holding `func` +# -------------------------------------------------------------------------- +class Acc: + def __init__(self, name): + self.name = name + self.reset() + + def reset(self): + self.ms = 0.0 + self.calls = 0 + self.extra = 0 # generic counter (bytes for _write) + + def report(self): + return f"{self.name:<28} {self.calls:>8} calls {self.ms:>9.2f} ms" + + +class Instrument: + """Wrap target functions with timing accumulators; restore afterwards.""" + + def __init__(self): + self.accs = {} + self.patched = [] # (holder, attrname, original) + self.in_payload = False # toggled by _write_raw_buffer spy + self.payload_bytes = 0 + self.meta_bytes = 0 + + def _make_spy(self, name, func): + acc = self.accs[name] + inst = self + + if name == "_write": + def spy(fid, data, kind, data_size, FIFFT_TYPE, dtype): + t0 = time.perf_counter_ns() + r = func(fid, data, kind, data_size, FIFFT_TYPE, dtype) + acc.ms += (time.perf_counter_ns() - t0) / 1e6 + acc.calls += 1 + try: + nb = int(np.asarray(data).nbytes) + 16 + except Exception: + nb = 16 + acc.extra += nb + if inst.in_payload: + inst.payload_bytes += nb + else: + inst.meta_bytes += nb + return r + elif name == "_write_raw_buffer": + def spy(fid, buf, cals, fmt): + inst.in_payload = True + try: + t0 = time.perf_counter_ns() + r = func(fid, buf, cals, fmt) + acc.ms += (time.perf_counter_ns() - t0) / 1e6 + acc.calls += 1 + finally: + inst.in_payload = False + return r + elif name == "deepcopy": + real = func + + def spy(x=None, memo=None, *a, **kw): + t0 = time.perf_counter_ns() + r = real(x, memo) if memo is not None else real(x) + acc.ms += (time.perf_counter_ns() - t0) / 1e6 + acc.calls += 1 + return r + else: + def spy(*args, **kwargs): + t0 = time.perf_counter_ns() + r = func(*args, **kwargs) + acc.ms += (time.perf_counter_ns() - t0) / 1e6 + acc.calls += 1 + return r + + spy.__name__ = f"spy_{name}" + return spy + + def __enter__(self): + import mne.annotations as ann_mod + import mne._fiff.meas_info as meas_info + import mne._fiff.write as fiff_write + import mne.io.base as io_base + import mne.utils.check as check_mod + + targets = { + "_check_fname": check_mod._check_fname, + "check_fname": check_mod.check_fname, + "_write": fiff_write._write, + "write_meas_info": meas_info.write_meas_info, + "_annotations_starts_stops": ann_mod._annotations_starts_stops, + "_write_annotations": ann_mod._write_annotations, + "_write_raw_buffer": io_base._write_raw_buffer, + "deepcopy": _copy.deepcopy, + } + for name, func in targets.items(): + self.accs[name] = Acc(name) + spy = self._make_spy(name, func) + # rebind in every loaded module that references the function object + n = 0 + for modname, mod in list(sys.modules.items()): + if mod is None or modname.split(".")[0] != "mne": + continue + try: + attrs = vars(mod) + except TypeError: + continue + for k, v in list(attrs.items()): + if v is func: + attrs[k] = spy + self.patched.append((mod, k, func)) + n += 1 + # deepcopy lives in stdlib too + if name == "deepcopy": + import copy as c2 + + if c2.deepcopy is func: + c2.deepcopy = spy + self.patched.append((c2, "deepcopy", func)) + return self + + def __exit__(self, *exc): + for holder, k, orig in reversed(self.patched): + try: + vars(holder)[k] = orig + except TypeError: + setattr(holder, k, orig) + self.patched.clear() + + def dump(self, total_ms, payload_mb): + print(f" {'component':<28} {'calls':>8} {'ms':>9} {'%total':>7}") + for a in self.accs.values(): + pct = 100 * a.ms / total_ms if total_ms else 0 + print(f" {a.name:<28} {a.calls:>8} {a.ms:>9.2f} {pct:>6.1f}%") + print(f" tag bytes: payload={mb(self.payload_bytes):.1f} MB, " + f"meta={mb(self.meta_bytes):.3f} MB") + other = total_ms - sum(a.ms for a in self.accs.values()) + print(f" {'(unattributed glue/loop)':<28} {'':>8} {other:>9.2f} " + f"{100 * other / total_ms:>6.1f}%") + + +def show(prof, n=18, sort="tottime"): + s = _io.StringIO() + st = pstats.Stats(prof, stream=s) + st.sort_stats(sort).print_stats(n) + txt = s.getvalue() + lines = txt.splitlines() + start = next(i for i, ln in enumerate(lines) if "ncalls" in ln) - 1 + print("\n".join(lines[start : start + n + 2])) + + +# -------------------------------------------------------------------------- +# FIF save decomposition +# -------------------------------------------------------------------------- +def bench_fif(): + from mne.io import read_raw_edf + + meta = json.loads((DATA / "meta.json").read_text()) + ch_names = [f"EEG{i:03d}" for i in range(meta["n_ch"])] + + raw_small = read_raw_edf(DATA / "bench.edf", preload=True) + out_small = TMP / "w7_small_raw.fif" + n_rep = 7 + med, ts, _ = median_ms( + lambda: raw_small.save(out_small, fmt="single", overwrite=True), n_rep + ) + size_small = out_small.stat().st_size + for p in TMP.glob("w7_small*"): + p.unlink() + rate = size_small / (med / 1e3) / 1e6 + RESULTS["fif_small"] = dict(file_mb=mb(size_small), ms=med, + mbps=rate, reps_ms=ts) + print(f"\n[FIF small] {mb(size_small):.1f} MB in {med:.1f} ms " + f"(median of {n_rep}) -> {rate:.0f} MB/s") + print(" reps ms:", ", ".join(f"{t:.1f}" for t in ts)) + + # proj=True variant (exercises info deepcopy path) + med_proj, _, _ = median_ms( + lambda: raw_small.save(TMP / "w7p_raw.fif", fmt="single", + proj=True, overwrite=True), 3 + ) + (TMP / "w7p_raw.fif").unlink(missing_ok=True) + print(f"[FIF small proj=True] {med_proj:.1f} ms (vs {med:.1f} default)") + RESULTS["fif_small_proj"] = dict(ms=med_proj) + + # instrumented small save (one rep; overhead ~us/call) + with Instrument() as ins: + t0 = time.perf_counter() + raw_small.save(out_small, fmt="single", overwrite=True) + tot = (time.perf_counter() - t0) * 1e3 + out_small.unlink() + print(f"\n[FIF small decomposed] total {tot:.1f} ms (instrumented)") + ins.dump(tot, mb(size_small)) + RESULTS["fif_small_decomp"] = { + a.name: dict(calls=a.calls, ms=a.ms) for a in ins.accs.values() + } | {"payload_bytes": ins.payload_bytes, "meta_bytes": ins.meta_bytes} + + # cProfile cross-check + prof = cProfile.Profile() + prof.enable() + raw_small.save(out_small, fmt="single", overwrite=True) + prof.disable() + out_small.unlink() + print("\n[FIF small cProfile top tottime]") + show(prof, n=16) + + # ---- BIG ---- + print("\nloading BIG edf (preload) ...", flush=True) + t0 = time.perf_counter() + raw_big = read_raw_edf(DATA_BIG / "bench_big.edf", preload=True) + print(f" loaded in {time.perf_counter() - t0:.1f} s") + out_big = TMP / "w7_big_raw.fif" + med_b, ts_b, _ = median_ms( + lambda: raw_big.save(out_big, fmt="single", overwrite=True), 3 + ) + size_big = out_big.stat().st_size + out_big.unlink() + rate_b = size_big / (med_b / 1e3) / 1e6 + RESULTS["fif_big"] = dict(file_mb=mb(size_big), ms=med_b, mbps=rate_b, + reps_ms=ts_b) + print(f"[FIF big] {mb(size_big):.1f} MB in {med_b:.1f} ms " + f"(median of 3) -> {rate_b:.0f} MB/s") + + with Instrument() as ins: + t0 = time.perf_counter() + raw_big.save(out_big, fmt="single", overwrite=True) + tot_b = (time.perf_counter() - t0) * 1e3 + out_big.unlink() + print(f"[FIF big decomposed] total {tot_b:.1f} ms (instrumented)") + ins.dump(tot_b, mb(size_big)) + RESULTS["fif_big_decomp"] = { + a.name: dict(calls=a.calls, ms=a.ms) for a in ins.accs.values() + } + + del raw_big + + # two-point fixed/marginal fit: t(bytes) = fixed + bytes/marginal_rate + b = np.array([size_small, size_big], dtype=float) + t = np.array([med, med_b]) + A = np.vstack([np.ones_like(b), b]).T + coef, *_ = np.linalg.lstsq(A, t, rcond=None) + fixed_ms, per_byte = coef + RESULTS["fif_fit"] = dict(fixed_ms=float(fixed_ms), + marginal_mbps=float(1e3 / per_byte / 1e6)) + print(f"\n[FIF linear fit] fixed ≈ {fixed_ms:.1f} ms/file; marginal ≈ " + f"{1e3 / per_byte / 1e6:.0f} MB/s") + pct_small = 100 * fixed_ms / med + print(f"[FIF small fixed overhead] ≈ {pct_small:.0f}% of the " + f"{med:.1f} ms small save") + + +def bench_fif_extra(): + """Non-preloaded save (the historical ~103ms case?), per-buffer fetch, + cold-start single save in a fresh interpreter.""" + from mne.io import read_raw_edf + + meta = json.loads((DATA / "meta.json").read_text()) + + # (a) non-preloaded source -> save reads through the EDF reader per buffer + raw_nop = read_raw_edf(DATA / "bench.edf", preload=False) + out = TMP / "w7n_raw.fif" + med, ts, _ = median_ms( + lambda: raw_nop.save(out, fmt="single", overwrite=True), 5 + ) + size = out.stat().st_size + out.unlink() + print(f"\n[FIF small NON-preloaded] {mb(size):.1f} MB in {med:.1f} ms " + f"(median of 5) -> {size / (med / 1e3) / 1e6:.0f} MB/s") + print(" reps ms:", ", ".join(f"{t:.1f}" for t in ts)) + RESULTS["fif_small_nopreload"] = dict(ms=med, + mbps=size / (med / 1e3) / 1e6) + + with Instrument() as ins: + t0 = time.perf_counter() + raw_nop.save(out, fmt="single", overwrite=True) + tot = (time.perf_counter() - t0) * 1e3 + out.unlink() + print(f"[FIF small NON-preloaded decomposed] total {tot:.1f} ms") + ins.dump(tot, mb(size)) + + # (b) instrumented preloaded save incl. per-buffer fetch (__getitem__) + raw_small = read_raw_edf(DATA / "bench.edf", preload=True) + import mne.io.base as io_base + + orig_gi = io_base.BaseRaw.__getitem__ + acc = Acc("BaseRaw.__getitem__") + + def gi_spy(*a, **kw): + t0 = time.perf_counter_ns() + r = orig_gi(*a, **kw) + acc.ms += (time.perf_counter_ns() - t0) / 1e6 + acc.calls += 1 + return r + + io_base.BaseRaw.__getitem__ = gi_spy + try: + with Instrument() as ins: + t0 = time.perf_counter() + raw_small.save(out, fmt="single", overwrite=True) + tot = (time.perf_counter() - t0) * 1e3 + finally: + io_base.BaseRaw.__getitem__ = orig_gi + out.unlink() + print(f"\n[FIF small preloaded +fetch] total {tot:.1f} ms; " + f"__getitem__: {acc.calls} calls {acc.ms:.2f} ms") + ins.dump(tot, mb(size)) + RESULTS["fif_small_fetch"] = dict(total_ms=tot, getitem_calls=acc.calls, + getitem_ms=acc.ms) + + # (c) cold start: fresh process, one save, wall time end-to-end + import subprocess + + code = ( + "import sys, time; sys.path.insert(0, %r);" + "import mne; mne.set_log_level('ERROR');" + "from mne.io import read_raw_edf;" + "t0=time.perf_counter();" + "raw=read_raw_edf(%r, preload=True); t1=time.perf_counter();" + "raw.save(%r, fmt='single', overwrite=True); t2=time.perf_counter();" + "print(f'{(t1-t0)*1e3:.1f} {(t2-t1)*1e3:.1f}')" + ) % (str(TREE), str(DATA / "bench.edf"), str(TMP / "w7c_raw.fif")) + r = subprocess.run([sys.executable, "-c", code], capture_output=True, + text=True, timeout=120) + load_ms, save_ms = r.stdout.split() + (TMP / "w7c_raw.fif").unlink(missing_ok=True) + print(f"\n[FIF cold process] load {load_ms} ms + save {save_ms} ms " + f"(single-shot, includes imports/page-cache warm-up)") + RESULTS["fif_cold"] = dict(load_ms=float(load_ms), save_ms=float(save_ms)) + + +def bench_bv(): + import pybv + + meta = json.loads((DATA_BIG / "meta.json").read_text()) + n_ch, sfreq = meta["n_ch"], int(meta["sfreq"]) + ch_names = [f"EEG{i:03d}" for i in range(n_ch)] + # pybv expects VOLTS; fresh copy per call because the locally-patched + # installed pybv scales `data` IN PLACE (caller-visible mutation). + base_uv = np.fromfile(DATA_BIG / "bench_big.bin", dtype=np.float32).reshape( + n_ch, -1 + ) + payload = base_uv.nbytes + print(f"\n[BV] payload {mb(payload):.1f} MB f32, pybv " + f"{getattr(pybv, '__version__', '?')}") + + def make_w(): + volts = base_uv * 1e-6 # prepare once, outside timed region + + def w(): + # fresh copy per call: installed pybv scales `data` IN PLACE + pybv.write_brainvision(data=volts.copy(), folder_out=str(TMP), + fname_base="w7bv", sfreq=sfreq, + ch_names=ch_names, fmt="binary_float32", + events=None, resolution=1e-7, unit="µV", + overwrite=True) + return w + + w = make_w() + + # warm once (imports etc.), then time; delete between runs + w() + for p in TMP.glob("w7bv.*"): + p.unlink() + + n_rep = 5 + med, ts, _ = median_ms(w, n_rep) + eeg = (TMP / "w7bv.eeg").stat().st_size + for p in TMP.glob("w7bv.*"): + p.unlink() + rate = payload / (med / 1e3) / 1e6 + RESULTS["bv_big"] = dict(payload_mb=mb(payload), ms=med, mbps=rate, + eeg_mb=mb(eeg)) + print(f"[BV big] wrote .eeg {mb(eeg):.1f} MB in {med:.1f} ms " + f"(median of {n_rep}) -> {rate:.0f} MB/s (of {mb(payload):.0f} MB payload)") + + # decompose: time _write_bveeg_file alone vs header writing + import pybv.io as pio + + acc = {"bveeg": 0.0, "calls": 0} + orig = pio._write_bveeg_file + + def spy(*a, **kw): + t0 = time.perf_counter_ns() + r = orig(*a, **kw) + acc["bveeg"] += (time.perf_counter_ns() - t0) / 1e6 + acc["calls"] += 1 + return r + + pio._write_bveeg_file = spy + try: + med2, _, _ = median_ms(w, n_rep) + finally: + pio._write_bveeg_file = orig + for p in TMP.glob("w7bv.*"): + p.unlink() + bveeg_per_call = acc["bveeg"] / n_rep + print(f"[BV decomposed] binary .eeg write: {bveeg_per_call:.1f} ms of " + f"{med2:.1f} ms total ({100 * bveeg_per_call / med2:.0f}%); " + f"headers/vmrk outside: {med2 - bveeg_per_call:.1f} ms") + RESULTS["bv_decomp"] = dict(total_ms=med2, bveeg_ms=bveeg_per_call) + + # cost of the per-call defensive copy we must make (installed pybv mutates) + volts = (base_uv * 1e-6) + copy_med, _, _ = median_ms(lambda: volts.copy(), 5) + print(f"[BV caller-side copy] {copy_med:.1f} ms (included in totals above; " + f"pybv mutates its input)") + RESULTS["bv_copy_ms"] = copy_med + + # BV-layout floor: multiplexed (time-major) write incl. transposition + floor_f, _, _ = median_ms(lambda: volts.ravel(order="F").tofile( + TMP / "w7floor.eeg"), 3) + (TMP / "w7floor.eeg").unlink() + print(f"[BV layout floor] ravel(F)+tofile: {floor_f:.1f} ms -> " + f"{mb(payload) / (floor_f / 1e3):.0f} MB/s") + + # scaling-pass cost estimate: the two in-place multiplies + range checks + scaled = base_uv.copy() + scales = np.ones((n_ch, 1)) + t0 = time.perf_counter() + scaled *= scales + t_mult1 = (time.perf_counter() - t0) * 1e3 + t0 = time.perf_counter() + scaled *= 1e-6 + t_mult2 = (time.perf_counter() - t0) * 1e3 + t0 = time.perf_counter() + ok = np.all(scaled >= np.finfo(np.float32).min) and np.all( + scaled <= np.finfo(np.float32).max) + t_range = (time.perf_counter() - t0) * 1e3 + del scaled, ok + print(f"[BV internal passes] mult1 {t_mult1:.1f} ms + mult2 {t_mult2:.1f}" + f" ms + range checks {t_range:.1f} ms " + f"(≈{t_mult1 + t_mult2 + t_range:.0f} ms of pure memory passes)") + RESULTS["bv_passes"] = dict(mult1=t_mult1, mult2=t_mult2, range=t_range) + + # pristine upstream pybv (3 internal full-array copies) via subprocess + import subprocess + + up = Path("/tmp/pybv_up/x") + if up.exists(): + code = ( + "import sys, time, tempfile, shutil, gc, numpy as np;" + "sys.path.insert(0, %r); import pybv;" + "print('using', pybv.__file__);" + "d=np.fromfile(%r, dtype=np.float32).reshape(128, -1);" + "ch=[f'EEG{i:03d}' for i in range(128)];" + "tmp=tempfile.mkdtemp();" + "def w():" + " d2=(d*1e-6).astype('f4');" + " pybv.write_brainvision(data=d2, folder_out=tmp, fname_base='p'," + "sfreq=512, ch_names=ch, fmt='binary_float32', resolution=1e-7," + "overwrite=True);" + "w(); ts=[];" + "for i in range(3):" + " shutil.rmtree(tmp, ignore_errors=True); tmp=tempfile.mkdtemp();" + " gc.collect(); gc.disable(); t0=time.perf_counter(); w();" + " ts.append(time.perf_counter()-t0); gc.enable()" + "print('pristine ms:', sorted(ts)[1] * 1e3)" + ) % (str(up), str(DATA_BIG / "bench_big.bin")) + r = subprocess.run([sys.executable, "-c", code], capture_output=True, + text=True, timeout=300) + out = r.stdout.strip().splitlines() + for ln in out: + if "using" in ln or "pristine" in ln: + print(f"[BV {ln.split(':')[0].strip()}]") + if "ms" in ln: + ms = float(ln.rsplit(":", 1)[1]) + RESULTS["bv_pristine_ms"] = ms + print(f" -> {mb(payload) / (ms / 1e3):.0f} MB/s") + + +def bench_edf(): + from edfio import Edf, EdfSignal + + meta = json.loads((DATA_BIG / "meta.json").read_text()) + n_ch, sfreq = meta["n_ch"], meta["sfreq"] + ch_names = [f"EEG{i:03d}" for i in range(n_ch)] + f32 = np.fromfile(DATA_BIG / "bench_big.bin", dtype=np.float32).reshape( + n_ch, -1 + ) + data = f32.astype(np.float64) # what generate_data feeds edfio + pmin, pmax = float(data.min()) - 1.0, float(data.max()) + 1.0 + payload_i16 = data.size * 2 + + def mk(): + return Edf(signals=[ + EdfSignal(data=data[i].copy(), sampling_frequency=sfreq, + physical_range=(pmin, pmax), label=ch_names[i], + physical_dimension="uV") for i in range(n_ch) + ]) + + def w(edf): + edf.write(TMP / "w7.edf") + + edf_obj = mk() + w(edf_obj) # warm + (TMP / "w7.edf").unlink() + + n_rep = 5 + med, ts, _ = median_ms(lambda: w(edf_obj), n_rep) + size = (TMP / "w7.edf").stat().st_size + (TMP / "w7.edf").unlink() + rate = size / (med / 1e3) / 1e6 + RESULTS["edf_big"] = dict(file_mb=mb(size), ms=med, mbps=rate, + payload_i16_mb=mb(payload_i16)) + print(f"\n[EDF big] write-only {mb(size):.1f} MB (int16) in {med:.1f} ms " + f"(median of {n_rep}) -> {rate:.0f} MB/s") + print(" note: digital int16 conversion happens at EdfSignal CONSTRUCTION," + " not in write(); conversion cost measured separately below") + + def full(): + w(mk()) + + med_e, _, _ = median_ms(full, n_rep) + (TMP / "w7.edf").unlink() + RESULTS["edf_big_e2e"] = dict(ms=med_e, mbps=size / (med_e / 1e3) / 1e6) + print(f"[EDF big end-to-end (construct+convert+write)] {med_e:.1f} ms -> " + f"{mb(size) / (med_e / 1e3):.0f} MB/s") + + # digital-conversion pass analog (what edfio must do before disk) + res = (pmax - pmin) / 32767.0 + t0 = time.perf_counter() + dig = ((data - pmin) / res).round().astype(np.int16) + t_conv = (time.perf_counter() - t0) * 1e3 + del dig + print(f"[EDF reference conversion pass] f64->i16 vectorized: " + f"{t_conv:.1f} ms (memory-bound)") + + +def bench_floor(): + small = np.fromfile(DATA / "bench.bin", dtype=np.float32) + big = np.fromfile(DATA_BIG / "bench_big.bin", dtype=np.float32) + out = TMP / "floor.bin" + + med_s, _, _ = median_ms(lambda: small.tofile(out), 7) + med_b, _, _ = median_ms(lambda: big.tofile(out), 5) + out.unlink() + print(f"\n[floor tofile] small {mb(small.nbytes):.1f} MB: {med_s:.1f} ms " + f"-> {small.nbytes / (med_s / 1e3) / 1e6:.0f} MB/s") + print(f"[floor tofile] big {mb(big.nbytes):.1f} MB: {med_b:.1f} ms " + f"-> {big.nbytes / (med_b / 1e3) / 1e6:.0f} MB/s") + RESULTS["floor"] = dict( + small_ms=med_s, small_mbps=small.nbytes / (med_s / 1e3) / 1e6, + big_ms=med_b, big_mbps=big.nbytes / (med_b / 1e3) / 1e6, + ) + + # open/close syscall cost (per-save fixed floor) + t0 = time.perf_counter() + for i in range(200): + with open(TMP / f"f{i}.bin", "wb"): + pass + t_open = (time.perf_counter() - t0) / 200 * 1e3 + for i in range(200): + (TMP / f"f{i}.bin").unlink() + print(f"[floor] bare open+close: {t_open * 1000:.0f} us") + RESULTS["floor"]["open_us"] = t_open * 1000 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--only", choices=["fif", "fifx", "bv", "edf", "floor"]) + args = ap.parse_args() + try: + if args.only in (None, "floor"): + bench_floor() + if args.only in (None, "fif"): + bench_fif() + if args.only in (None, "fifx"): + bench_fif_extra() + if args.only in (None, "bv"): + bench_bv() + if args.only in (None, "edf"): + bench_edf() + finally: + shutil.rmtree(TMP, ignore_errors=True) + (HERE / "results.json").write_text(json.dumps(RESULTS, indent=2)) + print(f"\nresults -> {HERE / 'results.json'}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/io_dl/workspaces/W7/consolidated.json b/benchmarks/io_dl/workspaces/W7/consolidated.json new file mode 100644 index 00000000000..913895f5394 --- /dev/null +++ b/benchmarks/io_dl/workspaces/W7/consolidated.json @@ -0,0 +1,11 @@ +{ + "floor_small": 1.8324169795960188, + "floor_big": 70.51770799444057, + "fif_small": 16.183792002266273, + "fif_big": 319.24983399221674, + "bv_copy_only": 34.75150000303984, + "bv_installed": 2424.0037919953465, + "edf_write_only": 65.58979200781323, + "edf_e2e": 263.01145800971426, + "floor_bv_layout": 211.66029199957848 +} \ No newline at end of file diff --git a/benchmarks/io_dl/workspaces/W7/findings.md b/benchmarks/io_dl/workspaces/W7/findings.md new file mode 100644 index 00000000000..4eb2dfaabaa --- /dev/null +++ b/benchmarks/io_dl/workspaces/W7/findings.md @@ -0,0 +1,174 @@ +# W7 — Write-path decomposition for DL caching pipelines + +Date: 2026-08-24 · Agent W7 · Measurement-only (no library edits). +Machine: Apple Silicon (arm64), macOS 15, py 3.12.12, numpy 2.4.6, MNE 1.13.dev0 +(this working tree, incl. sessions 1–4 reader/access patches), edfio 0.4.13, +pybv 0.7.6 (locally patched in site-packages — see §3). + +Method: medians over repeated saves to /tmp (deleted after), GC off, floors and +writers measured back-to-back in single-session passes (`bench_w7_consolidated.py`, +3 passes; spread reported). Decomposition via accumulator-wrapping of mne +internals (module-attribute rebinding, no cProfile inflation) + one cProfile run. +Prototypes monkeypatched in-process, outputs md5-verified byte-identical. + +## 1. Headline table + +| writer | payload | MB/s | % fixed overhead | top hotspot function / file:line | +|---|---|---|---|---| +| **FIF save, small, warm preloaded** | 19.7 MB f32 | **1220** (16.2 ms) | **~19 %** true logic (fit: 2–6 ms/file); ~87 % if first-save-in-process (see §2.1) | payload cast+transpose+`tobytes` copies: `write_float`/`_write` mne/_fiff/write.py:94,27; per-buffer fetch `BaseRaw.__getitem__` mne/io/base.py:973 | +| **FIF save, small, NON-preloaded** | 19.7 MB | 472 (41.7 ms) | ~60 % = per-buffer source reads through EDF reader (~87 µs × 300 buffers) | `_read_segment_file` buffer fetch per 1 s output buffer, called from `_write_raw_data` loop mne/io/base.py:3197 | +| FIF save, big (128 ch × 1800 s) | 471.9 MB f32 | **1478–1600** (295–320 ms) | ≈0 (marginal regime) | same as small; CPU side = `buf / cals` copy + `np.array(">f4")` strided cast + `.tobytes()` extra copy (`_write_raw_buffer` mne/io/base.py:3324 → `_write`) | +| floor: `np.tofile` big f32 | 471.9 MB | 4400–5700 (71–107 ms) | – | – | +| **BrainVision (installed pybv)** | 471.9 MB f32 | **199** (2420 ms) | ~90 % vs its own layout floor | chunked `chunk.T.tofile(fout)` pybv/io.py:827 — non-contiguous transposed write | +| BrainVision (pristine upstream 0.7.6) | 471.9 MB f32 | **958** (493 ms) | ~40 % | 3 internal full-array copies: `data * scales`, `astype(dtype)`, `ravel(order="F")` pybv/io.py:683,799,811 | +| floor: multiplexed layout (`ravel(F)`+tofile) | 471.9 MB | 1929–2406 (197–245 ms) | – | – | +| prototype fix (chunked `ascontiguousarray(chunk.T)`) | 471.9 MB | **2628** (191 ms) | ~0 | one fused transpose-copy per ≤32 MB block | +| **EDF write (edfio), serialization only** | 236.0 MB i16 | **3156–4000** (59–75 ms) | ≈0 | single contiguous `data_record.tofile()` edfio/_edfio.py `Edf.write` | +| EDF end-to-end (construct + int16 convert + write) | 236.0 MB i16 | **772–897** (263–306 ms) | conversion-dominated | f64→i16 digital quantization at `EdfSignal` construction (memory-bound pass, ~330 ms measured standalone) | + +Notes on the table: +- "% fixed overhead" for FIF small = two-point fit `t(bytes) = fixed + bytes/rate` + over {small, big} within-session: fixed = 2.3–5.7 ms/file across 3 passes, + marginal rate = 1490–1632 MB/s. The earlier "~190 MB/s ⇒ huge fixed cost" flag + is explained by §2.1/§2.2, not by tag machinery. +- BV "its own layout floor" = writing the same multiplexed (time-major) byte + layout via `volts.ravel(order="F").tofile()`. +- edfio converts float→digital-int16 at signal construction, so its `write()` + alone is nearly free; the honest end-to-end number is the last row. + +## 2. FIF save decomposition (mission item 1) + +Instrumented component breakdown (accumulator wrapping, ms): + +| component | small warm preloaded (17.6 ms total) | small NON-preloaded (41.7 ms) | big preloaded (281 ms instr.) | +|---|---|---|---| +| `_check_fname` + `check_fname` regexes | 0.13 (0.7 %) | 0.13 | 0.15 | +| info deepcopy | 0.15 (35 calls) | 0.14 | 0.24 | +| `write_meas_info` (metadata tags) | 0.68 | 0.67 | 1.27 | +| annotations handling | 0.01 (empty) | 0.01 | 0.01 | +| split-size logic (`fid.tell` checks) | <0.5 (in glue) | <0.5 | <1 | +| `_write_raw_buffer` (incl. payload `_write` inside) | 11.54 | 9.94 | 214.31 | +| per-buffer source read (`raw[picks, first:last]` via EDF reader) | 3.44 (11.5 µs/call) | **≈25** (glue row, 60 %) | ~21 (est. from 11.5 µs × 1800) | +| unattributed Python glue (loop, asserts, logger) | ~1 | 24.75 | ~45 | + +Tag counts: 319 tags for small (300 payload + 19 metadata), 1819 for big; +metadata bytes = 0.0003 MB. Disk I/O is only 4.6 ms (small) / 71–107 ms (big) +of the totals — everything else is CPU-side array traffic. + +### 2.1 The historical "~103 ms" mystery — SOLVED + +A fresh process's first `raw.save()` costs **121 ms**; with `mne_bids` already +imported it is **16 ms**. Cause: `_check_fname(..., check_bids_split=True)` +(mne/utils/check.py:270, called from BaseRaw.save at mne/io/base.py:1967 and +Epochs.save at mne/epochs.py:2357) executes `from mne_bids import BIDSPath` +inside a try/except. When mne-bids is installed this lazily imports mne_bids → +mne.viz → jinja2 → scipy.sparse/scipy.special/scipy.io (hundreds of modules; +standalone `from mne_bids import BIDSPath` = 265–279 ms; ≈105 ms after +`import mne`). Every caching pipeline pays this once per process (per DataLoader +worker!), and it was misattributed to tag serialization. + +Secondary factor: the old benchmark saved a NON-preloaded raw +(`bench_io.py::bench_micro`, `profile_io.profile_save`), adding per-buffer EDF +reader costs (now 42 ms; was ~100 ms before the session-1 reader patches). + +### 2.2 Rejected hypotheses (measured, not guessed) + +- **info deepcopy**: NOT executed by default (`proj=False` passes `self.info` by + reference, base.py:2006). With `proj=True`: +3.9 ms total, almost all + `setup_proj` SVD work, not the copy; `deepcopy(info)` alone ≈ 0.05 ms @64ch. + "Skip deepcopy when projs/comps empty" would save ~nothing on DL paths. +- **Tag serialization loop**: 19 metadata tags = 0.68 ms. Negligible. +- **Annotations**: empty case = 0.01 ms. +- **Split-size logic**: two `fid.tell()` + comparisons per buffer; µs-scale. +- **Bigger `buffer_size_sec`**: tested 4/16/60 s on BIG save — all SLOWER than + default 1 s (400/425/382 vs 297 ms; larger temps blow L2/L3). Reject. + +## 3. BrainVision writer (mission item 2) + +The site-packages pybv 0.7.6 has been locally patched (vs pristine PyPI wheel: +`data * scales` → in-place `*=`, full `ravel(order="F").tofile()` → chunked +`chunk.T.to_file`). Two defects introduced: + +1. **Performance pessimization**: `ndarray.T.tofile()` on a (128, n)-slice view + writes ~10× slower than necessary (numpy iterates the transposed view with + poor buffering): microbench same layout — `T.tofile` 2244–2341 ms vs + `ravel(F).tofile` 224 ms. Result: 199 MB/s installed vs 958 MB/s pristine + vs 2628 MB/s achievable (§1 prototype B: `np.ascontiguousarray(chunk.T)` + per ≤32 MB block, then scale+write — bounded memory AND fastest). +2. **Caller-visible mutation**: in-place scaling multiplies the user's array by + unit/resolution factors every call (verified: input grows ×1e13 per call, + overflow→inf→ValueError on 3rd call). Upstream returns copies; any pipeline + re-using a data array across `write_brainvision` calls silently corrupts it. + +Cost structure otherwise: scaling passes ≈108 ms of pure memory traffic +(2 multiplies + range check) out of ~2400 ms — i.e., I/O-pattern cost, not +copies, dominate. `raw.export(fmt="brainvision")` routes through pybv, so MNE +inherits both defects. + +## 4. EDF writer (mission item 3) + +Confirmed and refined the earlier "~590 MB/s": that was end-to-end. edfio's +`Edf.write` itself moves bytes at **~3200–4000 MB/s** (one contiguous +`tofile` of a pre-packed uint8 record matrix); the f64→int16 digital +quantization happens earlier, at `EdfSignal` construction, and costs ~330 ms +for 118 M samples (a pure memory-bound pass; measured standalone). End-to-end +construct+convert+write = 263–306 ms (**772–897 MB/s**, matches the earlier +observation). Headroom: the conversion could stream/fuse with generation, or +accept f32 input directly (halving one pass); disk-side nothing left. + +## 5. TOP-3 fix proposals (ranked by expected gain) + +### #1 — Fix the BrainVision binary write path (pybv patch or upstream PR) +Replace `chunk.T.tofile(fout)` with +`np.ascontiguousarray(data[:, s:e].T).tofile(fout)` blocks (≤32 MB), folding +the two scaling multiplies into one multiply on each transposed block (also +removes the caller-mutation bug since scaling then happens on the private +block copy). Expected: **2420 → ~200 ms on a 472 MB file (≈12×)**; MNE +`raw.export(fmt="brainvision")` inherits the win. Risk: none to output bytes +(layout identical; md5-verifiable against upstream writer). + +### #2 — De-guard the lazy `mne_bids` import in `_check_fname` +Only attempt `from mne_bids import BIDSPath` when `"mne_bids" in sys.modules` +(a real BIDSPath cannot reach these call sites unless mne-bids was imported), +or duck-type on `hasattr(fname, "split")`. Expected: **first save per process +121 → 16 ms (−105 ms)** for every pipeline with mne-bids installed; hits +`Raw.save` and `Epochs.save` (the two DL caching entry points) plus read-path +callers of `_check_fname`. Behavior-identical (import succeeds iff mne_bids +already loaded ⇒ identical validation outcomes). Trivially upstreamable. + +### #3 — Fuse the FIF payload buffer write into one pass +In `_write_raw_buffer` (single fmt): divide once into a fresh `(n_t, n_ch)` +`">f4"` C-contig output and `fid.write(out)` directly (buffer protocol), +eliminating the `buf/cals` temp, the separate strided cast, and the redundant +`tobytes()` bytes-object copy (4 payload passes → 2). Prototype verified +byte-identical md5 on the 472 MB save; measured gain modest but real: +**≈5–15 % on warm saves** (big 306→293 ms; small ~16→~14 ms), larger on +memory-constrained machines. Complementary free win: avoid saving from +NON-preloaded raws when possible (document/batch source reads) — that state +costs 42 vs 16 ms per small file (2.7×) today. + +Rejected for lack of measured support: skip-info-deepcopy fast path (§2.2), +default `buffer_size_sec` changes (§2.2), split-size/tag-loop micro-opts. + +## 6. Repro + +```bash +cd /Users/braristimunha/Projects/libraries/mne_python/mne_python_more_io_speed +python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only floor # np.tofile floors +python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only fif # small+BIG FIF decomp, cProfile +python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only fifx # non-preloaded, fetch, cold-process A/B +python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only bv # installed vs pristine pybv + prototypes +python benchmarks/io_dl/workspaces/W7/bench_w7_write.py --only edf # write-only vs end-to-end +python benchmarks/io_dl/workspaces/W7/bench_w7_consolidated.py # single-session pass (run 2-3x) +# cold-save mne_bids A/B: +python -c "import time,sys; sys.path.insert(0,'.'); import mne; mne.set_log_level('ERROR'); \ + from mne.io import read_raw_edf; r=read_raw_edf('benchmarks/io_dl/data/bench.edf',preload=True); \ + t=time.perf_counter(); r.save('/tmp/x.fif',overwrite=True); print((time.perf_counter()-t)*1e3)" +# compare with 'import mne_bids' inserted before the timer. +``` + +JSON artifacts: `results.json`, `consolidated.json` next to this file. +Requires `pip download pybv==0.7.6 --no-deps -d /tmp/pybv_up && unzip +/tmp/pybv_up/*.whl -d /tmp/pybv_up/x` for the pristine-pybv arm. + +COMPLETE: W7 - The FIF "fixed overhead" is a lazy `from mne_bids import BIDSPath` inside `_check_fname(check_bids_split=True)` costing ~105 ms of every process's first save (121→16 ms measured); worst writer is the locally-patched pybv at 199 MB/s (fixable to ~2600 MB/s, 12×), best is edfio serialization at ~3600 MB/s. diff --git a/benchmarks/io_dl/workspaces/W7/results.json b/benchmarks/io_dl/workspaces/W7/results.json new file mode 100644 index 00000000000..04e0098b60c --- /dev/null +++ b/benchmarks/io_dl/workspaces/W7/results.json @@ -0,0 +1,9 @@ +{ + "floor": { + "small_ms": 1.7277080041822046, + "small_mbps": 11379.700708920584, + "big_ms": 84.28358301171102, + "big_mbps": 5598.471056153797, + "open_us": 42.57041509845294 + } +} \ No newline at end of file diff --git a/doc/changes/dev/14202.bugfix.rst b/doc/changes/dev/14202.bugfix.rst new file mode 100644 index 00000000000..81b6bf816a2 --- /dev/null +++ b/doc/changes/dev/14202.bugfix.rst @@ -0,0 +1,11 @@ +Sped up raw data access for workloads making many small reads (e.g., deep-learning +training loops): ``Raw.get_data`` no longer materializes the full time axis or +re-resolves channel picks on every call; EDF/BDF readers gained vectorized and +optional-numba fast paths (bit-identical output, including stim channels); the FIF +reader now reads tag payloads through a memory map; batched internal window reads +avoid per-window allocations. + +Memory-mapped caches passed to ``Raw.load_data(memmap=...)`` are now reused +when valid instead of being re-decoded, and are no longer deleted when the +:class:`~mne.io.Raw` object is garbage-collected — reopening large recordings +becomes near-instantaneous and cache lifetime is managed by the caller. diff --git a/mne/_fiff/_mmap_cache.py b/mne/_fiff/_mmap_cache.py new file mode 100644 index 00000000000..e96a2e46a5a --- /dev/null +++ b/mne/_fiff/_mmap_cache.py @@ -0,0 +1,38 @@ +"""PID-keyed memmap caching for direct byte-offset reads. + +Used by readers that need random access into raw data files (currently the +FIF raw reader). Keyed by PID so forked worker processes (e.g., PyTorch +DataLoader workers) create their own mapping instead of sharing a parent's, +and validated against file size/mtime so stale mappings are never reused. +""" + +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors + +import os + +import numpy as np + +_MAX_CACHE = 16 +_cache = {} + + +def get_u8_memmap(path): + """Return a uint8 memmap of *path* (PID-keyed), or None on any failure.""" + try: + st = os.stat(path) + key = (os.getpid(), str(path)) + hit = _cache.get(key) + if hit is not None: + mm, mtime_ns, size = hit + if mtime_ns == st.st_mtime_ns and size == st.st_size: + return mm + _cache.pop(key, None) + mm = np.memmap(str(path), dtype=np.uint8, mode="r") + except Exception: + return None + _cache[key] = (mm, st.st_mtime_ns, st.st_size) + while len(_cache) > _MAX_CACHE: + _cache.pop(next(iter(_cache))) + return mm diff --git a/mne/_fiff/pick.py b/mne/_fiff/pick.py index 7f62644c254..b908fba22f5 100644 --- a/mne/_fiff/pick.py +++ b/mne/_fiff/pick.py @@ -1334,6 +1334,18 @@ def _picks_to_idx( ) raise TypeError(msg) del extra_repr + # Fast path: an integer ndarray with all values already in range needs no + # copy or further checks. This matters for callers resolving picks on + # every access (e.g., Raw.get_data in deep-learning training loops). + if ( + picks.dtype.kind == "i" + and picks.size + and picks.min() >= 0 + and picks.max() < n_chan + ): + if return_kind: + return picks, picked_ch_type_or_generic + return picks picks = picks.astype(int) # diff --git a/mne/_fiff/utils.py b/mne/_fiff/utils.py index b158914bb88..df28c582596 100644 --- a/mne/_fiff/utils.py +++ b/mne/_fiff/utils.py @@ -73,22 +73,28 @@ def _find_channels(ch_names, ch_type="EOG"): def _mult_cal_one(data_view, one, idx, cals, mult): """Take a chunk of raw data, multiply by mult or cals, and store.""" - one = np.asarray(one, dtype=data_view.dtype) assert data_view.shape[1] == one.shape[1], ( data_view.shape[1], one.shape[1], ) # noqa: E501 if mult is not None: + one = np.asarray(one, dtype=data_view.dtype) assert mult.ndim == one.ndim == 2 data_view[:] = mult @ one[idx] else: assert cals is not None if isinstance(idx, slice): - data_view[:] = one[idx] + # Hot path: gather + type-cast + calibration in a single pass, + # without materializing an intermediate float64 copy of `one` + # (`one[idx]` is a view for basic slices). Numerically identical + # to cast-then-scale because both are elementwise. + np.multiply(one[idx], cals.reshape(-1, 1), out=data_view, + casting="unsafe") else: + one = np.asarray(one, dtype=data_view.dtype) # faster than doing one = one[idx] np.take(one, idx, axis=0, out=data_view) - data_view *= cals + data_view *= cals def _blk_read_lims(start, stop, buf_len): @@ -215,6 +221,8 @@ def _read_segments_file( if n_channels is None: n_channels = raw._raw_extras[fi]["orig_nchan"] + import os as _os + n_bytes = np.dtype(dtype).itemsize # data_offset and data_left count data samples (channels x time points), # not bytes. @@ -224,6 +232,49 @@ def _read_segments_file( # Read up to 100 MB of data at a time, block_size is in data samples block_size = ((int(100e6) // n_bytes) // n_channels) * n_channels block_size = min(data_left, block_size) + + # Reuse a memory map across calls (keyed by PID so forked processes -- + # e.g., PyTorch DataLoader workers -- create their own mapping instead of + # sharing one). This removes the per-call open/seek/syscall overhead. + ex = raw._raw_extras[fi] if fi < len(raw._raw_extras) else {} + mm = ex.get("_mm") if isinstance(ex, dict) else None + if mm is not None and ex.get("_mm_pid") != _os.getpid(): + mm = None + if mm is not None and ( + mm.dtype != np.dtype(dtype) or mm.size * n_bytes < data_offset + data_left * n_bytes + ): + mm = None + if mm is None and isinstance(ex, dict): + try: + mm = np.memmap(raw.filenames[fi], dtype=dtype, mode="r") + ex["_mm"] = mm + ex["_mm_pid"] = _os.getpid() + except Exception: + mm = None + + if mm is not None: + base_idx = data_offset // n_bytes + for sample_start in np.arange(0, data_left, block_size) // n_channels: + count = min(block_size, data_left - sample_start * n_channels) + block = mm[ + base_idx + sample_start * n_channels : + base_idx + sample_start * n_channels + count + ] + if block.size != count: + raise RuntimeError( + f"Incorrect number of samples ({block.size} != {count}), " + "please report this error to MNE-Python developers" + ) + block = block.reshape(n_channels, -1, order="F") + n_samples = block.shape[1] + sample_stop = sample_start + n_samples + if trigger_ch is not None: + stim_ch = trigger_ch[start:stop][sample_start:sample_stop] + block = np.vstack((block, stim_ch)) + data_view = data[:, sample_start:sample_stop] + _mult_cal_one(data_view, block, idx, cals, mult) + return + with open(raw.filenames[fi], "rb", buffering=0) as fid: fid.seek(data_offset) # extract data in chunks diff --git a/mne/io/base.py b/mne/io/base.py index 79096cafaa3..b5ed1a07cdf 100644 --- a/mne/io/base.py +++ b/mne/io/base.py @@ -2,7 +2,7 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. -import os +import hashlib import shutil from collections import defaultdict from collections.abc import Callable, Sequence @@ -97,6 +97,7 @@ copy_doc, copy_function_doc_to_method_doc, fill_doc, + get_config, logger, repr_html, sizeof_fmt, @@ -621,16 +622,65 @@ def load_data( .. versionadded:: 0.10.0 """ if not self.preload: + if memmap == "memmap": + memmap = self._auto_memmap_path() if memmap is not None: _validate_type(memmap, "path-like", "memmap") self._preload_data(memmap if memmap is not None else True) return self + def _auto_memmap_path(self) -> Path: + """Return the managed memmap-cache path for a single-file Raw.""" + if len(self.filenames) != 1 or self.filenames[0] is None: + raise ValueError( + "preload='memmap' requires a Raw backed by exactly one file" + ) + src = Path(self.filenames[0]) + st = src.stat() + key = hashlib.sha256(f"{src}|{st.st_mtime_ns}|{st.st_size}".encode()) + cache_dir = Path( + get_config("MNE_MEMMAP_DIR", Path.home() / ".cache" / "mne" / "memmap") + ) + cache_dir.mkdir(parents=True, exist_ok=True) + return cache_dir / f"{src.stem}-{key.hexdigest()[:24]}.f64" + def _preload_data(self, preload): """Actually preload the data.""" data_buffer = preload + if preload == "memmap": + # sentinel: managed persistent cache keyed by source + # path/mtime/size, so staleness is handled by the key itself + preload = data_buffer = self._auto_memmap_path() + logger.info(f"Using memory-map cache at {preload}") if isinstance(preload, bool | np.bool_) and not preload: data_buffer = None + elif isinstance(preload, str | Path): + # A memory-map cache file: when one already exists with the exact + # expected size, mmap it and skip decoding entirely. This makes + # reopening huge recordings nearly instantaneous (first run pays + # the decode, subsequent runs reuse the OS page cache). Note the + # cache is validated by size only; delete it if the source file + # changed. + path = Path(preload) + n_bytes = ( + self.info["nchan"] + * (self._last_samps[-1] - self._first_samps[0] + 1) + * np.dtype(self._dtype).itemsize + ) + if path.exists() and path.stat().st_size == n_bytes: + self._data = np.memmap( + str(path), + mode="r+", + dtype=self._dtype, + shape=( + self.info["nchan"], + n_bytes + // (self.info["nchan"] * np.dtype(self._dtype).itemsize), + ), + ) + self.preload = True + logger.info(f"Reusing memory-mapped cache at {path}") + return t = self.times logger.info( f"Reading 0 ... {len(t) - 1} = {0.0:9.3f} ... {t[-1]:9.3f} secs..." @@ -813,16 +863,12 @@ def set_annotations( return self def __del__(self): # noqa: D105 - # remove file for memmap + # Memmap-backed data (load_data(memmap=...)) is owned by the caller: + # the file persists after the Raw object is deleted so that subsequent + # sessions can mmap it directly instead of re-decoding. fname = getattr(getattr(self, "_data", None), "filename", None) if fname is not None: - # First, close the file out; happens automatically on del del self._data - # Now file can be removed - try: - os.remove(fname) - except OSError: - pass # ignore file that no longer exists def __enter__(self): """Entering with block.""" @@ -877,6 +923,60 @@ def _parse_get_set_params(self, item): return sel, start, stop + def _get_windows(self, starts, width, *, out=None, sel=None): + """Read many equal-width windows with setup shared across them. + + This amortizes channel-pick resolution and output allocation over all + windows, which matters for workloads issuing many small reads, such as + deep-learning training loops. + + Parameters + ---------- + starts : array of int + First sample of each window. + width : int + Number of samples per window. + out : ndarray | None + Optional ``(len(starts), n_channels, width)`` float64 array to + fill; allocated when None. + sel : ndarray | slice | None + Channel selection resolved once for all windows. + + Returns + ------- + ndarray : ``(len(starts), n_channels, width)`` window stack. + """ + starts = np.atleast_1d(np.asarray(starts, dtype=np.int64)).ravel() + width = int(width) + if width <= 0: + raise ValueError(f"width must be positive, got {width}") + n_times = self.n_times + if len(starts) == 0: + starts = np.zeros(1, dtype=np.int64) + bad = (starts < 0) | (starts + width > n_times) + if bad.any(): + raise ValueError( + f"window out of bounds at index {int(np.flatnonzero(bad)[0])}" + ) + # sel=None lets _read_segment treat this as all channels + n_out = self.info["nchan"] if sel is None else len(sel) + if out is None: + out = np.empty((len(starts), n_out, width), dtype=self._dtype) + elif out.shape != (len(starts), n_out, width): + raise ValueError( + f"out has shape {out.shape}, need {(len(starts), n_out, width)}" + ) + elif out.dtype not in (np.float64, np.float32): + raise ValueError(f"out dtype must be float64 or float32, got {out.dtype}") + for j, s0 in enumerate(starts): + self._read_segment( + start=int(s0), + stop=int(s0) + width, + sel=sel if sel is not None else None, + data_buffer=out[j], + ) + return out + def __getitem__(self, item): """Get raw data and times. @@ -1001,7 +1101,14 @@ def get_data( stop, types=("int-like", None), item_name="stop", type_name="int, None" ) - picks = _picks_to_idx(self.info, picks, "all", exclude=()) + if picks is None: + # fast path: equivalent to _picks_to_idx(info, None, "all", + # exclude=()) but avoids channel-name resolution on every call, + # which matters for workloads making many small reads (e.g., + # deep-learning training loops) + picks = np.arange(self.info["nchan"]) + else: + picks = _picks_to_idx(self.info, picks, "all", exclude=()) # Get channel factors for conversion into specified unit # (vector of ones if no conversion needed) diff --git a/mne/io/edf/_bdf_numba.py b/mne/io/edf/_bdf_numba.py new file mode 100644 index 00000000000..8315915bace --- /dev/null +++ b/mne/io/edf/_bdf_numba.py @@ -0,0 +1,33 @@ +"""Numba-accelerated BDF (24-bit little-endian) sample decoding. + +Optional acceleration: falls back to the vectorized-numpy path in +``mne.io.edf.edf._read_ch`` when numba is unavailable. +""" + +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors + +import numpy as np + +from ..._numba import jit + + +@jit() +def decode_int24(buf): # pragma: no cover + """Decode packed 24-bit little-endian samples to int32. + + ``buf`` is a ``(n_samples, 3)`` uint8 array whose rows hold the low, + middle, and high bytes of each signed sample. + """ + n = buf.shape[0] + out = np.empty(n, dtype=np.int32) + for i in range(n): + # plain-Python integer arithmetic so the non-numba fallback follows + # the same semantics as the jitted version (values stay within + # [-2**23, 2**23) after the sign fix, so int32 stores never overflow) + v = int(buf[i, 0]) | (int(buf[i, 1]) << 8) | (int(buf[i, 2]) << 16) + if v >= (1 << 23): + v -= 1 << 24 + out[i] = v + return out diff --git a/mne/io/edf/_edf_numba.py b/mne/io/edf/_edf_numba.py new file mode 100644 index 00000000000..7910185f491 --- /dev/null +++ b/mne/io/edf/_edf_numba.py @@ -0,0 +1,69 @@ +"""Numba-accelerated EDF/BDF digital-to-physical window decoding. + +Optional acceleration: falls back to the vectorized-numpy path in +``mne.io.edf.edf._read_segment_file`` when numba is unavailable. +""" + +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors + +from ..._numba import jit + + +@jit(fastmath=False) +def decode_window(digital, cal_v, off_v, gain_v, out): # pragma: no cover + """Decode a block of digital samples to physical units. + + ``digital`` is a ``(k, n_blocks, buf_len)`` (possibly strided) integer + view of raw digital samples; ``cal_v``, ``off_v``, and ``gain_v`` are + length-``k`` float64 vectors; ``out`` is a ``(k, n_blocks * buf_len)`` + float64 array whose rows are filled with blocks concatenated along the + sample axis such that:: + + out[i, b * buf_len + j] = ((digital[i, b, j] * cal[i]) + off[i]) * gain[i] + + replicating exactly the operation order of the vectorized-numpy fallback + (hence ``fastmath=False``: no FMA contraction or reassociation is + allowed, so results are bit-identical to separate multiply/add rounding). + """ + k = digital.shape[0] + n_blk = digital.shape[1] + n_smp = digital.shape[2] + for i in range(k): + cal = cal_v[i] + off = off_v[i] + gain = gain_v[i] + out_i = out[i] + for b in range(n_blk): + base = b * n_smp + for j in range(n_smp): + out_i[base + j] = ((digital[i, b, j] * cal) + off) * gain + + +@jit(fastmath=False) +def decode_window_into( + dst, digital, cal_v, off_v, gain_v, s0, w +): # pragma: no cover + """Decode into a possibly-strided 2-D destination. + + ``dst`` is ``(k, w)`` with arbitrary strides (e.g., a column slice of the + caller's output buffer); ``digital`` is the ``(k, n_blocks, buf_len)`` + strided integer view covering whole data records; ``s0`` is the first + sample to take from that view and ``w`` the number of samples to write, + so edge records at window boundaries are handled without temporaries. + Elementwise op order is identical to :func:`decode_window`. + """ + k = digital.shape[0] + n_smp = digital.shape[2] + for i in range(k): + cal = cal_v[i] + off = off_v[i] + gain = gain_v[i] + dst_i = dst[i] + dig_i = digital[i] + for t in range(w): + g = s0 + t + b = g // n_smp + j = g - b * n_smp + dst_i[t] = ((dig_i[b, j] * cal) + off) * gain diff --git a/mne/io/edf/_edfio_backend.py b/mne/io/edf/_edfio_backend.py new file mode 100644 index 00000000000..dcacada0f60 --- /dev/null +++ b/mne/io/edf/_edfio_backend.py @@ -0,0 +1,162 @@ +"""Optional edfio-backed reader engine for EDF files. + +This module implements an alternative ``engine="edfio"`` for +:func:`mne.io.read_raw_edf` that parses the file with the +`edfio `_ package instead of the +native reader. It is faster on uniform-sampling-rate recordings and always +returns preloaded data. + +Scope (kept deliberately minimal): + +- uniform sampling rates only (the native engine handles mixed rates); +- all channels are typed ``eeg``; +- ``meas_date`` is not set; +- data is returned in volts, scaled from the header's physical dimension + using the same unit mapping as the native reader. +""" + +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors + +import numpy as np + +from ..._fiff.meas_info import _unique_channel_names +from ...annotations import Annotations +from ...utils import _check_fname, fill_doc, verbose +from ..base import BaseRaw + +_UNIT_MULT = { + "\u03bcV": 1e-6, # greek mu + "\u00b5V": 1e-6, # micro symbol + "uV": 1e-6, + "mV": 1e-3, +} + + +class _RawEdfio(BaseRaw): + """Raw from edfio-parsed EDF (always preloaded).""" + + _extra_attributes = () + + def __init__(self, info, data, annotations, *, verbose=None): + super().__init__( + info, + preload=data, + last_samps=[data.shape[1] - 1], + filenames=None, + orig_format="double", + verbose=verbose, + ) + if len(annotations): + self.set_annotations(annotations) + + +@fill_doc +@verbose +def read_raw_edf_edfio( + input_fname, + *, + preload=True, + exclude=(), + include=None, + verbose=None, +) -> _RawEdfio: + """Read an EDF file using the edfio parser. + + Parameters + ---------- + input_fname : path-like + Path to the EDF/EDF+ file. + %(preload)s + The edfio engine currently supports only preloaded reads; ``True`` + (or a truthy string) is required. + exclude : list of str + Channel names to exclude. + include : list of str | None + Restrict channels to these names (after ``exclude``). + %(verbose)s + + Returns + ------- + raw : instance of Raw + Preloaded raw data in volts. + + Notes + ----- + Uniform sampling rates only; all channels are typed ``eeg``; + ``info['meas_date']`` is not populated. + """ + from edfio import read_edf as _read_edf + + input_fname = str(_check_fname(input_fname, "read", True, "input_fname")) + if not preload: + raise NotImplementedError( + 'The "edfio" engine currently always loads data into memory; ' + 'use preload=True.' + ) + edf = _read_edf(input_fname) + + signals = edf.signals + ch_names = [sig.label for sig in signals] + sfreqs = {float(sig.sampling_frequency) for sig in signals} + if len(sfreqs) != 1: + raise NotImplementedError( + "The edfio engine requires a uniform sampling rate; this file has " + f"{len(sfreqs)} distinct rates. Use the default engine instead." + ) + sfreq = sfreqs.pop() + + keep = np.arange(len(signals)) + if include is not None: + keep = [i for i in keep if ch_names[i] in set(include)] + if len(exclude): + excluded = set(exclude) + keep = [i for i in keep if ch_names[i] not in excluded] + keep = np.asarray(keep, dtype=int) + if keep.size == 0: + raise ValueError("No channels selected") + + ch_names = list(np.array(ch_names)[keep]) + ch_names = _unique_channel_names(ch_names) + unit_mults = np.array( + [ + _UNIT_MULT.get(str(signals[i].physical_dimension).strip(), 1.0) + for i in keep + ], + dtype=float, + ) + # Stack digital samples once, then decode all channels in two fused + # passes: physical = (digital + offset) * (gain * unit_mult), matching + # edfio's calibration op order. + n_times = min(len(signals[i].digital) for i in keep) + dig = np.empty((len(keep), n_times), dtype=np.int16) + gains = np.empty(len(keep)) + offsets = np.empty(len(keep)) + for row_i, sig_i in enumerate(keep): + digital = signals[sig_i].digital + dig[row_i] = digital[:n_times] + sig = signals[sig_i] + gains[row_i] = (sig.physical_max - sig.physical_min) / ( + sig.digital_max - sig.digital_min + ) + offsets[row_i] = sig.physical_max / gains[row_i] - sig.digital_max + + info = _make_info_edfio(ch_names, sfreq) + data = np.empty((len(keep), n_times), dtype=np.float64) + np.add(dig, offsets[:, np.newaxis], out=data, casting="unsafe") + data *= (gains * unit_mults)[:, np.newaxis] + + annots = edf.annotations + mne_annots = Annotations( + onset=[a.onset for a in annots], + duration=[a.duration for a in annots], + description=[str(a.text) for a in annots], + ) + return _RawEdfio(info, data, mne_annots, verbose=verbose) + + +def _make_info_edfio(ch_names, sfreq): + import mne + + return mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types="eeg") diff --git a/mne/io/edf/_open.py b/mne/io/edf/_open.py index 38f8be4113d..1d8f1106dfe 100644 --- a/mne/io/edf/_open.py +++ b/mne/io/edf/_open.py @@ -1,12 +1,55 @@ # Authors: The MNE-Python contributors. # License: BSD-3-Clause -# Copyright the MNE-Python contributors. +# Copyright the MNE-Python contributors +import os from pathlib import Path from ..._fiff.open import _NoCloseRead from ...utils import _file_like, _validate_type, logger +# Persistent read handles for EDF/BDF/GDF files. Readers seek before every +# read, so a shared handle is safe; keying by PID keeps forked worker +# processes (e.g., PyTorch DataLoader workers) from sharing file-offset state +# through an inherited descriptor. +_HANDLE_CACHE = {} +_MAX_HANDLES = 8 + + +class _NoCloseCached(_NoCloseRead): + """A file object whose context manager detaches instead of closing. + + Used for handles shared through the per-process LRU cache: leaving the + reader's ``with`` block must not close a descriptor other reads may still + use. + """ + + def close(self): # noqa: D102 + pass + + def __exit__(self, *args): # noqa: D105 + # detach rather than close; the cache owns the lifetime + return False + + +def _get_cached_fid(fname): + """Return a persistent binary handle for *fname* (per process).""" + key = (os.getpid(), str(fname)) + hit = _HANDLE_CACHE.get(key) + if hit is not None: + hit.seek(0) # match fresh-open semantics + return hit + fid = open(fname, "rb") + cached = _NoCloseCached(fid) + _HANDLE_CACHE[key] = cached + while len(_HANDLE_CACHE) > _MAX_HANDLES: + old_key = next(iter(_HANDLE_CACHE)) + try: + _HANDLE_CACHE.pop(old_key).fid.close() + except Exception: + pass + return cached + def _gdf_edf_get_fid(fname, **kwargs): """Open a EDF/BDF/GDF file with no additional parsing.""" @@ -14,8 +57,8 @@ def _gdf_edf_get_fid(fname, **kwargs): logger.debug("Using file-like I/O") fid = _NoCloseRead(fname) fid.seek(0) - else: - _validate_type(fname, [Path, str], "fname", extra="or file-like") - logger.debug("Using normal I/O") - fid = open(fname, "rb", **kwargs) # Open in binary mode - return fid + return fid + _validate_type(fname, [Path, str], "fname", extra="or file-like") + logger.debug("Using normal I/O") + kwargs.pop("buffering", None) # cached handle manages its own buffering + return _get_cached_fid(Path(fname)) diff --git a/mne/io/edf/edf.py b/mne/io/edf/edf.py index 390289ae6bb..0a1fdef16e7 100644 --- a/mne/io/edf/edf.py +++ b/mne/io/edf/edf.py @@ -590,12 +590,36 @@ def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): ) +_decode_int24 = None + + +def _get_int24_decoder(): + """Return the numba int24 decoder if available, else False.""" + global _decode_int24 + if _decode_int24 is None: + dec = False + try: + from mne._numba import has_numba + + # only use the jitted decoder when numba is actually enabled, + # otherwise the decorated function would run as slow pure Python + if has_numba: + from mne.io.edf._bdf_numba import decode_int24 as dec + except Exception: + dec = False + _decode_int24 = dec + return _decode_int24 + + def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): """Read a number of samples for a single channel.""" assert dtype is not None # BDF if subtype == "bdf": ch_data = read_from_file_or_buffer(fid, dtype=dtype, count=samp * dtype_byte) + dec = _get_int24_decoder() + if dec is not False: + return dec(ch_data.reshape(-1, 3)) ch_data = ch_data.reshape(-1, 3).astype(INT32) ch_data = (ch_data[:, 0]) + (ch_data[:, 1] << 8) + (ch_data[:, 2] << 16) # 24th bit determines the sign @@ -610,8 +634,6 @@ def _read_ch(fid, subtype, samp, dtype_byte, dtype=None): def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, mult): """Read a chunk of raw data.""" - from scipy.interpolate import interp1d - n_samps = raw_extras["n_samps"] buf_len = int(raw_extras["max_samp"]) dtype = raw_extras["dtype_np"] @@ -640,6 +662,190 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, # Let's do ~10 MB chunks: n_per = max(10 * 1024 * 1024 // (ch_offsets[-1] * dtype_byte), 1) + # Fast path: uniform sampling rate among all requested channels, no + # annotations (TAL) channel and no stim-channel special casing. This is + # the common case for large DL corpora; it replaces the per-channel + # Python loop with a handful of vectorized operations. + # The vectorized path removes per-channel Python overhead, which dominates + # small reads (deep-learning window access); for very large outputs the + # legacy loop's cache-friendly working set is slightly faster, so gate on + # decoded size. + _fast_path_output_limit = int(32e6) + uni_mask = [n_samps[ci] == buf_len for ci in read_sel] + slow_ii = [ii for ii, u in enumerate(uni_mask) if not u] + if ( + subtype in ("edf", "bdf") + and len(tal_idx) == 0 + and all(uni_mask) # mixed-sfreq partial decode is future work + and (stop - start) * len(idx_arr) * 8 <= _fast_path_output_limit + ): + with _gdf_edf_get_fid(filenames, buffering=0) as fid: + start_offset = ( + data_offset + block_start_idx * ch_offsets[-1] * dtype_byte + ) + k = len(idx_arr) + uni_list = [ii for ii in range(k) if uni_mask[ii]] + sel_uni = [read_sel[ii] for ii in uni_list] + idx_arr_uni = idx_arr[uni_list] + k_u = len(uni_list) + cal_v = cal[idx_arr][:, np.newaxis, np.newaxis] + off_v = offsets[idx_arr][:, np.newaxis, np.newaxis] + gain_v = gains[idx_arr][:, np.newaxis, np.newaxis] + cal_u = cal[idx_arr_uni][:, np.newaxis, np.newaxis] + off_u = offsets[idx_arr_uni][:, np.newaxis, np.newaxis] + gain_u = gains[idx_arr_uni][:, np.newaxis, np.newaxis] + off_v = offsets[idx_arr][:, np.newaxis, np.newaxis] + gain_v = gains[idx_arr][:, np.newaxis, np.newaxis] + # With no projector/compensation and unit cals (always true for + # EDF/BDF), decoded samples can go straight into the output. + write_direct = ( + mult is None and bool(np.all(cals == 1)) and not slow_ii + ) + ones = None if write_direct else np.zeros( + (len(orig_sel), data.shape[-1]), dtype=data.dtype + ) + dec = _get_window_decoder() + dec_into = None + if dec is not False and k_u: + from ._edf_numba import decode_window_into + + dec_into = decode_window_into + cal_1d = cal_u.ravel() + off_1d = off_u.ravel() + gain_1d = gain_u.ravel() + # Uniform stim channels can stay on the fast path: legacy applies + # a truncating bitmask (float -> int cast, mask 2**17-1) AFTER + # calibration, which we replicate per row below. + stim_rows = [ + ii + for ii, orig in enumerate(idx_arr) + if int(orig) in stim_channel_idxs and uni_mask[ii] + ] + pos = 0 + for ai in range(0, len(r_lims), n_per): + block_offset = ai * ch_offsets[-1] * dtype_byte + n_read = min(len(r_lims) - ai, n_per) + fid.seek(start_offset + block_offset, 0) + many_chunk = _read_ch( + fid, subtype, ch_offsets[-1] * n_read, dtype_byte, dtype + ) + arr3 = many_chunk.reshape(n_read, len(n_samps), buf_len) + r_sidx = r_lims[ai][0] + r_eidx = buf_len * (n_read - 1) + r_lims[ai + n_read - 1][1] + # gather this call's channels as a strided view + # (k, n_read, buf_len), each row one requested channel; + # when all channels are requested this is a zero-copy + # transpose + if not slow_ii and k == len(n_samps): + view = arr3.transpose(1, 0, 2) + elif slow_ii: + view = arr3[:, sel_uni, :].transpose(1, 0, 2) + else: + view = arr3[:, read_sel, :].transpose(1, 0, 2) + # digital -> physical, preserving the legacy op order; + # when possible decode straight into the destination slice so + # the block never round-trips through a temporary + width = r_eidx - r_sidx + if dec_into is not None and write_direct: + if not view.dtype.isnative: + # numba only types native byteorder; real-world EDF is + # big-endian, so swap into a native copy per chunk + view = view.astype(view.dtype.newbyteorder("=")) + dst_full = ( + data[:, pos : pos + width] + if write_direct + else ones[idx_arr_uni, pos : pos + width] + ) + dec_into(dst_full, view, cal_1d, off_1d, gain_1d, + r_sidx, width) + block = None # values already in place + elif ones is None: + # no kernel and direct write: decode to a temp then copy + # once into the destination slice + one = np.empty((k, n_read, buf_len), dtype=np.float64) + np.multiply(view, cal_v, out=one) + one += off_v + one *= gain_v + dst_full = data[:, pos : pos + width] + dst_full[...] = one.reshape(k, -1)[:, r_sidx:r_eidx] + block = None # values already in place + else: + one = np.empty((k, n_read, buf_len), dtype=np.float64) + np.multiply(view, cal_v, out=one) + one += off_v + one *= gain_v + block = one.reshape(k, -1)[:, r_sidx:r_eidx] + if stim_rows: + for ii in stim_rows: + src_row = block[ii] if block is not None else dst_full[ii] + row_i = src_row.astype(int) + np.bitwise_and(row_i, 2**17 - 1, out=row_i) + if block is not None: + block[ii] = row_i + else: + dst_full[ii] = row_i + if block is not None: + if write_direct: + data[:, pos : pos + width] = block + elif slow_ii: + ones[idx_arr_uni, pos : pos + width] = block + else: + ones[idx_arr, pos : pos + width] = block + # legacy per-channel treatment for non-uniform rows: their + # columns live inside the same records we just read + for ii in slow_ii: + ci = read_sel[ii] + orig_idx = idx_arr[ii] + ch_data = many_chunk[ + :, ch_offsets[ci] : ch_offsets[ci + 1] + ].copy() + o_i = orig_idx + ch_data = ch_data * cal[o_i] + ch_data += offsets[o_i] + ch_data *= gains[o_i] + if int(orig_idx) in stim_channel_idxs: + from scipy.interpolate import interp1d + + s_n = n_samps[ci] + oldg = np.linspace(0, 1, s_n + 1, True) + newg = np.linspace(0, 1, buf_len, False) + ch_data = np.append( + ch_data, np.zeros((len(ch_data), 1)), -1 + ) + ch_data = interp1d(oldg, ch_data, kind="zero", axis=-1)( + newg + ) + one_i = ch_data.ravel()[r_sidx:r_eidx] + w0 = pos - width # first output column of this chunk + ones[o_i, w0 : w0 + len(one_i)] = one_i + pos += width + if slow_ii: + smp_exp = data.shape[-1] + resampled = False + for ii in slow_ii: + row = int(idx_arr[ii]) + ci = read_sel[ii] + if n_samps[ci] != buf_len and width != smp_exp: + resampled = True + ones[row, :] = resample( + ones[row, :width].astype(np.float64), + smp_exp, + width, + npad=0, + axis=-1, + ) + if resampled and raw_extras["nsamples"] != (stop - start): + warn( + "Loading an EDF with mixed sampling frequencies and " + "preload=False will result in edge artifacts. " + "It is recommended to use preload=True." + "See also " + "https://github.com/mne-tools/mne-python/issues/10635" + ) + if not write_direct: + _mult_cal_one(data[:, :], ones, idx, cals, mult) + return tal_data + with _gdf_edf_get_fid(filenames, buffering=0) as fid: # Extract data start_offset = data_offset + block_start_idx * ch_offsets[-1] * dtype_byte @@ -683,6 +889,8 @@ def _read_segment_file(data, idx, fi, start, stop, raw_extras, filenames, cals, if n_samps[ci] != buf_len: if orig_idx in stim_channel_idxs: # Stim channel will be interpolated + from scipy.interpolate import interp1d + old = np.linspace(0, 1, n_samps[ci] + 1, True) new = np.linspace(0, 1, buf_len, False) ch_data = np.append(ch_data, np.zeros((len(ch_data), 1)), -1) @@ -1902,9 +2110,10 @@ def read_raw_edf( units: dict | str | None = None, encoding: str = "utf8", exclude_after_unique: bool = False, + engine: Literal["mne", "edfio"] = "mne", *, verbose: bool | str | int | None = None, -) -> RawEDF: +) -> RawEDF | BaseRaw: """Reader function for EDF and EDF+ files. Parameters @@ -1951,6 +2160,13 @@ def read_raw_edf( %(units_edf_bdf_io)s %(encoding_edf)s %(exclude_after_unique)s + engine : ``'mne'`` | ``'edfio'`` + Parser backend. ``'mne'`` (default) uses the native reader; + ``'edfio'`` parses via the optional edfio package, which is faster on + uniform-sampling-rate recordings but always preloads, types all + channels as EEG, and does not set ``info['meas_date']``. + + .. versionadded:: 1.13 %(verbose)s Returns @@ -2010,6 +2226,19 @@ def read_raw_edf( The EDF specification allows storage of subseconds in measurement date. However, this reader currently sets subseconds to 0 by default. """ + if engine == "edfio": + from ._edfio_backend import read_raw_edf_edfio + + return read_raw_edf_edfio( + input_fname, + preload=preload, + exclude=exclude, + include=include, + verbose=verbose, + ) + if engine != "mne": + raise ValueError(f"Unknown engine {engine!r}; use 'mne' or 'edfio'.") + _check_args(input_fname, preload, "edf") return RawEDF( @@ -2361,3 +2590,22 @@ def _get_annotations_gdf(edf_info, sfreq): desc = events[2] return onset, duration, desc + + +_decode_window = None + + +def _get_window_decoder(): + """Return the numba fused window decoder if available, else False.""" + global _decode_window + if _decode_window is None: + dec = False + try: + from mne._numba import has_numba + + if has_numba: + from mne.io.edf._edf_numba import decode_window as dec + except Exception: + dec = False + _decode_window = dec + return _decode_window diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 211641dc726..612f3913129 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -8,6 +8,7 @@ from io import BytesIO from pathlib import Path +import mne import numpy as np import pytest from numpy.testing import ( @@ -1262,3 +1263,71 @@ def test_edf_read_from_file_like(): ] assert raw.ch_names == channels + + + +def requires_edfio(func): + import pytest + + return pytest.mark.skipif( + __import__("importlib.util", fromlist=["util"]).find_spec("edfio") is None, + reason="Requires edfio", + )(func) + + +@requires_edfio +def test_engine_edfio(tmp_path): + """Compare the optional edfio engine against the native one.""" + pytest.importorskip("edfio") + rng = np.random.default_rng(11) + info = mne.create_info(["EEG A", "EEG B"], sfreq=128.0, ch_types="eeg") + raw = mne.io.RawArray(rng.standard_normal((2, 512)) * 30e-6, info) + fname = tmp_path / "engine_test.edf" + raw.export(fname, verbose="error") + base = read_raw_edf(fname, preload=True, verbose="error").get_data() + alt = read_raw_edf(fname, preload=True, engine="edfio", + verbose="error").get_data() + assert base.shape == alt.shape + assert_allclose(base, alt, rtol=0, atol=1e-15) + + +def test_memmap_cache_reuse(tmp_path): + """load_data(memmap=...) reuses a valid cache without re-decoding.""" + pytest.importorskip("edfio") + rng = np.random.default_rng(3) + info = mne.create_info(["EEG A"], sfreq=64.0, ch_types="eeg") + raw = mne.io.RawArray(rng.standard_normal((1, 2048)) * 30e-6, info) + fname = tmp_path / "mm_cache.edf" + raw.export(fname, verbose="error") + cache = tmp_path / "cache.f64" + kwargs = dict(preload=False, verbose="error") + r1 = read_raw_edf(fname, **kwargs).load_data(memmap=str(cache)) + assert cache.exists() + for _ in range(3): # replay open→use cycles + r2 = read_raw_edf(fname, **kwargs).load_data(memmap=str(cache)) + assert r2.preload + d = r2.get_data() + assert_allclose(d, r1.get_data(), rtol=0, atol=0) + # stale-size guard: wrong-sized cache is rebuilt, not mmap-reused + cache.write_bytes(b"x" * 16) + r3 = read_raw_edf(fname, **kwargs).load_data(memmap=str(cache)) + assert_allclose(r3.get_data(), r1.get_data(), rtol=0, atol=0) + + +def test_preload_memmap_sentinel(tmp_path, monkeypatch): + """preload="memmap" manages a persistent cache automatically.""" + monkeypatch.setenv("MNE_MEMMAP_DIR", str(tmp_path)) + rng = np.random.default_rng(5) + info = mne.create_info(["EEG A"], sfreq=128.0, ch_types="eeg") + raw = mne.io.RawArray(rng.standard_normal((1, 1024)) * 30e-6, info) + fname = tmp_path / "sentinel.edf" + raw.export(fname, verbose="error") + kw = dict(preload="memmap", verbose="error") + r1 = read_raw_edf(fname, **kw) + caches = list(tmp_path.rglob("*.f64")) + assert len(caches) == 1 + d1 = r1.get_data() + del r1 + r2 = read_raw_edf(fname, **kw) + assert isinstance(r2._data, np.memmap) + assert_allclose(r2.get_data(), d1, rtol=0, atol=0) diff --git a/mne/io/fiff/raw.py b/mne/io/fiff/raw.py index 95c6db5dbec..87102f77544 100644 --- a/mne/io/fiff/raw.py +++ b/mne/io/fiff/raw.py @@ -9,10 +9,11 @@ import numpy as np +from ..._fiff._mmap_cache import get_u8_memmap from ..._fiff.constants import FIFF from ..._fiff.meas_info import read_meas_info from ..._fiff.open import _fiff_get_fid, _get_next_fname, fiff_open -from ..._fiff.tag import _call_dict, read_tag +from ..._fiff.tag import _call_dict, _simple_dict, read_tag from ..._fiff.tree import dir_tree_find from ..._fiff.utils import _mult_cal_one from ...annotations import Annotations, _read_annotations_fif @@ -403,13 +404,84 @@ def _dtype(self): def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): """Read a segment of data from a file.""" n_bad = 0 - with _fiff_get_fid(self._raw_extras[fi]["filename"]) as fid: - bounds = self._raw_extras[fi]["bounds"] - ents = self._raw_extras[fi]["ent"] - nchan = self._raw_extras[fi]["orig_nchan"] - use = (stop > bounds[:-1]) & (start < bounds[1:]) + bounds = self._raw_extras[fi]["bounds"] + ents = self._raw_extras[fi]["ent"] + nchan = self._raw_extras[fi]["orig_nchan"] + fname = self._raw_extras[fi]["filename"] + # Entries overlapping [start, stop) via binary search on sorted bounds + # (O(log n) instead of a mask over every entry; matters for long + # recordings with thousands of buffer entries). + eis = range( + max(np.searchsorted(bounds, start, side="right") - 1, 0), + min(np.searchsorted(bounds, stop, side="left"), len(bounds) - 1), + ) + + # Fast path: read tag payloads directly through a PID-keyed memory map, + # skipping per-call open/seek/read syscalls. Only taken for uncompressed + # real files whose touched tags are simple numeric types with the + # expected sizes; everything else falls back to the legacy loop below. + mm = None + if ( + isinstance(fname, Path) + and len(fname.suffixes) > 0 + and fname.suffixes[-1] != ".gz" + ): + mm = get_u8_memmap(fname) + if mm is not None: + for ei in eis: + ent = ents[ei] + if ent is None or ent.type not in _simple_dict: + mm = None + break + nsamp_ei = bounds[ei + 1] - bounds[ei] + itemsize = np.dtype(_simple_dict[ent.type]).itemsize + if getattr(ent, "size", None) != nsamp_ei * nchan * itemsize: + mm = None + break + if mm is not None: + offset = 0 + for ei in eis: + first = bounds[ei] + last = bounds[ei + 1] + nsamp = last - first + ent = ents[ei] + first_pick = max(start - first, 0) + last_pick = min(nsamp, stop - first) + picksamp = last_pick - first_pick + this_start = offset + offset += picksamp + this_stop = offset + if ent is None: + continue # gaps were zero-initialized by the caller + dtype_s = _simple_dict[ent.type] + itemsize = np.dtype(dtype_s).itemsize + nbytes = picksamp * nchan * itemsize + base = ent.pos + 16 + first_pick * nchan * itemsize + one = np.frombuffer( + mm[base : base + nbytes], dtype=dtype_s, count=picksamp * nchan + ) + if one.size != picksamp * nchan: + n_bad += picksamp + continue + one = one.reshape(picksamp, nchan) + _mult_cal_one( + data[:, this_start:this_stop], + one.T, + idx, + cals, + mult, + ) + if n_bad: + warn( + f"FIF raw buffer could not be read, acquisition error " + f"likely: {n_bad} samples set to zero" + ) + assert offset == stop - start + return + + with _fiff_get_fid(fname) as fid: offset = 0 - for ei in np.where(use)[0]: + for ei in eis: first = bounds[ei] last = bounds[ei + 1] nsamp = last - first diff --git a/mne/utils/docs.py b/mne/utils/docs.py index 868ee4fe0e4..40155ecf469 100644 --- a/mne/utils/docs.py +++ b/mne/utils/docs.py @@ -3686,7 +3686,12 @@ def _reflow_param_docstring(docstring, has_first_line=True, width=75): If True, the data will be preloaded into memory (fast, requires large amount of memory). If preload is a string, preload is the file name of a memory-mapped file which is used to store the data - on the hard drive (slower, requires less memory).""" + on the hard drive (slower, requires less memory). The special string + ``"memmap"`` selects an automatically managed memory-mapped cache: + the first read decodes into the cache, and subsequent reads mmap it + directly instead of decoding again. + + .. versionadded:: 1.13""" docdict["preload_concatenate"] = """ preload : bool | str | None diff --git a/mne/utils/mixin.py b/mne/utils/mixin.py index 04c55c62034..3addd688797 100644 --- a/mne/utils/mixin.py +++ b/mne/utils/mixin.py @@ -575,8 +575,13 @@ def _handle_tmin_tmax(self, tmin, tmax): type_name="int, float, None", ) - # handle tmin/tmax as start and stop indices into data array - n_times = self.times.size + # handle tmin/tmax as start and stop indices into data array. + # Prefer an integer n_times (available on Raw); falling back to + # times.size there would materialize the full time vector on every + # call, which dominates the cost of many small get_data() reads. + n_times = getattr(self, "n_times", None) + if n_times is None: + n_times = self.times.size start = 0 if tmin is None else self.time_as_index(tmin)[0] stop = n_times if tmax is None else self.time_as_index(tmax)[0]