Braindecode 1.3.0.dev175368549__tar.gz → 1.3.0.dev175435903__tar.gz

{braindecode-1.3.0.dev175368549/braindecode.egg-info → braindecode-1.3.0.dev175435903}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.0.dev175368549
+Version: 1.3.0.dev175435903
 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>
@@ -52,6 +52,9 @@ Requires-Dist: pytest-cov; extra == "tests"
 Requires-Dist: codecov; extra == "tests"
 Requires-Dist: pytest_cases; extra == "tests"
 Requires-Dist: mypy; extra == "tests"
+Provides-Extra: typing
+Requires-Dist: pydantic<3.0,>=2.0; extra == "typing"
+Requires-Dist: numpydantic>=1.7; extra == "typing"
 Provides-Extra: docs
 Requires-Dist: sphinx_gallery; extra == "docs"
 Requires-Dist: sphinx_rtd_theme; extra == "docs"
@@ -76,6 +79,7 @@ Requires-Dist: braindecode[hub]; extra == "all"
 Requires-Dist: braindecode[tests]; extra == "all"
 Requires-Dist: braindecode[docs]; extra == "all"
 Requires-Dist: braindecode[eegprep]; extra == "all"
+Requires-Dist: braindecode[typing]; extra == "all"
 Dynamic: license-file
 
 .. image:: https://badges.gitter.im/braindecodechat/community.svg

{braindecode-1.3.0.dev175368549 → braindecode-1.3.0.dev175435903}/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.dev175368549 → braindecode-1.3.0.dev175435903}/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.dev175368549 → braindecode-1.3.0.dev175435903}/braindecode/datautil/serialization.py

@@ -138,7 +138,7 @@ def _load_signals(fif_file, preload, is_raw):
         with open(pkl_file, "rb") as f:
             signals = pickle.load(f)
 
-        if all(f.exists() for f in signals.filenames):
+        if all(Path(f).exists() for f in signals.filenames):
             if preload:
                 signals.load_data()
             return signals

{braindecode-1.3.0.dev175368549 → braindecode-1.3.0.dev175435903}/braindecode/models/atcnet.py

@@ -141,7 +141,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
     - Pool sizes ``P1,P2`` trade temporal resolution for stability/compute; they set
       ``T_c = T/(P1·P2)`` and thus window width ``T_w``.
     - ``n_windows`` controls the ensemble over shifts (compute ∝ windows).
-    - ``att_num_heads``, ``att_head_dim`` set attention capacity; keep ``H·d_h ≈ F2``.
+    - ``num_heads``, ``head_dim`` set attention capacity; keep ``H·d_h ≈ F2``.
     - ``tcn_depth``, ``tcn_kernel_size`` govern receptive field; larger values demand
       longer inputs (see minimum length above). The implementation warns and *rescales*
       kernels/pools/windows if inputs are too short.
@@ -194,10 +194,10 @@ class ATCNet(EEGModuleMixin, nn.Module):
         table 1 of the paper [1]_. Defaults to 0.3 as in [1]_.
     n_windows : int
         Number of sliding windows, denoted n in [1]_. Defaults to 5 as in [1]_.
-    att_head_dim : int
+    head_dim : int
         Embedding dimension used in each self-attention head, denoted dh in
         table 1 of the paper [1]_. Defaults to 8 as in [1]_.
-    att_num_heads : int
+    num_heads : int
         Number of attention heads, denoted H in table 1 of the paper [1]_.
         Defaults to 2 as in [1]_.
     att_dropout : float
@@ -248,13 +248,13 @@ class ATCNet(EEGModuleMixin, nn.Module):
         conv_block_depth_mult=2,
         conv_block_dropout=0.3,
         n_windows=5,
-        att_head_dim=8,
-        att_num_heads=2,
+        head_dim=8,
+        num_heads=2,
         att_drop_prob=0.5,
         tcn_depth=2,
         tcn_kernel_size=4,
         tcn_drop_prob=0.3,
-        tcn_activation: nn.Module = nn.ELU,
+        tcn_activation: type[nn.Module] = nn.ELU,
         concat=False,
         max_norm_const=0.25,
         chs_info=None,
@@ -316,8 +316,8 @@ class ATCNet(EEGModuleMixin, nn.Module):
         self.conv_block_depth_mult = conv_block_depth_mult
         self.conv_block_dropout = conv_block_dropout
         self.n_windows = n_windows
-        self.att_head_dim = att_head_dim
-        self.att_num_heads = att_num_heads
+        self.head_dim = head_dim
+        self.num_heads = num_heads
         self.att_dropout = att_drop_prob
         self.tcn_depth = tcn_depth
         self.tcn_kernel_size = tcn_kernel_size
