Persistent managed memmap caches + preload="memmap" sentinel

doc/changes/dev/14216.bugfix.rst

@@ -0,0 +1 @@
+Fix bug where :func:`mne.io.read_raw_ant` and :func:`mne.io.read_raw_curry` ignored projections and non-boolean ``preload`` values respectively when reading data lazily, by `Bruno Aristimunha`_.

doc/changes/dev/14216.newfeature.rst

@@ -0,0 +1 @@
+Speed up repeated preloading of any file-backed :class:`~mne.io.Raw` with ``preload="auto"``, which persists decoded data in a copy-on-write cache below :func:`mne.set_cache_dir`, by `Bruno Aristimunha`_.

mne/channels/channels.py

@@ -748,6 +748,7 @@ def add_channels(self, add_list, force_update_info=False):
         # Now update the attributes
         if (
             isinstance(self._data, np.memmap)
+            and self._data.mode != "c"
             and con_axis == 0
             and sys.platform != "darwin"
         ):  # resizing not available--no mremap

mne/io/_preload_cache.py

@@ -0,0 +1,114 @@
+"""Persistent decoded-data cache for Raw readers."""
+
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+import hashlib
+import os
+import pickle
+from pathlib import Path
+
+import numpy as np
+
+from .. import __version__ as MNE_VERSION  # ty: ignore[unresolved-import]
+from ..utils import get_config, logger
+
+_RAW_PRELOAD_CACHE_VERSION = 1
+
+
+def _raw_preload_cache_info(raw):
+    """Return the cache path and decoded array description."""
+    cache_root = get_config("MNE_CACHE_DIR", None)
+    if cache_root is None:
+        raise ValueError(
+            'preload="auto" requires a configured cache directory; use '
+            "mne.set_cache_dir(...) first"
+        )
+    cache_dir = Path(cache_root).expanduser().resolve()
+    cache_dir = cache_dir / f"raw-preload-v{_RAW_PRELOAD_CACHE_VERSION}"
+    cache_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
+
+    sources = []
+    for filename in raw.filenames:
+        if filename is None:
+            raise ValueError(
+                'preload="auto" requires stable source files; use preload=True '
+                "or an explicit memory-map path"
+            )
+        path = Path(filename).resolve(strict=True)
+        # some formats (e.g., CTF) name a directory rather than a single file
+        members = sorted(path.rglob("*")) if path.is_dir() else [path]
+        for member in members:
+            if not member.is_file():
+                continue
+            result = member.stat()
+            sources.append((str(member), int(result.st_size), int(result.st_mtime_ns)))
+
+    dtype = np.dtype(raw._dtype)
+    shape = (int(raw.info["nchan"]), int(raw.n_times))
+    identity = (
+        _RAW_PRELOAD_CACHE_VERSION,
+        MNE_VERSION,
+        type(raw).__module__,
+        type(raw).__qualname__,
+        sources,
+        raw._raw_extras,
+        raw._cals,
+        dtype.str,
+        shape,
+    )
+    try:
+        key = hashlib.sha256(pickle.dumps(identity, protocol=5)).hexdigest()
+    except Exception as exc:
+        raise ValueError(
+            f'preload="auto" cannot identify this {type(raw).__name__} source'
+        ) from exc
+    return cache_dir / f"{key}.data", sources, shape, dtype
+
+
+def _raw_preload_cache_read(path, shape, dtype):
+    """Map a complete decoded-data cache entry."""
+    try:
+        nbytes = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize
+        if path.stat().st_size != nbytes:
+            return None
+        return np.memmap(path, mode="c", dtype=dtype, shape=shape)
+    except OSError:
+        return None
+
+
+def _raw_preload_auto(raw):
+    """Reuse or create an automatic decoded-data cache entry."""
+    path, sources, shape, dtype = _raw_preload_cache_info(raw)
+    data = _raw_preload_cache_read(path, shape, dtype)
+    if data is not None:
+        logger.info(f"Reusing decoded data from {path}")
+        return data
+
+    # The temporary is per-process and os.replace is atomic, so concurrent
+    # misses need no lock; they at worst decode the same entry twice.
+    logger.info(f"Creating decoded data cache in {path.parent}")
+    temporary = path.with_suffix(f".{os.getpid()}.tmp")
+    try:
+        data = np.memmap(temporary, mode="w+", dtype=dtype, shape=shape)
+        try:
+            raw._read_segment(data_buffer=data)
+            data.flush()
+        finally:
+            data._mmap.close()  # ty: ignore[unresolved-attribute]
+        if _raw_preload_cache_info(raw)[1] != sources:
+            raise RuntimeError(
+                "Source data changed while decoded cache was created; retry"
+            )
+        try:
+            os.replace(temporary, path)
+        except OSError:
+            # Windows refuses to replace an entry another process already mapped.
+            pass
+    finally:
+        temporary.unlink(missing_ok=True)
+    data = _raw_preload_cache_read(path, shape, dtype)
+    if data is None:
+        raise RuntimeError(f"Could not read back the decoded data cache at {path}")
+    return data

