braindecode 1.8.1.dev169692045__tar.gz → 1.8.1.dev169754522__tar.gz

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169754522}/NOTICE.txt

@@ -47,7 +47,6 @@ The following components are licensed under the MIT License:
 - `braindecode/models/mirepnet.py`
 - `braindecode/models/popt.py`
 - `braindecode/models/tfm_tokenizer.py`
-- `braindecode/modules/popt_modules.py`
 - `braindecode/models/brainomni.py`
 - `braindecode/modules/quantization.py`
 - `braindecode/modules/linear_attention.py`

{braindecode-1.8.1.dev169692045/braindecode.egg-info → braindecode-1.8.1.dev169754522}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev169692045
+Version: 1.8.1.dev169754522
 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>
@@ -22,8 +22,8 @@ Requires-Python: >=3.11
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
 License-File: NOTICE.txt
-Requires-Dist: torch>=2.4
-Requires-Dist: torchaudio>=2.4
+Requires-Dist: torch>=2.6
+Requires-Dist: torchaudio>=2.6
 Requires-Dist: mne>=1.11.0
 Requires-Dist: mne_bids>=0.18
 Requires-Dist: h5py
@@ -32,8 +32,8 @@ Requires-Dist: joblib
 Requires-Dist: torchinfo
 Requires-Dist: pandas
 Requires-Dist: wfdb>=4.3.1
-Requires-Dist: linear_attention_transformer
 Requires-Dist: docstring_inheritance
+Requires-Dist: einops
 Requires-Dist: pydantic>=2.0
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.4.3; extra == "moabb"
@@ -50,6 +50,7 @@ Provides-Extra: tests
 Requires-Dist: pytest<9.1; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
 Requires-Dist: pytest-xdist>=3.8.0; extra == "tests"
+Requires-Dist: pytest-timeout; extra == "tests"
 Requires-Dist: codecov; extra == "tests"
 Requires-Dist: pytest_cases; extra == "tests"
 Requires-Dist: mypy; extra == "tests"

braindecode-1.8.1.dev169754522/braindecode/__init__.py

@@ -0,0 +1,53 @@
+import importlib
+from typing import TYPE_CHECKING
+
+from .version import __version__
+
+if TYPE_CHECKING:  # type checkers and IDEs see the classes; runtime uses __getattr__
+    from .classifier import EEGClassifier
+    from .regressor import EEGRegressor
+
+__all__ = [
+    "__version__",
+    "EEGClassifier",
+    "EEGRegressor",
+]
+
+# Subpackages reachable as ``braindecode.<name>`` after a bare ``import braindecode``.
+_SUBMODULES = frozenset(
+    {
+        "augmentation",
+        "classifier",
+        "datasets",
+        "datautil",
+        "eegneuralnet",
+        "functional",
+        "models",
+        "modules",
+        "preprocessing",
+        "regressor",
+        "samplers",
+        "training",
+        "util",
+        "visualization",
+    }
+)
+
+if not TYPE_CHECKING:  # hidden from checkers so misspelled names stay errors
+
+    def __getattr__(name):
+        # PEP 562: skorch, the datasets and the models load on first use.
+        if name == "EEGClassifier":
+            from .classifier import EEGClassifier
+
+            return EEGClassifier
+        if name == "EEGRegressor":
+            from .regressor import EEGRegressor
+
+            return EEGRegressor
+        if name in _SUBMODULES:
+            return importlib.import_module(f".{name}", __name__)
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    def __dir__():
+        return sorted(set(globals()) | set(__all__) | _SUBMODULES)

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169754522}/braindecode/augmentation/base.py

@@ -9,10 +9,11 @@
 from numbers import Real
 from typing import Any, Callable, Optional, Union
 
+import numpy as np
 import torch
 from sklearn.utils import check_random_state
 from torch import Tensor, nn
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, get_worker_info
 from torch.utils.data._utils.collate import default_collate
 
 from .functional import identity
@@ -191,6 +192,59 @@ def _as_mixed_target(y, lam_dtype):
     return y, y, lam
 
 
