Braindecode 1.3.0.dev177551367__tar.gz → 1.3.0.dev177628147__tar.gz

{braindecode-1.3.0.dev177551367/braindecode.egg-info → braindecode-1.3.0.dev177628147}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.0.dev177551367
+Version: 1.3.0.dev177628147
 Summary: Deep learning software to decode EEG, ECG or MEG signals
-Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>
+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>
 License: BSD-3-Clause
 Project-URL: homepage, https://braindecode.org
@@ -38,12 +38,14 @@ Requires-Dist: wfdb
 Requires-Dist: h5py
 Requires-Dist: linear_attention_transformer
 Requires-Dist: docstring_inheritance
+Requires-Dist: rotary_embedding_torch
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.2.0; extra == "moabb"
 Provides-Extra: eegprep
 Requires-Dist: eegprep[eeglabio]>=0.1.1; extra == "eegprep"
-Provides-Extra: hug
-Requires-Dist: huggingface_hub[torch]>=0.20.0; extra == "hug"
+Provides-Extra: hub
+Requires-Dist: huggingface_hub[torch]>=0.20.0; extra == "hub"
+Requires-Dist: zarr<3.0,>=2.18; extra == "hub"
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
@@ -69,7 +71,11 @@ Requires-Dist: pre-commit; extra == "docs"
 Requires-Dist: openneuro-py; extra == "docs"
 Requires-Dist: plotly; extra == "docs"
 Provides-Extra: all
-Requires-Dist: braindecode[docs,eegprep,hug,moabb,tests]; extra == "all"
+Requires-Dist: braindecode[moabb]; extra == "all"
+Requires-Dist: braindecode[hub]; extra == "all"
+Requires-Dist: braindecode[tests]; extra == "all"
+Requires-Dist: braindecode[docs]; extra == "all"
+Requires-Dist: braindecode[eegprep]; extra == "all"
 Dynamic: license-file
 
 .. image:: https://badges.gitter.im/braindecodechat/community.svg

{braindecode-1.3.0.dev177551367 → braindecode-1.3.0.dev177628147}/braindecode/augmentation/functional.py

@@ -1,6 +1,7 @@
 # Authors: Cédric Rommel <cedric.rommel@inria.fr>
 #          Alexandre Gramfort <alexandre.gramfort@inria.fr>
 #          Gustavo Rodrigues <gustavenrique01@gmail.com>
+#          Bruna Lopes <brunajaflopes@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -1194,3 +1195,103 @@ def mask_encoding(
     X[mask] = 0
 
     return X, y  # Return the masked tensor and labels
