braindecode 1.2.0.dev184328194__py3-none-any.whl → 1.3.0.dev171178473__py3-none-any.whl

braindecode/datasets/base.py

@@ -16,14 +16,17 @@ import json
 import os
 import shutil
 import warnings
+from abc import abstractmethod
 from collections.abc import Callable
 from glob import glob
-from typing import Iterable, no_type_check
+from typing import Generic, Iterable, no_type_check
 
 import mne.io
 import numpy as np
 import pandas as pd
+from mne.utils.docs import deprecated
 from torch.utils.data import ConcatDataset, Dataset
+from typing_extensions import TypeVar
 
 
 def _create_description(description) -> pd.Series:
@@ -37,7 +40,64 @@ def _create_description(description) -> pd.Series:
     return description
 
 
-class BaseDataset(Dataset):
+class RecordDataset(Dataset[tuple[np.ndarray, int | str, tuple[int, int, int]]]):
+    def __init__(
+        self,
+        description: dict | pd.Series | None = None,
+        transform: Callable | None = None,
+    ):
+        self._description = _create_description(description)
+        self.transform = transform
+
+    @abstractmethod
+    def __len__(self) -> int:
+        pass
+
+    @property
+    def description(self) -> pd.Series:
+        return self._description
+
+    def set_description(self, description: dict | pd.Series, overwrite: bool = False):
+        """Update (add or overwrite) the dataset description.
+
+        Parameters
+        ----------
+        description: dict | pd.Series
+            Description in the form key: value.
+        overwrite: bool
+            Has to be True if a key in description already exists in the
+            dataset description.
+        """
+        description = _create_description(description)
+        if self.description is None:
+            self._description = description
+        else:
+            for key, value in description.items():
+                # if the key is already in the existing description, drop it
+                if key in self._description:
+                    assert overwrite, (
+                        f"'{key}' already in description. Please "
+                        f"rename or set overwrite to True."
+                    )
+                    self._description.pop(key)
+            self._description = pd.concat([self.description, description])
+
+    @property
+    def transform(self) -> Callable | None:
+        return self._transform
+
+    @transform.setter
+    def transform(self, value: Callable | None):
+        if value is not None and not callable(value):
+            raise ValueError("Transform needs to be a callable.")
+        self._transform = value
+
+
+# Type of the datasets contained in BaseConcatDataset
+T = TypeVar("T", bound=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.
@@ -64,9 +124,8 @@ class BaseDataset(Dataset):
         target_name: str | tuple[str, ...] | None = None,
         transform: Callable | None = None,
     ):
+        super().__init__(description, transform)
         self.raw = raw
-        self._description = _create_description(description)
-        self.transform = transform
 
         # save target name for load/save later
         self.target_name = self._target_name(target_name)
@@ -85,45 +144,6 @@ class BaseDataset(Dataset):
     def __len__(self):
         return len(self.raw)
 
-    @property
-    def transform(self):
-        return self._transform
-
-    @transform.setter
-    def transform(self, value):
-        if value is not None and not callable(value):
-            raise ValueError("Transform needs to be a callable.")
-        self._transform = value
-
-    @property
-    def description(self) -> pd.Series:
-        return self._description
-
-    def set_description(self, description: dict | pd.Series, overwrite: bool = False):
-        """Update (add or overwrite) the dataset description.
-
-        Parameters
-        ----------
-        description: dict | pd.Series
-            Description in the form key: value.
-        overwrite: bool
-            Has to be True if a key in description already exists in the
-            dataset description.
-        """
-        description = _create_description(description)
-        for key, value in description.items():
-            # if the key is already in the existing description, drop it
-            if self._description is not None and key in self._description:
-                assert overwrite, (
-                    f"'{key}' already in description. Please "
-                    f"rename or set overwrite to True."
-                )
-                self._description.pop(key)
-        if self._description is None:
-            self._description = description
-        else:
-            self._description = pd.concat([self.description, description])
-
     def _target_name(self, target_name):
         if target_name is not None and not isinstance(target_name, (str, tuple, list)):
             raise ValueError("target_name has to be None, str, tuple or list")
