braindecode 1.7.0.dev183981614__tar.gz → 1.8.0.dev1122__tar.gz

{braindecode-1.7.0.dev183981614/braindecode.egg-info → braindecode-1.8.0.dev1122}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.7.0.dev183981614
+Version: 1.8.0.dev1122
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/augmentation/base.py

@@ -3,6 +3,7 @@
 #          Bruno Aristimunha <b.aristimunha@gmail.com>
 #          Martin Wimpff <martin.wimpff@iss.uni-stuttgart.de>
 #          Valentin Iovene <val@too.gy>
+#          Sarthak Tayal <sarthaktayal2@gmail.com>
 # License: BSD (3-clause)
 
 from numbers import Real
@@ -176,6 +177,14 @@ class Compose(Transform):
         return X, y
 
 
+def _as_mixed_target(y, lam_dtype):
+    # an untouched target is the same as being mixed with itself with lam of one
+    if isinstance(y, (tuple, list)) and len(y) == 3:
+        return tuple(y)
+    lam = torch.ones(y.shape[0], device=y.device, dtype=lam_dtype)
+    return y, y, lam
+
+
 class _AugmentationCollate:
     """Collate that applies a transform to each batch, with optional expansion.
 
@@ -193,7 +202,10 @@ class _AugmentationCollate:
         ``0`` (default) applies the transform in place (batch size unchanged).
         ``> 0`` keeps the clean originals and appends ``n_augmentation``
         independently transformed copies, returning ``(X, y)`` of
-        ``(1 + n_augmentation)`` times the original size.
+        ``(1 + n_augmentation)`` times the original size. When the transform
+        mixes targets, as :class:`braindecode.augmentation.Mixup` does, ``y``
+        stays the ``(y_a, y_b, lam)`` triple and the clean originals get a
+        mixing coefficient of one.
     """
 
     def __init__(self, transform, device=None, n_augmentation=0):
@@ -216,6 +228,18 @@ class _AugmentationCollate:
             aug_X, aug_y = self.transform(X, y)
             xs.append(aug_X)
             ys.append(aug_y)
+        mixed_ys = [
+            aug_y
+            for aug_y in ys
+            if isinstance(aug_y, (tuple, list)) and len(aug_y) == 3
+        ]
+        if mixed_ys:
+            # a target-mixing transform such as Mixup returns (y_a, y_b, lam),
+            # so the parts are concatenated one by one and the untouched copies
+            # are given a mixing coefficient of one
+            lam_dtype = mixed_ys[0][2].dtype
+            ys = [_as_mixed_target(aug_y, lam_dtype) for aug_y in ys]
+            return torch.cat(xs), tuple(torch.cat(part) for part in zip(*ys))
         return torch.cat(xs), torch.cat(ys)
 
 

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/augmentation/functional.py

@@ -1065,13 +1065,11 @@ def mixup(
     batch_size, n_channels, n_times = X.shape
 
     X_mix = torch.zeros((batch_size, n_channels, n_times)).to(device)
-    y_a = torch.arange(batch_size).to(device)
-    y_b = torch.arange(batch_size).to(device)
+    y_a = y.clone()
+    y_b = y[idx_perm].clone()
 
     for idx in range(batch_size):
         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)
 

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/augmentation/transforms.py

@@ -1068,16 +1068,18 @@ class Mixup(Transform):
         device = X.device
         batch_size, _, _ = X.shape
 
+        # lam follows the dtype of X, numpy draws float64 and that would leak
+        # into the mixed signal and into the loss returned by mixup_criterion
         if self.alpha > 0:
             if self.beta_per_sample:
                 lam = torch.as_tensor(
                     self.rng.beta(self.alpha, self.alpha, batch_size)
-                ).to(device)
+                ).to(device=device, dtype=X.dtype)
             else:
-                lam = torch.ones(batch_size).to(device)
+                lam = torch.ones(batch_size, dtype=X.dtype).to(device)
                 lam *= self.rng.beta(self.alpha, self.alpha)
         else:
-            lam = torch.ones(batch_size).to(device)
+            lam = torch.ones(batch_size, dtype=X.dtype).to(device)
 
         idx_perm = torch.as_tensor(
             self.rng.permutation(

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/classifier.py

@@ -228,8 +228,9 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
             if return_targets:
                 return preds, X.get_metadata()["target"].to_numpy()
             return preds
+        self.check_is_fitted()
         return predict_trials(
-            module=self.module,
+            module=self.module_,
             dataset=X,
             return_targets=return_targets,
             batch_size=self.batch_size,

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/datasets/sleep_physionet.py

@@ -117,9 +117,15 @@ class SleepPhysionet(BaseConcatDataset):
             sleep_event_inds = np.where(mask)[0]
 
             # Crop raw
-            tmin = annots[int(sleep_event_inds[0])]["onset"] - crop_wake_mins * 60
-            tmax = annots[int(sleep_event_inds[-1])]["onset"] + crop_wake_mins * 60
-            raw.crop(tmin=max(tmin, raw.times[0]), tmax=min(tmax, raw.times[-1]))
+            a_tmin = annots[sleep_event_inds[0]]
+            a_tmax = annots[sleep_event_inds[-1]]
+            tmin = a_tmin["onset"] - crop_wake_mins * 60
+            tmax = a_tmax["onset"] + a_tmax["duration"] + crop_wake_mins * 60
+            raw.crop(
+                tmin=max(tmin, raw.times[0]),
+                tmax=min(tmax, raw.times[-1] + 1 / raw.info["sfreq"]),
+                include_tmax=False,
+            )
 
         # Rename EEG channels
         ch_names = {i: i.replace("EEG ", "") for i in raw.ch_names if "EEG" in i}

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/eegneuralnet.py

@@ -1,5 +1,6 @@
 # Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
 #          Pierre Guetschel <pierre.guetschel@gmail.com>
+#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -126,7 +127,8 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
                 self._last_window_inds_ = None
 
     def predict_with_window_inds_and_ys(self, dataset):
-        self.module.eval()
+        # self.module can still be a name or a class, self.module_ is the built one
+        self.module_.eval()
         preds = []
         i_window_in_trials = []
         i_window_stops = []
@@ -142,7 +144,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
             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))))
+                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)

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/models/__init__.py

@@ -71,6 +71,7 @@ from .util import (
     models_mandatory_parameters,
     positions_from_chs_info,
 )
+from .zuna import ZUNA
 
 # Call this last in order to make sure the dataset list is populated with
 # the models imported in this file.