mne/io/ant/ant.py

@@ -11,6 +11,7 @@
 
 from ..._fiff.constants import FIFF
 from ..._fiff.meas_info import create_info
+from ..._fiff.utils import _mult_cal_one
 from ...annotations import Annotations
 from ...utils import (
     _check_fname,
@@ -185,11 +186,7 @@ def _read_segment_file(self, data, idx, fi, start, stop, cals, mult):
             one = read_data(cnt, i_start, i_stop)
             _scale_data(one, ch_units)
             data_view = data[:, i_start - start : i_stop - start]
-            if isinstance(idx, slice):
-                data_view[:] = one[idx]
-            else:
-                # faster than doing one = one[idx]
-                np.take(one, idx, axis=0, out=data_view)
+            _mult_cal_one(data_view, one, idx, cals, mult)
 
 
 def _handle_bipolar_channels(

mne/io/ant/tests/test_ant.py

@@ -14,6 +14,7 @@
 from mne.datasets import testing
 from mne.io import BaseRaw, read_raw, read_raw_ant, read_raw_brainvision
 from mne.io.ant.ant import RawANT
+from mne.io.tests.test_raw import _test_raw_reader
 
 pytest.importorskip("antio", minversion="0.5.0")
 data_path = testing.data_path(download=False) / "antio"
@@ -428,6 +429,12 @@ def test_annotations_and_preload(ca_208: TypeDataset):
     assert raw_cnt.annotations.description[0] == "impedance"
 
 
+@testing.requires_testing_data
+def test_ant_raw_reader(ca_208: TypeDataset):
+    """Test the generic reader checks, including projected lazy reads."""
+    _test_raw_reader(read_raw_ant, fname=ca_208["cnt"]["short"])
+
+
 @testing.requires_testing_data
 def test_read_raw(ca_208: TypeDataset):
     """Test loading through read_raw."""

mne/io/base.py

@@ -103,6 +103,7 @@
     warn,
 )
 from ..utils._typing import Color, Self
+from ._preload_cache import _raw_preload_auto
 
 if TYPE_CHECKING:
     # Heavy/optional deps kept out of the runtime import path (see
@@ -143,8 +144,11 @@ class BaseRaw(
         freshly created memory-mapped file used to store the data on the hard
         drive (slower, requires less memory). An existing file is overwritten.
         The caller owns the file and is responsible for removing it after the
-        Raw object is no longer in use. If preload is an ndarray, the data are
-        taken from that array. If False, data are not read until save.
+        Raw object is no longer in use. For supported file readers, the exact
+        string ``"auto"`` instead reuses decoded data below the directory
+        configured by :func:`mne.set_cache_dir`. Use ``Path("auto")`` for a
+        literal filename. If preload is an ndarray, the data are taken from that
+        array. If False, data are not read until save.
     first_samps : sequence
         Sequence of the first sample number from each raw file. For unsplit raw
         files this should be a length-one list or tuple.
@@ -601,11 +605,14 @@ def load_data(
 
         Parameters
         ----------
-        memmap : path-like | None
+        memmap : path-like | str | None
             If not ``None``, preload data into a freshly created memory-mapped file
             at this path. An existing file is overwritten. The caller owns the file
             and is responsible for removing it after the Raw object is no longer in
-            use. If ``None`` (default), preload data into RAM.
+            use. The exact string ``"auto"`` instead means the same as
+            ``preload="auto"``: reuse decoded data below the directory configured by
+            :func:`mne.set_cache_dir`. Use ``Path("auto")`` for a literal filename.
+            If ``None`` (default), preload data into RAM.
 
             .. versionadded:: 1.13
         %(verbose)s
@@ -623,13 +630,23 @@ def load_data(
         .. versionadded:: 0.10.0
         """
         if not self.preload:
-            if memmap is not None:
+            if isinstance(memmap, str) and memmap == "auto":
+                pass  # sentinel, resolved in _preload_data
+            elif memmap is not None:
                 _validate_type(memmap, "path-like", "memmap")
+                memmap = Path(memmap)
             self._preload_data(memmap if memmap is not None else True)
         return self
 
     def _preload_data(self, preload):
         """Actually preload the data."""
+        if isinstance(preload, str) and preload == "auto":
+            self._data = _raw_preload_auto(self)
+            assert len(self._data) == self.info["nchan"]
+            self.preload = True
+            self._comp = None
+            self.close()
+            return
         data_buffer = preload
         if isinstance(preload, bool | np.bool_) and not preload:
             data_buffer = None
@@ -793,17 +810,18 @@ def set_annotations(
                     "of the raw object."
                 )
 
-            delta = 1.0 / self.info["sfreq"]
+            # This is algebraically ``self.times[-1] + 1 / sfreq`` without
+            # allocating the full time vector for large file-backed recordings.
+            sfreq = self.info["sfreq"]
+            annotation_end = (self.n_times - 1) / sfreq + 1.0 / sfreq
             new_annotations = annotations.copy()
             new_annotations._prune_ch_names(self.info, on_missing)
             if annotations.orig_time is None:
-                new_annotations.crop(
-                    0, self.times[-1] + delta, emit_warning=emit_warning
-                )
+                new_annotations.crop(0, annotation_end, emit_warning=emit_warning)
                 new_annotations.onset += self._first_time
             else:
                 tmin = meas_date + timedelta(0, self._first_time)
-                tmax = tmin + timedelta(seconds=self.times[-1] + delta)
+                tmax = tmin + timedelta(seconds=annotation_end)
                 new_annotations.crop(tmin=tmin, tmax=tmax, emit_warning=emit_warning)
                 new_annotations.onset -= (
                     meas_date - new_annotations.orig_time

mne/io/curry/curry.py

@@ -807,8 +807,9 @@ def __init__(self, fname, preload=False, on_bad_hpi_match="warn", verbose=None):
 
         # scale data to SI units
         self._cals = np.array(cals)
-        if isinstance(preload, bool | np.bool_) and preload:
-            self.load_data()
+        if not isinstance(preload, bool | np.bool_) or preload:
+            # preload can also be a memory-map path or the "auto" sentinel
+            self._preload_data(preload)
 
         # set events / annotations
         # format from curryreader: sample, etype, startsample, endsample

mne/io/fiff/raw.py

@@ -98,6 +98,12 @@ def __init__(
         on_split_missing: str = "raise",
         verbose: bool | str | int | None = None,
     ):
+        if isinstance(preload, str) and preload == "auto":
+            if _file_like(fname):
+                raise ValueError(
+                    'preload="auto" requires stable source files and is not '
+                    "supported for file-like FIF inputs"
+                )
         raws = []
         do_check_ext = not _file_like(fname)
         next_fname = fname
@@ -198,7 +204,9 @@ def _read_raw_file(
                 check_fname(fname, "raw", endings)
             # filename
             fname = _check_fname(fname, "read", True, "fname")
-            whole_file = preload if fname.suffix == ".gz" else False
+            whole_file = (
+                preload if preload != "auto" and fname.suffix == ".gz" else False
+            )
         else:
             # file-like
             if not preload:

mne/io/fiff/tests/test_raw_fiff.py

@@ -2120,6 +2120,15 @@ def test_file_like(kind, preload, split, tmp_path):
     assert file_fid.closed
 
 
+def test_file_like_auto_preload_rejected(tmp_path, monkeypatch):
+    """Test that automatic caching cannot misidentify a named stream."""
+    monkeypatch.setenv("MNE_CACHE_DIR", str(tmp_path))
+    stream = BytesIO(test_fif_fname.read_bytes())
+    stream.name = str(test_fif_fname)
+    with pytest.raises(ValueError, match="stable source files"):
+        read_raw_fif(stream, preload="auto")
+
+
 def test_str_like():
     """Test handling with str-like objects."""
     fname = pathlib.Path(test_fif_fname)

mne/io/fil/tests/test_fil.py

@@ -13,6 +13,7 @@
 from mne.datasets import testing
 from mne.io import read_raw_fil
 from mne.io.fil.sensors import _get_pos_units
+from mne.io.tests.test_raw import _test_raw_reader
 from mne.utils import copytree_rw
 
 fil_path = testing.data_path(download=False) / "FIL"
@@ -155,6 +156,7 @@ def test_fil_complete():
     _fil_megmag(raw, mat)
     _fil_stim(raw, mat)
     _fil_sensorpos(raw, mat)
+    _test_raw_reader(read_raw_fil, binfile=binname)
 
 
 @testing.requires_testing_data