braindecode 1.5.0.dev1010__tar.gz → 1.5.0.dev1013__tar.gz

{braindecode-1.5.0.dev1010/braindecode.egg-info → braindecode-1.5.0.dev1013}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.5.0.dev1010
+Version: 1.5.0.dev1013
 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.5.0.dev1010 → braindecode-1.5.0.dev1013}/braindecode/augmentation/__init__.py

@@ -4,6 +4,7 @@ from . import functional
 from .base import AugmentedDataLoader, Compose, IdentityTransform, Transform
 from .transforms import (
     AmplitudeScale,
+    BandRotation,
     BandstopFilter,
     ChannelsDropout,
     ChannelsReref,
@@ -47,6 +48,7 @@ __all__ = [
     "SegmentationReconstruction",
     "MaskEncoding",
     "AmplitudeScale",
+    "BandRotation",
     "ChannelsReref",
     "functional",
 ]

{braindecode-1.5.0.dev1010 → braindecode-1.5.0.dev1013}/braindecode/augmentation/functional.py

@@ -1298,3 +1298,130 @@ def amplitude_scale(
     X = s * X
 
     return X, y
+
+
+def band_rotation(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    num_bands: int = 2,
+    electrodes_per_band: int = 16,
+    band_offsets: tuple[int, ...] = (-1, 0, 1),
+    max_temporal_jitter: int = 0,
+    circular_jitter: bool = True,
+    random_state: int | np.random.RandomState | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-band electrode rotation + inter-band temporal jitter.
+
+    Models small wristband rotation between sessions and relative timing
+    noise between two arms.  Introduced in [Sivakumar2024]_ for the
+    emg2qwerty CTC keystroke decoding task: each electrode band gets its
+    own circular roll along the channel axis (``Uniform(band_offsets)``
+    positions), and band 1 also gets a sample-level temporal shift
+    (``Uniform(-max_temporal_jitter, +max_temporal_jitter)``) along the
+    time axis.
+
+    Channel layout assumes ``(B, num_bands * electrodes_per_band, T)`` with
+    bands contiguous along the channel axis.  Same offset / shift is
+    applied to every sample in the batch (one set of parameters per call).
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        EMG input batch of shape ``(B, C, T)`` with
+        ``C == num_bands * electrodes_per_band``.
+    y : torch.Tensor
+        Labels (returned unchanged).
+    num_bands : int, optional
+        Number of electrode bands (e.g. ``2`` for left + right wristband).
+        Must be ``>= 1``.  Defaults to 2.
+    electrodes_per_band : int, optional
+        Electrodes per band (e.g. ``16``).  Must be ``>= 1``.  Defaults
+        to 16.
+    band_offsets : tuple of int, optional
+        Per-band roll values to sample from uniformly.  ``(-1, 0, 1)``
+        covers ±1-electrode misalignment.  Must be non-empty.  Defaults
+        to ``(-1, 0, 1)``.
+    max_temporal_jitter : int, optional
+        Max ±-sample temporal shift applied to band 1 only when
+        ``num_bands >= 2``.  Defaults to 0 (disabled).  Must be ``>= 0``.
+    circular_jitter : bool, optional
+        If True (the default, paper-faithful), the temporal jitter is a
+        circular ``torch.roll`` — samples shifted off one edge wrap to
+        the other.  If False, the gap left by the shift is zero-padded
+        and the shifted-off samples are dropped, avoiding wrap-around
+        discontinuity at the cost of a small zeroed margin.  Has no
+        effect when ``max_temporal_jitter == 0``.
+    random_state : int | numpy.random.RandomState, optional
+        Seed / generator for sampling rotation + jitter values.
+
+    Returns
+    -------
+    torch.Tensor
+        Transformed inputs.
+    torch.Tensor
+        Labels (unchanged).
+
+    References
+    ----------
+    .. [Sivakumar2024] Sivakumar, V., Seely, J., Du, A., Bittner, S. R.,
+       Berenzweig, A., Bolarinwa, A., Gramfort, A., & Mandel, M. I. (2024).
+       "emg2qwerty: A Large Dataset with Baselines for Touch Typing using
+       Surface Electromyography." *NeurIPS Datasets and Benchmarks Track*.
+    """
+    if num_bands < 1:
+        raise ValueError(f"num_bands must be >= 1, got {num_bands}")
+    if electrodes_per_band < 1:
+        raise ValueError(f"electrodes_per_band must be >= 1, got {electrodes_per_band}")
+    # Normalise to a tuple before truth-testing so callers can pass any
+    # sequence-like (incl. ``np.ndarray``) without hitting numpy's
+    # ambiguous-truth-value error on ``if not band_offsets``.
+    band_offsets = tuple(band_offsets)
+    if not band_offsets:
+        raise ValueError("band_offsets must be non-empty")
+    if not all(isinstance(o, (int, np.integer)) for o in band_offsets):
+        raise ValueError(f"band_offsets must contain integers, got {band_offsets!r}")
+    if max_temporal_jitter < 0:
+        raise ValueError(f"max_temporal_jitter must be >= 0, got {max_temporal_jitter}")
+    expected_channels = num_bands * electrodes_per_band
+    if X.shape[1] != expected_channels:
+        raise ValueError(
+            f"X.shape[1]={X.shape[1]} != num_bands * electrodes_per_band="
+            f"{expected_channels}"
+        )
+
+    rng = check_random_state(random_state)
+    band_offsets_arr = np.asarray(band_offsets)
+    out = X.clone()
+
+    # Per-band channel-axis rolls.  A vectorized ``torch.gather`` was
+    # benchmarked and is ~16 % slower for the typical ``num_bands == 2``
+    # case on CPU (the index tensor is larger than what two contiguous
+    # rolls touch); the gather only wins past ``num_bands >= 8``.
+    for b in range(num_bands):
+        offset = int(rng.choice(band_offsets_arr))
+        if offset:
+            sl = slice(b * electrodes_per_band, (b + 1) * electrodes_per_band)
+            out[:, sl, :] = torch.roll(out[:, sl, :], offset, dims=1)
+
+    # Inter-band temporal jitter — paper recipe applies it to band 1 only.
+    if max_temporal_jitter > 0 and num_bands >= 2:
+        shift = int(rng.randint(-max_temporal_jitter, max_temporal_jitter + 1))
+        if shift:
+            sl = slice(electrodes_per_band, 2 * electrodes_per_band)
+            band1 = out[:, sl, :]
+            if circular_jitter:
+                # Paper-faithful circular shift; wraps end-of-window
+                # samples to the start (and vice versa).
+                out[:, sl, :] = torch.roll(band1, shift, dims=2)
+            else:
+                # Crop-and-pad shift: drop samples that fall off one end,
+                # zero-pad the gap on the other.  Avoids the wrap-around
+                # discontinuity at the cost of a ``|shift|``-sample margin.
+                shifted = torch.zeros_like(band1)
+                if shift > 0:
+                    shifted[:, :, shift:] = band1[:, :, :-shift]
+                else:  # shift < 0
+                    shifted[:, :, :shift] = band1[:, :, -shift:]
+                out[:, sl, :] = shifted
+
+    return out, y

{braindecode-1.5.0.dev1010 → braindecode-1.5.0.dev1013}/braindecode/augmentation/transforms.py

@@ -16,6 +16,7 @@ from mne.channels import make_standard_montage
 from .base import Transform
 from .functional import (
     amplitude_scale,
+    band_rotation,
     bandstop_filter,
     channels_dropout,
     channels_permute,
@@ -1356,3 +1357,99 @@ class AmplitudeScale(Transform):
     def get_augmentation_params(self, *batch):
         """Return transform parameters."""
         return {"random_state": self.rng, "scale": self.scale}
+
+
+class BandRotation(Transform):
+    """Per-band electrode rotation + inter-band temporal jitter.
+
+    Models small wristband rotation between sessions and relative timing
+    noise between two arms.  Introduced in [Sivakumar2024]_ for the
+    emg2qwerty surface-EMG keystroke decoding task: the channel axis is
+    laid out as ``(B, num_bands * electrodes_per_band, T)`` with bands
+    contiguous, each band gets a uniform circular roll along the channel
+    axis, and when ``num_bands >= 2``, band 1 also gets a sample-level
+    temporal shift.  The same offset / shift is applied to every sample
+    in a transformed sub-batch (one set of parameters per call).
+
+    Parameters
+    ----------
+    probability : float
+        Float setting the probability of applying the operation.
+    num_bands : int, optional
+        Number of electrode bands (e.g. ``2`` for left + right wristband).
+        Must be ``>= 1``.  Defaults to 2.
+    electrodes_per_band : int, optional
+        Electrodes per band (e.g. ``16``).  Must be ``>= 1``.  Defaults
+        to 16.
+    band_offsets : tuple of int, optional
+        Per-band roll values to sample from uniformly.  ``(-1, 0, 1)``
+        covers ±1-electrode misalignment.  Must be non-empty.  Defaults
+        to ``(-1, 0, 1)``.
+    max_temporal_jitter : int, optional
+        Max ±-sample temporal shift applied to band 1.  Defaults to 0
+        (jitter disabled).  Must be ``>= 0``.  The emg2qwerty paper uses
+        120 samples (60 ms at 2 kHz).
+    circular_jitter : bool, optional
+        If True (default, paper-faithful) the jitter is a circular roll;
+        if False the gap left by the shift is zero-padded.  See
+        :func:`band_rotation`.
+    random_state : int | numpy.random.RandomState, optional
+        Seed for the rotation / jitter sampler.  Defaults to None.
+
+    References
+    ----------
+    .. [Sivakumar2024] Sivakumar, V., Seely, J., Du, A., Bittner, S. R.,
+       Berenzweig, A., Bolarinwa, A., Gramfort, A., & Mandel, M. I. (2024).
+       "emg2qwerty: A Large Dataset with Baselines for Touch Typing using
+       Surface Electromyography." *NeurIPS Datasets and Benchmarks Track*.
+    """
+
+    operation = staticmethod(band_rotation)  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        probability,
+        num_bands=2,
+        electrodes_per_band=16,
+        band_offsets=(-1, 0, 1),
+        max_temporal_jitter=0,
+        circular_jitter=True,
+        random_state=None,
+    ):
+        super().__init__(probability=probability, random_state=random_state)
+        # Up-front parameter validation; the underlying ``band_rotation``
+        # also re-checks at call time, but raising here surfaces config
+        # mistakes when the Transform is built rather than on the first
+        # batch.
+        if num_bands < 1:
+            raise ValueError(f"num_bands must be >= 1, got {num_bands}")
+        if electrodes_per_band < 1:
+            raise ValueError(
+                f"electrodes_per_band must be >= 1, got {electrodes_per_band}"
+            )
+        band_offsets = tuple(band_offsets)
+        if not band_offsets:
+            raise ValueError("band_offsets must be non-empty")
+        if not all(isinstance(o, (int, np.integer)) for o in band_offsets):
+            raise ValueError(
+                f"band_offsets must contain integers, got {band_offsets!r}"
+            )
+        if max_temporal_jitter < 0:
+            raise ValueError(
+                f"max_temporal_jitter must be >= 0, got {max_temporal_jitter}"
+            )
+        self.num_bands = num_bands
+        self.electrodes_per_band = electrodes_per_band
+        self.band_offsets = band_offsets
+        self.max_temporal_jitter = max_temporal_jitter
+        self.circular_jitter = circular_jitter
+
+    def get_augmentation_params(self, *batch):
+        return {
+            "num_bands": self.num_bands,
+            "electrodes_per_band": self.electrodes_per_band,
+            "band_offsets": self.band_offsets,
+            "max_temporal_jitter": self.max_temporal_jitter,
+            "circular_jitter": self.circular_jitter,
+            "random_state": self.rng,
+        }

braindecode-1.5.0.dev1013/braindecode/version.py

@@ -0,0 +1 @@
+__version__ = "1.5.0.dev1013"

{braindecode-1.5.0.dev1010 → braindecode-1.5.0.dev1013/braindecode.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.5.0.dev1010
+Version: 1.5.0.dev1013
 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.5.0.dev1010 → braindecode-1.5.0.dev1013}/docs/whats_new.rst

@@ -28,6 +28,14 @@ Current 1.5.0 (stable)
 Enhancements
 ============
 
+- Add :class:`braindecode.augmentation.BandRotation` and
+  :func:`braindecode.augmentation.functional.band_rotation`: per-band
+  circular roll along the channel axis plus inter-band temporal jitter,
+  for surface-EMG inputs shaped ``(B, num_bands * electrodes_per_band, T)``.
+  Models small wristband rotation between sessions and relative timing
+  noise between two arms, from the emg2qwerty paper (Sivakumar et al.,
+  NeurIPS 2024).  By `Bruno Aristimunha`_.
+
 - Add :meth:`braindecode.datasets.BaseConcatDataset.set_target` to swap
   any per-window metadata column or per-record description field
   (e.g. a BIDS entity, a participants.tsv extra) into the dataset's

braindecode-1.5.0.dev1010/braindecode/version.py

@@ -1 +0,0 @@
-__version__ = "1.5.0.dev1010"