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

braindecode/augmentation/functional.py

@@ -1,17 +1,18 @@
 # Authors: Cédric Rommel <cedric.rommel@inria.fr>
 #          Alexandre Gramfort <alexandre.gramfort@inria.fr>
+#          Gustavo Rodrigues <gustavenrique01@gmail.com>
 #
 # License: BSD (3-clause)
 
 from numbers import Real
 
 import numpy as np
+import torch
+from mne.filter import notch_filter
 from scipy.interpolate import Rbf
 from sklearn.utils import check_random_state
-import torch
 from torch.fft import fft, ifft
-from torch.nn.functional import pad, one_hot
-from mne.filter import notch_filter
+from torch.nn.functional import one_hot, pad
 
 
 def identity(X, y):
@@ -79,11 +80,14 @@ def _new_random_fft_phase_odd(batch_size, c, n, device, random_state):
     random_phase = torch.from_numpy(
         2j * np.pi * rng.random((batch_size, c, (n - 1) // 2))
     ).to(device)
-    return torch.cat([
-        torch.zeros((batch_size, c, 1), device=device),
-        random_phase,
-        -torch.flip(random_phase, [-1])
-    ], dim=-1)
+    return torch.cat(
+        [
+            torch.zeros((batch_size, c, 1), device=device),
+            random_phase,
+            -torch.flip(random_phase, [-1]),
+        ],
+        dim=-1,
+    )
 
 
 def _new_random_fft_phase_even(batch_size, c, n, device, random_state):
@@ -91,27 +95,21 @@ def _new_random_fft_phase_even(batch_size, c, n, device, random_state):
     random_phase = torch.from_numpy(
         2j * np.pi * rng.random((batch_size, c, n // 2 - 1))
     ).to(device)
-    return torch.cat([
-        torch.zeros((batch_size, c, 1), device=device),
-        random_phase,
-        torch.zeros((batch_size, c, 1), device=device),
-        -torch.flip(random_phase, [-1])
-    ], dim=-1)
-
-
-_new_random_fft_phase = {
-    0: _new_random_fft_phase_even,
-    1: _new_random_fft_phase_odd
-}
-
-
-def ft_surrogate(
-    X,
-    y,
-    phase_noise_magnitude,
-    channel_indep,
-    random_state=None
-):
+    return torch.cat(
+        [
+            torch.zeros((batch_size, c, 1), device=device),
+            random_phase,
+            torch.zeros((batch_size, c, 1), device=device),
+            -torch.flip(random_phase, [-1]),
+        ],
+        dim=-1,
+    )
+
+
+_new_random_fft_phase = {0: _new_random_fft_phase_even, 1: _new_random_fft_phase_odd}
+
+
+def ft_surrogate(X, y, phase_noise_magnitude, channel_indep, random_state=None):
     """FT surrogate augmentation of a single EEG channel, as proposed in [1]_.
 
     Function copied from https://github.com/cliffordlab/sleep-convolutions-tf
@@ -148,11 +146,12 @@ def ft_surrogate(
        Problems of Noisy Signals by using Fourier Transform Surrogates. arXiv
        preprint arXiv:1806.08675.
     """
-    assert isinstance(
-        phase_noise_magnitude,
-        (Real, torch.FloatTensor, torch.cuda.FloatTensor)
-    ) and 0 <= phase_noise_magnitude <= 1, (
-        f"eps must be a float between 0 and 1. Got {phase_noise_magnitude}.")
+    assert (
+        isinstance(
+            phase_noise_magnitude, (Real, torch.FloatTensor, torch.cuda.FloatTensor)
+        )
+        and 0 <= phase_noise_magnitude <= 1
+    ), f"eps must be a float between 0 and 1. Got {phase_noise_magnitude}."
 
     f = fft(X.double(), dim=-1)
     device = X.device
@@ -163,7 +162,7 @@ def ft_surrogate(
         f.shape[-2] if channel_indep else 1,
         n,
         device=device,
-        random_state=random_state
+        random_state=random_state,
     )
     if not channel_indep:
         random_phase = torch.tile(random_phase, (1, f.shape[-2], 1))
@@ -186,7 +185,7 @@ def _pick_channels_randomly(X, p_pick, random_state):
         device=X.device,
     )
     # equivalent to a 0s and 1s mask
-    return torch.sigmoid(1000*(unif_samples - p_pick))
+    return torch.sigmoid(1000 * (unif_samples - p_pick))
 
 
 def channels_dropout(X, y, p_drop, random_state=None):
@@ -234,7 +233,8 @@ def _make_permutation_matrix(X, mask, random_state):
         channels_to_shuffle = torch.arange(n_channels, device=X.device)
         channels_to_shuffle = channels_to_shuffle[mask.bool()]
         reordered_channels = torch.tensor(
-            rng.permutation(channels_to_shuffle.cpu()), device=X.device)
+            rng.permutation(channels_to_shuffle.cpu()), device=X.device
+        )
         channels_permutation = torch.arange(n_channels, device=X.device)
         channels_permutation[channels_to_shuffle] = reordered_channels
         batch_permutations[b, ...] = one_hot(channels_permutation)
@@ -320,12 +320,14 @@ def gaussian_noise(X, y, std, random_state=None):
     rng = check_random_state(random_state)
     if isinstance(std, torch.Tensor):
         std = std.to(X.device)
-    noise = torch.from_numpy(
-        rng.normal(
-            loc=np.zeros(X.shape),
-            scale=1
-        ),
-    ).float().to(X.device) * std
+    noise = (
+        torch.from_numpy(
+            rng.normal(loc=np.zeros(X.shape), scale=1),
+        )
+        .float()
+        .to(X.device)
+        * std
+    )
     transformed_X = X + noise
     return transformed_X, y
 
@@ -399,9 +401,14 @@ def smooth_time_mask(X, y, mask_start_per_sample, mask_len_samples):
     t = t.repeat(batch_size, n_channels, 1)
     mask_start_per_sample = mask_start_per_sample.view(-1, 1, 1)
     s = 1000 / seq_len
-    mask = (torch.sigmoid(s * -(t - mask_start_per_sample)) +
-            torch.sigmoid(s * (t - mask_start_per_sample - mask_len_samples))
-            ).float().to(X.device)
+    mask = (
+        (
+            torch.sigmoid(s * -(t - mask_start_per_sample))
+            + torch.sigmoid(s * (t - mask_start_per_sample - mask_len_samples))
+        )
+        .float()
+        .to(X.device)
+    )
     return X * mask, y
 
 
@@ -447,17 +454,18 @@ def bandstop_filter(X, y, sfreq, bandwidth, freqs_to_notch):
     if bandwidth == 0:
         return X, y
     transformed_X = X.clone()
-    for c, (sample, notched_freq) in enumerate(
-            zip(transformed_X, freqs_to_notch)):
+    for c, (sample, notched_freq) in enumerate(zip(transformed_X, freqs_to_notch)):
         sample = sample.cpu().numpy().astype(np.float64)
-        transformed_X[c] = torch.as_tensor(notch_filter(
-            sample,
-            Fs=sfreq,
-            freqs=notched_freq,
-            method='fir',
-            notch_widths=bandwidth,
-            verbose=False
-        ))
+        transformed_X[c] = torch.as_tensor(
+            notch_filter(
+                sample,
+                Fs=sfreq,
+                freqs=notched_freq,
+                method="fir",
+                notch_widths=bandwidth,
+                verbose=False,
+            )
+        )
     return transformed_X, y
 
 
@@ -470,10 +478,10 @@ def _analytic_transform(x):
     h = torch.zeros_like(f)
     if N % 2 == 0:
         h[..., 0] = h[..., N // 2] = 1
-        h[..., 1:N // 2] = 2
+        h[..., 1 : N // 2] = 2
     else:
         h[..., 0] = 1
-        h[..., 1:(N + 1) // 2] = 2
+        h[..., 1 : (N + 1) // 2] = 2
 
     return ifft(f * h, dim=-1)
 
@@ -500,7 +508,8 @@ def _frequency_shift(X, fs, f_shift):
         f_shift = torch.as_tensor(f_shift).float()
     f_shift_stack = f_shift.repeat(N_padded, n_channels, 1)
     reshaped_f_shift = f_shift_stack.permute(
-        *torch.arange(f_shift_stack.ndim - 1, -1, -1))
+        *torch.arange(f_shift_stack.ndim - 1, -1, -1)
+    )
     shifted = analytical * torch.exp(2j * np.pi * reshaped_f_shift * t)
     return shifted[..., :N_orig].real.float()
 
@@ -539,7 +548,7 @@ def frequency_shift(X, y, delta_freq, sfreq):
 def _torch_normalize_vectors(rr):
     """Normalize surface vertices."""
     norm = torch.linalg.norm(rr, axis=1, keepdim=True)
-    mask = (norm > 0)
+    mask = norm > 0
     norm[~mask] = 1  # in case norm is zero, divide by 1
     new_rr = rr / norm
     return new_rr
@@ -631,7 +640,7 @@ def _torch_legval(x, c, tensor=True):
     if isinstance(x, (tuple, list)):
         x = torch.as_tensor(x)
     if isinstance(x, torch.Tensor) and tensor:
-        c = c.view(c.shape + (1,)*x.ndim)
+        c = c.view(c.shape + (1,) * x.ndim)
 
     c = c.to(x.device)
 
@@ -648,9 +657,9 @@ def _torch_legval(x, c, tensor=True):
         for i in range(3, len(c) + 1):
             tmp = c0
             nd = nd - 1
-            c0 = c[-i] - (c1*(nd - 1))/nd
-            c1 = tmp + (c1*x*(2*nd - 1))/nd
-    return c0 + c1*x
+            c0 = c[-i] - (c1 * (nd - 1)) / nd
+            c1 = tmp + (c1 * x * (2 * nd - 1)) / nd
+    return c0 + c1 * x
 
 
 def _torch_calc_g(cosang, stiffness=4, n_legendre_terms=50):
@@ -702,9 +711,10 @@ def _torch_calc_g(cosang, stiffness=4, n_legendre_terms=50):
     OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
     DAMAGE.
     """
-    factors = [(2 * n + 1) / (n ** stiffness * (n + 1) ** stiffness *
-                              4 * np.pi)
-               for n in range(1, n_legendre_terms + 1)]
+    factors = [
+        (2 * n + 1) / (n**stiffness * (n + 1) ** stiffness * 4 * np.pi)
+        for n in range(1, n_legendre_terms + 1)
+    ]
     return _torch_legval(cosang, [0] + factors)
 
 
@@ -783,15 +793,20 @@ def _torch_make_interpolation_matrix(pos_from, pos_to, alpha=1e-5):
     assert G_to_from.shape == (n_to, n_from)
 
     if alpha is not None:
-        G_from.flatten()[::len(G_from) + 1] += alpha
+        G_from.flatten()[:: len(G_from) + 1] += alpha
 
     device = G_from.device
-    C = torch.vstack([
+    C = torch.vstack(
+        [
             torch.hstack([G_from, torch.ones((n_from, 1), device=device)]),
-            torch.hstack([
-                torch.ones((1, n_from), device=device),
-                torch.as_tensor([[0]], device=device)])
-        ])
+            torch.hstack(
+                [
+                    torch.ones((1, n_from), device=device),
+                    torch.as_tensor([[0]], device=device),
+                ]
+            ),
+        ]
+    )
 
     try:
         C_inv = torch.linalg.inv(C)
@@ -800,10 +815,9 @@ def _torch_make_interpolation_matrix(pos_from, pos_to, alpha=1e-5):
         # see https://github.com/pytorch/pytorch/issues/75494
         C_inv = torch.linalg.pinv(C.cpu()).to(device)
 
-    interpolation = torch.hstack([
-        G_to_from,
-        torch.ones((n_to, 1), device=device)
-    ]).matmul(C_inv[:, :-1])
+    interpolation = torch.hstack(
+        [G_to_from, torch.ones((n_to, 1), device=device)]
+    ).matmul(C_inv[:, :-1])
     assert interpolation.shape == (n_to, n_from)
     return interpolation
 
@@ -815,11 +829,15 @@ def _rotate_signals(X, rotations, sensors_positions_matrix, spherical=True):
     ]
     if spherical:
         interpolation_matrix = torch.stack(
-            [torch.as_tensor(
-                _torch_make_interpolation_matrix(
-                    sensors_positions_matrix.T, rot_sensors_matrix.T
-                ), device=X.device
-            ).float() for rot_sensors_matrix in rot_sensors_matrices]
+            [
+                torch.as_tensor(
+                    _torch_make_interpolation_matrix(
+                        sensors_positions_matrix.T, rot_sensors_matrix.T
+                    ),
+                    device=X.device,
+                ).float()
+                for rot_sensors_matrix in rot_sensors_matrices
+            ]
         )
         return torch.matmul(interpolation_matrix, X)
     else:
@@ -836,7 +854,7 @@ def _rotate_signals(X, rotations, sensors_positions_matrix, spherical=True):
 
 
 def _make_rotation_matrix(axis, angle, degrees=True):
-    assert axis in ['x', 'y', 'z'], "axis should be either x, y or z."
+    assert axis in ["x", "y", "z"], "axis should be either x, y or z."
 
     if isinstance(angle, (float, int, np.ndarray, list)):
         angle = torch.as_tensor(angle)
@@ -846,11 +864,13 @@ def _make_rotation_matrix(axis, angle, degrees=True):
 
     device = angle.device
     zero = torch.zeros(1, device=device)
-    rot = torch.stack([
-        torch.as_tensor([1, 0, 0], device=device),
-        torch.hstack([zero, torch.cos(angle), -torch.sin(angle)]),
-        torch.hstack([zero, torch.sin(angle), torch.cos(angle)]),
-    ])
+    rot = torch.stack(
+        [
+            torch.as_tensor([1, 0, 0], device=device),
+            torch.hstack([zero, torch.cos(angle), -torch.sin(angle)]),
+            torch.hstack([zero, torch.sin(angle), torch.cos(angle)]),
+        ]
+    )
     if axis == "x":
         return rot
     elif axis == "y":
@@ -861,8 +881,7 @@ def _make_rotation_matrix(axis, angle, degrees=True):
         return rot[:, [1, 2, 0]]
 
 
-def sensors_rotation(X, y, sensors_positions_matrix, axis, angles,
-                     spherical_splines):
+def sensors_rotation(X, y, sensors_positions_matrix, axis, angles, spherical_splines):
     """Interpolates EEG signals over sensors rotated around the desired axis
     with the desired angle.
 
@@ -900,13 +919,8 @@ def sensors_rotation(X, y, sensors_positions_matrix, axis, angles,
        Conference of the IEEE Engineering in Medicine and Biology Society
        (EMBC) (pp. 471-474).
     """
-    rots = [
-        _make_rotation_matrix(axis, angle, degrees=True)
-        for angle in angles
-    ]
-    rotated_X = _rotate_signals(
-        X, rots, sensors_positions_matrix, spherical_splines
-    )
+    rots = [_make_rotation_matrix(axis, angle, degrees=True) for angle in angles]
+    rotated_X = _rotate_signals(X, rots, sensors_positions_matrix, spherical_splines)
     return rotated_X, y
 
 
@@ -942,7 +956,7 @@ def mixup(X, y, lam, idx_perm):
         International Conference on Learning Representations (ICLR)
         Online: https://arxiv.org/abs/1710.09412
     .. [2] https://github.com/facebookresearch/mixup-cifar10/blob/master/train.py
-     """
+    """
     device = X.device
     batch_size, n_channels, n_times = X.shape
 
@@ -951,9 +965,132 @@ def mixup(X, y, lam, idx_perm):
     y_b = torch.arange(batch_size).to(device)
 
     for idx in range(batch_size):
-        X_mix[idx] = lam[idx] * X[idx] \
-            + (1 - lam[idx]) * X[idx_perm[idx]]
+        X_mix[idx] = lam[idx] * X[idx] + (1 - lam[idx]) * X[idx_perm[idx]]
         y_a[idx] = y[idx]
         y_b[idx] = y[idx_perm[idx]]
 
     return X_mix, (y_a, y_b, lam)
+
+
+def segmentation_reconstruction(
+    X, y, n_segments, data_classes, rand_indices, idx_shuffle
+):
+    """Segment and reconstruct EEG data from [1]_.
+
+    See [1]_ for details.
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        EEG input example or batch.
+    y : torch.Tensor
+        EEG labels for the example or batch.
+    n_segments : int
+        Number of segments to use in the batch.
+    rand_indices: array-like
+        Array of indices that indicates which trial to use in each segment.
+    idx_shuffle: array-like
+        Array of indices to shuffle the new generated trials.
+    Returns
+    -------
+    torch.Tensor
+        Transformed inputs.
+    torch.Tensor
+        Transformed labels.
+    References
+    ----------
+    .. [1] Lotte, F. (2015). Signal processing approaches to minimize or
+        suppress calibration time in oscillatory activity-based brain–computer
+        interfaces. Proceedings of the IEEE, 103(6), 871-890.
+    """
+
+    # Initialize lists to store augmented data and corresponding labels
+    aug_data = []
+    aug_label = []
+
+    # Iterate through each class to separate and augment data
+    for class_index, X_class in data_classes:
+        # Determine class-specific dimensions
+        # Store the augmented data and the corresponding class labels
+        n_trials, n_channels, window_size = X_class.shape
+        # Segment Size
+        segment_size = window_size // n_segments
+        # Initialize an empty tensor for augmented data
+        X_aug = torch.zeros_like(X_class)
+        # Generate random indices within the class-specific dataset
+        rand_idx = rand_indices[class_index]
+        for idx_segment in range(n_segments):
+            start = idx_segment * segment_size
+            end = (idx_segment + 1) * segment_size
+
+            # Perform the data augmentation
+            X_aug[np.arange(n_trials), :, start:end] = X_class[
+                rand_idx[:, idx_segment], :, start:end
+            ]
+        aug_data.append(X_aug)
+        aug_label.append(torch.full((n_trials,), class_index))
+    # Concatenate the augmented data and labels
+    aug_data = torch.cat(aug_data, dim=0)
+    aug_data = aug_data.to(dtype=X.dtype, device=X.device)
+    aug_data = aug_data[idx_shuffle]
+
+    if y is not None:
+        aug_label = torch.cat(aug_label, dim=0)
+        aug_label = aug_label.to(dtype=y.dtype, device=y.device)
+        aug_label = aug_label[idx_shuffle]
+        return aug_data, aug_label
+
+    return aug_data, y
+
+
+def mask_encoding(X, y, time_start, segment_length, n_segments):
+    """Mark encoding from Ding et al. (2024) from [ding2024]_.
+
+    Replaces a contiguous part (or parts) of all channels by zeros
+    (if more than one segment, it may overlap).
+
+    Implementation based on [ding2024]_
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        EEG input example or batch.
+    y : torch.Tensor
+        EEG labels for the example or batch.
+    time_start : torch.Tensor
+        Tensor of integers containing the position (in last dimension) where to
+        start masking the signal. Should have "n_segments" times the size of the first
+        dimension of X (i.e. "n_segments" start positions per example in the batch).
+    segment_length : int
+        Length of each segment to zero out.
+    n_segments : int
+        Number of segments to zero out in each example.
+
+    Returns
+    -------
+    torch.Tensor
+        Transformed inputs.
+    torch.Tensor
+        Transformed labels.
+
+    References
+    ----------
+    .. [ding2024] Ding, Wenlong, et al. A Novel Data Augmentation Approach
+       Using Mask Encoding for Deep Learning-Based Asynchronous SSVEP-BCI.
+       IEEE Transactions on Neural Systems and Rehabilitation Engineering
+       32 (2024): 875-886.
+    """
+
+    batch_indices = torch.arange(X.shape[0]).repeat_interleave(n_segments)
+    start_indices = time_start.flatten()
+    mask_indices = start_indices[:, None] + torch.arange(segment_length)
+
+    # Create a boolean mask with the same shape as X
+    mask = torch.zeros_like(X, dtype=torch.bool)
+    for batch_index, grouped_mask_indices in zip(batch_indices, mask_indices):
+        mask[batch_index, :, grouped_mask_indices] = True
+
+    # Apply the mask to set the values to 0
+    X[mask] = 0
+
+    return X, y  # Return the masked tensor and labels