Braindecode 1.3.2.dev170805792__tar.gz → 1.3.2.dev171423126__tar.gz

{braindecode-1.3.2.dev170805792/braindecode.egg-info → braindecode-1.3.2.dev171423126}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev170805792
+Version: 1.3.2.dev171423126
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>

{braindecode-1.3.2.dev170805792 → braindecode-1.3.2.dev171423126}/braindecode/datasets/base.py

@@ -20,24 +20,18 @@ from abc import abstractmethod
 from collections import Counter
 from collections.abc import Callable
 from glob import glob
-from pathlib import Path
 from typing import Any, Generic, Iterable, no_type_check
 
-import mne
 import mne.io
 import numpy as np
 import pandas as pd
-from mne.utils import _soft_import
 from mne.utils.docs import deprecated
 from torch.utils.data import ConcatDataset, Dataset, IterableDataset
 from typing_extensions import TypeVar
 
 from .bids.hub import HubDatasetMixin
-from .bids.hub_io import _restore_nan_from_json
 from .registry import register_dataset
 
-zarr = _soft_import("zarr", purpose="lazy loading from Zarr", strict=False)
-
 
 def _create_description(description) -> pd.Series:
     if description is not None:
@@ -182,11 +176,14 @@ def _concat_metadata_summary(datasets):
 
 
 def _channel_info(mne_obj):
-    """Extract (n_ch, type_str, sfreq) from an mne object or mne.Info."""
-    info = mne_obj if isinstance(mne_obj, mne.Info) else mne_obj.info
+    """Extract channel summary from an mne object.
+
+    Returns (n_ch, type_str, sfreq).
+    """
+    info = mne_obj.info
     n_ch = info["nchan"]
     sfreq = info["sfreq"]
-    ch_types = info.get_channel_types()
+    ch_types = mne_obj.get_channel_types()
     type_counts = Counter(ch_types)
     type_str = ", ".join(f"{cnt} {t.upper()}" for t, cnt in sorted(type_counts.items()))
     return n_ch, type_str, sfreq