@@ -150,7 +170,17 @@ class BaseDataset(Dataset):
         return target_name if len(target_name) > 1 else target_name[0]
 
 
-class EEGWindowsDataset(BaseDataset):
+@deprecated(
+    "The BaseDataset class is deprecated. "
+    "If you want to instantiate a dataset containing raws, use RawDataset instead. "
+    "If you want to type a Braindecode dataset (i.e. RawDataset|EEGWindowsDataset|WindowsDataset), "
+    "use the RecordDataset class instead."
+)
+class BaseDataset(RawDataset):
+    pass
+
+
+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
@@ -161,12 +191,12 @@ class EEGWindowsDataset(BaseDataset):
     required to serve information about the windowing (e.g., useful for cropped
     training).
     See `braindecode.datautil.windowers` to directly create a `WindowsDataset`
-    from a `BaseDataset` object.
+    from a `RawDataset` object.
 
     Parameters
     ----------
     windows : mne.Raw or mne.Epochs (Epochs is outdated)
-        Windows obtained through the application of a windower to a BaseDataset
+        Windows obtained through the application of a windower to a ``RawDataset``
         (see `braindecode.datautil.windowers`).
     description : dict | pandas.Series | None
         Holds additional info about the windows.
@@ -185,18 +215,17 @@ class EEGWindowsDataset(BaseDataset):
 
     def __init__(
         self,
-        raw: mne.io.BaseRaw | mne.BaseEpochs,
+        raw: mne.io.BaseRaw,
         metadata: pd.DataFrame,
         description: dict | pd.Series | None = None,
         transform: Callable | None = None,
         targets_from: str = "metadata",
         last_target_only: bool = True,
     ):
+        super().__init__(description, transform)
         self.raw = raw
         self.metadata = metadata
-        self._description = _create_description(description)
 
-        self.transform = transform
         self.last_target_only = last_target_only
         if targets_from not in ("metadata", "channels"):
             raise ValueError("Wrong value for parameter `targets_from`.")
@@ -255,44 +284,8 @@ class EEGWindowsDataset(BaseDataset):
     def __len__(self):
         return len(self.crop_inds)
 
-    @property
-    def transform(self):
-        return self._transform
-
-    @transform.setter
-    def transform(self, value):
-        if value is not None and not callable(value):
-            raise ValueError("Transform needs to be a callable.")
-        self._transform = value
-
-    @property
-    def description(self) -> pd.Series:
-        return self._description
-
-    def set_description(self, description: dict | pd.Series, overwrite: bool = False):
-        """Update (add or overwrite) the dataset description.
-
-        Parameters
-        ----------
-        description: dict | pd.Series
-            Description in the form key: value.
-        overwrite: bool
-            Has to be True if a key in description already exists in the
-            dataset description.
-        """
-        description = _create_description(description)
-        for key, value in description.items():
-            # if they key is already in the existing description, drop it
-            if key in self._description:
-                assert overwrite, (
-                    f"'{key}' already in description. Please "
-                    f"rename or set overwrite to True."
-                )
-                self._description.pop(key)
-        self._description = pd.concat([self.description, description])
 
-
-class WindowsDataset(BaseDataset):
+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
@@ -303,12 +296,12 @@ class WindowsDataset(BaseDataset):
     required to serve information about the windowing (e.g., useful for cropped
     training).
     See `braindecode.datautil.windowers` to directly create a `WindowsDataset`
-    from a `BaseDataset` object.
+    from a ``RawDataset`` object.
 
     Parameters
     ----------
     windows : mne.Epochs
-        Windows obtained through the application of a windower to a BaseDataset
+        Windows obtained through the application of a windower to a RawDataset
         (see `braindecode.datautil.windowers`).
     description : dict | pandas.Series | None
         Holds additional info about the windows.
@@ -327,19 +320,20 @@ class WindowsDataset(BaseDataset):
         targets_from: str = "metadata",
         last_target_only: bool = True,
     ):
+        super().__init__(description, transform)
         self.windows = windows