+
+
+def channels_rereference(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    random_state: int | np.random.RandomState | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Randomly re-reference channels in EEG data matrix.
+
+    Part of the augmentations proposed in [1]_
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        EEG input example or batch.
+    y : torch.Tensor
+        EEG labels for the example or batch.
+    random_state: int | numpy.random.Generator, optional
+        Seed to be used to instantiate numpy random number generator instance.
+        Defaults to None.
+
+    Returns
+    -------
+    torch.Tensor
+        Transformed inputs.
+    torch.Tensor
+        Transformed labels.
+
+    References
+    ----------
+    .. [1] Mohsenvand, M.N., Izadi, M.R. &amp; Maes, P.. (2020). Contrastive
+        Representation Learning for Electroencephalogram Classification. Proceedings
+        of the Machine Learning for Health NeurIPS Workshop, in Proceedings of Machine
+        Learning Research 136:238-253
+
+    """
+
+    rng = check_random_state(random_state)
+    batch_size, n_channels, _ = X.shape
+
+    ch = rng.randint(0, n_channels, size=batch_size)
+
+    X_ch = X[torch.arange(batch_size), ch, :]
+    X = X - X_ch.unsqueeze(1)
+    X[torch.arange(batch_size), ch, :] = -X_ch
+
+    return X, y
+
+
+def amplitude_scale(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    scale: tuple,
+    random_state: int | np.random.RandomState | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Rescale amplitude of each channel based on a random sampled scaling value.
+
+    Part of the augmentations proposed in [1]_
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        EEG input example or batch.
+    y : torch.Tensor
+        EEG labels for the example or batch.
+    scale : tuple of floats
+        Interval from which ypu sample the scaling value
+    random_state: int | numpy.random.Generator, optional
+        Seed to be used to instantiate numpy random number generator instance.
+        Defaults to None.
+
+    Returns
+    -------
+    torch.Tensor
+        Transformed inputs.
+    torch.Tensor
+        Transformed labels.
+
+    References
+    ----------
+    .. [1] Mohsenvand, M.N., Izadi, M.R. &amp; Maes, P.. (2020). Contrastive
+        Representation Learning for Electroencephalogram Classification. Proceedings
+        of the Machine Learning for Health NeurIPS Workshop, in Proceedings of Machine
+        Learning Research 136:238-253
+
+    """
+
+    rng = torch.Generator()
+    rng.manual_seed(random_state)
+    batch_size, n_channels, _ = X.shape
+
+    # Parameter for scaling amplitude / channel / trial
+    l, h = scale
+    s = l + (h - l) * torch.rand(
+        batch_size, n_channels, 1, generator=rng, device=X.device, dtype=X.dtype
+    )
+
+    X = s * X
+
+    return X, y

{braindecode-1.3.0.dev177551367 → braindecode-1.3.0.dev177628147}/braindecode/augmentation/transforms.py

@@ -1,6 +1,7 @@
 # Authors: Cédric Rommel <cedric.rommel@inria.fr>
 #          Alexandre Gramfort <alexandre.gramfort@inria.fr>
 #          Gustavo Rodrigues <gustavenrique01@gmail.com>
+#          Bruna Lopes <brunajaflopes@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -13,9 +14,11 @@ from mne.channels import make_standard_montage
 
 from .base import Transform
 from .functional import (
+    amplitude_scale,
     bandstop_filter,
     channels_dropout,
     channels_permute,
+    channels_rereference,
     channels_shuffle,
     frequency_shift,
     ft_surrogate,
@@ -1271,3 +1274,74 @@ class MaskEncoding(Transform):
             "segment_length": segment_length,
             "n_segments": self.n_segments,
         }
+
+
+class ChannelsReref(Transform):
+    """Randomly re-reference channels in EEG data matrix.
+
+    Part of the augmentations proposed in [1]_
+
+    Parameters
+    ----------
+    probability: float
+        Float setting the probability of applying the operation.
+    random_state: int | numpy.random.Generator, optional
+        Seed to be used to instantiate numpy random number generator instance.
+        Used to decide whether or not to transform given the probability
+        argument, to sample which channels to shuffle and to carry the shuffle.
+        Defaults to None.
+
+    References
+    ----------
+    .. [1] Mohsenvand, M.N., Izadi, M.R. &amp; Maes, P.. (2020). Contrastive
+        Representation Learning for Electroencephalogram Classification. Proceedings
+        of the Machine Learning for Health NeurIPS Workshop, in Proceedings of Machine
+        Learning Research 136:238-253 Available from https://proceedings.mlr.press/v136/mohsenvand20a.html.
+
+    """
+
+    operation = staticmethod(channels_rereference)  # type: ignore[assignment]
+
+    def __init__(self, probability, random_state=None):
+        super().__init__(probability=probability, random_state=random_state)
+
+    def get_augmentation_params(self, *batch):
+        """Return transform parameters"""
+        return {
+            "random_state": self.rng,
+        }
+
+
+class AmplitudeScale(Transform):
+    """Rescale amplitude based on a random sampled scaling value.
+
+    Part of the augmentations proposed in [1]_
+
+    Parameters
+    ----------
+    probability: float
+        Float setting the probability of applying the operation.
+    random_state: int | numpy.random.Generator, optional
+        Seed to be used to instantiate numpy random number generator instance.
+        Used to decide whether or not to transform given the probability
+        argument, to sample which channels to shuffle and to carry the shuffle.
+        Defaults to None.
+
+    References
+    ----------
+    .. [1] Mohsenvand, M.N., Izadi, M.R. &amp; Maes, P.. (2020). Contrastive
+        Representation Learning for Electroencephalogram Classification. Proceedings
+        of the Machine Learning for Health NeurIPS Workshop, in Proceedings of Machine
+        Learning Research 136:238-253 Available from https://proceedings.mlr.press/v136/mohsenvand20a.html.
+
+    """
+
+    operation = staticmethod(amplitude_scale)  # type: ignore[assignment]
+
+    def __init__(self, probability, interval=(0.5, 2), random_state=None):
+        super().__init__(probability=probability, random_state=random_state)
+        self.scale = interval
+
+    def get_augmentation_params(self, *batch):
+        """Return transform parameters"""
+        return {"random_state": self.rng, "scale": self.scale}

{braindecode-1.3.0.dev177551367 → braindecode-1.3.0.dev177628147}/braindecode/datasets/base.py

@@ -19,7 +19,7 @@ import warnings
 from abc import abstractmethod
 from collections.abc import Callable
 from glob import glob
-from typing import Generic, Iterable, no_type_check
+from typing import Any, Generic, Iterable, no_type_check
 
 import mne.io
 import numpy as np
@@ -28,6 +28,9 @@ from mne.utils.docs import deprecated
 from torch.utils.data import ConcatDataset, Dataset
 from typing_extensions import TypeVar
 
+from .hub import HubDatasetMixin
+from .registry import register_dataset
+
 
 def _create_description(description) -> pd.Series:
     if description is not None:
@@ -97,6 +100,7 @@ class RecordDataset(Dataset[tuple[np.ndarray, int | str, tuple[int, int, int]]])
 T = TypeVar("T", bound=RecordDataset)
 
 
+@register_dataset
 class RawDataset(RecordDataset):
     """Returns samples from an mne.io.Raw object along with a target.
 
@@ -129,6 +133,7 @@ class RawDataset(RecordDataset):
 
         # save target name for load/save later
         self.target_name = self._target_name(target_name)
+        self.raw_preproc_kwargs: list[dict[str, Any]] = []
 
     def __getitem__(self, index):
         X = self.raw[:, index][0]
@@ -176,10 +181,12 @@ class RawDataset(RecordDataset):
     "If you want to type a Braindecode dataset (i.e. RawDataset|EEGWindowsDataset|WindowsDataset), "
     "use the RecordDataset class instead."
 )
+@register_dataset
 class BaseDataset(RawDataset):
     pass
 
 
+@register_dataset
 class EEGWindowsDataset(RecordDataset):
     """Returns windows from an mne.Raw object, its window indices, along with a target.
 
@@ -235,6 +242,7 @@ class EEGWindowsDataset(RecordDataset):
         ].to_numpy()
         if self.targets_from == "metadata":
             self.y = metadata.loc[:, "target"].to_list()