@@ -297,88 +294,6 @@ def _build_windowed_repr(
     return b
 
 
-def _zarr_to_memmap(zarr_path, group_name):
-    """Decompress a zarr array to a float64 ``.npy`` memmap (one-time, atomic).
-
-    Saved as float64 so MNE objects can wrap the memmap zero-copy
-    (MNE requires float64 internally).  Copy-on-write mode (``'c'``)
-    means preprocessing writes go to RAM, the file stays untouched.
-    """
-    zarr_p = Path(zarr_path)
-    cache_dir = zarr_p.parent / f".{zarr_p.name}_memmap"
-    npy_path = cache_dir / f"{group_name}.npy"
-
-    if not npy_path.exists():
-        arr = zarr.open(zarr_path, mode="r")[group_name]["data"]
-        cache_dir.mkdir(parents=True, exist_ok=True)
-        tmp_path = cache_dir / f"{group_name}.{os.getpid()}.tmp.npy"
-        # Write chunk-by-chunk to avoid materializing the full array in RAM
-        mm = np.lib.format.open_memmap(
-            tmp_path, mode="w+", dtype=np.float64, shape=arr.shape
-        )
-        chunk0 = arr.chunks[0] if hasattr(arr, "chunks") else arr.shape[0]
-        for start in range(0, arr.shape[0], chunk0):
-            end = min(start + chunk0, arr.shape[0])
-            mm[start:end] = np.asarray(arr[start:end], dtype=np.float64)
-        mm.flush()
-        del mm
-        try:
-            tmp_path.rename(npy_path)
-        except OSError:
-            tmp_path.unlink(missing_ok=True)
-
-    return npy_path
-
-
-class _ZarrMixin:
-    """Zarr-to-memmap lazy loading with MNE compatibility.
-
-    Decompresses zarr to float64 ``.npy`` memmaps on first access.
-    MNE objects (``raw``/``windows``) are reconstructed lazily on first
-    property access, backed by copy-on-write memmaps so preprocessing
-    writes go to RAM while the file stays untouched.
-    """
-
-    _zarr_data = None
-
-    def _open_zarr(self):
-        npy_path = _zarr_to_memmap(self._zarr_path, self._group_name)
-        # 'c' = copy-on-write: reads from file, writes to RAM
-        self._zarr_data = np.load(npy_path, mmap_mode="c")
-
-    def _make_mne_info(self):
-        """Reconstruct ``mne.Info`` from stored dict (no data I/O)."""
-        return mne.Info.from_json_dict(_restore_nan_from_json(self._info_dict))
-
-    @classmethod
-    def _init_zarr_base(cls, zarr_path, group_name, description, info_dict, transform):
-        """Shared ``_from_zarr`` preamble: create instance with zarr attrs."""
-        if zarr is False:
-            raise ImportError("zarr is required: pip install braindecode[hub]")
-        obj = object.__new__(cls)
-        RecordDataset.__init__(obj, description, transform)
-        obj._zarr_path = Path(zarr_path)
-        obj._group_name = group_name
-        obj._info_dict = info_dict
-        return obj
-
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        state.pop("_zarr_data", None)
-        # Drop lazily-constructed mne objects; they'll be rebuilt from memmap
-        if hasattr(self, "_zarr_path"):
-            state.pop("_raw", None)
-            state.pop("_windows", None)
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        if hasattr(self, "_zarr_path"):
-            self._open_zarr()
-        else:
-            self._zarr_data = None
-
-
 class RecordDataset(Dataset[tuple[np.ndarray, int | str, tuple[int, int, int]]]):
     def __init__(
         self,
@@ -446,7 +361,7 @@ T = TypeVar("T", bound=RecordDataset)
 
 
 @register_dataset
-class RawDataset(_ZarrMixin, RecordDataset):
+class RawDataset(RecordDataset):
     """Returns samples from an mne.io.Raw object along with a target.
 
     Dataset which serves samples from an mne.io.Raw object along with a target.
@@ -480,43 +395,8 @@ class RawDataset(_ZarrMixin, RecordDataset):
         self.target_name = self._target_name(target_name)
         self.raw_preproc_kwargs: list[dict[str, Any]] = []
 
-    @classmethod
-    def _from_zarr(
-        cls,
-        zarr_path,
-        group_name,
-        description,
-        info_dict,
-        target_name=None,
-        transform=None,
-    ):
-        """Create a lazy Zarr-backed RawDataset (internal API)."""
-        obj = cls._init_zarr_base(
-            zarr_path, group_name, description, info_dict, transform
-        )
-        obj._raw = None
-        obj.target_name = cls._target_name(obj, target_name)
-        obj.raw_preproc_kwargs = []
-        obj._open_zarr()
-        return obj
-
-    @property
-    def raw(self):
-        if self._raw is None and self._zarr_data is not None:
-            self._raw = mne.io.RawArray(self._zarr_data, self._make_mne_info())
-        return self._raw
-
-    @raw.setter
-    def raw(self, value):
-        self._raw = value
-
     def __getitem__(self, index):
-        if self._zarr_data is not None:
-            X = np.array(self._zarr_data[:, index], dtype="float32")
-            if X.ndim == 1:
-                X = X[:, np.newaxis]
-        else:
-            X = self._raw[:, index][0]
+        X = self.raw[:, index][0]
         y = None
         if self.target_name is not None:
             y = self.description[self.target_name]
@@ -527,19 +407,13 @@ class RawDataset(_ZarrMixin, RecordDataset):
         return X, y
 
     def __len__(self):
-        if self._raw is not None:
-            return len(self._raw)
-        return self._zarr_data.shape[1]
+        return len(self.raw)
 
     def _build_repr(self):
         b = _ReprBuilder(type(self).__name__)
-        if self._raw is not None:
-            n_ch, type_str, sfreq = _channel_info(self._raw)
-            n_times = len(self._raw.times)
-        else:
-            n_ch, type_str, sfreq = _channel_info(self._make_mne_info())
-            n_times = self._zarr_data.shape[1]
-        duration = n_times / sfreq if sfreq > 0 else 0
+        n_ch, type_str, sfreq = _channel_info(self.raw)
+        n_times = len(self.raw.times)
+        duration = n_times / sfreq
         b.add_header(f"{n_ch} ch ({type_str})", "Channels", f"{n_ch} ({type_str})")
         b.add_header(f"{sfreq:.1f} Hz", "Sfreq", f"{sfreq:.1f} Hz")
         b.add_header(
@@ -550,8 +424,6 @@ class RawDataset(_ZarrMixin, RecordDataset):
         if self.description is not None:
             desc_items = ", ".join(f"{k}={v}" for k, v in self.description.items())
             b.add_row("description", desc_items)
-        if self._zarr_data is not None:
-            b.add_footnote("Data loaded lazily from Zarr")
         return b
 
     def _target_name(self, target_name):
@@ -592,7 +464,7 @@ class BaseDataset(RawDataset):
 
 
 @register_dataset
-class EEGWindowsDataset(_ZarrMixin, RecordDataset):
+class EEGWindowsDataset(RecordDataset):
     """Returns windows from an mne.Raw object, its window indices, along with a target.
 
     Dataset which serves windows from an mne.Epochs object along with their
@@ -649,47 +521,6 @@ class EEGWindowsDataset(_ZarrMixin, RecordDataset):
             self.y = metadata.loc[:, "target"].to_list()
         self.raw_preproc_kwargs: list[dict[str, Any]] = []
 
-    @classmethod
-    def _from_zarr(
-        cls,
-        zarr_path,
-        group_name,
-        metadata,
-        description,
-        info_dict,
-        targets_from="metadata",
-        last_target_only=True,
-        transform=None,
-    ):
-        """Create a lazy Zarr-backed EEGWindowsDataset (internal API)."""
-        obj = cls._init_zarr_base(
-            zarr_path, group_name, description, info_dict, transform
-        )
-        obj._raw = None
-        obj.metadata = metadata
-        obj.last_target_only = last_target_only
-        if targets_from not in ("metadata", "channels"):
-            raise ValueError("Wrong value for parameter `targets_from`.")
-        obj.targets_from = targets_from
-        obj.crop_inds = metadata.loc[
-            :, ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"]
-        ].to_numpy()
-        if obj.targets_from == "metadata":
-            obj.y = metadata.loc[:, "target"].to_list()
-        obj.raw_preproc_kwargs = []
-        obj._open_zarr()
-        return obj
-
-    @property
-    def raw(self):
-        if self._raw is None and self._zarr_data is not None:
-            self._raw = mne.io.RawArray(self._zarr_data, self._make_mne_info())
-        return self._raw
-
-    @raw.setter
-    def raw(self, value):
-        self._raw = value
-
     def __getitem__(self, index: int):
         """Get a window and its target.
 
@@ -713,13 +544,11 @@ class EEGWindowsDataset(_ZarrMixin, RecordDataset):
         crop_inds = self.crop_inds[index].tolist()
 
         i_window_in_trial, i_start, i_stop = crop_inds
-        if self._zarr_data is not None:
-            X = self._zarr_data[:, i_start:i_stop]
-        else:
-            X = self.raw._getitem(
-                (slice(None), slice(i_start, i_stop)), return_times=False
-            )
-        X = np.array(X, dtype="float32")
+        X = self.raw._getitem((slice(None), slice(i_start, i_stop)), return_times=False)
+        X = X.astype("float32")
+        # ensure we don't give the user the option
+        # to accidentally modify the underlying array
+        X = X.copy()
         if self.transform is not None:
             X = self.transform(X)
         if self.targets_from == "metadata":
@@ -730,7 +559,10 @@ class EEGWindowsDataset(_ZarrMixin, RecordDataset):
                 y = X[misc_mask, -1]
             else:
                 y = X[misc_mask, :]
+            # ensure we don't give the user the option
+            # to accidentally modify the underlying array
             y = y.copy()
+            # remove the target channels from raw
             X = X[~misc_mask, :]
         return X, y, crop_inds
 
@@ -738,11 +570,10 @@ class EEGWindowsDataset(_ZarrMixin, RecordDataset):
         return len(self.crop_inds)
 
     def _build_repr(self):
-        mne_obj = self._raw if self._raw is not None else self._make_mne_info()
         return _build_windowed_repr(
             type(self).__name__,
             len(self),
-            mne_obj,
+            self.raw,
             self.crop_inds,
             self.description,
             self.metadata,
@@ -833,7 +664,7 @@ class EEGWindowsDataset(_ZarrMixin, RecordDataset):
 
 
 @register_dataset
-class WindowsDataset(_ZarrMixin, RecordDataset):
+class WindowsDataset(RecordDataset):
     """Returns windows from an mne.Epochs object along with a target.
 
     Dataset which serves windows from an mne.Epochs object along with their
@@ -886,7 +717,6 @@ class WindowsDataset(_ZarrMixin, RecordDataset):
 
         metadata = self.windows.metadata
         assert metadata is not None, "WindowsDataset requires windows with metadata."
-        self.metadata = metadata
         self.crop_inds = metadata.loc[
             :, ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"]
         ].to_numpy()
@@ -895,67 +725,6 @@ class WindowsDataset(_ZarrMixin, RecordDataset):
         self.raw_preproc_kwargs: list[dict[str, Any]] = []
         self.window_preproc_kwargs: list[dict[str, Any]] = []
 
-    @classmethod
-    def _from_zarr(
-        cls,
-        zarr_path,
-        group_name,
-        metadata,
-        description,
-        info_dict,
-        target_name=None,
-        transform=None,
-        targets_from="metadata",
-        last_target_only=True,
-    ):
-        """Create a lazy Zarr-backed WindowsDataset (internal API)."""
-        obj = cls._init_zarr_base(
-            zarr_path, group_name, description, info_dict, transform
-        )
-        obj._windows = None
-        obj.target_name = target_name
-        obj._fast_disk = False
-        obj.last_target_only = last_target_only
-        if targets_from not in ("metadata", "channels"):
-            raise ValueError("Wrong value for parameter `targets_from`.")
-        obj.targets_from = targets_from
-        obj.metadata = metadata
-        obj.crop_inds = metadata.loc[
-            :, ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"]
-        ].to_numpy()
-        if obj.targets_from == "metadata":
-            obj.y = metadata.loc[:, "target"].to_list()
-        obj.raw_preproc_kwargs = []
-        obj.window_preproc_kwargs = []
-        obj._open_zarr()
-        return obj
-
-    @property
-    def windows(self):
-        if self._windows is None and self._zarr_data is not None:
-            info = self._make_mne_info()
-            targets = self.metadata["target"].values
-            event_ids = (
-                targets
-                if np.issubdtype(targets.dtype, np.integer)
-                else np.ones(len(self.metadata), dtype=int)
-            )
-            events = np.column_stack(
-                [
-                    self.metadata["i_start_in_trial"].values.astype(int),
-                    np.zeros(len(self.metadata), dtype=int),
-                    event_ids,
-                ]
-            )
-            self._windows = mne.EpochsArray(
-                self._zarr_data, info, events=events, metadata=self.metadata
-            )
-        return self._windows
-
-    @windows.setter
-    def windows(self, value):
-        self._windows = value
-
     @staticmethod
     def _can_use_fast_get_epoch_from_raw(epochs: mne.BaseEpochs) -> bool:
         """Check if we can use the fast _get_epoch_from_raw method,
@@ -987,9 +756,7 @@ class WindowsDataset(_ZarrMixin, RecordDataset):
         np.ndarray
             Crop indices.
         """
-        if self._zarr_data is not None:
-            X = self._zarr_data[index].astype("float32")
-        elif self._fast_disk:
+        if self._fast_disk:
             X = self.windows._get_epoch_from_raw(index).astype("float32")
         else:
             X = self.windows.get_data(item=index)[0].astype("float32")
@@ -1003,6 +770,7 @@ class WindowsDataset(_ZarrMixin, RecordDataset):
                 y = X[misc_mask, -1]
             else:
                 y = X[misc_mask, :]
+            # remove the target channels from raw
             X = X[~misc_mask, :]
         # necessary to cast as list to get list of three tensors from batch,
         # otherwise get single 2d-tensor...
@@ -1010,20 +778,16 @@ class WindowsDataset(_ZarrMixin, RecordDataset):
         return X, y, crop_inds
 
     def __len__(self) -> int:
-        if self._windows is not None:
-            return len(self._windows.events)
-        return len(self.crop_inds)
+        return len(self.windows.events)
 
     def _build_repr(self):
-        mne_obj = self._windows if self._windows is not None else self._make_mne_info()
-        md = self._windows.metadata if self._windows is not None else self.metadata
         return _build_windowed_repr(
             type(self).__name__,
             len(self),
-            mne_obj,
+            self.windows,
             self.crop_inds,
             self.description,
-            md,
+            self.windows.metadata,
         )
 
 
@@ -1250,8 +1014,8 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
 
         all_dfs = list()
         for ds in self.datasets:
-            if hasattr(ds, "_windows") and ds._windows is not None:
-                df = ds._windows.metadata
+            if hasattr(ds, "windows"):
+                df = ds.windows.metadata
             else:
                 df = ds.metadata
             for k, v in ds.description.items():

{braindecode-1.3.2.dev170805792 → braindecode-1.3.2.dev171423126}/braindecode/datasets/bids/hub.py

@@ -39,6 +39,7 @@ import mne
 import numpy as np
 import pandas as pd
 import scipy
+from mne._fiff.meas_info import Info
 from mne.utils import _soft_import
 
 if TYPE_CHECKING:
@@ -732,6 +733,7 @@ class HubDatasetMixin:
                     info_dict,
                     compressor,
                     target_name,
+                    chunk_size,
                 )
 
             elif dataset_type == "EEGWindowsDataset":
@@ -834,7 +836,6 @@ class HubDatasetMixin:
         datasets = []
         for i_ds in range(n_datasets):
             grp = root[f"recording_{i_ds}"]
-            group_name = f"recording_{i_ds}"
 
             if dataset_type == "WindowsDataset":
                 # Load using inlined function
@@ -842,36 +843,24 @@ class HubDatasetMixin:
                     _load_windows_from_zarr(grp, preload)
                 )
 
