braindecode 1.3.0.dev175955015__py3-none-any.whl → 1.3.0.dev176481332__py3-none-any.whl

braindecode/augmentation/base.py

@@ -189,7 +189,7 @@ class AugmentedDataLoader(DataLoader):
 
     Parameters
     ----------
-    dataset : BaseDataset
+    dataset : RecordDataset
         The dataset containing the signals.
     transforms : list | Transform, optional
         Transform or sequence of Transform to be applied to each batch.

braindecode/datasets/__init__.py

@@ -2,7 +2,13 @@
 Loader code for some datasets.
 """
 
-from .base import BaseConcatDataset, BaseDataset, WindowsDataset
+from .base import (
+    BaseConcatDataset,
+    EEGWindowsDataset,
+    RawDataset,
+    RecordDataset,
+    WindowsDataset,
+)
 from .bcicomp import BCICompetitionIVDataset4
 from .bids import BIDSDataset, BIDSEpochsDataset
 from .mne import create_from_mne_epochs, create_from_mne_raw
@@ -15,7 +21,9 @@ from .xy import create_from_X_y
 
 __all__ = [
     "WindowsDataset",
-    "BaseDataset",
+    "EEGWindowsDataset",
+    "RecordDataset",
+    "RawDataset",
     "BaseConcatDataset",
     "BIDSDataset",
     "BIDSEpochsDataset",

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
         ----------

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

@@ -25,14 +25,14 @@ class BIDSIterableDataset(IterableDataset):
 
     Examples
     --------
-    >>> from braindecode.datasets import BaseDataset, BaseConcatDataset
+    >>> 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 = BaseDataset(raw, description=desc)
+    ...     ds = RawDataset(raw, description=desc)
     ...     windows_ds = create_fixed_length_windows(
     ...         BaseConcatDataset([ds]),
     ...         window_size_samples=400,

braindecode/datasets/mne.py

@@ -9,7 +9,7 @@ import mne
 import numpy as np
 import pandas as pd
 
-from .base import BaseConcatDataset, BaseDataset, WindowsDataset
+from .base import BaseConcatDataset, RawDataset, WindowsDataset
 
 
 def create_from_mne_raw(
@@ -75,11 +75,9 @@ def create_from_mne_raw(
                 f"length of 'raws' ({len(raws)}) and 'description' "
                 f"({len(descriptions)}) has to match"
             )
-        base_datasets = [
-            BaseDataset(raw, desc) for raw, desc in zip(raws, descriptions)
-        ]
+        base_datasets = [RawDataset(raw, desc) for raw, desc in zip(raws, descriptions)]
     else:
-        base_datasets = [BaseDataset(raw) for raw in raws]
+        base_datasets = [RawDataset(raw) for raw in raws]
 
     base_datasets = BaseConcatDataset(base_datasets)
     windows_datasets = create_windows_from_events(

braindecode/datasets/moabb.py

@@ -18,7 +18,7 @@ import pandas as pd
 
 from braindecode.util import _update_moabb_docstring
 
-from .base import BaseConcatDataset, BaseDataset
+from .base import BaseConcatDataset, RawDataset
 
 
 def _find_dataset_in_moabb(dataset_name, dataset_kwargs=None):
@@ -164,7 +164,7 @@ class MOABBDataset(BaseConcatDataset):
             dataset_load_kwargs=dataset_load_kwargs,
         )
         all_base_ds = [
-            BaseDataset(raw, row) for raw, (_, row) in zip(raws, description.iterrows())
+            RawDataset(raw, row) for raw, (_, row) in zip(raws, description.iterrows())
         ]
         super().__init__(all_base_ds)
 

braindecode/datasets/nmt.py

@@ -31,7 +31,7 @@ import pandas as pd
 from joblib import Parallel, delayed
 from mne.datasets import fetch_dataset
 
-from braindecode.datasets.base import BaseConcatDataset, BaseDataset
+from braindecode.datasets.base import BaseConcatDataset, RawDataset
 
 NMT_URL = "https://zenodo.org/record/10909103/files/NMT.zip"
 NMT_archive_name = "NMT.zip"
@@ -172,7 +172,7 @@ class NMT(BaseConcatDataset):
         d["n_samples"] = raw.n_times
         d["sfreq"] = raw.info["sfreq"]
         d["train"] = "train" in d.path.split(os.sep)
-        base_dataset = BaseDataset(raw, d, target_name)
+        base_dataset = RawDataset(raw, d, target_name)
         return base_dataset
 
 

braindecode/datasets/sleep_physio_challe_18.py

@@ -21,7 +21,7 @@ from mne.datasets.sleep_physionet._utils import _fetch_one
 from mne.datasets.utils import _get_path
 from mne.utils import warn
 
-from braindecode.datasets import BaseConcatDataset, BaseDataset
+from braindecode.datasets import BaseConcatDataset, RawDataset
 
 PC18_DIR = op.join(op.dirname(__file__), "data", "pc18")
 PC18_RECORDS = op.join(PC18_DIR, "sleep_records.csv")
@@ -403,7 +403,7 @@ class SleepPhysionetChallenge2018(BaseConcatDataset):
             },
             name="",
         )
-        base_dataset = BaseDataset(raw_file, desc)
+        base_dataset = RawDataset(raw_file, desc)
 
         if preproc is not None:
             from braindecode.preprocessing.preprocess import _preprocess

braindecode/datasets/sleep_physionet.py

@@ -12,7 +12,7 @@ import numpy as np
 import pandas as pd
 from mne.datasets.sleep_physionet.age import fetch_data
 
-from .base import BaseConcatDataset, BaseDataset
+from .base import BaseConcatDataset, RawDataset
 
 
 class SleepPhysionet(BaseConcatDataset):
@@ -71,7 +71,7 @@ class SleepPhysionet(BaseConcatDataset):
                 crop_wake_mins=crop_wake_mins,
                 crop=crop,
             )
-            base_ds = BaseDataset(raw, desc)
+            base_ds = RawDataset(raw, desc)
             all_base_ds.append(base_ds)
         super().__init__(all_base_ds)
 

braindecode/datasets/tuh.py

@@ -22,7 +22,7 @@ import numpy as np
 import pandas as pd
 from joblib import Parallel, delayed
 
-from .base import BaseConcatDataset, BaseDataset
+from .base import BaseConcatDataset, RawDataset
 
 
 class TUH(BaseConcatDataset):
@@ -214,7 +214,7 @@ class TUH(BaseConcatDataset):
             d["report"] = physician_report
         additional_description = pd.Series(d)
         description = pd.concat([description, additional_description])
-        base_dataset = BaseDataset(raw, description, target_name=target_name)
+        base_dataset = RawDataset(raw, description, target_name=target_name)
         return base_dataset
 
 

braindecode/datasets/xy.py

@@ -12,7 +12,7 @@ import numpy as np
 import pandas as pd
 from numpy.typing import ArrayLike, NDArray
 
-from .base import BaseConcatDataset, BaseDataset
+from .base import BaseConcatDataset, RawDataset
 
 log = logging.getLogger(__name__)
 
@@ -69,7 +69,7 @@ def create_from_X_y(
         n_samples_per_x.append(x.shape[1])
         info = mne.create_info(ch_names=ch_names, sfreq=sfreq)
         raw = mne.io.RawArray(x, info)
-        base_dataset = BaseDataset(
+        base_dataset = RawDataset(
             raw, pd.Series({"target": target}), target_name="target"
         )
         base_datasets.append(base_dataset)