braindecode 0.8.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

braindecode/modules/layers.py

@@ -0,0 +1,133 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#
+# License: BSD (3-clause)
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops.layers.torch import Rearrange
+from torch import Tensor, nn
+
+from braindecode.functional import drop_path, safe_log
+
+
+class Ensure4d(nn.Module):
+    def forward(self, x):
+        while len(x.shape) < 4:
+            x = x.unsqueeze(-1)
+        return x
+
+
+class Chomp1d(nn.Module):
+    def __init__(self, chomp_size):
+        super().__init__()
+        self.chomp_size = chomp_size
+
+    def extra_repr(self):
+        return "chomp_size={}".format(self.chomp_size)
+
+    def forward(self, x):
+        return x[:, :, : -self.chomp_size].contiguous()
+
+
+class TimeDistributed(nn.Module):
+    """Apply module on multiple windows.
+
+    Apply the provided module on a sequence of windows and return their
+    concatenation.
+    Useful with sequence-to-prediction models (e.g. sleep stager which must map
+    a sequence of consecutive windows to the label of the middle window in the
+    sequence).
+
+    Parameters
+    ----------
+    module : nn.Module
+        Module to be applied to the input windows. Must accept an input of
+        shape (batch_size, n_channels, n_times).
+    """
+
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+
+    def forward(self, x):
+        """
+        Parameters
+        ----------
+        x : torch.Tensor
+            Sequence of windows, of shape (batch_size, seq_len, n_channels,
+            n_times).
+
+        Returns
+        -------
+        torch.Tensor
+            Shape (batch_size, seq_len, output_size).
+        """
+        b, s, c, t = x.shape
+        out = self.module(x.view(b * s, c, t))
+        return out.view(b, s, -1)
+
+
+class DropPath(nn.Module):
+    """Drop paths, also known as Stochastic Depth, per sample.
+
+        When applied in main path of residual blocks.
+
+        Parameters:
+        -----------
+        drop_prob: float (default=None)
+            Drop path probability (should be in range 0-1).
+
+        Notes
+        -----
+        Code copied and modified from VISSL facebookresearch:
+    https://github.com/facebookresearch/vissl/blob/0b5d6a94437bc00baed112ca90c9d78c6ccfbafb/vissl/models/model_helpers.py#L676
+        All rights reserved.
+
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+    # Utility function to print DropPath module
+    def extra_repr(self) -> str:
+        return f"p={self.drop_prob}"
+
+
+class SqueezeFinalOutput(nn.Module):
+    """
+
+    Removes empty dimension at end and potentially removes empty time
+    dimension. It does  not just use squeeze as we never want to remove
+    first dimension.
+
+    Returns
+    -------
+    x: torch.Tensor
+        squeezed tensor
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self.squeeze = Rearrange("b c t 1 -> b c t")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 1) drop feature dim
+        x = self.squeeze(x)
+        # 2) drop time dim if singleton
+        if x.shape[-1] == 1:
+            x = x.squeeze(-1)
+        return x

braindecode/modules/linear.py

@@ -0,0 +1,50 @@
+import torch
+from torch import Tensor, nn
+from torch.nn.utils.parametrize import register_parametrization
+
+from braindecode.modules.parametrization import MaxNorm, MaxNormParametrize
+
+
+class MaxNormLinear(nn.Linear):
+    """Linear layer with MaxNorm constraining on weights.
+
+    Equivalent of Keras tf.keras.Dense(..., kernel_constraint=max_norm())
+    [1]_ and [2]_. Implemented as advised in [3]_.
+
+    Parameters
+    ----------
+    in_features: int
+        Size of each input sample.
+    out_features: int
+        Size of each output sample.
+    bias: bool, optional
+        If set to ``False``, the layer will not learn an additive bias.
+        Default: ``True``.
+
+    References
+    ----------
+    .. [1] https://keras.io/api/layers/core_layers/dense/#dense-class
+    .. [2] https://www.tensorflow.org/api_docs/python/tf/keras/constraints/
+           MaxNorm
+    .. [3] https://discuss.pytorch.org/t/how-to-correctly-implement-in-place-
+           max-norm-constraint/96769
+    """
+
+    def __init__(
+        self, in_features, out_features, bias=True, max_norm_val=2, eps=1e-5, **kwargs
+    ):
+        super().__init__(
+            in_features=in_features, out_features=out_features, bias=bias, **kwargs
+        )
+        self._max_norm_val = max_norm_val
+        self._eps = eps
+        register_parametrization(self, "weight", MaxNorm(self._max_norm_val, self._eps))
+
+
+class LinearWithConstraint(nn.Linear):
+    """Linear layer with max-norm constraint on the weights."""
+
+    def __init__(self, *args, max_norm=1.0, **kwargs):
+        super(LinearWithConstraint, self).__init__(*args, **kwargs)
+        self.max_norm = max_norm
+        register_parametrization(self, "weight", MaxNormParametrize(self.max_norm))