-                if preload:
-                    # Convert to MNE objects and create dataset
-                    info = mne.Info.from_json_dict(info_dict)
-                    targets = metadata["target"].values
-                    if np.issubdtype(targets.dtype, np.integer):
-                        event_ids = targets
-                    else:
-                        event_ids = np.ones(len(metadata), dtype=int)
-                    events = np.column_stack(
-                        [
-                            metadata["i_start_in_trial"].values.astype(int),
-                            np.zeros(len(metadata), dtype=int),
-                            event_ids,
-                        ]
-                    )
-                    epochs = mne.EpochsArray(
-                        data, info, events=events, metadata=metadata
-                    )
-                    ds = WindowsDataset(epochs, description)
-                    if target_name is not None:
-                        ds.target_name = target_name
+                # Convert to MNE objects and create dataset
+                info = Info.from_json_dict(info_dict)
+                targets = metadata["target"].values
+                if np.issubdtype(targets.dtype, np.integer):
+                    event_ids = targets
                 else:
-                    ds = WindowsDataset._from_zarr(
-                        zarr_path=input_path,
-                        group_name=group_name,
-                        metadata=metadata,
-                        description=description,
-                        info_dict=info_dict,
-                        target_name=target_name,
-                    )
+                    event_ids = np.ones(len(metadata), dtype=int)
+                events = np.column_stack(
+                    [
+                        metadata["i_start_in_trial"].values.astype(int),
+                        np.zeros(len(metadata), dtype=int),
+                        event_ids,
+                    ]
+                )
+                epochs = mne.EpochsArray(data, info, events=events, metadata=metadata)
+                ds = WindowsDataset(epochs, description)
+                if target_name is not None:
+                    ds.target_name = target_name
 
             elif dataset_type == "EEGWindowsDataset":
                 # Load using inlined function