@@ -145,6 +146,7 @@ __all__ = [
     "TIDNet",
     "TSception",
     "USleep",
+    "ZUNA",
     "build_model_config",
     "_init_models_dict",
     "models_mandatory_parameters",

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/models/fblightconvnet.py

@@ -1,3 +1,7 @@
+# Authors: Sarthak Tayal <sarthaktayal2@gmail.com>
+#
+# License: BSD-3
+
 from __future__ import annotations
 
 from typing import Optional
@@ -61,15 +65,18 @@ class FBLightConvNet(EEGModuleMixin, nn.Module):
 
     Parameters
     ----------
-    n_bands : int or None or list of tuple of int, default=8
+    n_bands : int or None or list of tuple of int, default=9
         Number of frequency bands or a list of frequency band tuples. If a list of tuples is provided,
         each tuple defines the lower and upper bounds of a frequency band.
     n_filters_spat : int, default=32
         Number of spatial filters in the depthwise convolutional layer.
     n_dim : int, default=3
         Number of dimensions for the temporal reduction layer.
-    stride_factor : int, default=4
-        Stride factor used for reshaping the temporal dimension.
+    win_len : int, default=250
+        Length in samples of the non-overlapping temporal windows the signal is
+        split into before the variance based feature extraction. The number of
+        windows passed to the attention module is ``n_times // win_len``, so
+        ``n_times`` has to be at least ``win_len``.
     activation : nn.Module, default=nn.ELU
         Activation function class to apply after convolutional layers.
     verbose : bool, default=False
@@ -82,6 +89,9 @@ class FBLightConvNet(EEGModuleMixin, nn.Module):
         If True, applies softmax to the attention weights.
     bias : bool, default=False
         If True, includes a bias term in the convolutional layers.
+    stride_factor : int or None, default=None
+        Deprecated and ignored, it will be removed in a future release. The
+        temporal segmentation of this model is controlled by ``win_len``.
 
     References
     ----------
@@ -106,7 +116,7 @@ class FBLightConvNet(EEGModuleMixin, nn.Module):
         n_bands=9,
         n_filters_spat: int = 32,
         n_dim: int = 3,
-        stride_factor: int = 4,
+        stride_factor: Optional[int] = None,
         win_len: int = 250,
         heads: int = 8,
         weight_softmax: bool = True,
@@ -125,6 +135,14 @@ class FBLightConvNet(EEGModuleMixin, nn.Module):
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
 
+        if stride_factor is not None:
+            warn(
+                "The parameter `stride_factor` is deprecated and ignored, it "
+                "will be removed in a future release. The temporal "
+                "segmentation of FBLightConvNet is set by `win_len`.",
+                DeprecationWarning,
+            )
+
         # Parameters
         self.n_bands = n_bands
         self.n_filters_spat = n_filters_spat
@@ -138,6 +156,13 @@ class FBLightConvNet(EEGModuleMixin, nn.Module):
         self.filter_parameters = filter_parameters or {}
 
         # Checkers
+        if self.n_times < self.win_len:
+            raise ValueError(
+                f"Time dimension ({self.n_times}) is shorter than win_len "
+                f"({self.win_len}), so the model cannot build a single "
+                f"temporal window. Pass a longer input or lower `win_len`."
+            )
+
         self.n_times_truncated = self.n_times
         if self.n_times % self.win_len != 0:
             warn(
@@ -194,10 +219,12 @@ class FBLightConvNet(EEGModuleMixin, nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass of the FBLightConvNet model.
+
         Parameters
         ----------
         x : torch.Tensor
             Input tensor with shape (batch_size, n_chans, n_times).
+
         Returns
         -------
         torch.Tensor

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/models/summary.csv

@@ -61,3 +61,4 @@ EEGDINO,General,Classification,200,"n_chans, n_outputs, n_times",4539698,"EEGDIN
 DANCE,General,Event Detection,200,"n_chans, n_outputs, n_times, sfreq, chs_info",29813542,"DANCE(n_chans=19, n_outputs=4, n_times=6400, sfreq=200, input_window_seconds=32.0, chs_info=[{'ch_name': f'E{i+1}', 'kind': 'eeg', 'loc': numpy.random.default_rng(2024).random(12)} for i in range(19)])","Convolution,Attention/Transformer,Channel",EEG
 STEEGFormer,General,Classification,250,"n_chans, n_outputs, n_times",25305604,"STEEGFormer(n_chans=22, n_outputs=4, n_times=1000)","Attention/Transformer,Foundation Model",EEG
 MVPFormer,General,"Prediction, Embedding",512,"n_chans, n_outputs, n_times, sfreq",3356978,"MVPFormer(n_chans=22, n_outputs=4, n_times=2000, sfreq=100)","Foundation Model,Attention/Transformer",EEG
+ZUNA,General,"Classification,Embedding",256,"chs_info, n_outputs, n_times",172104100,"ZUNA(chs_info=<user>, n_outputs=4, n_times=1280)","Foundation Model,Channel,Attention/Transformer",EEG

{braindecode-1.7.0.dev183981614 → braindecode-1.8.0.dev1122}/braindecode/models/util.py

@@ -503,6 +503,15 @@ models_mandatory_parameters: list[
             ],
         },
     ),
+    (
+        "ZUNA",
+        ["chs_info", "n_outputs", "n_times"],
+        {
+            "n_times": 1280,
+            "sfreq": 256.0,
+            "input_window_seconds": 5.0,
+        },
+    ),
 ]
 
 ################################################################