braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/eegneuralnet.py

@@ -0,0 +1,372 @@
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#          Pierre Guetschel <pierre.guetschel@gmail.com>
+#
+# License: BSD (3-clause)
+
+
+import abc
+import inspect
+import logging
+from typing import Literal
+
+import mne
+import numpy as np
+import torch
+from sklearn.metrics import get_scorer
+from skorch import NeuralNet
+from skorch.callbacks import BatchScoring, EpochScoring, EpochTimer, PrintLog
+from skorch.utils import noop, to_numpy, train_loss_score, valid_loss_score
+
+from braindecode.datautil import infer_signal_properties
+
+from .models.util import models_dict
+from .training.scoring import (
+    CroppedTimeSeriesEpochScoring,
+    CroppedTrialEpochScoring,
+    PostEpochTrainScoring,
+)
+
+log = logging.getLogger(__name__)
+
+
+def _get_model(model: str):
+    """Returns the corresponding class in case the model passed is a string."""
+    if isinstance(model, str):
+        if model in models_dict:
+            model = models_dict[model]
+        else:
+            raise ValueError(f"Unknown model name {model!r}.")
+    return model
+
+
+class _EEGNeuralNet(NeuralNet, abc.ABC):
+    signal_args_set_ = False
+
+    @property
+    def log(self):
+        return log.getChild(self.__class__.__name__)
+
+    def initialize_module(self):
+        """Initializes the module.
+
+        A Braindecode model name can also be passed as module argument.
+
+        If the module is already initialized and no parameter was changed, it
+        will be left as is.
+        """
+        kwargs = self.get_params_for("module")
+        module = _get_model(self.module)
+        module = self.initialized_instance(module, kwargs)
+        # pylint: disable=attribute-defined-outside-init
+        self.module_ = module
+        return self
+
+    def _yield_callbacks(self):
+        # Here we parse the callbacks supplied as strings,
+        # e.g. 'accuracy', to the callbacks skorch expects
+        for name, cb, named_by_user in super()._yield_callbacks():
+            if name == "str":
+                train_cb, valid_cb = self._parse_str_callback(cb)
+                yield train_cb
+                if self.train_split is not None:
+                    yield valid_cb
+            else:
+                yield name, cb, named_by_user
+
+    def _parse_str_callback(self, cb_supplied_name):
+        scoring = get_scorer(cb_supplied_name)
+        scoring_name = scoring._score_func.__name__
+        assert scoring_name.endswith(("_score", "_error", "_deviance", "_loss"))
+        if scoring_name.endswith("_score") or cb_supplied_name.startswith("neg_"):
+            lower_is_better = False
+        else:
+            lower_is_better = True
+        train_name = f"train_{cb_supplied_name}"
+        valid_name = f"valid_{cb_supplied_name}"
+        if self.cropped:
+            train_scoring = CroppedTrialEpochScoring(
+                cb_supplied_name, lower_is_better, on_train=True, name=train_name
+            )
+            valid_scoring = CroppedTrialEpochScoring(
+                cb_supplied_name, lower_is_better, on_train=False, name=valid_name
+            )
+        else:
+            train_scoring = PostEpochTrainScoring(
+                cb_supplied_name, lower_is_better, name=train_name
+            )
+            valid_scoring = EpochScoring(
+                cb_supplied_name, lower_is_better, on_train=False, name=valid_name
+            )
+        named_by_user = True
+        train_valid_callbacks = [
+            (train_name, train_scoring, named_by_user),
+            (valid_name, valid_scoring, named_by_user),
+        ]
+        return train_valid_callbacks
+
+    def on_batch_end(self, net, *batch, training=False, **kwargs):
+        # If training is false, assume that our loader has indices for this
+        # batch
+        if not training:
+            epoch_cbs = []
+            for name, cb in self.callbacks_:
+                if (
+                    isinstance(
+                        cb, (CroppedTrialEpochScoring, CroppedTimeSeriesEpochScoring)
+                    )
+                    and (hasattr(cb, "window_inds_"))
+                    and (not cb.on_train)
+                ):
+                    epoch_cbs.append(cb)
+            # for trialwise decoding stuffs it might also be we don't have
+            # cropped loader, so no indices there
+            if len(epoch_cbs) > 0:
+                assert self._last_window_inds_ is not None
+                for cb in epoch_cbs:
+                    cb.window_inds_.append(self._last_window_inds_)
+                self._last_window_inds_ = None
+
+    def predict_with_window_inds_and_ys(self, dataset):
+        self.module.eval()
+        preds = []
+        i_window_in_trials = []
+        i_window_stops = []
+        window_ys = []
+        for X, y, i in self.get_iterator(dataset, drop_index=False):
+            i_window_in_trials.append(i[0].cpu().numpy())
+            i_window_stops.append(i[2].cpu().numpy())
+            with torch.no_grad():
+                preds.append(to_numpy(self.module.forward(X.to(self.device))))
+            window_ys.append(y.cpu().numpy())
+        preds = np.concatenate(preds)
+        i_window_in_trials = np.concatenate(i_window_in_trials)
+        i_window_stops = np.concatenate(i_window_stops)
+        window_ys = np.concatenate(window_ys)
+        return dict(
+            preds=preds,
+            i_window_in_trials=i_window_in_trials,
+            i_window_stops=i_window_stops,
+            window_ys=window_ys,
+        )
+
+    # Changes the default target extractor to noop
+    @property
+    def _default_callbacks(self):
+        return [
+            ("epoch_timer", EpochTimer()),
+            (
+                "train_loss",
+                BatchScoring(
+                    train_loss_score,
+                    name="train_loss",
+                    on_train=True,
+                    target_extractor=noop,
+                ),
+            ),
+            (
+                "valid_loss",
+                BatchScoring(
+                    valid_loss_score,
+                    name="valid_loss",
+                    target_extractor=noop,
+                ),
+            ),
+            ("print_log", PrintLog()),
+        ]
+
+    @property
+    @abc.abstractmethod
+    def mode(self) -> Literal["classification", "regression"]:
+        pass
+
+    def _set_signal_args(self, X, y, classes):
+        is_init = isinstance(self.module, torch.nn.Module)
+        if is_init:
+            self.log.info(
+                "The module passed is already initialized which is not recommended. "
+                "Instead, you can pass the module class and its parameters separately.\n"
+                "For more details, see "
+                "https://skorch.readthedocs.io/en/stable/user/neuralnet.html#module \n"
+                "Skipping setting signal-related parameters from data."
+            )
+            return
+        if classes is None:
+            classes = getattr(self, "classes", None)
+        signal_kwargs = infer_signal_properties(X, y, mode=self.mode, classes=classes)
+        if not signal_kwargs:
+            return
+
+        # kick out missing kwargs:
+        module_kwargs = dict()
+        module = _get_model(self.module)
+        all_module_kwargs = inspect.signature(module.__init__).parameters.keys()
+        for k, v in signal_kwargs.items():
+            if v is None:
+                continue
+            if k in all_module_kwargs:
+                module_kwargs[k] = v
+            else:
+                self.log.warning(f"Module {self.module!r} is missing parameter {k!r}.")
+
+        # kick out inferred signal kwargs if user specifies kwargs:
+        user_specified_kwargs = self.get_params_for("module").items()
+        if len(user_specified_kwargs) > 0:
+            self.log.info(
+                f"Overriding inferred parameters with user "
+                f"specified parameters{user_specified_kwargs!r}."
+            )
+            for k, v in self.get_params_for("module").items():
+                if k in module_kwargs:
+                    module_kwargs.pop(k)
+                    module_kwargs[k] = v
+
+        # save kwargs to self:
+        self.log.info(
+            f"Passing additional parameters {module_kwargs!r} "
+            f"to module {self.module!r}."
+        )
+        module_kwargs = {f"module__{k}": v for k, v in module_kwargs.items()}
+        self.set_params(**module_kwargs)
+
+    def get_dataset(self, X, y=None):
+        """Get a dataset that contains the input data and is passed to.
+
+        the iterator.
+
+        Override this if you want to initialize your dataset
+        differently.
+
+        Parameters
+        ----------
+        X : input data, compatible with skorch.dataset.Dataset
+          By default, you should be able to pass:
+
+            * mne.Epochs
+            * numpy arrays
+            * torch tensors
+            * pandas DataFrame or Series
+            * scipy sparse CSR matrices
+            * a dictionary of the former three
+            * a list/tuple of the former three
+            * a Dataset
+
+          If this doesn't work with your data, you have to pass a
+          ``Dataset`` that can deal with the data.
+
+        y : target data, compatible with skorch.dataset.Dataset
+          The same data types as for ``X`` are supported. If your X is
+          a Dataset that contains the target, ``y`` may be set to
+          None.
+
+        Returns
+        -------
+        dataset
+          The initialized dataset.
+        """
+        if isinstance(X, mne.BaseEpochs):
+            X = X.get_data(units="uV")
+        return super().get_dataset(X, y)
+
+    def partial_fit(self, X, y=None, classes=None, **fit_params):
+        """Fit the module.
+
+        If the module is initialized, it is not re-initialized, which
+        means that this method should be used if you want to continue
+        training a model (warm start).
+        If possible, signal-related parameters are inferred from the
+        data and passed to the module at initialisation.
+        Depending on the type of input passed, the following parameters
+        are inferred:
+
+          * mne.Epochs: ``n_times``, ``n_chans``, ``n_outputs``, ``chs_info``,
+            ``sfreq``, ``input_window_seconds``
+          * array-like: ``n_times``, ``n_chans``, ``n_outputs``
+          * WindowsDataset with ``targets_from='metadata'``
+            (or BaseConcatDataset of such datasets): ``n_times``, ``n_chans``, ``n_outputs``
+          * other Dataset: ``n_times``, ``n_chans``
+          * other types: no parameters are inferred.
+
+        Parameters
+        ----------
+        X : input data, compatible with skorch.dataset.Dataset
+          By default, you should be able to pass:
+
+            * mne.Epochs
+            * numpy arrays
+            * torch tensors
+            * pandas DataFrame or Series
+            * scipy sparse CSR matrices
+            * a dictionary of the former three
+            * a list/tuple of the former three
+            * a Dataset
+
+          If this doesn't work with your data, you have to pass a
+          ``Dataset`` that can deal with the data.
+
+        y : target data, compatible with skorch.dataset.Dataset
+          The same data types as for ``X`` are supported. If your X is
+          a Dataset that contains the target, ``y`` may be set to
+          None.
+
+        classes : array, sahpe (n_classes,)
+          Solely for sklearn compatibility, currently unused.
+
+        **fit_params : dict
+          Additional parameters passed to the ``forward`` method of
+          the module and to the ``self.train_split`` call.
+        """
+        # this needs to be executed before the net is initialized:
+        if not self.signal_args_set_:
+            self._set_signal_args(X, y, classes)
+            self.signal_args_set_ = True
+        return super().partial_fit(X=X, y=y, classes=classes, **fit_params)
+
+    def fit(self, X, y=None, **fit_params):
+        """Initialize and fit the module.
+
+        If the module was already initialized, by calling fit, the
+        module will be re-initialized (unless ``warm_start`` is True).
+        If possible, signal-related parameters are inferred from the
+        data and passed to the module at initialisation.
+        Depending on the type of input passed, the following parameters
+        are inferred:
+
+          * mne.Epochs: ``n_times``, ``n_chans``, ``n_outputs``, ``chs_info``,
+            ``sfreq``, ``input_window_seconds``
+          * array-like: ``n_times``, ``n_chans``, ``n_outputs``
+          * WindowsDataset with ``targets_from='metadata'``
+            (or BaseConcatDataset of such datasets): ``n_times``, ``n_chans``, ``n_outputs``
+          * other Dataset: ``n_times``, ``n_chans``
+          * other types: no parameters are inferred.
+
+        Parameters
+        ----------
+        X : input data, compatible with skorch.dataset.Dataset
+          By default, you should be able to pass:
+
+            * mne.Epochs
+            * numpy arrays
+            * torch tensors
+            * pandas DataFrame or Series
+            * scipy sparse CSR matrices
+            * a dictionary of the former three
+            * a list/tuple of the former three
+            * a Dataset
+
+          If this doesn't work with your data, you have to pass a
+          ``Dataset`` that can deal with the data.
+
+        y : target data, compatible with skorch.dataset.Dataset
+          The same data types as for ``X`` are supported. If your X is
+          a Dataset that contains the target, ``y`` may be set to
+          None.
+
+        **fit_params : dict
+          Additional parameters passed to the ``forward`` method of
+          the module and to the ``self.train_split`` call.
+        """
+        # this needs to be executed before the net is initialized:
+        if not self.signal_args_set_:
+            self._set_signal_args(X, y, classes=None)
+            self.signal_args_set_ = True
+        return super().fit(X=X, y=y, **fit_params)