@@ -884,28 +873,17 @@ class HubDatasetMixin:
                     last_target_only,
                 ) = _load_eegwindows_from_zarr(grp, preload)
 
-                if preload:
-                    # Convert to MNE objects and create dataset
-                    # Data is already in continuous format [n_channels, n_timepoints]
-                    info = mne.Info.from_json_dict(info_dict)
-                    raw = mne.io.RawArray(data, info)
-                    ds = EEGWindowsDataset(
-                        raw=raw,
-                        metadata=metadata,
-                        description=description,
-                        targets_from=targets_from,
-                        last_target_only=last_target_only,
-                    )
-                else:
-                    ds = EEGWindowsDataset._from_zarr(
-                        zarr_path=input_path,
-                        group_name=group_name,
-                        metadata=metadata,
-                        description=description,
-                        info_dict=info_dict,
-                        targets_from=targets_from,
-                        last_target_only=last_target_only,
-                    )
+                # Convert to MNE objects and create dataset
+                # Data is already in continuous format [n_channels, n_timepoints]
+                info = Info.from_json_dict(info_dict)
+                raw = mne.io.RawArray(data, info)
+                ds = EEGWindowsDataset(
+                    raw=raw,
+                    metadata=metadata,
+                    description=description,
+                    targets_from=targets_from,
+                    last_target_only=last_target_only,
+                )
 
             elif dataset_type == "RawDataset":
                 # Load using inlined function
@@ -913,22 +891,13 @@ class HubDatasetMixin:
                     grp, preload
                 )
 
-                if preload:
-                    # Convert to MNE objects and create dataset
-                    # Data is in continuous format [n_channels, n_timepoints]
-                    info = mne.Info.from_json_dict(info_dict)
-                    raw = mne.io.RawArray(data, info)
-                    ds = RawDataset(raw, description)
-                    if target_name is not None:
-                        ds.target_name = target_name
-                else:
-                    ds = RawDataset._from_zarr(
-                        zarr_path=input_path,
-                        group_name=group_name,
-                        description=description,
-                        info_dict=info_dict,
-                        target_name=target_name,
-                    )
+                # Convert to MNE objects and create dataset
+                # Data is in continuous format [n_channels, n_timepoints]
+                info = Info.from_json_dict(info_dict)
+                raw = mne.io.RawArray(data, info)
+                ds = RawDataset(raw, description)
+                if target_name is not None:
+                    ds.target_name = target_name
 
             else:
                 raise ValueError(f"Unsupported dataset_type: {dataset_type}")