+def _transform_rngs(transform):
+    """Return the ``RandomState`` objects of a transform, in a fixed order.
+
+    Walks the transform, its submodules and the ``transforms`` of a
+    :class:`Compose`, and lists each ``rng`` once (pooled transforms of
+    :class:`braindecode.augmentation.TrivialAugment` share theirs). The global
+    numpy generator, used with ``random_state=None``, is included: a spawned
+    worker receives a pickled copy of it, which torch does not reseed.
+    """
+    rngs, seen, stack = [], set(), [transform]
+    while stack:
+        tr = stack.pop(0)
+        if id(tr) in seen:
+            continue
+        seen.add(id(tr))
+        # Compose forwards to its children without using its inherited RNG.
+        rng = None if type(tr) is Compose else getattr(tr, "rng", None)
+        if isinstance(rng, np.random.RandomState) and id(rng) not in seen:
+            seen.add(id(rng))
+            rngs.append(rng)
+        children = list(getattr(tr, "transforms", None) or [])
+        if isinstance(tr, nn.Module):
+            children += list(tr.children())
+        stack.extend(children)
+    return rngs
+
+
+def _assert_stackable(items, name):
+    # Fixed expansion concatenates the clean originals with the transformed
+    # copies, so every copy must keep the original's shape except along the
+    # batch dimension. A transform that reshapes X or y (e.g. one returning
+    # class distributions instead of class indices) would otherwise fail deep
+    # inside torch.cat with a bare dimension-mismatch RuntimeError.
+    reference = items[0]
+    for item in items[1:]:
+        pairs = (
+            zip(reference, item)
+            if isinstance(reference, (tuple, list))
+            else [(reference, item)]
+        )
+        for ref_part, part in pairs:
+            if part.dim() != ref_part.dim() or part.shape[1:] != ref_part.shape[1:]:
+                raise ValueError(
+                    f"The transform returned {name} with shape "
+                    f"{tuple(part.shape)} instead of {tuple(ref_part.shape)}: "
+                    f"the fixed-expansion mode (n_augmentation > 0) keeps the "
+                    f"clean originals and concatenates the transformed "
+                    f"copies, so it cannot accommodate a transform that "
+                    f"changes the shape of {name}. Use n_augmentation = 0 "
+                    f"(in-place augmentation) instead."
+                )
+
+
 class _AugmentationCollate:
     """Collate that applies a transform to each batch, with optional expansion.
 
@@ -211,15 +265,44 @@ class _AugmentationCollate:
         ``(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.
+        mixing coefficient of one. In this mode a transform that changes the
+        shape of ``X`` or ``y`` is rejected with a clear
+        :class:`ValueError`, as its copies cannot be concatenated with the
+        clean originals.
     """
 
     def __init__(self, transform, device=None, n_augmentation=0):
         self.transform = transform
         self.device = device
         self.n_augmentation = n_augmentation
+        self.worker_seeds = None
+        self._worker_rngs_seeded = False
+
+    def draw_worker_seeds(self):
+        """Draw, from each generator of the transform, a seed for the workers.
+
+        Called in the main process when a multi-process iteration starts. A
+        worker receives a copy of the transform whose generators are in the
+        state of the main process; without a new seed, every worker, and every
+        epoch, would draw the same augmentations.
+        """
+        self.worker_seeds = [
+            int(rng.randint(2**31)) for rng in _transform_rngs(self.transform)
+        ]
+
+    def _seed_worker_rngs(self):
+        if self._worker_rngs_seeded or self.worker_seeds is None:
+            return
+        worker_info = get_worker_info()
+        if worker_info is None:
+            return
+        self._worker_rngs_seeded = True
+        rngs = _transform_rngs(self.transform)
+        for rng, seed in zip(rngs, self.worker_seeds):
+            rng.seed(np.random.SeedSequence([seed, worker_info.id]).generate_state(1))
 
     def __call__(self, batch):
+        self._seed_worker_rngs()
         collated = default_collate(batch)
         X, y = collated[:2]
         if self.device is not None:
@@ -239,13 +322,16 @@ class _AugmentationCollate:
             for aug_y in ys
             if isinstance(aug_y, (tuple, list)) and len(aug_y) == 3
         ]
+        _assert_stackable(xs, "X")
         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]