braindecode/modules/parametrization.py

@@ -0,0 +1,38 @@
+import torch
+from torch import nn
+
+
+class MaxNorm(nn.Module):
+    def __init__(self, max_norm_val=2.0, eps=1e-5):
+        super().__init__()
+        self.max_norm_val = max_norm_val
+        self.eps = eps
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        norm = X.norm(2, dim=0, keepdim=True)
+        denom = norm.clamp(min=self.max_norm_val / 2)
+        number = denom.clamp(max=self.max_norm_val)
+        return X * (number / (denom + self.eps))
+
+    def right_inverse(self, X: torch.Tensor) -> torch.Tensor:
+        # Assuming the forward scales X by a factor s,
+        # the right inverse would scale it back by 1/s.
+        norm = X.norm(2, dim=0, keepdim=True)
+        denom = norm.clamp(min=self.max_norm_val / 2)
+        number = denom.clamp(max=self.max_norm_val)
+        scale = number / (denom + self.eps)
+        return X / scale
+
+
+class MaxNormParametrize(nn.Module):
+    """
+    Enforce a max‑norm constraint on the rows of a weight tensor via parametrization.
+    """
+
+    def __init__(self, max_norm: float = 1.0):
+        super().__init__()
+        self.max_norm = max_norm
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        # Renormalize each "row" (dim=0 slice) to have at most self.max_norm L2-norm
+        return X.renorm(p=2, dim=0, maxnorm=self.max_norm)

braindecode/modules/stats.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+
+
+class StatLayer(nn.Module):
+    """
+    Generic layer to compute a statistical function along a specified dimension.
+    Parameters
+    ----------
+    stat_fn : Callable
+        A function like torch.mean, torch.std, etc.
+    dim : int
+        Dimension along which to apply the function.
+    keepdim : bool, default=True
+        Whether to keep the reduced dimension.
+    clamp_range : tuple(float, float), optional
+        Used only for functions requiring clamping (e.g., log variance).
+    apply_log : bool, default=False
+        Whether to apply log after computation (used for LogVarLayer).
+    """
+
+    def __init__(
+        self,
+        stat_fn: Callable[..., torch.Tensor],
+        dim: int,
+        keepdim: bool = True,
+        clamp_range: Optional[tuple[float, float]] = None,
+        apply_log: bool = False,
+    ) -> None:
+        super().__init__()
+        self.stat_fn = stat_fn
+        self.dim = dim
+        self.keepdim = keepdim
+        self.clamp_range = clamp_range
+        self.apply_log = apply_log
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.stat_fn(x, dim=self.dim, keepdim=self.keepdim)
+        if self.clamp_range is not None:
+            out = torch.clamp(out, min=self.clamp_range[0], max=self.clamp_range[1])
+        if self.apply_log:
+            out = torch.log(out)
+        return out
+
+
+# make things more simple
+def _max_fn(x: torch.Tensor, dim: int, keepdim: bool) -> torch.Tensor:
+    return x.max(dim=dim, keepdim=keepdim)[0]
+
+
+def _power_fn(x: torch.Tensor, dim: int, keepdim: bool) -> torch.Tensor:
+    # compute mean of squared values along `dim`
+    return torch.mean(x**2, dim=dim, keepdim=keepdim)
+
+
+MeanLayer: Callable[[int, bool], StatLayer] = partial(StatLayer, torch.mean)
+MaxLayer: Callable[[int, bool], StatLayer] = partial(StatLayer, _max_fn)
+VarLayer: Callable[[int, bool], StatLayer] = partial(StatLayer, torch.var)
+StdLayer: Callable[[int, bool], StatLayer] = partial(StatLayer, torch.std)
+LogVarLayer: Callable[[int, bool], StatLayer] = partial(
+    StatLayer,
+    torch.var,
+    clamp_range=(1e-6, 1e6),
+    apply_log=True,
+)
+
+LogPowerLayer: Callable[[int, bool], StatLayer] = partial(
+    StatLayer,
+    _power_fn,
+    clamp_range=(1e-4, 1e4),
+    apply_log=True,
+)