@@ -356,8 +356,8 @@ class ATCNet(EEGModuleMixin, nn.Module):
             [
                 _AttentionBlock(
                     in_shape=self.F2,
-                    head_dim=self.att_head_dim,
-                    num_heads=att_num_heads,
+                    head_dim=self.head_dim,
+                    num_heads=num_heads,
                     dropout=att_drop_prob,
                 )
                 for _ in range(self.n_windows)
@@ -656,7 +656,7 @@ class _TCNResidualBlock(nn.Module):
         kernel_size=4,
         n_filters=32,
         dropout=0.3,
-        activation: nn.Module = nn.ELU,
+        activation: type[nn.Module] = nn.ELU,
         dilation=1,
     ):
         super().__init__()

{braindecode-1.3.0.dev175368549 → braindecode-1.3.0.dev175435903}/braindecode/models/attentionbasenet.py

@@ -235,7 +235,7 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
     kernel_size : int, default=9
         The kernel size used in certain types of attention mechanisms for convolution
         operations.
-    activation: nn.Module, default=nn.ELU
+    activation: type[nn.Module] = nn.ELU,
         Activation function class to apply. Should be a PyTorch activation
         module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
     extra_params : bool, default=False
@@ -277,7 +277,7 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         freq_idx: int = 0,
         n_codewords: int = 4,
         kernel_size: int = 9,
-        activation: nn.Module = nn.ELU,
+        activation: type[nn.Module] = nn.ELU,
         extra_params: bool = False,
     ):
         super(AttentionBaseNet, self).__init__()
@@ -453,7 +453,7 @@ class _FeatureExtractor(nn.Module):
         pool_length: int = 75,
         pool_stride: int = 15,
         drop_prob: float = 0.5,
-        activation: nn.Module = nn.ELU,
+        activation: type[nn.Module] = nn.ELU,
     ):
         super().__init__()
 
@@ -592,7 +592,7 @@ class _ChannelAttentionBlock(nn.Module):
         n_codewords: int = 4,
         kernel_size: int = 9,
         extra_params: bool = False,
-        activation: nn.Module = nn.ELU,
+        activation: type[nn.Module] = nn.ELU,
     ):
         super().__init__()
         self.conv = nn.Sequential(

{braindecode-1.3.0.dev175368549 → braindecode-1.3.0.dev175435903}/braindecode/models/attn_sleep.py

@@ -90,8 +90,8 @@ class AttnSleep(EEGModuleMixin, nn.Module):
         d_ff=120,
         n_attn_heads=5,
         drop_prob=0.1,
-        activation_mrcnn: nn.Module = nn.GELU,
-        activation: nn.Module = nn.ReLU,
+        activation_mrcnn: type[nn.Module] = nn.GELU,
+        activation: type[nn.Module] = nn.ReLU,
         input_window_seconds=None,
         n_outputs=None,
         after_reduced_cnn_size=30,
@@ -230,7 +230,7 @@ class _SEBasicBlock(nn.Module):
         planes,
         stride=1,
         downsample=None,
-        activation: nn.Module = nn.ReLU,
+        activation: type[nn.Module] = nn.ReLU,
         *,
         reduction=16,
     ):
@@ -278,8 +278,8 @@ class _MRCNN(nn.Module):
         self,
         after_reduced_cnn_size,
         kernel_size=7,
-        activation: nn.Module = nn.GELU,
-        activation_se: nn.Module = nn.ReLU,
+        activation: type[nn.Module] = nn.GELU,
+        activation_se: type[nn.Module] = nn.ReLU,
     ):
         super(_MRCNN, self).__init__()
         drate = 0.5
@@ -325,7 +325,7 @@ class _MRCNN(nn.Module):
         )
 
     def _make_layer(
-        self, block, planes, blocks, stride=1, activate: nn.Module = nn.ReLU
+        self, block, planes, blocks, stride=1, activate: type[nn.Module] = nn.ReLU
     ):  # makes residual SE block
         downsample = None
         if stride != 1 or self.inplanes != planes * block.expansion:
@@ -526,7 +526,9 @@ class _EncoderLayer(nn.Module):
 class _PositionwiseFeedForward(nn.Module):
     """Positionwise feed-forward network."""
 
-    def __init__(self, d_model, d_ff, dropout=0.1, activation: nn.Module = nn.ReLU):
+    def __init__(
+        self, d_model, d_ff, dropout=0.1, activation: type[nn.Module] = nn.ReLU
+    ):
         super().__init__()
         self.w_1 = nn.Linear(d_model, d_ff)
         self.w_2 = nn.Linear(d_ff, d_model)

{braindecode-1.3.0.dev175368549 → braindecode-1.3.0.dev175435903}/braindecode/models/bendr.py

@@ -176,7 +176,7 @@ class BENDR(EEGModuleMixin, nn.Module):
         projection_head=False,  # Whether encoder should project back to input feature size (unused in original fine-tuning)
         drop_prob=0.1,  # General dropout probability (paper: 0.15 for pretraining, 0.0 for fine-tuning)
         layer_drop=0.0,  # Probability of dropping transformer layers during training (paper: 0.01 for pretraining)
-        activation=nn.GELU,  # Activation function
+        activation: type[nn.Module] = nn.GELU,  # Activation function
         # Transformer specific parameters
         transformer_layers=8,
         transformer_heads=8,