+            _assert_stackable(ys, "y")
             return torch.cat(xs), tuple(torch.cat(part) for part in zip(*ys))
+        _assert_stackable(ys, "y")
         return torch.cat(xs), torch.cat(ys)
 
 
@@ -268,7 +354,10 @@ class AugmentedDataLoader(DataLoader):
         the clean originals are kept and ``n_augmentation`` independently
         transformed copies are appended. This expresses augmentations defined as
         a fixed set-expansion (e.g. the EEG-Inception MI 6x training set,
-        ``n_augmentation=5``). In this mode batches are returned as ``(X, y)``.
+        ``n_augmentation=5``). In this mode batches are returned as ``(X, y)``
+        and transforms that change the shape of ``X`` or ``y`` are rejected
+        with a :class:`ValueError`, as their copies cannot be concatenated
+        with the clean originals.
     **kwargs : dict, optional
         keyword arguments to pass to standard DataLoader class.
     """
@@ -301,3 +390,10 @@ class AugmentedDataLoader(DataLoader):
             transform, device=device, n_augmentation=n_augmentation
         )
         super().__init__(dataset, collate_fn=self.collated_tr, **kwargs)
+
+    def __iter__(self):
+        # the workers get a copy of the transform; give them new seeds so that
+        # seeded transforms differ across workers and epochs
+        if self.num_workers > 0:
+            self.collated_tr.draw_worker_seeds()
+        return super().__iter__()

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169754522}/braindecode/augmentation/functional.py

@@ -1089,10 +1089,12 @@ def mixup(
         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
 
-    X_mix = torch.zeros((batch_size, n_channels, n_times)).to(device)
+    X_mix = X.new_zeros(
+        (batch_size, n_channels, n_times),
+        dtype=X.dtype if X.is_floating_point() else torch.get_default_dtype(),
+    )
     y_a = y.clone()
     y_b = y[idx_perm].clone()
 

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169754522}/braindecode/augmentation/transforms.py

@@ -1173,7 +1173,8 @@ class LabelSmoothing(Transform):
     fixed-expansion mode (``n_augmentation > 0``) of
     :class:`braindecode.augmentation.AugmentedDataLoader` is not supported:
     it cannot concatenate the unsmoothed original targets with the reshaped
-    smoothed copies.
+    smoothed copies and rejects the combination with a :class:`ValueError`
+    pointing to ``n_augmentation = 0``.
 
     Parameters
     ----------

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169754522}/braindecode/classifier.py

@@ -33,7 +33,12 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
 
     cropped: bool (default=False)
         Defines whether torch model passed to this class is cropped or not.
-        Currently used for callbacks definition.
+        Currently used for callbacks definition. For input without window
+        indices (such as numpy arrays), scoring callbacks treat each example
+        as a separate trial and average its temporal predictions. To aggregate
+        multiple windows of the same trial, pass a window dataset retaining
+        its window indices. Cross-validation splits must keep windows from
+        the same trial together.
 
     callbacks: None or list of strings or list of Callback instances (default=None)
         More callbacks, in addition to those returned by
@@ -195,7 +200,7 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
         """
         return self.predict_proba(X).argmax(1)
 
-    def predict_trials(self, X, return_targets=True):
+    def predict_trials(self, X, return_targets=True, *, loader_kwargs=None):
         """Create trialwise predictions from a cropped dataset.
 
         Parameters
@@ -204,6 +209,9 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
             A braindecode dataset to be predicted.
         return_targets : bool
             If True, additionally returns the trial targets.
+        loader_kwargs : dict | None
+            Cropped mode only: DataLoader overrides, see
+            :func:`braindecode.training.iter_trial_predictions`.
 
         Returns
         -------
@@ -220,6 +228,8 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
             trial lengths differ, a list with one target array per trial.
         """
         if not self.cropped:
+            if loader_kwargs is not None:
+                raise ValueError("loader_kwargs requires cropped=True")
             warnings.warn(
                 "This method was designed to predict trials in cropped mode. "
                 "Calling it when cropped is False will give the same result as "
@@ -237,6 +247,7 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
             return_targets=return_targets,
             batch_size=self.batch_size,
             num_workers=self.get_iterator(X, training=False).loader.num_workers,
+            loader_kwargs=loader_kwargs,
         )
 
     @property

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169754522}/braindecode/datasets/base.py