braindecode/functional/__init__.py

@@ -0,0 +1,22 @@
+from .functions import (
+    _get_gaussian_kernel1d,
+    drop_path,
+    hilbert_freq,
+    identity,
+    plv_time,
+    safe_log,
+    square,
+)
+from .initialization import glorot_weight_zero_bias, rescale_parameter
+
+__all__ = [
+    "_get_gaussian_kernel1d",
+    "drop_path",
+    "hilbert_freq",
+    "identity",
+    "plv_time",
+    "safe_log",
+    "square",
+    "glorot_weight_zero_bias",
+    "rescale_parameter",
+]

braindecode/functional/functions.py

@@ -0,0 +1,251 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#
+# License: BSD (3-clause)
+
+import torch
+import torch.nn.functional as F
+
+
+def square(x):
+    return x * x
+
+
+def safe_log(x, eps: float = 1e-6) -> torch.Tensor:
+    """Prevents :math:`log(0)` by using :math:`log(max(x, eps))`."""
+    return torch.log(torch.clamp(x, min=eps))
+
+
+def identity(x):
+    return x
+
+
+def drop_path(
+    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
+):
+    """Drop paths (Stochastic Depth) per sample.
+
+    Notes: This implementation is taken from timm library.
+
+    All credit goes to Ross Wightman.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        input tensor
+    drop_prob : float, optional
+        survival rate (i.e. probability of being kept), by default 0.0
+    training : bool, optional
+        whether the model is in training mode, by default False
+    scale_by_keep : bool, optional
+        whether to scale output by (1/keep_prob) during training, by default True
+
+    Returns
+    -------
+    torch.Tensor
+        output tensor
+
+    Notes from Ross Wightman:
+    (when applied in main path of residual blocks)
+    This is the same as the DropConnect impl I created for EfficientNet,
+    etc. networks, however,
+    the original name is misleading as 'Drop Connect' is a different form
+    of dropout in a separate paper...
+    See discussion : https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956
+    ... I've opted for changing the layer and argument names to 'drop path'
+    rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+def _get_gaussian_kernel1d(kernel_size: int, sigma: float) -> torch.Tensor:
+    """
+    Generates a 1-dimensional Gaussian kernel based on the specified kernel.
+
+    size and standard deviation (sigma).
+    This kernel is useful for Gaussian smoothing or filtering operations in
+    image processing. The function calculates a range limit to ensure the kernel
+    effectively covers the Gaussian distribution. It generates a tensor of
+    specified size and type, filled with values distributed according to a
+    Gaussian curve, normalized using a softmax function
+    to ensure all weights sum to 1.
+
+    Parameters
+    ----------
+    kernel_size : int
+    sigma : float
+
+    Returns
+    -------
+    kernel1d : torch.Tensor
+
+    Notes
+    -----
+    Code copied and modified from TorchVision:
+    https://github.com/pytorch/vision/blob/main/torchvision/transforms/_functional_tensor.py#L725-L732
+    All rights reserved.
+
+    LICENSE in https://github.com/pytorch/vision/blob/main/LICENSE
+    """
+    ksize_half = (kernel_size - 1) * 0.5
+    x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
+    pdf = torch.exp(-0.5 * (x / sigma).pow(2))
+    kernel1d = pdf / pdf.sum()
+    return kernel1d
+
+
+def hilbert_freq(x, forward_fourier=True):
+    r"""
+    Compute the Hilbert transform using PyTorch, separating the real and
+    imaginary parts.
+
+    The analytic signal :math:`x_a(t)` of a real-valued signal :math:`x(t)`
+    is defined as:
+
+    .. math::
+
+        x_a(t) = x(t) + i y(t) = \mathcal{F}^{-1} \{ U(f) \mathcal{F}\{x(t)\} \}
+
+    where:
+    - :math:`\mathcal{F}` is the Fourier transform,
+    - :math:`U(f)` is the unit step function,
+    - :math:`y(t)` is the Hilbert transform of :math:`x(t)`.
+
+
+    Parameters
+    ----------
+    input : torch.Tensor
+        Input tensor. The expected shape depends on the `forward_fourier` parameter:
+
+        - If `forward_fourier` is True:
+            (..., seq_len)
+        - If `forward_fourier` is False:
+            (..., seq_len / 2 + 1, 2)
+
+    forward_fourier : bool, optional
+        Determines the format of the input tensor.
+        - If True, the input is in the forward Fourier domain.
+        - If False, the input contains separate real and imaginary parts.
+        Default is True.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor with shape (..., seq_len, 2), where the last dimension represents
+        the real and imaginary parts of the Hilbert transform.
+
+    Examples
+    --------
+    >>> import torch
+    >>> input = torch.randn(10, 100)  # Example input tensor
+    >>> output = hilbert_transform(input)
+    >>> print(output.shape)
+    torch.Size([10, 100, 2])
+
+    Notes
+    -----
+    The implementation is matching scipy implementation, but using torch.
+    https://github.com/scipy/scipy/blob/v1.14.1/scipy/signal/_signaltools.py#L2287-L2394
+
+    """
+    if forward_fourier:
+        x = torch.fft.rfft(x, norm=None, dim=-1)
+        x = torch.view_as_real(x)
+    x = x * 2.0
+    x[..., 0, :] = x[..., 0, :] / 2.0  # Don't multiply the DC-term by 2
+    x = F.pad(
+        x, [0, 0, 0, x.shape[-2] - 2]
+    )  # Fill Fourier coefficients to retain shape
+    x = torch.view_as_complex(x)
+    x = torch.fft.ifft(x, norm=None, dim=-1)  # returns complex signal
+    x = torch.view_as_real(x)
+
+    return x
+
+
+def plv_time(x, forward_fourier=True, epsilon: float = 1e-6):
+    """Compute the Phase Locking Value (PLV) metric in the time domain.
+
+    The Phase Locking Value (PLV) is a measure of the synchronization between
+    different channels by evaluating the consistency of phase differences
+    over time. It ranges from 0 (no synchronization) to 1 (perfect
+    synchronization) [Lachaux1999]_.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor containing the signal data.
+
+        - If `forward_fourier` is `True`, the shape should be `(..., channels, time)`.
+        - If `forward_fourier` is `False`, the shape should be `(..., channels, freqs, 2)`,
+          where the last dimension represents the real and imaginary parts.
+
+    forward_fourier : bool, optional
+        Specifies the format of the input tensor `x`.
+
+        - If `True`, `x` is assumed to be in the time domain.
+        - If `False`, `x` is assumed to be in the Fourier domain with separate real and
+          imaginary components.
+
+        Default is `True`.
+    epsilon : float, default 1e-6
+        Small numerical value to ensure positivity constraint on the complex part
+
+    Returns
+    -------
+    plv : torch.Tensor
+        The Phase Locking Value matrix with shape `(..., channels, channels)`. Each
+        element `[i, j]` represents the PLV between channel `i` and channel `j`.
+
+    References
+    ----------
+    .. [Lachaux1999] Lachaux, J. P., Rodriguez, E., Martinerie, J., & Varela, F. J. (1999).
+        Measuring phase synchrony in brain signals. Human brain mapping,
+        8(4), 194-208.
+    """
+    # Compute the analytic signal using the Hilbert transform.
+    # x_a has separate real and imaginary parts.
+    analytic_signal = hilbert_freq(x, forward_fourier)
+    # Calculate the amplitude (magnitude) of the analytic signal.
+    # Adding a small epsilon (1e-6) to avoid division by zero.
+    amplitude = torch.sqrt(
+        analytic_signal[..., 0] ** 2 + analytic_signal[..., 1] ** 2 + 1e-6
+    )
+    # Normalize the analytic signal to obtain unit vectors (phasors).
+    unit_phasor = analytic_signal / amplitude.unsqueeze(-1)
+
+    # Compute the real part of the outer product between phasors of
+    # different channels.
+    real_real = torch.matmul(unit_phasor[..., 0], unit_phasor[..., 0].transpose(-2, -1))
+
+    # Compute the imaginary part of the outer product between phasors of
+    # different channels.
+    imag_imag = torch.matmul(unit_phasor[..., 1], unit_phasor[..., 1].transpose(-2, -1))
+
+    # Compute the cross-terms for the real and imaginary parts.
+    real_imag = torch.matmul(unit_phasor[..., 0], unit_phasor[..., 1].transpose(-2, -1))
+    imag_real = torch.matmul(unit_phasor[..., 1], unit_phasor[..., 0].transpose(-2, -1))
+
+    # Combine the real and imaginary parts to form the complex correlation.
+    correlation_real = real_real + imag_imag
+    correlation_imag = real_imag - imag_real
+
+    # Determine the number of time points (or frequency bins if in Fourier domain).
+    time = amplitude.shape[-1]
+
+    # Calculate the PLV by averaging the magnitude of the complex correlation over time.
+    # epsilon is small numerical value to ensure positivity constraint on the complex part
+    plv_matrix = (
+        1 / time * torch.sqrt(correlation_real**2 + correlation_imag**2 + epsilon)
+    )
+
+    return plv_matrix

braindecode/functional/initialization.py

@@ -0,0 +1,47 @@
+import math
+
+from torch import nn
+
+
+def glorot_weight_zero_bias(model):
+    """Initialize parameters of all modules by initializing weights with.
+
+    glorot uniform/xavier initialization, and setting biases to zero. Weights from
+    batch norm layers are set to 1.
+
+    Parameters
+    ----------
+    model : Module
+    """
+    for module in model.modules():
+        if hasattr(module, "weight"):
+            if "BatchNorm" in module.__class__.__name__:
+                nn.init.constant_(module.weight, 1)
+        if hasattr(module, "bias"):
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+def rescale_parameter(param, layer_id):
+    r"""Recaling the l-th transformer layer.
+
+    Rescales the parameter tensor by the inverse square root of the layer id.
+    Made inplace. :math:`\frac{1}{\sqrt{2 \cdot \text{layer\_id}}}` [Beit2022]
+
+    In the labram, this is used to rescale the output matrices
+    (i.e., the last linear projection within each sub-layer) of the
+    self-attention module.
+
+    Parameters
+    ----------
+    param: :class:`torch.Tensor`
+        tensor to be rescaled
+    layer_id: int
+        layer id in the neural network
+
+    References
+    ----------
+    [Beit2022] Hangbo Bao, Li Dong, Songhao Piao, Furu We (2022). BEIT: BERT
+    Pre-Training of Image Transformers.
+    """
+    param.div_(math.sqrt(2.0 * layer_id))