braindecode/modules/util.py

@@ -0,0 +1,77 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#          Hubert Banville <hubert.jbanville@gmail.com>
+#
+# License: BSD (3-clause)
+
+import numpy as np
+import torch
+from scipy.special import log_softmax
+
+
+def _pad_shift_array(x, stride=1):
+    """Zero-pad and shift rows of a 3D array.
+
+    E.g., used to align predictions of corresponding windows in
+    sequence-to-sequence models.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Array of shape (n_rows, n_classes, n_windows).
+    stride : int
+        Number of non-overlapping elements between two consecutive sequences.
+
+    Returns
+    -------
+    np.ndarray :
+        Array of shape (n_rows, n_classes, (n_rows - 1) * stride + n_windows)
+        where each row is obtained by zero-padding the corresponding row in
+        ``x`` before and after in the last dimension.
+    """
+    if x.ndim != 3:
+        raise NotImplementedError(
+            f"x must be of shape (n_rows, n_classes, n_windows), got {x.shape}"
+        )
+    x_padded = np.pad(x, ((0, 0), (0, 0), (0, (x.shape[0] - 1) * stride)))
+    orig_strides = x_padded.strides
+    new_strides = (
+        orig_strides[0] - stride * orig_strides[2],
+        orig_strides[1],
+        orig_strides[2],
+    )
+    return np.lib.stride_tricks.as_strided(x_padded, strides=new_strides)
+
+
+def aggregate_probas(logits, n_windows_stride=1):
+    """Aggregate predicted probabilities with self-ensembling.
+
+    Aggregate window-wise predicted probabilities obtained on overlapping
+    sequences of windows using multiplicative voting as described in
+    [Phan2018]_.
+
+    Parameters
+    ----------
+    logits : np.ndarray
+        Array of shape (n_sequences, n_classes, n_windows) containing the
+        logits (i.e. the raw unnormalized scores for each class) for each
+        window of each sequence.
+    n_windows_stride : int
+        Number of windows between two consecutive sequences. Default is 1
+        (maximally overlapping sequences).
+
+    Returns
+    -------
+    np.ndarray :
+        Array of shape ((n_rows - 1) * stride + n_windows, n_classes)
+        containing the aggregated predicted probabilities for each window
+        contained in the input sequences.
+
+    References
+    ----------
+    .. [Phan2018] Phan, H., Andreotti, F., Cooray, N., Chén, O. Y., &
+        De Vos, M. (2018). Joint classification and prediction CNN framework
+        for automatic sleep stage classification. IEEE Transactions on
+        Biomedical Engineering, 66(5), 1285-1296.
+    """
+    log_probas = log_softmax(logits, axis=1)
+    return _pad_shift_array(log_probas, stride=n_windows_stride).sum(axis=0).T

braindecode/modules/wrapper.py

@@ -0,0 +1,75 @@
+from typing import List, Tuple
+
+import torch
+from torch import nn
+
+
+class Expression(nn.Module):
+    """Compute given expression on forward pass.
+
+    Parameters
+    ----------
+    expression_fn : callable
+        Should accept variable number of objects of type
+        `torch.autograd.Variable` to compute its output.
+    """
+
+    def __init__(self, expression_fn):
+        super().__init__()
+        self.expression_fn = expression_fn
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.expression_fn(x)
+
+    def __repr__(self):
+        if hasattr(self.expression_fn, "func") and hasattr(
+            self.expression_fn, "kwargs"
+        ):
+            expression_str = "{:s} {:s}".format(
+                self.expression_fn.func.__name__, str(self.expression_fn.kwargs)
+            )
+        elif hasattr(self.expression_fn, "__name__"):
+            expression_str = self.expression_fn.__name__
+        else:
+            expression_str = repr(self.expression_fn)
+        return self.__class__.__name__ + "(expression=%s) " % expression_str
+
+
+class IntermediateOutputWrapper(nn.Module):
+    """Wraps network model such that outputs of intermediate layers can be returned.
+    forward() returns list of intermediate activations in a network during forward pass.
+
+    Parameters
+    ----------
+    to_select : list
+        list of module names for which activation should be returned
+    model : model object
+        network model
+
+    Examples
+    --------
+    >>> model = Deep4Net()
+    >>> select_modules = ['conv_spat','conv_2','conv_3','conv_4'] # Specify intermediate outputs
+    >>> model_pert = IntermediateOutputWrapper(select_modules,model) # Wrap model
+    """
+
+    def __init__(self, to_select, model):
+        if not len(list(model.children())) == len(list(model.named_children())):
+            raise Exception("All modules in model need to have names!")
+
+        super().__init__()
+
+        modules_list = model.named_children()
+        for key, module in modules_list:
+            self.add_module(key, module)
+            self._modules[key].load_state_dict(module.state_dict())
+        self._to_select = to_select
+
+    def forward(self, x):
+        # Call modules individually and append activation to output if module is in to_select
+        o = []
+        for name, module in self._modules.items():
+            x = module(x)
+            if name in self._to_select:
+                o.append(x)
+        return o