@@ -26,6 +26,7 @@ from itertools import chain
 from operator import attrgetter
 from pathlib import Path
 from typing import Any, Generic, Iterable, no_type_check
+from uuid import uuid4
 
 import mne
 import mne.io
@@ -343,6 +344,9 @@ def _read_fif_window(raw, start, stop, willneed=False):
     return np.multiply(picked, raw._cals[:, np.newaxis], order="C")
 
 
+_COALESCED_MAX_BYTES = 2 * 1024 * 1024
+
+
 def _window_info(crop_inds, sfreq):
     """Extract window size from crop indices.
 
@@ -479,14 +483,14 @@ def _zarr_to_memmap(zarr_path, group_name):
     ``SIGBUS`` when workers page-faulted on the unlinked inode.
 
     Under ``N`` concurrent writers, every racer materialises the data
-    into its own per-pid temp file and then tries to publish it; at
+    into its own per-call temp file and then tries to publish it; at
     most one ``os.link`` wins, and the losers silently ``unlink``
     their temp file in a ``finally`` block.  The wasted I/O scales
     with the number of racers; for the workloads braindecode targets
     (a handful of concurrent workers per group) it is negligible.
 
     Crash recovery: if a writer dies mid-run the only debris is a
-    ``<group>.<pid>.tmp.npy`` file.  Such files are namespaced by pid
+    ``<group>.<uuid>.tmp.npy`` file.  Such files are unique per call
     so they can never collide with a live writer, they are never
     read by the fast path, and they do not prevent subsequent
     callers from producing a correct result.  Caches written by
@@ -504,7 +508,10 @@ def _zarr_to_memmap(zarr_path, group_name):
 
     cache_dir.mkdir(parents=True, exist_ok=True)
     arr = zarr.open(zarr_path, mode="r")[group_name]["data"]
-    tmp_path = cache_dir / f"{group_name}.{os.getpid()}.tmp.npy"
+    # Unique per call, not per pid: threads of one process must not share it.
+    # Exclusive creation, not tempfile.mkstemp, which would make it 0600.
+    tmp_path = cache_dir / f"{group_name}.{uuid4().hex}.tmp.npy"
+    tmp_path.open("xb").close()
     try:
         # Write chunk-by-chunk to avoid materialising the whole zarr
         # array in RAM.
@@ -531,9 +538,9 @@ def _zarr_to_memmap(zarr_path, group_name):
             # canonical one; ours is discarded in the finally block.
             pass
     finally:
-        # Always remove our private per-pid temp file.  It is
-        # namespaced by pid so this cannot clobber a concurrent
-        # writer's in-progress tmp file.
+        # Always remove our private per-call temp file.  Its name is
+        # unique per call so this cannot clobber a concurrent writer's
+        # in-progress tmp file.
         try:
             os.unlink(tmp_path)
         except FileNotFoundError:
@@ -953,6 +960,10 @@ class EEGWindowsDataset(_ZarrMixin, RecordDataset):
                 X = self.raw._getitem(
                     (slice(None), slice(i_start, i_stop)), return_times=False
                 )
+        return self._finish_window(index, crop_inds, X)
+
+    def _finish_window(self, index, crop_inds, X):
+        """Item ``index`` from its raw samples ``X``, copied before the transform."""
         X = np.array(X, dtype="float32")
         if self.transform is not None:
             X = self.transform(X)
@@ -1344,11 +1355,24 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
         Optional function to call on targets before returning them.
     lazy : bool, default False
         If True, defer computing cumulative sizes until length or item access.
+    coalesce_reads : bool, default False
+        Opt in to reading a :class:`~torch.utils.data.DataLoader` batch of
+        consecutive, overlapping windows of one lazy (``preload=False``) FIF
+        recording as one union of samples instead of one read per window.
+        Only a batch of consecutive ascending indices into one
+        :class:`~braindecode.datasets.EEGWindowsDataset` (not a subclass)
+        whose union is at most a third of the summed window samples and at
+        most 2 MiB as float64 is read this way; every other batch takes the
+        default path. Outputs, order and random number use equal the
+        default's. It was faster only for heavily overlapping windows and
+        slower for some other layouts (numbers in :ref:`whats_new`), so it is off
+        by default. Containers returned by ``split`` default to ``False``.
     """
 
     plot = _viewer_plot  # eegdash-viewer embed; a direct member so the API docs list it
 
     datasets: list[T]
