braindecode 1.2.0.dev184328194__py3-none-any.whl → 1.3.0.dev168011974__py3-none-any.whl

braindecode/augmentation/functional.py

@@ -4,9 +4,13 @@
 #
 # License: BSD (3-clause)
 
+from __future__ import annotations
+
 from numbers import Real
+from typing import Literal
 
 import numpy as np
+import numpy.typing as npt
 import torch
 from mne.filter import notch_filter
 from scipy.interpolate import Rbf
@@ -15,7 +19,7 @@ from torch.fft import fft, ifft
 from torch.nn.functional import one_hot, pad
 
 
-def identity(X, y):
+def identity(X: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """Identity operation.
 
     Parameters
@@ -35,7 +39,7 @@ def identity(X, y):
     return X, y
 
 
-def time_reverse(X, y):
+def time_reverse(X: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """Flip the time axis of each input.
 
     Parameters
@@ -55,7 +59,7 @@ def time_reverse(X, y):
     return torch.flip(X, [-1]), y
 
 
-def sign_flip(X, y):
+def sign_flip(X: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """Flip the sign axis of each input.
 
     Parameters
@@ -75,7 +79,13 @@ def sign_flip(X, y):
     return -X, y
 
 
-def _new_random_fft_phase_odd(batch_size, c, n, device, random_state):
+def _new_random_fft_phase_odd(
+    batch_size: int,
+    c: int,
+    n: int,
+    device: torch.device,
+    random_state: int | np.random.RandomState | None,
+) -> torch.Tensor:
     rng = check_random_state(random_state)
     random_phase = torch.from_numpy(
         2j * np.pi * rng.random((batch_size, c, (n - 1) // 2))
@@ -90,7 +100,13 @@ def _new_random_fft_phase_odd(batch_size, c, n, device, random_state):
     )
 
 
-def _new_random_fft_phase_even(batch_size, c, n, device, random_state):
+def _new_random_fft_phase_even(
+    batch_size: int,
+    c: int,
+    n: int,
+    device: torch.device,
+    random_state: int | np.random.RandomState | None,
+) -> torch.Tensor:
     rng = check_random_state(random_state)
     random_phase = torch.from_numpy(
         2j * np.pi * rng.random((batch_size, c, n // 2 - 1))
@@ -109,7 +125,13 @@ def _new_random_fft_phase_even(batch_size, c, n, device, random_state):
 _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):
+def ft_surrogate(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    phase_noise_magnitude: float,
+    channel_indep: bool,
+    random_state: int | np.random.RandomState | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """FT surrogate augmentation of a single EEG channel, as proposed in [1]_.
 
     Function copied from https://github.com/cliffordlab/sleep-convolutions-tf
@@ -175,7 +197,9 @@ def ft_surrogate(X, y, phase_noise_magnitude, channel_indep, random_state=None):
     return transformed_X, y
 
 
-def _pick_channels_randomly(X, p_pick, random_state):
+def _pick_channels_randomly(
+    X: torch.Tensor, p_pick: float, random_state: int | np.random.RandomState | None
+) -> torch.Tensor:
     rng = check_random_state(random_state)
     batch_size, n_channels, _ = X.shape
     # allows to use the same RNG
@@ -188,7 +212,12 @@ def _pick_channels_randomly(X, p_pick, random_state):
     return torch.sigmoid(1000 * (unif_samples - p_pick))
 
 
-def channels_dropout(X, y, p_drop, random_state=None):
+def channels_dropout(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    p_drop: float,
+    random_state: int | np.random.RandomState | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Randomly set channels to flat signal.
 
     Part of the CMSAugment policy proposed in [1]_
@@ -222,7 +251,9 @@ def channels_dropout(X, y, p_drop, random_state=None):
     return X * mask.unsqueeze(-1), y
 
 
-def _make_permutation_matrix(X, mask, random_state):
+def _make_permutation_matrix(
+    X: torch.Tensor, mask: torch.Tensor, random_state: int | np.random.Generator | None
+) -> torch.Tensor:
     rng = check_random_state(random_state)
     batch_size, n_channels, _ = X.shape
     hard_mask = mask.round()
@@ -241,7 +272,12 @@ def _make_permutation_matrix(X, mask, random_state):
     return batch_permutations
 
 
-def channels_shuffle(X, y, p_shuffle, random_state=None):
+def channels_shuffle(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    p_shuffle: float,
+    random_state: int | np.random.RandomState | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Randomly shuffle channels in EEG data matrix.
 
     Part of the CMSAugment policy proposed in [1]_
@@ -280,7 +316,12 @@ def channels_shuffle(X, y, p_shuffle, random_state=None):
     return torch.matmul(batch_permutations, X), y
 
 
-def gaussian_noise(X, y, std, random_state=None):
+def gaussian_noise(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    std: float,
+    random_state: int | np.random.RandomState | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Randomly add white Gaussian noise to all channels.
 
     Suggested e.g. in [1]_, [2]_ and [3]_
@@ -332,7 +373,9 @@ def gaussian_noise(X, y, std, random_state=None):
     return transformed_X, y
 
 
-def channels_permute(X, y, permutation):
+def channels_permute(
+    X: torch.Tensor, y: torch.Tensor, permutation: list[int]
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Permute EEG channels according to fixed permutation matrix.
 
     Suggested e.g. in [1]_
@@ -362,7 +405,12 @@ def channels_permute(X, y, permutation):
     return X[..., permutation, :], y
 
 
-def smooth_time_mask(X, y, mask_start_per_sample, mask_len_samples):
+def smooth_time_mask(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    mask_start_per_sample: torch.Tensor,
+    mask_len_samples: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Smoothly replace a contiguous part of all channels by zeros.
 
     Originally proposed in [1]_ and [2]_
@@ -412,7 +460,13 @@ def smooth_time_mask(X, y, mask_start_per_sample, mask_len_samples):
     return X * mask, y
 
 
-def bandstop_filter(X, y, sfreq, bandwidth, freqs_to_notch):
+def bandstop_filter(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    sfreq: float,
+    bandwidth: float,
+    freqs_to_notch: npt.ArrayLike | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Apply a band-stop filter with desired bandwidth at the desired frequency
     position.
 
@@ -451,7 +505,7 @@ def bandstop_filter(X, y, sfreq, bandwidth, freqs_to_notch):
        Representation Learning for Electroencephalogram Classification. In
        Machine Learning for Health (pp. 238-253). PMLR.
     """
-    if bandwidth == 0:
+    if bandwidth == 0 or freqs_to_notch is None:
         return X, y
     transformed_X = X.clone()
     for c, (sample, notched_freq) in enumerate(zip(transformed_X, freqs_to_notch)):
@@ -469,7 +523,7 @@ def bandstop_filter(X, y, sfreq, bandwidth, freqs_to_notch):
     return transformed_X, y
 
 
-def _analytic_transform(x):
+def _analytic_transform(x: torch.Tensor) -> torch.Tensor:
     if torch.is_complex(x):
         raise ValueError("x must be real.")
 
@@ -486,12 +540,12 @@ def _analytic_transform(x):
     return ifft(f * h, dim=-1)
 
 
-def _nextpow2(n):
+def _nextpow2(n: int) -> int:
     """Return the first integer N such that 2**N >= abs(n)."""
     return int(np.ceil(np.log2(np.abs(n))))
 
 
-def _frequency_shift(X, fs, f_shift):
+def _frequency_shift(X: torch.Tensor, fs: float, f_shift: float) -> torch.Tensor:
     """
     Shift the specified signal by the specified frequency.
 
@@ -504,9 +558,13 @@ def _frequency_shift(X, fs, f_shift):
     t = torch.arange(N_padded, device=X.device) / fs
     padded = pad(X, (0, N_padded - N_orig))
     analytical = _analytic_transform(padded)
-    if isinstance(f_shift, (float, int, np.ndarray, list)):
-        f_shift = torch.as_tensor(f_shift).float()
-    f_shift_stack = f_shift.repeat(N_padded, n_channels, 1)
+    if isinstance(f_shift, torch.Tensor):
+        _f_shift = f_shift
+    elif isinstance(f_shift, (float, int, np.ndarray, list)):
+        _f_shift = torch.as_tensor(f_shift).float()
+    else:
+        raise ValueError(f"Invalid f_shift type: {type(f_shift)}")
+    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)
     )
@@ -514,7 +572,9 @@ def _frequency_shift(X, fs, f_shift):
     return shifted[..., :N_orig].real.float()
 
 
-def frequency_shift(X, y, delta_freq, sfreq):
+def frequency_shift(
+    X: torch.Tensor, y: torch.Tensor, delta_freq: float, sfreq: float
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Adds a shift in the frequency domain to all channels.
 
     Note that here, the shift is the same for all channels of a single example.
@@ -545,7 +605,7 @@ def frequency_shift(X, y, delta_freq, sfreq):
     return transformed_X, y
 
 
-def _torch_normalize_vectors(rr):
+def _torch_normalize_vectors(rr: torch.Tensor) -> torch.Tensor:
     """Normalize surface vertices."""
     norm = torch.linalg.norm(rr, axis=1, keepdim=True)
     mask = norm > 0
@@ -554,7 +614,9 @@ def _torch_normalize_vectors(rr):
     return new_rr
 
 
-def _torch_legval(x, c, tensor=True):
+def _torch_legval(
+    x: torch.Tensor, c: torch.Tensor, tensor: bool = True
+) -> torch.Tensor:
     """
     Evaluate a Legendre series at points x.
     If `c` is of length `n + 1`, this function returns the value:
@@ -662,7 +724,9 @@ def _torch_legval(x, c, tensor=True):
     return c0 + c1 * x
 
 
-def _torch_calc_g(cosang, stiffness=4, n_legendre_terms=50):
+def _torch_calc_g(
+    cosang: torch.Tensor, stiffness: float = 4, n_legendre_terms: int = 50
+) -> torch.Tensor:
     """Calculate spherical spline g function between points on a sphere.
 
     Parameters
@@ -718,23 +782,25 @@ def _torch_calc_g(cosang, stiffness=4, n_legendre_terms=50):
     return _torch_legval(cosang, [0] + factors)
 
 
-def _torch_make_interpolation_matrix(pos_from, pos_to, alpha=1e-5):
+def _torch_make_interpolation_matrix(
+    pos_from: torch.Tensor, pos_to: torch.Tensor, alpha: float = 1e-5
+) -> torch.Tensor:
     """Compute interpolation matrix based on spherical splines.
 
     Implementation based on [1]_
 
     Parameters
     ----------
-    pos_from : np.ndarray of float, shape(n_good_sensors, 3)
+    pos_from : torch.Tensor of float, shape(n_good_sensors, 3)
         The positions to interpolate from.
-    pos_to : np.ndarray of float, shape(n_bad_sensors, 3)
+    pos_to : torch.Tensor of float, shape(n_bad_sensors, 3)
         The positions to interpolate.
     alpha : float
         Regularization parameter. Defaults to 1e-5.
 
     Returns
     -------
-    interpolation : np.ndarray of float, shape(len(pos_from), len(pos_to))
+    interpolation : torch.Tensor of float, shape(len(pos_from), len(pos_to))
         The interpolation matrix that maps good signals to the location
         of bad signals.
 
@@ -822,7 +888,12 @@ def _torch_make_interpolation_matrix(pos_from, pos_to, alpha=1e-5):
     return interpolation
 
 
-def _rotate_signals(X, rotations, sensors_positions_matrix, spherical=True):
+def _rotate_signals(
+    X: torch.Tensor,
+    rotations: list[torch.Tensor],
+    sensors_positions_matrix: torch.Tensor,
+    spherical: bool = True,
+) -> torch.Tensor:
     sensors_positions_matrix = sensors_positions_matrix.to(X.device)
     rot_sensors_matrices = [
         rotation.matmul(sensors_positions_matrix) for rotation in rotations
@@ -853,22 +924,29 @@ def _rotate_signals(X, rotations, sensors_positions_matrix, spherical=True):
         return transformed_X
 
 
-def _make_rotation_matrix(axis, angle, degrees=True):
+def _make_rotation_matrix(
+    axis: Literal["x", "y", "z"],
+    angle: float | int | np.ndarray | list | torch.Tensor,
+    degrees: bool = True,
+) -> torch.Tensor:
     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)
+    if isinstance(angle, torch.Tensor):
+        _angle = angle
+    elif isinstance(angle, (float, int, np.ndarray, list)):
+        _angle = torch.as_tensor(angle)
+    else:
+        raise ValueError(f"Invalid angle type: {type(angle)}")
 
     if degrees:
-        angle = angle * np.pi / 180
+        _angle = _angle * np.pi / 180
 
-    device = angle.device
+    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)]),
+            torch.hstack([zero, torch.cos(_angle), -torch.sin(_angle)]),
+            torch.hstack([zero, torch.sin(_angle), torch.cos(_angle)]),
         ]
     )
     if axis == "x":
@@ -881,7 +959,14 @@ 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: torch.Tensor,
+    y: torch.Tensor,
+    sensors_positions_matrix: torch.Tensor,
+    axis: Literal["x", "y", "z"],
+    angles: npt.ArrayLike,
+    spherical_splines: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Interpolates EEG signals over sensors rotated around the desired axis
     with the desired angle.
 
@@ -893,7 +978,7 @@ def sensors_rotation(X, y, sensors_positions_matrix, axis, angles, spherical_spl
         EEG input example or batch.
     y : torch.Tensor
         EEG labels for the example or batch.
-    sensors_positions_matrix : numpy.ndarray
+    sensors_positions_matrix : torch.Tensor
         Matrix giving the positions of each sensor in a 3D cartesian coordinate
         system. Should have shape (3, n_channels), where n_channels is the
         number of channels. Standard 10-20 positions can be obtained from
@@ -924,7 +1009,9 @@ def sensors_rotation(X, y, sensors_positions_matrix, axis, angles, spherical_spl
     return rotated_X, y
 
 
-def mixup(X, y, lam, idx_perm):
+def mixup(
+    X: torch.Tensor, y: torch.Tensor, lam: torch.Tensor, idx_perm: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Mixes two channels of EEG data.
 
     See [1]_ for details.
@@ -973,8 +1060,13 @@ def mixup(X, y, lam, idx_perm):
 
 
 def segmentation_reconstruction(
-    X, y, n_segments, data_classes, rand_indices, idx_shuffle
-):
+    X: torch.Tensor,
+    y: torch.Tensor,
+    n_segments: int,
+    data_classes: list[tuple[int, torch.Tensor]],
+    rand_indices: npt.ArrayLike,
+    idx_shuffle: npt.ArrayLike,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Segment and reconstruct EEG data from [1]_.
 
     See [1]_ for details.
@@ -987,6 +1079,8 @@ def segmentation_reconstruction(
         EEG labels for the example or batch.
     n_segments : int
         Number of segments to use in the batch.
+    data_classes: list[tuple[int, torch.Tensor]]
+        List of tuples. Each tuple contains the class index and the corresponding EEG data.
     rand_indices: array-like
         Array of indices that indicates which trial to use in each segment.
     idx_shuffle: array-like
@@ -1005,8 +1099,8 @@ def segmentation_reconstruction(
     """
 
     # Initialize lists to store augmented data and corresponding labels
-    aug_data = []
-    aug_label = []
+    aug_data: list[torch.Tensor] = []
+    aug_label: list[torch.Tensor] = []
 
     # Iterate through each class to separate and augment data
     for class_index, X_class in data_classes:
@@ -1030,20 +1124,26 @@ def segmentation_reconstruction(
         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]
+    concat_aug_data = torch.cat(aug_data, dim=0)
+    concat_aug_data = concat_aug_data.to(dtype=X.dtype, device=X.device)
+    concat_aug_data = concat_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
+        concat_label = torch.cat(aug_label, dim=0)
+        concat_label = concat_label.to(dtype=y.dtype, device=y.device)
+        concat_label = concat_label[idx_shuffle]
+        return concat_aug_data, concat_label
 
-    return aug_data, y
+    return concat_aug_data, None
 
 
-def mask_encoding(X, y, time_start, segment_length, n_segments):
+def mask_encoding(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    time_start: torch.Tensor,
+    segment_length: int,
+    n_segments: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Mark encoding from Ding et al. (2024) from [ding2024]_.
 
     Replaces a contiguous part (or parts) of all channels by zeros

braindecode/augmentation/transforms.py

@@ -161,10 +161,10 @@ class ChannelsDropout(Transform):
     ----------
     probability: float
         Float setting the probability of applying the operation.
-    proba_drop: float | None, optional
+    p_drop: float | None, optional
         Float between 0 and 1 setting the probability of dropping each channel.
         Defaults to 0.2.
-    random_state: int | numpy.random.Generator, optional
+    random_state: int | numpy.random.RandomState, optional
         Seed to be used to instantiate numpy random number generator instance.
         Used to decide whether or not to transform given the probability
         argument and to sample channels to erase. Defaults to None.

braindecode/datasets/experimental.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import Callable, Sequence
+
+import mne_bids
+from torch.utils.data import IterableDataset, get_worker_info
+
+
+class BIDSIterableDataset(IterableDataset):
+    """Dataset for loading BIDS.
+
+    .. warning::
+        This class is experimental and may change in the future.
+
+    .. warning::
+        This dataset is not consistent with the Braindecode API.
+
+    This class has the same parameters as the :func:`mne_bids.find_matching_paths` function
+    as it will be used to find the files to load. The default ``extensions`` parameter was changed.
+
+    More information on BIDS (Brain Imaging Data Structure)
+    can be found at https://bids.neuroimaging.io
+
+    Examples
+    --------
+    >>> from braindecode.datasets import BaseDataset, BaseConcatDataset
+    >>> from braindecode.datasets.bids import BIDSIterableDataset, _description_from_bids_path
+    >>> from braindecode.preprocessing import create_fixed_length_windows
+    >>>
+    >>> def my_reader_fn(path):
+    ...     raw = mne_bids.read_raw_bids(path)
+    ...     desc = _description_from_bids_path(path)
+    ...     ds = BaseDataset(raw, description=desc)
+    ...     windows_ds = create_fixed_length_windows(
+    ...         BaseConcatDataset([ds]),
+    ...         window_size_samples=400,
+    ...         window_stride_samples=200,
+    ...     )
+    ...     return windows_ds
+    >>>
+    >>> dataset = BIDSIterableDataset(
+    ...     reader_fn=my_reader_fn,
+    ...     root="root/of/my/bids/dataset/",
+    ... )
+
+    Parameters
+    ----------
+    reader_fn : Callable[[mne_bids.BIDSPath], Sequence]
+        A function that takes a BIDSPath and returns a dataset.
+    pool_size : int
+        The number of recordings to read and sample from.
+    bids_paths : list[mne_bids.BIDSPath] | None
+        A list of BIDSPaths to load. If None, will use the paths found by
+        :func:`mne_bids.find_matching_paths` and the arguments below.
+    root : pathlib.Path | str
+        The root of the BIDS path.
+    subjects : str | array-like of str | None
+        The subject ID. Corresponds to "sub".
+    sessions : str | array-like of str | None
+        The acquisition session. Corresponds to "ses".
+    tasks : str | array-like of str | None
+        The experimental task. Corresponds to "task".
+    acquisitions: str | array-like of str | None
+        The acquisition parameters. Corresponds to "acq".
+    runs : str | array-like of str | None
+        The run number. Corresponds to "run".
+    processings : str | array-like of str | None
+        The processing label. Corresponds to "proc".
+    recordings : str | array-like of str | None
+        The recording name. Corresponds to "rec".
+    spaces : str | array-like of str | None
+        The coordinate space for anatomical and sensor location
+        files (e.g., ``*_electrodes.tsv``, ``*_markers.mrk``).
+        Corresponds to "space".
+        Note that valid values for ``space`` must come from a list
+        of BIDS keywords as described in the BIDS specification.
+    splits : str | array-like of str | None
+        The split of the continuous recording file for ``.fif`` data.
+        Corresponds to "split".
+    descriptions : str | array-like of str | None
+        This corresponds to the BIDS entity ``desc``. It is used to provide
+        additional information for derivative data, e.g., preprocessed data
+        may be assigned ``description='cleaned'``.
+    suffixes : str | array-like of str | None
+        The filename suffix. This is the entity after the
+        last ``_`` before the extension. E.g., ``'channels'``.
+        The following filename suffix's are accepted:
+        'meg', 'markers', 'eeg', 'ieeg', 'T1w',
+        'participants', 'scans', 'electrodes', 'coordsystem',
+        'channels', 'events', 'headshape', 'digitizer',
+        'beh', 'physio', 'stim'
+    extensions : str | array-like of str | None
+        The extension of the filename. E.g., ``'.json'``.
+        By default, uses the ones accepted by :func:`mne_bids.read_raw_bids`.
+    datatypes : str | array-like of str | None
+        The BIDS data type, e.g., ``'anat'``, ``'func'``, ``'eeg'``, ``'meg'``,
+        ``'ieeg'``.
+    check : bool
+        If ``True``, only returns paths that conform to BIDS. If ``False``
+        (default), the ``.check`` attribute of the returned
+        :class:`mne_bids.BIDSPath` object will be set to ``True`` for paths that
+        do conform to BIDS, and to ``False`` for those that don't.
+    preload : bool
+        If True, preload the data. Defaults to False.
+    n_jobs : int
+        Number of jobs to run in parallel. Defaults to 1.
+    """
+
+    def __init__(
+        self,
+        reader_fn: Callable[[mne_bids.BIDSPath], Sequence],
+        pool_size: int = 4,
+        bids_paths: list[mne_bids.BIDSPath] | None = None,
+        root: Path | str | None = None,
+        subjects: str | list[str] | None = None,
+        sessions: str | list[str] | None = None,
+        tasks: str | list[str] | None = None,
+        acquisitions: str | list[str] | None = None,
+        runs: str | list[str] | None = None,
+        processings: str | list[str] | None = None,
+        recordings: str | list[str] | None = None,
+        spaces: str | list[str] | None = None,
+        splits: str | list[str] | None = None,
+        descriptions: str | list[str] | None = None,
+        suffixes: str | list[str] | None = None,
+        extensions: str | list[str] | None = [
+            ".con",
+            ".sqd",
+            ".pdf",
+            ".fif",
+            ".ds",
+            ".vhdr",
+            ".set",
+            ".edf",
+            ".bdf",
+            ".EDF",
+            ".snirf",
+            ".cdt",
+            ".mef",
+            ".nwb",
+        ],
+        datatypes: str | list[str] | None = None,
+        check: bool = False,
+    ):
+        if bids_paths is None:
+            bids_paths = mne_bids.find_matching_paths(
+                root=root,
+                subjects=subjects,
+                sessions=sessions,
+                tasks=tasks,
+                acquisitions=acquisitions,
+                runs=runs,
+                processings=processings,
+                recordings=recordings,
+                spaces=spaces,
+                splits=splits,
+                descriptions=descriptions,
+                suffixes=suffixes,
+                extensions=extensions,
+                datatypes=datatypes,
+                check=check,
+                ignore_json=True,
+            )
+            # Filter out _epo.fif files:
+            bids_paths = [
+                bids_path
+                for bids_path in bids_paths
+                if not (bids_path.suffix == "epo" and bids_path.extension == ".fif")
+            ]
+        self.bids_paths = bids_paths
+        self.reader_fn = reader_fn
+        self.pool_size = pool_size
+
+    def __add__(self, other):
+        assert isinstance(other, BIDSIterableDataset)
+        return BIDSIterableDataset(
+            reader_fn=self.reader_fn,
+            bids_paths=self.bids_paths + other.bids_paths,
+            pool_size=self.pool_size,
+        )
+
+    def __iadd__(self, other):
+        assert isinstance(other, BIDSIterableDataset)
+        self.bids_paths += other.bids_paths
+        return self
+
+    def __iter__(self):
+        worker_info = get_worker_info()
+        if worker_info is None:  # single-process data loading, return the full iterator
+            bids_paths = self.bids_paths
+        else:  # in a worker process
+            # split workload
+            bids_paths = self.bids_paths[worker_info.id :: worker_info.num_workers]
+
+        pool = []
+        end = False
+        paths_it = iter(random.sample(bids_paths, k=len(bids_paths)))
+        while not (end and len(pool) == 0):
+            while not end and len(pool) < self.pool_size:
+                try:
+                    bids_path = next(paths_it)
+                    ds = self.reader_fn(bids_path)
+                    if ds is None:
+                        print(f"Skipping {bids_path} as it is too short.")
+                        continue
+                    idx = iter(random.sample(range(len(ds)), k=len(ds)))
+                    pool.append((ds, idx))
+                except StopIteration:
+                    end = True
+            i_pool = random.randint(0, len(pool) - 1)
+            ds, idx = pool[i_pool]
+            try:
+                i_ds = next(idx)
+                yield ds[i_ds]
+            except StopIteration:
+                pool.pop(i_pool)