braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/models/config.py

@@ -0,0 +1,233 @@
+from collections.abc import Callable
+from inspect import signature
+from types import UnionType
+from typing import Annotated, Any, Literal, Union, get_args, get_origin
+
+import numpy as np
+from mne.utils import _soft_import
+from typing_extensions import TypedDict
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.models.util import SigArgName, models_dict, models_mandatory_parameters
+
+pydantic = _soft_import(name="pydantic", purpose="model configuration", strict=False)
+
+try:
+    from numpydantic import NDArray, Shape
+except ImportError:
+    # we can't use soft import for numpydantic because numpydantic does not define its version in __init__
+    NDArray = Any  # type: ignore
+    Shape = Any  # type: ignore
+
+
+class ChsInfoType(TypedDict, total=False, closed=True):  # type: ignore[call-arg]
+    cal: float
+    ch_name: str
+    coil_type: int
+    coord_frame: int
+    kind: str
+    loc: NDArray[Shape["12"], np.float64]  # type: ignore[misc]
+    logno: int
+    range: float
+    scanno: int
+    unit: int
+    unit_mul: int
+
+
+def _replace_type_hints(type_hint: Any) -> Any:
+    origin = get_origin(type_hint)
+    args = get_args(type_hint)
+    if origin is type or origin is Callable or type_hint is Callable:
+        return pydantic.ImportString
+    if origin is None:
+        return type_hint
+    replaced_args = tuple(_replace_type_hints(arg) for arg in args)
+    if origin is UnionType:
+        origin = Union
+    return origin[replaced_args]
+
+
+SIGNAL_ARGS_TYPES = {
+    "n_chans": int,
+    "n_times": int,
+    "sfreq": float,
+    "input_window_seconds": float,
+    "n_outputs": int,
+    "chs_info": list[ChsInfoType],
+}
+
+
+class BaseBraindecodeModelConfig(pydantic.BaseModel):  # type: ignore
+    def create_instance(self) -> EEGModuleMixin:
+        model_cls = models_dict[self.model_name_]
+        kwargs = self.model_dump(mode="python", exclude={"model_name_"})
+        if kwargs.get("n_chans") is not None and kwargs.get("chs_info") is not None:
+            kwargs.pop("n_chans")
+        if (
+            kwargs.get("n_times") is not None
+            and kwargs.get("input_window_seconds") is not None
+            and kwargs.get("sfreq") is not None
+        ):
+            kwargs.pop("n_times")
+        return model_cls(**kwargs)
+
+
+def make_model_config(
+    model_class: type[EEGModuleMixin],
+    required: list[SigArgName],
+) -> type[BaseBraindecodeModelConfig]:
+    """Create a pydantic model config for a given model class.
+
+    Parameters
+    ----------
+    model_class : type[EEGModuleMixin]
+        The model class for which to create the config.
+    required : list of SigArgName
+        The required signal arguments for the model.
+
+    Returns
+    -------
+    type
+        A pydantic BaseModel subclass representing the model config.
+    """
+    if not pydantic:
+        raise ImportError(
+            "pydantic is required to use make_model_config. "
+            "Please install braindecode[typing]."
+        )
+
+    # ironically, we need to ignore the type here to have the soft dependency.
+
+    @pydantic.model_validator(mode="before")
+    def validate_signal_params(cls, data: Any):
+        n_outputs = data.get("n_outputs")
+        n_chans = data.get("n_chans")
+        chs_info = data.get("chs_info")
+        n_times = data.get("n_times")
+        input_window_seconds = data.get("input_window_seconds")
+        sfreq = data.get("sfreq")
+
+        # Check that required parameters are provided or can be inferred
+        if "n_outputs" in required and n_outputs is None:
+            raise ValueError("n_outputs is a required parameter but was not provided.")
+        if "n_chans" in required and n_chans is None and chs_info is None:
+            raise ValueError(
+                "n_chans is required and could not be inferred. Either specify n_chans or chs_info."
+            )
+        if "chs_info" in required and chs_info is None:
+            raise ValueError("chs_info is a required parameter but was not provided.")
+        if "n_times" in required and (
+            n_times is None and (sfreq is None or input_window_seconds is None)
+        ):
+            raise ValueError(
+                "n_times is required and could not be inferred."
+                "Either specify n_times or input_window_seconds and sfreq."
+            )
+        if "sfreq" in required and (
+            sfreq is None and (n_times is None or input_window_seconds is None)
+        ):
+            raise ValueError(
+                "sfreq is required and could not be inferred."
+                "Either specify sfreq or input_window_seconds and n_times."
+            )
+        if "input_window_seconds" in required and (
+            input_window_seconds is None and (n_times is None or sfreq is None)
+        ):
+            raise ValueError(
+                "input_window_seconds is required and could not be inferred."
+                "Either specify input_window_seconds or n_times and sfreq."
+            )
+
+        # Infer missing parameters if possible, and check consistency
+        if chs_info is not None:
+            if n_chans is None:
+                data["n_chans"] = len(chs_info)
+            elif n_chans != len(chs_info):
+                raise ValueError(
+                    f"Provided {n_chans=} does not match length of chs_info: {len(chs_info)}."
+                )
+        if (
+            n_times is not None
+            and sfreq is not None
+            and input_window_seconds is not None
+        ):
+            if n_times != round(input_window_seconds * sfreq):
+                raise ValueError(
+                    f"Provided {n_times=} does not match {input_window_seconds=} * {sfreq=}."
+                )
+        elif n_times is None and sfreq is not None and input_window_seconds is not None:
+            data["n_times"] = round(input_window_seconds * sfreq)
+        elif sfreq is None and n_times is not None and input_window_seconds is not None:
+            data["sfreq"] = n_times / input_window_seconds
+        elif input_window_seconds is None and n_times is not None and sfreq is not None:
+            data["input_window_seconds"] = n_times / sfreq
+        return data
+
+    signature_params = signature(model_class.__init__, eval_str=True).parameters
+    has_args = any(p.kind == p.VAR_POSITIONAL for p in signature_params.values())
+    has_kwargs = any(p.kind == p.VAR_KEYWORD for p in signature_params.values())
+    if has_args:
+        raise ValueError("Model __init__ methods cannot have *args")
+
+    extra = "allow" if has_kwargs else "forbid"
+    fields = {}
+    for name, p in signature_params.items():
+        if name == "self" or p.kind == p.VAR_KEYWORD:
+            continue
+
+        annot = p.annotation
+        if annot is p.empty:
+            annot = Any
+        # case with type[nn.Module] or callable
+        else:
+            annot = _replace_type_hints(annot)
+        # Most models did not specify types for signal args, so we add them here
+        if name in SIGNAL_ARGS_TYPES:
+            annot = SIGNAL_ARGS_TYPES[name] | None
+
+        fields[name] = (annot, p.default) if p.default is not p.empty else annot
+
+    name = model_class.__name__
+    model_config = pydantic.create_model(
+        f"{name}Config",
+        model_name_=(Literal[name], name),
+        __config__=pydantic.ConfigDict(
+            arbitrary_types_allowed=True, extra=extra, validate_default=True
+        ),
+        __doc__=f"Pydantic config of model {model_class.__name__}\n\n{model_class.__doc__}",
+        __base__=BaseBraindecodeModelConfig,
+        __module__="braindecode.models.config",
+        __validators__={"validate_signal_params": validate_signal_params},
+        **fields,
+    )
+    return model_config
+
+
+# Automatically generate and add classes to the global namespace
+# and define __all__ based on generated classes
+__all__ = ["make_model_config"]
+
+if not pydantic:
+    pass
+else:
+    models_configs: list[type[BaseBraindecodeModelConfig]] = []
+    for model_name, req, _ in models_mandatory_parameters:
+        model_cls = models_dict[model_name]
+        model_cfg = make_model_config(model_cls, req)
+        globals()[model_cfg.__name__] = model_cfg
+        __all__.append(model_cfg.__name__)
+        models_configs.append(model_cfg)
+
+    BraindecodeModelConfig = Annotated[  # type: ignore
+        Union[tuple(models_configs)],
+        pydantic.Field(
+            discriminator="model_name_", description="Braindecode model configuration"
+        ),
+    ]
+
+# # Example usage:
+#
+# class DummyConfigWithModel(pydantic.BaseModel):
+#     model: BraindecodeModelConfig
+#
+# DummyConfigWithModel.model_validate({'model': dict(model_name_='EEGNet', n_chans=16, n_outputs=1, n_times=200)})