+    coalesce_reads: bool = False  # default for pickles made before this option
 
     def __init__(
         self,
@@ -1358,6 +1382,7 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
         target_transform: Callable | None = None,
         *,
         lazy: bool = False,
+        coalesce_reads: bool = False,
     ):
         # Adapted from torch.utils.data.ConcatDataset:
         # added lazy init to defer cumulative size computation until access,
@@ -1387,6 +1412,9 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
 
         self.target_transform = target_transform
         self._lazy = lazy
+        if type(coalesce_reads) is not bool:
+            raise TypeError("coalesce_reads must be a bool.")
+        self.coalesce_reads = coalesce_reads
 
     @property
     def cumulative_sizes(self) -> list[int]:
@@ -1454,6 +1482,10 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
         another: on a network file system a batch read from cold files then
         costs about one round trip, not one per window.
         """
+        if self.coalesce_reads:
+            coalesced = self._get_coalesced_items(indices)
+            if coalesced is not None:
+                return coalesced
         # ponytail: advises the whole batch at once; chunk it if a batch's
         # windows ever outgrow the page cache
         if hasattr(os, "posix_fadvise"):
@@ -1466,6 +1498,52 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
                     _read_fif_window(ds.raw, start, stop, willneed=True)
         return [self[idx] for idx in indices]
 
+    def _get_coalesced_items(self, indices):
+        """The items cut from one read of their union, or None if not eligible."""
+        n = len(indices)
+        if (
+            type(self) is not BaseConcatDataset  # subclasses may override items
+            or n < 2
+            or type(indices[0]) is not int
+            or indices[0] < 0
+            or indices != list(range(indices[0], indices[0] + n))
+        ):
+            return None
+        sizes = self._ensure_cumulative_sizes()
+        k = bisect.bisect_right(sizes, indices[0])
+        if k == len(sizes) or indices[-1] >= sizes[k]:
+            return None  # past the end, or across recordings
+        ds = self.datasets[k]
+        if (
+            type(ds) is not EEGWindowsDataset
+            or ds._zarr_data is not None
+            or not isinstance(ds.crop_inds, np.ndarray)  # lazy metadata
+            or ds.raw.preload
+        ):
+            return None
+        first = indices[0] - (sizes[k - 1] if k else 0)
+        crops = ds.crop_inds[first : first + n]
+        start, stop = int(crops[:, 1].min()), int(crops[:, 2].max())
+        if (
+            start < 0
+            or stop > ds.raw.n_times
+            or (crops[:, 2] <= crops[:, 1]).any()
+            or 3 * (stop - start) > (crops[:, 2] - crops[:, 1]).sum()
+            or (stop - start) * ds.raw.info["nchan"] * 8 > _COALESCED_MAX_BYTES
+        ):
+            return None
+        union = _read_fif_window(ds.raw, start, stop)
+        if union is None:
+            return None
+        items = []
+        for i, crop in enumerate(crops.tolist(), first):
+            X = union[:, crop[1] - start : crop[2] - start]
+            item = ds._finish_window(i, crop, X)
+            if self.target_transform is not None:
+                item = item[:1] + (self.target_transform(item[1]),) + item[2:]
+            items.append(item)
+        return items
+
     def __getitem__(self, idx: int | list):
         """
         ---

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169754522}/braindecode/datasets/bids/datasets.py

@@ -224,15 +224,39 @@ class BIDSEpochsDataset(BIDSDataset):
         )
 
     def _set_metadata(self, epochs: mne.BaseEpochs) -> None:
-        # events = mne.events_from_annotations(epochs
         n_times = epochs.times.shape[0]
