braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/models/sinc_shallow.py

@@ -0,0 +1,339 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops.layers.torch import Rearrange
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class SincShallowNet(EEGModuleMixin, nn.Module):
+    r"""Sinc-ShallowNet from Borra, D et al (2020) [borra2020]_.
+
+    :bdg-success:`Convolution` :bdg-warning:`Interpretability`
+
+    .. figure:: https://ars.els-cdn.com/content/image/1-s2.0-S0893608020302021-gr2_lrg.jpg
+        :align: center
+        :alt: SincShallowNet Architecture
+
+    The Sinc-ShallowNet architecture has these fundamental blocks:
+
+    1. **Block 1: Spectral and Spatial Feature Extraction**
+
+        - *Temporal Sinc-Convolutional Layer*: Uses parametrized sinc functions to learn band-pass filters,
+          significantly reducing the number of trainable parameters by only
+          learning the lower and upper cutoff frequencies for each filter.
+        - *Spatial Depthwise Convolutional Layer*: Applies depthwise convolutions to learn spatial filters for
+          each temporal feature map independently, further reducing
+          parameters and enhancing interpretability.
+        - *Batch Normalization*
+
+    2. **Block 2: Temporal Aggregation**
+
+        - *Activation Function*: ELU
+        - *Average Pooling Layer*: Aggregation by averaging spatial dim
+        - *Dropout Layer*
+        - *Flatten Layer*
+
+    3. **Block 3: Classification**
+
+        - *Fully Connected Layer*: Maps the feature vector to n_outputs.
+
+    **Implementation Notes:**
+
+    - The sinc-convolutional layer initializes cutoff frequencies uniformly
+        within the desired frequency range and updates them during training while
+        ensuring the lower cutoff is less than the upper cutoff.
+
+    Parameters
+    ----------
+    num_time_filters : int
+        Number of temporal filters in the SincFilter layer.
+    time_filter_len : int
+        Size of the temporal filters.
+    depth_multiplier : int
+        Depth multiplier for spatial filtering.
+    activation : nn.Module, optional
+        Activation function to use. Default is nn.ELU().
+    drop_prob : float, optional
+        Dropout probability. Default is 0.5.
+    first_freq : float, optional
+        The starting frequency for the first Sinc filter. Default is 5.0.
+    min_freq : float, optional
+        Minimum frequency allowed for the low frequencies of the filters. Default is 1.0.
+    freq_stride : float, optional
+        Frequency stride for the Sinc filters. Controls the spacing between the filter frequencies.
+        Default is 1.0.
+    padding : str, optional
+        Padding mode for convolution, either 'same' or 'valid'. Default is 'same'.
+    bandwidth : float, optional
+        Initial bandwidth for each Sinc filter. Default is 4.0.
+    pool_size : int, optional
+        Size of the pooling window for the average pooling layer. Default is 55.
+    pool_stride : int, optional
+        Stride of the pooling operation. Default is 12.
+
+    Notes
+    -----
+    This implementation is based on the implementation from [sincshallowcode]_.
+
+    References
+    ----------
+    .. [borra2020] Borra, D., Fantozzi, S., & Magosso, E. (2020). Interpretable
+       and lightweight convolutional neural network for EEG decoding: Application
+       to movement execution and imagination. Neural Networks, 129, 55-74.
+    .. [sincshallowcode] Sinc-ShallowNet re-implementation source code:
+       https://github.com/marcellosicbaldi/SincNet-Tensorflow
+    """
+
+    def __init__(
+        self,
+        num_time_filters: int = 32,
+        time_filter_len: int = 33,
+        depth_multiplier: int = 2,
+        activation: type[nn.Module] | None = nn.ELU,
+        drop_prob: float = 0.5,
+        first_freq: float = 5.0,
+        min_freq: float = 1.0,
+        freq_stride: float = 1.0,
+        padding: str = "same",
+        bandwidth: float = 4.0,
+        pool_size: int = 55,
+        pool_stride: int = 12,
+        # braindecode parameters
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        chs_info=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, input_window_seconds, sfreq
+
+        # Checkers and creating variables
+        if activation is None:
+            activation = nn.ELU()
+
+        # Define low frequencies for the SincFilter
+        low_freqs = torch.arange(
+            first_freq,
+            first_freq + num_time_filters * freq_stride,
+            freq_stride,
+            dtype=torch.float32,
+        )
+        self.n_filters = len(low_freqs)
+
+        if padding.lower() == "valid":
+            n_times_after_sinc_filter = self.n_times - time_filter_len + 1
+        elif padding.lower() == "same":
+            n_times_after_sinc_filter = self.n_times
+        else:
+            raise ValueError("Padding must be 'valid' or 'same'.")
+
+        size_after_pooling = (
+            (n_times_after_sinc_filter - pool_size) // pool_stride
+        ) + 1
+        flattened_size = num_time_filters * depth_multiplier * size_after_pooling
+
+        # Layers
+        self.ensuredims = Rearrange("batch chans times -> batch chans times 1")
+
+        # Block 1: Sinc filter
+        self.sinc_filter_layer = _SincFilter(
+            low_freqs=low_freqs,
+            kernel_size=time_filter_len,
+            sfreq=self.sfreq,
+            padding=padding,
+            bandwidth=bandwidth,
+            min_freq=min_freq,
+        )
+
+        self.depthwiseconv = nn.Sequential(
+            # Matching dim to depth wise conv!
+            Rearrange("batch timefil time nfilter -> batch nfilter timefil time"),
+            nn.BatchNorm2d(
+                self.n_filters, momentum=0.99
+            ),  # To match keras implementation
+            nn.Conv2d(
+                in_channels=self.n_filters,
+                out_channels=depth_multiplier * self.n_filters,
+                kernel_size=(self.n_chans, 1),
+                groups=self.n_filters,
+                bias=False,
+            ),
+        )
+
+        # Block 2: Batch norm, activation, pooling, dropout
+        self.temporal_aggregation = nn.Sequential(
+            nn.BatchNorm2d(depth_multiplier * self.n_filters, momentum=0.99),
+            activation(),
+            nn.AvgPool2d(kernel_size=(1, pool_size), stride=(1, pool_stride)),
+            nn.Dropout(p=drop_prob),
+            nn.Flatten(),
+        )
+
+        # Final classification layer
+        self.final_layer = nn.Linear(
+            flattened_size,
+            self.n_outputs,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the model.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape [batch_size, num_channels, num_samples].
+
+        Returns
+        -------
+        torch.Tensor
+            Output logits of shape [batch_size, num_classes].
+        """
+        x = self.ensuredims(x)
+        x = self.sinc_filter_layer(x)
+        x = self.depthwiseconv(x)
+        x = self.temporal_aggregation(x)
+
+        return self.final_layer(x)
+
+
+class _SincFilter(nn.Module):
+    r"""Sinc-Based Convolutional Layer for Band-Pass Filtering from Ravanelli and Bengio (2018) [ravanelli]_.
+
+    The `SincFilter` layer implements a convolutional layer where each kernel is
+    defined using a parametrized sinc function.
+    This design enforces each kernel to represent a band-pass filter,
+    reducing the number of trainable parameters.
+
+    Parameters
+    ----------
+    low_freqs : torch.Tensor
+        Initial low cutoff frequencies for each filter.
+    kernel_size : int
+        Size of the convolutional kernels (filters). Must be odd.
+    sfreq : float
+        Sampling rate of the input signal.
+    bandwidth : float, optional
+        Initial bandwidth for each filter. Default is 4.0.
+    min_freq : float, optional
+        Minimum frequency allowed for low frequencies. Default is 1.0.
+    padding : str, optional
+        Padding mode, either 'same' or 'valid'. Default is 'same'.
+
+    References
+    ----------
+    .. [ravanelli] Ravanelli, M., & Bengio, Y. (2018, December). Speaker
+       recognition from raw waveform with sincnet. In 2018 IEEE spoken language
+       technology workshop (SLT) (pp. 1021-1028). IEEE.
+    """
+
+    def __init__(
+        self,
+        low_freqs: torch.Tensor,
+        kernel_size: int,
+        sfreq: float,
+        bandwidth: float = 4.0,
+        min_freq: float = 1.0,
+        padding: str = "same",
+    ):
+        super().__init__()
+        if kernel_size % 2 == 0:
+            raise ValueError("Kernel size must be odd.")
+
+        self.num_filters = low_freqs.numel()
+        self.kernel_size = kernel_size
+        self.sfreq = sfreq
+        self.min_freq = min_freq
+        self.padding = padding.lower()
+
+        # Precompute constants
+        window = torch.hamming_window(kernel_size, periodic=False)
+
+        self.register_buffer("window", window[: kernel_size // 2].unsqueeze(-1))
+
+        n_pi = (
+            torch.arange(-(kernel_size // 2), 0, dtype=torch.float32)
+            / sfreq
+            * 2
+            * math.pi
+        )
+        self.register_buffer("n_pi", n_pi.unsqueeze(-1))
+
+        # Initialize learnable parameters
+        bandwidths = torch.full((1, self.num_filters), bandwidth)
+        self.bandwidths = nn.Parameter(bandwidths)
+        self.low_freqs = nn.Parameter(low_freqs.unsqueeze(0))
+
+        # Constant tensor of ones for filter construction
+        self.register_buffer("ones", torch.ones(1, 1, 1, self.num_filters))
+
+    def build_sinc_filters(self) -> torch.Tensor:
+        """Builds the sinc filters based on current parameters."""
+        # Computing the low frequencies of the filters
+        low_freqs = self.min_freq + torch.abs(self.low_freqs)
+        # Setting a minimum band and minimum freq
+        high_freqs = torch.clamp(
+            low_freqs + torch.abs(self.bandwidths),
+            min=self.min_freq,
+            max=self.sfreq / 2.0,
+        )
+        bandwidths = high_freqs - low_freqs
+
+        # Passing from n_ to the corresponding f_times_t domain
+        low = self.n_pi * low_freqs  # [kernel_size // 2, num_filters]
+        high = self.n_pi * high_freqs  # [kernel_size // 2, num_filters]
+
+        filters_left = (torch.sin(high) - torch.sin(low)) / (self.n_pi / 2.0)
+        filters_left *= self.window
+        filters_left /= 2.0 * bandwidths
+
+        # [1, kernel_size // 2, 1, num_filters]
+        filters_left = filters_left.unsqueeze(0).unsqueeze(2)
+        filters_right = torch.flip(filters_left, dims=[1])
+
+        filters = torch.cat(
+            [filters_left, self.ones, filters_right], dim=1
+        )  # [1, kernel_size, 1, num_filters]
+        filters = filters / torch.std(filters)
+        return filters
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        Apply sinc filters to the input signal.
+
+        Parameters
+        ----------
+        inputs : torch.Tensor
+            Input tensor of shape [batch_size, num_channels, num_samples, 1].
+
+        Returns
+        -------
+        torch.Tensor
+            Filtered output tensor of shape [batch_size, num_channels, num_samples, num_filters].
+        """
+        filters = self.build_sinc_filters().to(
+            inputs.device
+        )  # [1, kernel_size, 1, num_filters]
+
+        # Convert from channels_last to channels_first format
+        inputs = inputs.permute(0, 3, 1, 2)
+        # Permuting to match conv:
+        filters = filters.permute(3, 2, 0, 1)
+        # Apply convolution
+        outputs = F.conv2d(inputs, filters, padding=self.padding)
+        # Changing the dimensional
+        outputs = outputs.permute(0, 2, 3, 1)
+
+        return outputs

braindecode/models/sleep_stager_blanco_2020.py

@@ -0,0 +1,169 @@
+# Authors: Divyesh Narayanan <divyesh.narayanan@gmail.com>
+#
+# License: BSD (3-clause)
+
+import torch
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class SleepStagerBlanco2020(EEGModuleMixin, nn.Module):
+    r"""Sleep staging architecture from Blanco et al (2020) from [Blanco2020]_
+
+    :bdg-success:`Convolution`
+
+    .. figure:: https://media.springernature.com/full/springer-static/image/art%3A10.1007%2Fs00500-019-04174-1/MediaObjects/500_2019_4174_Fig2_HTML.png
+        :align: center
+        :alt: SleepStagerBlanco2020 Architecture
+
+    Convolutional neural network for sleep staging described in [Blanco2020]_.
+    A series of seven convolutional layers with kernel sizes running down from 7 to 3,
+    in an attempt to extract more general features at the beginning, while more specific
+    and complex features were extracted in the final stages.
+
+    Parameters
+    ----------
+    n_conv_chans : int
+        Number of convolutional channels. Set to 20 in [Blanco2020]_.
+    n_groups : int
+        Number of groups for the convolution. Set to 2 in [Blanco2020]_ for 2 Channel EEG.
+        controls the connections between inputs and outputs. n_channels and n_conv_chans must be
+        divisible by n_groups.
+    drop_prob : float
+        Dropout rate before the output dense layer.
+    apply_batch_norm : bool
+        If True, apply batch normalization after both temporal convolutional
+        layers.
+    return_feats : bool
+        If True, return the features, i.e. the output of the feature extractor
+        (before the final linear layer). If False, pass the features through
+        the final linear layer.
+    n_channels : int
+        Alias for `n_chans`.
+    n_classes : int
+        Alias for `n_outputs`.
+    input_size_s : float
+        Alias for `input_window_seconds`.
+    activation: nn.Module, default=nn.ReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+
+    References
+    ----------
+    .. [Blanco2020] Fernandez-Blanco, E., Rivero, D. & Pazos, A. Convolutional
+        neural networks for sleep stage scoring on a two-channel EEG signal.
+        Soft Comput 24, 4067–4079 (2020). https://doi.org/10.1007/s00500-019-04174-1
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        sfreq=None,
+        n_conv_chans=20,
+        input_window_seconds=None,
+        n_outputs=5,
+        n_groups=2,
+        max_pool_size=2,
+        drop_prob=0.5,
+        apply_batch_norm=False,
+        return_feats=False,
+        activation: type[nn.Module] = nn.ReLU,
+        chs_info=None,
+        n_times=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, input_window_seconds, sfreq
+
+        self.mapping = {
+            "fc.1.weight": "final_layer.1.weight",
+            "fc.1.bias": "final_layer.1.bias",
+        }
+
+        batch_norm = nn.BatchNorm2d if apply_batch_norm else nn.Identity
+
+        self.feature_extractor = nn.Sequential(
+            nn.Conv2d(self.n_chans, n_conv_chans, (1, 7), groups=n_groups, padding=0),
+            batch_norm(n_conv_chans),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 7), groups=n_conv_chans, padding=0
+            ),
+            batch_norm(n_conv_chans),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0
+            ),
+            batch_norm(n_conv_chans),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0
+            ),
+            batch_norm(n_conv_chans),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 5), groups=n_conv_chans, padding=0
+            ),
+            batch_norm(n_conv_chans),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 3), groups=n_conv_chans, padding=0
+            ),
+            batch_norm(n_conv_chans),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+            nn.Conv2d(
+                n_conv_chans, n_conv_chans, (1, 3), groups=n_conv_chans, padding=0
+            ),
+            batch_norm(n_conv_chans),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+        )
+
+        self.len_last_layer = self._len_last_layer(self.n_chans, self.n_times)
+        self.return_feats = return_feats
+
+        # TODO: Add new way to handle return_features == True
+        if not return_feats:
+            self.final_layer = nn.Sequential(
+                nn.Dropout(drop_prob),
+                nn.Linear(self.len_last_layer, self.n_outputs),
+                nn.Identity(),
+            )
+
+    def _len_last_layer(self, n_channels, input_size):
+        self.feature_extractor.eval()
+        with torch.no_grad():
+            out = self.feature_extractor(
+                torch.Tensor(1, n_channels, 1, input_size)
+            )  # batch_size,n_channels,height,width
+        self.feature_extractor.train()
+        return len(out.flatten())
+
+    def forward(self, x):
+        """Forward pass.
+        Parameters
+        ----------
+        x: torch.Tensor
+            Batch of EEG windows of shape (batch_size, n_channels, n_times).
+        """
+        if x.ndim == 3:
+            x = x.unsqueeze(2)
+
+        feats = self.feature_extractor(x).flatten(start_dim=1)
+        if self.return_feats:
+            return feats
+        else:
+            return self.final_layer(feats)

braindecode/models/sleep_stager_chambon_2018.py

@@ -0,0 +1,159 @@
+# Authors: Hubert Banville <hubert.jbanville@gmail.com>
+#
+# License: BSD (3-clause)
+
+import math
+
+import torch
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class SleepStagerChambon2018(EEGModuleMixin, nn.Module):
+    r"""Sleep staging architecture from Chambon et al. (2018) [Chambon2018]_.
+
+    :bdg-success:`Convolution`
+
+    .. figure:: https://braindecode.org/dev/_static/model/SleepStagerChambon2018.jpg
+        :align: center
+        :alt: SleepStagerChambon2018 Architecture
+
+    Convolutional neural network for sleep staging described in [Chambon2018]_.
+
+    Parameters
+    ----------
+    n_conv_chs : int
+        Number of convolutional channels. Set to 8 in [Chambon2018]_.
+    time_conv_size_s : float
+        Size of filters in temporal convolution layers, in seconds. Set to 0.5
+        in [Chambon2018]_ (64 samples at sfreq=128).
+    max_pool_size_s : float
+        Max pooling size, in seconds. Set to 0.125 in [Chambon2018]_ (16
+        samples at sfreq=128).
+    pad_size_s : float
+        Padding size, in seconds. Set to 0.25 in [Chambon2018]_ (half the
+        temporal convolution kernel size).
+    drop_prob : float
+        Dropout rate before the output dense layer.
+    apply_batch_norm : bool
+        If True, apply batch normalization after both temporal convolutional
+        layers.
+    return_feats : bool
+        If True, return the features, i.e. the output of the feature extractor
+        (before the final linear layer). If False, pass the features through
+        the final linear layer.
+    n_channels : int
+        Alias for `n_chans`.
+    input_size_s:
+        Alias for `input_window_seconds`.
+    n_classes:
+        Alias for `n_outputs`.
+    activation: nn.Module, default=nn.ReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+
+    References
+    ----------
+    .. [Chambon2018] Chambon, S., Galtier, M. N., Arnal, P. J., Wainrib, G., &
+           Gramfort, A. (2018). A deep learning architecture for temporal sleep
+           stage classification using multivariate and multimodal time series.
+           IEEE Transactions on Neural Systems and Rehabilitation Engineering,
+           26(4), 758-769.
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        sfreq=None,
+        n_conv_chs=8,
+        time_conv_size_s=0.5,
+        max_pool_size_s=0.125,
+        pad_size_s=0.25,
+        activation: type[nn.Module] = nn.ReLU,
+        input_window_seconds=None,
+        n_outputs=5,
+        drop_prob=0.25,
+        apply_batch_norm=False,
+        return_feats=False,
+        chs_info=None,
+        n_times=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, input_window_seconds, sfreq
+
+        self.mapping = {
+            "fc.1.weight": "final_layer.1.weight",
+            "fc.1.bias": "final_layer.1.bias",
+        }
+
+        time_conv_size = math.ceil(time_conv_size_s * self.sfreq)
+        max_pool_size = math.ceil(max_pool_size_s * self.sfreq)
+        pad_size = math.ceil(pad_size_s * self.sfreq)
+
+        if self.n_chans > 1:
+            self.spatial_conv = nn.Conv2d(1, self.n_chans, (self.n_chans, 1))
+        else:
+            self.spatial_conv = nn.Identity()
+
+        batch_norm = nn.BatchNorm2d if apply_batch_norm else nn.Identity
+
+        self.feature_extractor = nn.Sequential(
+            nn.Conv2d(1, n_conv_chs, (1, time_conv_size), padding=(0, pad_size)),
+            batch_norm(n_conv_chs),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+            nn.Conv2d(
+                n_conv_chs, n_conv_chs, (1, time_conv_size), padding=(0, pad_size)
+            ),
+            batch_norm(n_conv_chs),
+            activation(),
+            nn.MaxPool2d((1, max_pool_size)),
+        )
+        self.return_feats = return_feats
+
+        dim_conv_1 = (
+            self.n_times + 2 * pad_size - (time_conv_size - 1)
+        ) // max_pool_size
+        dim_after_conv = (
+            dim_conv_1 + 2 * pad_size - (time_conv_size - 1)
+        ) // max_pool_size
+
+        self.len_last_layer = n_conv_chs * self.n_chans * dim_after_conv
+
+        # TODO: Add new way to handle return_features == True
+        if not return_feats:
+            self.final_layer = nn.Sequential(
+                nn.Dropout(p=drop_prob),
+                nn.Linear(in_features=self.len_last_layer, out_features=self.n_outputs),
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Batch of EEG windows of shape (batch_size, n_channels, n_times).
+        """
+        if x.ndim == 3:
+            x = x.unsqueeze(1)
+
+        if self.n_chans > 1:
+            x = self.spatial_conv(x)
+            x = x.transpose(1, 2)
+
+        feats = self.feature_extractor(x).flatten(start_dim=1)
+
+        if self.return_feats:
+            return feats
+
+        return self.final_layer(feats)