braindecode/models/contrawr.py

@@ -0,0 +1,319 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from mne.utils import warn
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class ContraWR(EEGModuleMixin, nn.Module):
+    r"""Contrast with the World Representation ContraWR from Yang et al (2021) [Yang2021]_.
+
+    :bdg-success:`Convolution`
+
+    This model is a convolutional neural network that uses a spectral
+    representation with a series of convolutional layers and residual blocks.
+    The model is designed to learn a representation of the EEG signal that can
+    be used for sleep staging.
+
+    Parameters
+    ----------
+    steps : int, optional
+        Number of steps to take the frequency decomposition `hop_length`
+        parameters by default 20.
+    emb_size : int, optional
+        Embedding size for the final layer, by default 256.
+    res_channels : list[int], optional
+        Number of channels for each residual block, by default [32, 64, 128].
+    activation: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+    drop_prob : float, default=0.5
+        The dropout rate for regularization. Values should be between 0 and 1.
+
+    .. versionadded:: 0.9
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors. The modifications are minimal and the model is expected
+    to work as intended. the original code from [Code2023]_.
+
+    References
+    ----------
+    .. [Yang2021] Yang, C., Xiao, C., Westover, M. B., & Sun, J. (2023).
+       Self-supervised electroencephalogram representation learning for automatic
+       sleep staging: model development and evaluation study. JMIR AI, 2(1), e46769.
+    .. [Code2023] Yang, C., Westover, M.B. and Sun, J., 2023. BIOT
+       Biosignal Transformer for Cross-data Learning in the Wild.
+       GitHub https://github.com/ycq091044/BIOT (accessed 2024-02-13)
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_outputs=None,
+        sfreq=None,
+        emb_size: int = 256,
+        res_channels: list[int] = [32, 64, 128],
+        steps=20,
+        activation: type[nn.Module] = nn.ELU,
+        drop_prob: float = 0.5,
+        stride_res: int = 2,
+        kernel_size_res: int = 3,
+        padding_res: int = 1,
+        # Another way to pass the EEG parameters
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=None,
+    ):
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            chs_info=chs_info,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, chs_info, n_times, sfreq, input_window_seconds
+        if not isinstance(res_channels, list):
+            raise ValueError("res_channels must be a list of integers.")
+
+        if self.input_window_seconds < 1.0:
+            warning_msg = (
+                "The input window is less than 1 second, which may not be "
+                "sufficient for the model to learn meaningful representations."
+                "changing the `n_fft` to `n_times`."
+            )
+            warn(warning_msg, UserWarning)
+            self.n_fft = self.n_times
+        else:
+            self.n_fft = int(self.sfreq)
+
+        self.steps = steps
+
+        res_channels = [self.n_chans] + res_channels + [emb_size]
+
+        self.torch_stft = _STFTModule(
+            n_fft=self.n_fft,
+            hop_length=int(self.n_fft // self.steps),
+        )
+
+        self.convs = nn.ModuleList(
+            [
+                _ResBlock(
+                    in_channels=res_channels[i],
+                    out_channels=res_channels[i + 1],
+                    stride=stride_res,
+                    use_downsampling=True,
+                    pooling=True,
+                    drop_prob=drop_prob,
+                    kernel_size=kernel_size_res,
+                    padding=padding_res,
+                    activation=activation,
+                )
+                for i in range(len(res_channels) - 1)
+            ]
+        )
+        self.adaptative_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.flatten_layer = nn.Flatten()
+
+        self.activation_layer = activation()
+        self.final_layer = nn.Linear(emb_size, self.n_outputs)
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        X: Tensor
+            Input tensor of shape (batch_size, n_channels, n_times).
+        Returns
+        -------
+        Tensor
+            Output tensor of shape (batch_size, n_outputs).
+        """
+        X = self.torch_stft(X)
+
+        for conv in self.convs:
+            X = conv.forward(X)
+
+        emb = self.adaptative_pool(X)
+        emb = self.flatten_layer(emb)
+        emb = self.activation_layer(emb)
+
+        return self.final_layer(emb)
+
+
+class _ResBlock(nn.Module):
+    r"""Convolutional Residual Block 2D.
+
+    This block stacks two convolutional layers with batch normalization,
+    max pooling, dropout, and residual connection.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    stride : int (default=1)
+        Stride of the convolutional layers.
+    use_downsampling : bool (default=True)
+        Whether to use a downsampling residual connection.
+    pooling : bool (default=True)
+        Whether to use max pooling.
+    kernel_size : int (default=3)
+        Kernel size of the convolutional layers.
+    padding : int (default=1)
+        Padding of the convolutional layers.
+    activation: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+    drop_prob : float, default=0.5
+        The dropout rate for regularization. Values should be between 0 and 1.
+
+    Examples
+    --------
+    >>> import torch
+    >>> model = ResBlock2D(6, 16, 1, True, True)
+    >>> input_ = torch.randn((16, 6, 28, 150))  # (batch, channel, height, width)
+    >>> output = model(input_)
+    >>> output.shape
+    torch.Size([16, 16, 14, 75])
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride=1,
+        use_downsampling=True,
+        pooling=True,
+        kernel_size=3,
+        padding=1,
+        drop_prob=0.5,
+        activation: type[nn.Module] = nn.ReLU,
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+        )
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.relu = activation()
+        self.conv2 = nn.Conv2d(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            padding=padding,
+        )
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.maxpool = nn.MaxPool2d(
+            kernel_size=kernel_size, stride=stride, padding=padding
+        )
+        self.downsample = nn.Sequential(
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+            ),
+            nn.BatchNorm2d(out_channels),
+        )
+        self.use_downsampling = use_downsampling
+        self.pooling = pooling
+        self.dropout = nn.Dropout(drop_prob)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+
+        Parameters
+        ----------
+        X: Tensor
+            Input tensor of shape (batch_size, n_channels, n_freqs, n_times).
+
+        Returns
+        -------
+        Tensor
+            Output tensor of shape (batch_size, n_channels, n_freqs, n_times).
+        """
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        if self.use_downsampling:
+            residual = self.downsample(x)
+            out += residual
+        if self.pooling:
+            out = self.maxpool(out)
+        out = self.dropout(out)
+        return out
+
+
+class _STFTModule(nn.Module):
+    r"""
+    A PyTorch module that computes the Short-Time Fourier Transform (STFT)
+    of an EEG batch tensor.
+
+    Expects input of shape (batch_size, n_channels, n_times) and returns
+    (batch_size, n_channels, n_freqs, n_times).
+    """
+
+    def __init__(
+        self,
+        n_fft: int,
+        hop_length: int,
+        center: bool = True,
+        onesided: bool = True,
+        return_complex: bool = True,
+        normalized: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        n_fft : int
+            Number of FFT points (window size).
+        steps : int
+            Number of hops per window (i.e. hop_length = n_fft // steps).
+        """
+        super().__init__()
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.center = center
+        self.one_sided = onesided
+        self.return_complex = return_complex
+        self.normalized = normalized
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        window = torch.ones(self.n_fft, device=x.device)
+
+        # x: (B, C, T)
+        B, C, T = x.shape
+        # flatten batch & channel into one dim
+        x_flat = x.reshape(B * C, T)
+
+        # compute stft on 2D tensor
+        spec_flat = torch.stft(
+            x_flat,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.n_fft,
+            window=window,
+            normalized=self.normalized,
+            center=self.center,
+            onesided=self.one_sided,
+            return_complex=self.return_complex,
+        )
+
+        F, L = spec_flat.shape[-2], spec_flat.shape[-1]
+        spec = spec_flat.view(B, C, F, L)
+
+        return torch.abs(spec)