-        # id_event = {v: k for k, v in epochs.event_id.items()}
-        annotations = epochs.annotations
-        if annotations is not None:
-            target = annotations.description
-        else:
-            id_events = {v: k for k, v in epochs.event_id.items()}
-            target = [id_events[event_id] for event_id in epochs.events[:, -1]]
+        id_events = {v: k for k, v in epochs.event_id.items()}
+        target = [id_events[event_id] for event_id in epochs.events[:, -1]]
+        if epochs.annotations is not None and len(epochs.annotations) > 0:
+            # Match triggers on the original recording's sample clock, even if
+            # epochs have been cropped past time zero, decimated or resampled.
+            onset_samples = np.round(
+                epochs.annotations.onset * epochs._raw_sfreq
+            ).astype(int)
+            at_zero = []
+            for sample in epochs.events[:, 0]:
+                matches = np.flatnonzero(onset_samples == sample)
+                if len(matches) == 0:
+                    # FIF stores onsets as float32. Only the quantization bins
+                    # containing this sample's half-sample boundaries can move
+                    # a valid onset to an adjacent rounded sample on readback.
+                    boundaries = np.asarray(
+                        [
+                            (sample - 0.5) / epochs._raw_sfreq,
+                            (sample + 0.5) / epochs._raw_sfreq,
+                        ],
+                        dtype=np.float32,
+                    ).astype(float)
+                    matches = np.flatnonzero(
+                        np.isin(epochs.annotations.onset, boundaries)
+                        & (np.abs(onset_samples - sample) == 1)
+                    )
+                at_zero.append(matches)
+            if all(len(matches) == 1 for matches in at_zero):
+                indices = [matches[0] for matches in at_zero]
+                # A boundary annotation cannot label two neighboring events.
+                if len(set(indices)) == len(set(epochs.events[:, 0])):
+                    target = epochs.annotations.description[indices].tolist()
         metadata_dict = {
             "i_window_in_trial": np.zeros(len(epochs)),
             "i_start_in_trial": np.zeros(len(epochs)),

{braindecode-1.8.1.dev169692045 → braindecode-1.8.1.dev169754522}/braindecode/datasets/chb_mit.py

@@ -16,12 +16,7 @@ It reorganizes the file structure to comply with the BIDS specification. To this
 # License: BSD (3-clause)
 from __future__ import annotations
 
-from pathlib import Path
-
-from mne.datasets import fetch_dataset
-
-from braindecode.datasets import BIDSDataset
-from braindecode.datasets.utils import _correct_dataset_path
+from braindecode.datasets.utils import _ZenodoBIDSDataset
 
 CHB_MIT_URL = "https://zenodo.org/records/10259996/files/BIDS_CHB-MIT.zip"
 CHB_MIT_archive_name = "chb_mit_bids.zip"
@@ -38,7 +33,7 @@ CHB_MIT_dataset_params = {
 }
 
 
-class CHBMIT(BIDSDataset):
+class CHBMIT(_ZenodoBIDSDataset):
     """The Children's Hospital Boston EEG Dataset.
 
     This database, collected at the Children's Hospital Boston, consists of EEG recordings
@@ -129,35 +124,6 @@ class CHBMIT(BIDSDataset):
     .. footbibliography::
     """
 
-    def __init__(self, root=None, *args, **kwargs):
-        # Download dataset if not present
-        if root is None:
-            path_root = fetch_dataset(
-                dataset_params=CHB_MIT_dataset_params,
-                path=None,
-                processor="unzip",
-                force_update=False,
-            )
-            # First time we fetch the dataset, we need to move the files to the
-            # correct directory.
-            path_root = _correct_dataset_path(
-                path_root, CHB_MIT_archive_name, "BIDS_CHB-MIT"
-            )
-        else:
-            # Validate that the provided root is a valid BIDS dataset
-            if not Path(f"{root}/participants.tsv").exists():
-                raise ValueError(
-                    f"The provided root directory {root} does not contain a valid "
-                    "BIDS dataset (missing participants.tsv). Please ensure the "
-                    "root points directly to the BIDS dataset directory."
-                )
-            path_root = root
-
-        kwargs["root"] = path_root
-
-        super().__init__(
-            *args,
-            extensions=".edf",
-            check=False,
-            **kwargs,
-        )
+    _dataset_params = CHB_MIT_dataset_params
+    _archive_name = CHB_MIT_archive_name
+    _unzipped_dir = "BIDS_CHB-MIT"