{braindecode-1.3.0.dev175368549 → braindecode-1.3.0.dev175435903}/braindecode/models/biot.py

@@ -45,11 +45,11 @@ class BIOT(EEGModuleMixin, nn.Module):
 
     Parameters
     ----------
-    emb_size : int, optional
+    embed_dim : int, optional
         The size of the embedding layer, by default 256
-    att_num_heads : int, optional
+    num_heads : int, optional
         The number of attention heads, by default 8
-    n_layers : int, optional
+    num_layers : int, optional
         The number of transformer layers, by default 4
     activation: nn.Module, default=nn.ELU
         Activation function class to apply. Should be a PyTorch activation
@@ -76,9 +76,9 @@ class BIOT(EEGModuleMixin, nn.Module):
 
     def __init__(
         self,
-        emb_size=256,
-        att_num_heads=8,
-        n_layers=4,
+        embed_dim=256,
+        num_heads=8,
+        num_layers=4,
         sfreq=200,
         hop_length=100,
         return_feature=False,
@@ -87,12 +87,12 @@ class BIOT(EEGModuleMixin, nn.Module):
         chs_info=None,
         n_times=None,
         input_window_seconds=None,
-        activation: nn.Module = nn.ELU,
+        activation: type[nn.Module] = nn.ELU,
         drop_prob: float = 0.5,
         # Parameters for the encoder
         max_seq_len: int = 1024,
-        attn_dropout=0.2,
-        attn_layer_dropout=0.2,
+        att_drop_prob=0.2,
+        att_layer_drop_prob=0.2,
     ):
         super().__init__(
             n_outputs=n_outputs,
@@ -103,10 +103,10 @@ class BIOT(EEGModuleMixin, nn.Module):
             sfreq=sfreq,
         )
         del n_outputs, n_chans, chs_info, n_times, sfreq
-        self.emb_size = emb_size
+        self.embed_dim = embed_dim
         self.hop_length = hop_length
-        self.att_num_heads = att_num_heads
-        self.n_layers = n_layers
+        self.num_heads = num_heads
+        self.num_layers = num_layers
         self.return_feature = return_feature
         if (self.sfreq != 200) & (self.sfreq is not None):
             warn(
@@ -114,7 +114,7 @@ class BIOT(EEGModuleMixin, nn.Module):
                 + "no guarantee to generalize well with the default parameters",
                 UserWarning,
             )
-        if self.n_chans > emb_size:
+        if self.n_chans > embed_dim:
             warn(
                 "The number of channels is larger than the embedding size. "
                 + "This may cause overfitting. Consider using a larger "
@@ -142,20 +142,20 @@ class BIOT(EEGModuleMixin, nn.Module):
             self.n_fft = int(self.sfreq)
 
         self.encoder = _BIOTEncoder(
-            emb_size=emb_size,
-            att_num_heads=att_num_heads,
-            n_layers=n_layers,
+            emb_size=self.embed_dim,
+            num_heads=self.num_heads,
+            n_layers=self.num_layers,
             n_chans=self.n_chans,
             n_fft=self.n_fft,
             hop_length=hop_length,
             drop_prob=drop_prob,
             max_seq_len=max_seq_len,
-            attn_dropout=attn_dropout,
-            attn_layer_dropout=attn_layer_dropout,
+            attn_dropout=att_drop_prob,
+            attn_layer_dropout=att_layer_drop_prob,
         )
 
         self.final_layer = _ClassificationHead(
-            emb_size=emb_size,
+            emb_size=self.embed_dim,
             n_outputs=self.n_outputs,
             activation=activation,
         )
@@ -250,7 +250,9 @@ class _ClassificationHead(nn.Sequential):
         (batch, n_outputs)
     """
 
-    def __init__(self, emb_size: int, n_outputs: int, activation: nn.Module = nn.ELU):
+    def __init__(
+        self, emb_size: int, n_outputs: int, activation: type[nn.Module] = nn.ELU
+    ):
         super().__init__()
         self.activation_layer = activation()
         self.classification_head = nn.Linear(emb_size, n_outputs)
@@ -345,7 +347,7 @@ class _BIOTEncoder(nn.Module):
         The number of channels
     emb_size: int
         The size of the embedding layer
-    att_num_heads: int
+    num_heads: int
         The number of attention heads
     n_layers: int
         The number of transformer layers
@@ -358,7 +360,7 @@ class _BIOTEncoder(nn.Module):
     def __init__(
         self,
         emb_size=256,  # The size of the embedding layer
-        att_num_heads=8,  # The number of attention heads
+        num_heads=8,  # The number of attention heads
         n_chans=16,  # The number of channels
         n_layers=4,  # The number of transformer layers
         n_fft=200,  # Related with the frequency resolution
@@ -378,7 +380,7 @@ class _BIOTEncoder(nn.Module):
         )
         self.transformer = LinearAttentionTransformer(
             dim=emb_size,
-            heads=att_num_heads,
+            heads=num_heads,
             depth=n_layers,
             max_seq_len=max_seq_len,
             attn_layer_dropout=attn_layer_dropout,