-        self._description = _create_description(description)
-        self.transform = transform
         self.last_target_only = last_target_only
         if targets_from not in ("metadata", "channels"):
             raise ValueError("Wrong value for parameter `targets_from`.")
         self.targets_from = targets_from
 
-        self.crop_inds = self.windows.metadata.loc[
+        metadata = self.windows.metadata
+        assert metadata is not None, "WindowsDataset requires windows with metadata."
+        self.crop_inds = metadata.loc[
             :, ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"]
         ].to_numpy()
         if self.targets_from == "metadata":
-            self.y = self.windows.metadata.loc[:, "target"].to_list()
+            self.y = metadata.loc[:, "target"].to_list()
 
     def __getitem__(self, index: int):
         """Get a window and its target.
@@ -379,44 +373,8 @@ class WindowsDataset(BaseDataset):
     def __len__(self) -> int:
         return len(self.windows.events)
 
-    @property
-    def transform(self):
-        return self._transform
-
-    @transform.setter
-    def transform(self, value):
-        if value is not None and not callable(value):
-            raise ValueError("Transform needs to be a callable.")
-        self._transform = value
 
-    @property
-    def description(self) -> pd.Series:
-        return self._description
-
-    def set_description(self, description: dict | pd.Series, overwrite: bool = False):
-        """Update (add or overwrite) the dataset description.
-
-        Parameters
-        ----------
-        description: dict | pd.Series
-            Description in the form key: value.
-        overwrite: bool
-            Has to be True if a key in description already exists in the
-            dataset description.
-        """
-        description = _create_description(description)
-        for key, value in description.items():
-            # if they key is already in the existing description, drop it
-            if key in self._description:
-                assert overwrite, (
-                    f"'{key}' already in description. Please "
-                    f"rename or set overwrite to True."
-                )
-                self._description.pop(key)
-        self._description = pd.concat([self.description, description])
-
-
-class BaseConcatDataset(ConcatDataset):
+class BaseConcatDataset(ConcatDataset, Generic[T]):
     """A base class for concatenated datasets.
 
     Holds either mne.Raw or mne.Epoch in self.datasets and has
@@ -425,22 +383,27 @@ class BaseConcatDataset(ConcatDataset):
     Parameters
     ----------
     list_of_ds : list
-        list of BaseDataset, BaseConcatDataset or WindowsDataset
+        list of RecordDataset
     target_transform : callable | None
         Optional function to call on targets before returning them.
 
     """
 
+    datasets: list[T]
+
     def __init__(
         self,
-        list_of_ds: list[BaseDataset | BaseConcatDataset | WindowsDataset]
-        | None = None,
+        list_of_ds: list[T | BaseConcatDataset[T]],
         target_transform: Callable | None = None,
     ):
         # if we get a list of BaseConcatDataset, get all the individual datasets
-        if list_of_ds and isinstance(list_of_ds[0], BaseConcatDataset):
-            list_of_ds = [d for ds in list_of_ds for d in ds.datasets]
-        super().__init__(list_of_ds)
+        flattened_list_of_ds: list[T] = []
+        for ds in list_of_ds:
+            if isinstance(ds, BaseConcatDataset):
+                flattened_list_of_ds.extend(ds.datasets)
+            else:
+                flattened_list_of_ds.append(ds)
+        super().__init__(flattened_list_of_ds)
 
         self.target_transform = target_transform
 
@@ -703,22 +666,23 @@ class BaseConcatDataset(ConcatDataset):
                 ds.set_description({key: value_}, overwrite=overwrite)
 
     def save(self, path: str, overwrite: bool = False, offset: int = 0):
-        """Save datasets to files by creating one subdirectory for each dataset:
-        path/
-            0/
-                0-raw.fif | 0-epo.fif
-                description.json
-                raw_preproc_kwargs.json (if raws were preprocessed)
-                window_kwargs.json (if this is a windowed dataset)
-                window_preproc_kwargs.json  (if windows were preprocessed)
-                target_name.json (if target_name is not None and dataset is raw)
-            1/
-                1-raw.fif | 1-epo.fif
-                description.json
-                raw_preproc_kwargs.json (if raws were preprocessed)
-                window_kwargs.json (if this is a windowed dataset)
-                window_preproc_kwargs.json  (if windows were preprocessed)
-                target_name.json (if target_name is not None and dataset is raw)
+        """Save datasets to files by creating one subdirectory for each dataset::
+
+            path/
+                0/
+                    0-raw.fif | 0-epo.fif
+                    description.json
+                    raw_preproc_kwargs.json (if raws were preprocessed)
+                    window_kwargs.json (if this is a windowed dataset)
+                    window_preproc_kwargs.json  (if windows were preprocessed)
+                    target_name.json (if target_name is not None and dataset is raw)
+                1/
+                    1-raw.fif | 1-epo.fif
+                    description.json
+                    raw_preproc_kwargs.json (if raws were preprocessed)
+                    window_kwargs.json (if this is a windowed dataset)
+                    window_preproc_kwargs.json  (if windows were preprocessed)
+                    target_name.json (if target_name is not None and dataset is raw)
 
         Parameters
         ----------
@@ -815,7 +779,7 @@ class BaseConcatDataset(ConcatDataset):
     @staticmethod
     def _save_description(sub_dir, description):
         description_file_path = os.path.join(sub_dir, "description.json")
-        description.to_json(description_file_path)
+        description.to_json(description_file_path, default_handler=str)
 
     @staticmethod
     def _save_kwargs(sub_dir, ds):

braindecode/datasets/bcicomp.py

@@ -16,7 +16,7 @@ import numpy as np
 from mne.utils import verbose
 from scipy.io import loadmat
 
-from braindecode.datasets import BaseConcatDataset, BaseDataset
+from braindecode.datasets import BaseConcatDataset, RawDataset
 
 DATASET_URL = (
     "https://stacks.stanford.edu/file/druid:zk881ps0522/"
@@ -73,8 +73,8 @@ class BCICompetitionIVDataset4(BaseConcatDataset):
                 file_name=file_path.split("/")[-1],
                 session="test",
             )
-            datasets.append(BaseDataset(raw_train, description=desc_train))
-            datasets.append(BaseDataset(raw_test, description=desc_test))
+            datasets.append(RawDataset(raw_train, description=desc_train))
+            datasets.append(RawDataset(raw_test, description=desc_test))
         super().__init__(datasets)
 
     @staticmethod
@@ -85,7 +85,7 @@ class BCICompetitionIVDataset4(BaseConcatDataset):
         ----------
         path  (None | str) – Location of where to look for the data storing location.
         If None, the environment variable or config parameter
-        MNE_DATASETS_(dataset)_PATH is used. If it doesn’t exist, the “~/mne_data”
+        MNE_DATASETS_(dataset)_PATH is used. If it doesn't exist, the “~/mne_data”
         directory is used. If the dataset is not found under the given path, the data
         will be automatically downloaded to the specified folder.
         force_update (bool) – Force update of the dataset even if a local copy exists.

braindecode/datasets/bids.py

@@ -19,7 +19,7 @@ import numpy as np
 import pandas as pd
 from joblib import Parallel, delayed
 
-from .base import BaseConcatDataset, BaseDataset, WindowsDataset
+from .base import BaseConcatDataset, RawDataset, WindowsDataset
 
 
 def _description_from_bids_path(bids_path: mne_bids.BIDSPath) -> dict[str, Any]:
@@ -186,12 +186,12 @@ class BIDSDataset(BaseConcatDataset):
         )
         super().__init__(all_base_ds)
 
-    def _get_dataset(self, bids_path: mne_bids.BIDSPath) -> BaseDataset:
+    def _get_dataset(self, bids_path: mne_bids.BIDSPath) -> RawDataset:
         description = _description_from_bids_path(bids_path)
         raw = mne_bids.read_raw_bids(bids_path, verbose=False)
         if self.preload:
             raw.load_data()
-        return BaseDataset(raw, description)
+        return RawDataset(raw, description)
 
 
 class BIDSEpochsDataset(BIDSDataset):

braindecode/datasets/experimental.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import Callable, Sequence
+
+import mne_bids
+from torch.utils.data import IterableDataset, get_worker_info
+
+
+class BIDSIterableDataset(IterableDataset):
+    """Dataset for loading BIDS.
+
+    .. warning::
+        This class is experimental and may change in the future.
+
+    .. warning::
+        This dataset is not consistent with the Braindecode API.
+
+    This class has the same parameters as the :func:`mne_bids.find_matching_paths` function
+    as it will be used to find the files to load. The default ``extensions`` parameter was changed.
+
+    More information on BIDS (Brain Imaging Data Structure)
+    can be found at https://bids.neuroimaging.io
+
+    Examples
+    --------
+    >>> from braindecode.datasets import RecordDataset, BaseConcatDataset
+    >>> from braindecode.datasets.bids import BIDSIterableDataset, _description_from_bids_path
+    >>> from braindecode.preprocessing import create_fixed_length_windows
+    >>>
+    >>> def my_reader_fn(path):
+    ...     raw = mne_bids.read_raw_bids(path)
+    ...     desc = _description_from_bids_path(path)
+    ...     ds = RawDataset(raw, description=desc)
+    ...     windows_ds = create_fixed_length_windows(
+    ...         BaseConcatDataset([ds]),
+    ...         window_size_samples=400,
+    ...         window_stride_samples=200,
+    ...     )
+    ...     return windows_ds
+    >>>
+    >>> dataset = BIDSIterableDataset(
+    ...     reader_fn=my_reader_fn,
+    ...     root="root/of/my/bids/dataset/",
+    ... )
+
+    Parameters
+    ----------
+    reader_fn : Callable[[mne_bids.BIDSPath], Sequence]
+        A function that takes a BIDSPath and returns a dataset.
+    pool_size : int
+        The number of recordings to read and sample from.
+    bids_paths : list[mne_bids.BIDSPath] | None
+        A list of BIDSPaths to load. If None, will use the paths found by
+        :func:`mne_bids.find_matching_paths` and the arguments below.
+    root : pathlib.Path | str
+        The root of the BIDS path.
+    subjects : str | array-like of str | None
+        The subject ID. Corresponds to "sub".
+    sessions : str | array-like of str | None
+        The acquisition session. Corresponds to "ses".
+    tasks : str | array-like of str | None
+        The experimental task. Corresponds to "task".
+    acquisitions: str | array-like of str | None
+        The acquisition parameters. Corresponds to "acq".
+    runs : str | array-like of str | None
+        The run number. Corresponds to "run".
+    processings : str | array-like of str | None
+        The processing label. Corresponds to "proc".
+    recordings : str | array-like of str | None
+        The recording name. Corresponds to "rec".
+    spaces : str | array-like of str | None
+        The coordinate space for anatomical and sensor location
+        files (e.g., ``*_electrodes.tsv``, ``*_markers.mrk``).
+        Corresponds to "space".
+        Note that valid values for ``space`` must come from a list
+        of BIDS keywords as described in the BIDS specification.
+    splits : str | array-like of str | None
+        The split of the continuous recording file for ``.fif`` data.
+        Corresponds to "split".
+    descriptions : str | array-like of str | None
+        This corresponds to the BIDS entity ``desc``. It is used to provide
+        additional information for derivative data, e.g., preprocessed data
+        may be assigned ``description='cleaned'``.
+    suffixes : str | array-like of str | None
+        The filename suffix. This is the entity after the
+        last ``_`` before the extension. E.g., ``'channels'``.
+        The following filename suffix's are accepted:
+        'meg', 'markers', 'eeg', 'ieeg', 'T1w',
+        'participants', 'scans', 'electrodes', 'coordsystem',
+        'channels', 'events', 'headshape', 'digitizer',
+        'beh', 'physio', 'stim'
+    extensions : str | array-like of str | None
+        The extension of the filename. E.g., ``'.json'``.
+        By default, uses the ones accepted by :func:`mne_bids.read_raw_bids`.
+    datatypes : str | array-like of str | None
+        The BIDS data type, e.g., ``'anat'``, ``'func'``, ``'eeg'``, ``'meg'``,
+        ``'ieeg'``.
+    check : bool
+        If ``True``, only returns paths that conform to BIDS. If ``False``
+        (default), the ``.check`` attribute of the returned
+        :class:`mne_bids.BIDSPath` object will be set to ``True`` for paths that
+        do conform to BIDS, and to ``False`` for those that don't.
+    preload : bool
+        If True, preload the data. Defaults to False.
+    n_jobs : int
+        Number of jobs to run in parallel. Defaults to 1.
+    """
+
+    def __init__(
+        self,
+        reader_fn: Callable[[mne_bids.BIDSPath], Sequence],
+        pool_size: int = 4,
+        bids_paths: list[mne_bids.BIDSPath] | None = None,
+        root: Path | str | None = None,
+        subjects: str | list[str] | None = None,
+        sessions: str | list[str] | None = None,
+        tasks: str | list[str] | None = None,
+        acquisitions: str | list[str] | None = None,
+        runs: str | list[str] | None = None,
+        processings: str | list[str] | None = None,
+        recordings: str | list[str] | None = None,
+        spaces: str | list[str] | None = None,
+        splits: str | list[str] | None = None,
+        descriptions: str | list[str] | None = None,
+        suffixes: str | list[str] | None = None,
+        extensions: str | list[str] | None = [
+            ".con",
+            ".sqd",
+            ".pdf",
+            ".fif",
+            ".ds",
+            ".vhdr",
+            ".set",
+            ".edf",
+            ".bdf",
+            ".EDF",
+            ".snirf",
+            ".cdt",
+            ".mef",
+            ".nwb",
+        ],
+        datatypes: str | list[str] | None = None,
+        check: bool = False,
+    ):
+        if bids_paths is None:
+            bids_paths = mne_bids.find_matching_paths(
+                root=root,
+                subjects=subjects,
+                sessions=sessions,
+                tasks=tasks,
+                acquisitions=acquisitions,
+                runs=runs,
+                processings=processings,
+                recordings=recordings,
+                spaces=spaces,
+                splits=splits,
+                descriptions=descriptions,
+                suffixes=suffixes,
+                extensions=extensions,
+                datatypes=datatypes,
+                check=check,
+                ignore_json=True,
+            )
+            # Filter out _epo.fif files:
+            bids_paths = [
+                bids_path
+                for bids_path in bids_paths
+                if not (bids_path.suffix == "epo" and bids_path.extension == ".fif")
+            ]
+        self.bids_paths = bids_paths
+        self.reader_fn = reader_fn
+        self.pool_size = pool_size
+
+    def __add__(self, other):
+        assert isinstance(other, BIDSIterableDataset)
+        return BIDSIterableDataset(
+            reader_fn=self.reader_fn,
+            bids_paths=self.bids_paths + other.bids_paths,
+            pool_size=self.pool_size,
+        )
+
+    def __iadd__(self, other):
+        assert isinstance(other, BIDSIterableDataset)
+        self.bids_paths += other.bids_paths
+        return self
+
+    def __iter__(self):
+        worker_info = get_worker_info()
+        if worker_info is None:  # single-process data loading, return the full iterator
+            bids_paths = self.bids_paths
+        else:  # in a worker process
+            # split workload
+            bids_paths = self.bids_paths[worker_info.id :: worker_info.num_workers]
+
+        pool = []
+        end = False
+        paths_it = iter(random.sample(bids_paths, k=len(bids_paths)))
+        while not (end and len(pool) == 0):
+            while not end and len(pool) < self.pool_size:
+                try:
+                    bids_path = next(paths_it)
+                    ds = self.reader_fn(bids_path)
+                    if ds is None:
+                        print(f"Skipping {bids_path} as it is too short.")
+                        continue
+                    idx = iter(random.sample(range(len(ds)), k=len(ds)))
+                    pool.append((ds, idx))
+                except StopIteration:
+                    end = True
+            i_pool = random.randint(0, len(pool) - 1)
+            ds, idx = pool[i_pool]
+            try:
+                i_ds = next(idx)
+                yield ds[i_ds]
+            except StopIteration:
+                pool.pop(i_pool)