braindecode/preprocessing/__init__.py

@@ -1,12 +1,37 @@
-from .preprocess import (exponential_moving_demean,
-                         exponential_moving_standardize, filterbank,
-                         preprocess, Preprocessor)
-from .mne_preprocess import (Resample, DropChannels, SetEEGReference, Filter, Pick, Crop)
-from .windowers import (create_windows_from_events, create_fixed_length_windows,
-                        create_windows_from_target_channels)
+from .mne_preprocess import (  # type: ignore[attr-defined]
+    Crop,
+    DropChannels,
+    Filter,
+    Pick,
+    Resample,
+    SetEEGReference,
+)
+from .preprocess import (
+    Preprocessor,
+    exponential_moving_demean,
+    exponential_moving_standardize,
+    filterbank,
+    preprocess,
+)
+from .windowers import (
+    create_fixed_length_windows,
+    create_windows_from_events,
+    create_windows_from_target_channels,
+)
 
-__all__ = ["exponential_moving_demean", "exponential_moving_standardize",
-           "filterbank", "preprocess", "Preprocessor", "Resample", "DropChannels",
-           "SetEEGReference", "Filter", "Pick", "Crop",
-           "create_windows_from_events", "create_fixed_length_windows",
-           "create_windows_from_target_channels"]
+__all__ = [
+    "exponential_moving_demean",
+    "exponential_moving_standardize",
+    "filterbank",
+    "preprocess",
+    "Preprocessor",
+    "Resample",
+    "DropChannels",
+    "SetEEGReference",
+    "Filter",
+    "Pick",
+    "Crop",
+    "create_windows_from_events",
+    "create_fixed_length_windows",
+    "create_windows_from_target_channels",
+]

braindecode/preprocessing/mne_preprocess.py

@@ -1,11 +1,14 @@
 """Preprocessor objects based on mne methods."""
+
 # Authors: Bruna Lopes <brunajaflopes@gmail.com>
 #          Bruno Aristimunha <b.aristimunha@gmail.com>
 #
 # License: BSD-3
 import inspect
+
 import mne.io
-from braindecode.preprocessing import Preprocessor
+
+from braindecode.preprocessing.preprocess import Preprocessor
 from braindecode.util import _update_moabb_docstring
 
 
@@ -31,9 +34,9 @@ def _generate_mne_pre_processor(function):
     """
     Generate a class based on an MNE function for preprocessing.
     """
-    class_name = ''.join(
-        word.title() for word in function.__name__.split('_')).replace('Eeg',
-                                                                       'EEG')
+    class_name = "".join(word.title() for word in function.__name__.split("_")).replace(
+        "Eeg", "EEG"
+    )
     import_path = f"{function.__module__}.{function.__name__}"
     doc = f" See more details in {import_path}"
 
@@ -55,7 +58,7 @@ mne_functions = [
     mne.io.Raw.filter,
     mne.io.Raw.crop,
     mne.io.Raw.pick,
-    mne.io.Raw.set_eeg_reference
+    mne.io.Raw.set_eeg_reference,
 ]
 
 # Automatically generate and add classes to the global namespace
@@ -64,8 +67,11 @@ for function in mne_functions:
     globals()[class_obj.__name__] = class_obj
 
 # Define __all__ based on the generated class names
-__all__ = [class_obj.__name__ for class_obj in globals().values() if
-           isinstance(class_obj, type)]
+__all__ = [
+    class_obj.__name__
+    for class_obj in globals().values()
+    if isinstance(class_obj, type)
+]
 
 # Clean up unnecessary variables
 del mne_functions, function, class_obj