+        self.raw_preproc_kwargs: list[dict[str, Any]] = []
 
     def __getitem__(self, index: int):
         """Get a window and its target.
@@ -285,6 +293,7 @@ class EEGWindowsDataset(RecordDataset):
         return len(self.crop_inds)
 
 
+@register_dataset
 class WindowsDataset(RecordDataset):
     """Returns windows from an mne.Epochs object along with a target.
 
@@ -334,6 +343,8 @@ class WindowsDataset(RecordDataset):
         ].to_numpy()
         if self.targets_from == "metadata":
             self.y = metadata.loc[:, "target"].to_list()
+        self.raw_preproc_kwargs: list[dict[str, Any]] = []
+        self.window_preproc_kwargs: list[dict[str, Any]] = []
 
     def __getitem__(self, index: int):
         """Get a window and its target.
@@ -374,12 +385,16 @@ class WindowsDataset(RecordDataset):
         return len(self.windows.events)
 
 
-class BaseConcatDataset(ConcatDataset, Generic[T]):
+@register_dataset
+class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
     """A base class for concatenated datasets.
 
     Holds either mne.Raw or mne.Epoch in self.datasets and has
     a pandas DataFrame with additional description.
 
+    Includes Hugging Face Hub integration via HubDatasetMixin for
+    uploading and downloading datasets.
+
     Parameters
     ----------
     list_of_ds : list
@@ -794,7 +809,7 @@ class BaseConcatDataset(ConcatDataset, Generic[T]):
                 kwargs = getattr(ds, kwargs_name)
                 if kwargs is not None:
                     with open(kwargs_file_path, "w") as f:
-                        json.dump(kwargs, f)
+                        json.dump(kwargs, f, indent=2)
 
     @staticmethod
     def _save_target_name(sub_dir, ds):