eegdash 0.3.3.dev61__py3-none-any.whl → 0.5.0.dev180784713__py3-none-any.whl

eegdash/const.py

@@ -0,0 +1,349 @@
+# Authors: The EEGDash contributors.
+# License: BSD-3-Clause
+# Copyright the EEGDash contributors.
+
+"""Configuration constants and mappings for EEGDash.
+
+This module contains global configuration settings, allowed query fields, and mapping
+constants used throughout the EEGDash package. It defines the interface between EEGDash
+releases and OpenNeuro dataset identifiers, as well as validation rules for database queries.
+"""
+
+__all__ = [
+    "config",
+    "ALLOWED_QUERY_FIELDS",
+    "RELEASE_TO_OPENNEURO_DATASET_MAP",
+    "SUBJECT_MINI_RELEASE_MAP",
+]
+
+ALLOWED_QUERY_FIELDS = {
+    "data_name",
+    "dataset",
+    "subject",
+    "task",
+    "session",
+    "run",
+    "modality",
+    "sampling_frequency",
+    "nchans",
+    "ntimes",
+}
+"""set: A set of field names that are permitted in database queries constructed
+via :func:`~eegdash.api.EEGDash.find` with keyword arguments."""
+
+RELEASE_TO_OPENNEURO_DATASET_MAP = {
+    "R11": "ds005516",
+    "R10": "ds005515",
+    "R9": "ds005514",
+    "R8": "ds005512",
+    "R7": "ds005511",
+    "R6": "ds005510",
+    "R4": "ds005508",
+    "R5": "ds005509",
+    "R3": "ds005507",
+    "R2": "ds005506",
+    "R1": "ds005505",
+}
+"""dict: A mapping from Healthy Brain Network (HBN) release identifiers (e.g., "R11")
+to their corresponding OpenNeuro dataset identifiers (e.g., "ds005516")."""
+
+SUBJECT_MINI_RELEASE_MAP = {
+    "R11": [
+        "NDARAB678VYW",
+        "NDARAG788YV9",
+        "NDARAM946HJE",
+        "NDARAY977BZT",
+        "NDARAZ532KK0",
+        "NDARCE912ZXW",
+        "NDARCM214WFE",
+        "NDARDL033XRG",
+        "NDARDT889RT9",
+        "NDARDZ794ZVP",
+        "NDAREV869CPW",
+        "NDARFN221WW5",
+        "NDARFV289RKB",
+        "NDARFY623ZTE",
+        "NDARGA890MKA",
+        "NDARHN206XY3",
+        "NDARHP518FUR",
+        "NDARJL292RYV",
+        "NDARKM199DXW",
+        "NDARKW236TN7",
+    ],
+    "R10": [
+        "NDARAR935TGZ",
+        "NDARAV474ADJ",
+        "NDARCB869VM8",
+        "NDARCJ667UPL",
+        "NDARCM677TC1",
+        "NDARET671FTC",
+        "NDARKM061NHZ",
+        "NDARLD501HDK",
+        "NDARLL176DJR",
+        "NDARMT791WDH",
+        "NDARMW299ZAB",
+        "NDARNC405WJA",
+        "NDARNP962TJK",
+        "NDARPB967KU7",
+        "NDARRU560AGK",
+        "NDARTB173LY2",
+        "NDARUW377KAE",
+        "NDARVH565FX9",
+        "NDARVP799KGY",
+        "NDARVY962GB5",
+    ],
+    "R9": [
+        "NDARAC589YMB",
+        "NDARAC853CR6",
+        "NDARAH239PGG",
+        "NDARAL897CYV",
+        "NDARAN160GUF",
+        "NDARAP049KXJ",
+        "NDARAP457WB5",
+        "NDARAW216PM7",
+        "NDARBA004KBT",
+        "NDARBD328NUQ",
+        "NDARBF042LDM",
+        "NDARBH019KPD",
+        "NDARBH728DFK",
+        "NDARBM370JCB",
+        "NDARBU183TDJ",
+        "NDARBW971DCW",
+        "NDARBZ444ZHK",
+        "NDARCC620ZFT",
+        "NDARCD182XT1",
+        "NDARCK113CJM",
+    ],
+    "R8": [
+        "NDARAB514MAJ",
+        "NDARAD571FLB",
+        "NDARAF003VCL",
+        "NDARAG191AE8",
+        "NDARAJ977PRJ",
+        "NDARAP912JK3",
+        "NDARAV454VF0",
+        "NDARAY298THW",
+        "NDARBJ375VP4",
+        "NDARBT436PMT",
+        "NDARBV630BK6",
+        "NDARCB627KDN",
+        "NDARCC059WTH",
+        "NDARCM953HKD",
+        "NDARCN681CXW",
+        "NDARCT889DMB",
+        "NDARDJ204EPU",
+        "NDARDJ544BU5",
+        "NDARDP292DVC",
+        "NDARDW178AC6",
+    ],
+    "R7": [
+        "NDARAY475AKD",
+        "NDARBW026UGE",
+        "NDARCK162REX",
+        "NDARCK481KRH",
+        "NDARCV378MMX",
+        "NDARCX462NVA",
+        "NDARDJ970ELG",
+        "NDARDU617ZW1",
+        "NDAREM609ZXW",
+        "NDAREW074ZM2",
+        "NDARFE555KXB",
+        "NDARFT176NJP",
+        "NDARGK442YHH",
+        "NDARGM439FZD",
+        "NDARGT634DUJ",
+        "NDARHE283KZN",
+        "NDARHG260BM9",
+        "NDARHL684WYU",
+        "NDARHN224TPA",
+        "NDARHP841RMR",
+    ],
+    "R6": [
+        "NDARAD224CRB",
+        "NDARAE301XTM",
+        "NDARAT680GJA",
+        "NDARCA578CEB",
+        "NDARDZ147ETZ",
+        "NDARFL793LDE",
+        "NDARFX710UZA",
+        "NDARGE994BMX",
+        "NDARGP191YHN",
+        "NDARGV436PFT",
+        "NDARHF545HFW",
+        "NDARHP039DBU",
+        "NDARHT774ZK1",
+        "NDARJA830BYV",
+        "NDARKB614KGY",
+        "NDARKM250ET5",
+        "NDARKZ085UKQ",
+        "NDARLB581AXF",
+        "NDARNJ899HW7",
+        "NDARRZ606EDP",
+    ],
+    "R4": [
+        "NDARAC350BZ0",
+        "NDARAD615WLJ",
+        "NDARAG584XLU",
+        "NDARAH503YG1",
+        "NDARAX272ZJL",
+        "NDARAY461TZZ",
+        "NDARBC734UVY",
+        "NDARBL444FBA",
+        "NDARBT640EBN",
+        "NDARBU098PJT",
+        "NDARBU928LV0",
+        "NDARBV059CGE",
+        "NDARCG037CX4",
+        "NDARCG947ZC0",
+        "NDARCH001CN2",
+        "NDARCU001ZN7",
+        "NDARCW497XW2",
+        "NDARCX053GU5",
+        "NDARDF568GL5",
+        "NDARDJ092YKH",
+    ],
+    "R5": [
+        "NDARAH793FBF",
+        "NDARAJ689BVN",
+        "NDARAP785CTE",
+        "NDARAU708TL8",
+        "NDARBE091BGD",
+        "NDARBE103DHM",
+        "NDARBF851NH6",
+        "NDARBH228RDW",
+        "NDARBJ674TVU",
+        "NDARBM433VER",
+        "NDARCA740UC8",
+        "NDARCU633GCZ",
+        "NDARCU736GZ1",
+        "NDARCU744XWL",
+        "NDARDC843HHM",
+        "NDARDH086ZKK",
+        "NDARDL305BT8",
+        "NDARDU853XZ6",
+        "NDARDV245WJG",
+        "NDAREC480KFA",
+    ],
+    "R3": [
+        "NDARAA948VFH",
+        "NDARAD774HAZ",
+        "NDARAE828CML",
+        "NDARAG340ERT",
+        "NDARBA839HLG",
+        "NDARBE641DGZ",
+        "NDARBG574KF4",
+        "NDARBM642JFT",
+        "NDARCL016NHB",
+        "NDARCV944JA6",
+        "NDARCY178KJP",
+        "NDARDY150ZP9",
+        "NDAREC542MH3",
+        "NDAREK549XUQ",
+        "NDAREM887YY8",
+        "NDARFA815FXE",
+        "NDARFF644ZGD",
+        "NDARFV557XAA",
+        "NDARFV780ABD",
+        "NDARGB102NWJ",
+    ],
+    "R2": [
+        "NDARAB793GL3",
+        "NDARAM675UR8",
+        "NDARBM839WR5",
+        "NDARBU730PN8",
+        "NDARCT974NAJ",
+        "NDARCW933FD5",
+        "NDARCZ770BRG",
+        "NDARDW741HCF",
+        "NDARDZ058NZN",
+        "NDAREC377AU2",
+        "NDAREM500WWH",
+        "NDAREV527ZRF",
+        "NDAREV601CE7",
+        "NDARFF070XHV",
+        "NDARFR108JNB",
+        "NDARFT305CG1",
+        "NDARGA056TMW",
+        "NDARGH775KF5",
+        "NDARGJ878ZP4",
+        "NDARHA387FPM",
+    ],
+    "R1": [
+        "NDARAC904DMU",
+        "NDARAM704GKZ",
+        "NDARAP359UM6",
+        "NDARBD879MBX",
+        "NDARBH024NH2",
+        "NDARBK082PDD",
+        "NDARCA153NKE",
+        "NDARCE721YB5",
+        "NDARCJ594BWQ",
+        "NDARCN669XPR",
+        "NDARCW094JCG",
+        "NDARCZ947WU5",
+        "NDARDH670PXH",
+        "NDARDL511UND",
+        "NDARDU986RBM",
+        "NDAREM731BYM",
+        "NDAREN519BLJ",
+        "NDARFK610GY5",
+        "NDARFT581ZW5",
+        "NDARFW972KFQ",
+    ],
+}
+"""dict: A mapping from HBN release identifiers to a list of subject IDs.
+This is used to select a small, representative subset of subjects for creating
+"mini" datasets for testing and demonstration purposes."""
+
+config = {
+    "required_fields": ["data_name"],
+    # Default set of user-facing primary record attributes expected in the database. Records
+    # where any of these are missing will be loaded with the respective attribute set to None.
+    # Additional fields may be returned if they are present in the database, notably bidsdependencies.
+    "attributes": {
+        "data_name": "str",
+        "dataset": "str",
+        "bidspath": "str",
+        "subject": "str",
+        "task": "str",
+        "session": "str",
+        "run": "str",
+        "sampling_frequency": "float",
+        "modality": "str",
+        "nchans": "int",
+        "ntimes": "int",  # note: this is really the number of seconds in the data, rounded down
+    },
+    # queryable descriptive fields for a given recording
+    "description_fields": ["subject", "session", "run", "task", "age", "gender", "sex"],
+    # list of filenames that may be present in the BIDS dataset directory that are used
+    # to load and interpret a given BIDS recording.
+    "bids_dependencies_files": [
+        "dataset_description.json",
+        "participants.tsv",
+        "events.tsv",
+        "events.json",
+        "eeg.json",
+        "electrodes.tsv",
+        "channels.tsv",
+        "coordsystem.json",
+    ],
+    "accepted_query_fields": ["data_name", "dataset"],
+}
+"""dict: A global configuration dictionary for the EEGDash package.
+
+Keys
+----
+required_fields : list
+    Fields that must be present in every database record.
+attributes : dict
+    A schema defining the expected primary attributes and their types for a
+    database record.
+description_fields : list
+    A list of fields considered to be descriptive metadata for a recording,
+    which can be used for filtering and display.
+bids_dependencies_files : list
+    A list of BIDS metadata filenames that are relevant for interpreting an
+    EEG recording.
+accepted_query_fields : list
+    Fields that are accepted for lightweight existence checks in the database.
+"""

eegdash/dataset/__init__.py

@@ -0,0 +1,28 @@
+"""Public API for dataset helpers and dynamically generated datasets."""
+
+from . import dataset as _dataset_mod  # triggers dynamic class registration
+from .bids_dataset import EEGBIDSDataset
+from .dataset import EEGChallengeDataset, EEGDashDataset
+from .registry import register_openneuro_datasets
+
+# Re-export dynamically generated dataset classes at the package level so that
+# ``eegdash.dataset`` shows them in the API docs and users can import as
+# ``from eegdash.dataset import DSXXXXX``.
+_dyn_names = []
+for _name in getattr(_dataset_mod, "__all__", []):
+    if _name == "EEGChallengeDataset":
+        # Already imported explicitly above
+        continue
+    _obj = getattr(_dataset_mod, _name, None)
+    if _obj is not None:
+        globals()[_name] = _obj
+        _dyn_names.append(_name)
+
+__all__ = [
+    "EEGBIDSDataset",
+    "EEGDashDataset",
+    "EEGChallengeDataset",
+    "register_openneuro_datasets",
+] + _dyn_names
+
+del _dataset_mod, _name, _obj, _dyn_names

eegdash/dataset/base.py

@@ -0,0 +1,311 @@
+# Authors: The EEGDash contributors.
+# License: BSD-3-Clause
+# Copyright the EEGDash contributors.
+
+"""Data utilities and dataset classes for EEG data handling.
+
+This module provides core dataset classes for working with EEG data in the EEGDash ecosystem,
+including classes for individual recordings and collections of datasets. It integrates with
+braindecode for machine learning workflows and handles data loading from both local and remote sources.
+"""
+
+import io
+import os
+import traceback
+from contextlib import redirect_stderr
+from pathlib import Path
+from typing import Any
+
+import mne
+import mne_bids
+from mne._fiff.utils import _read_segments_file
+from mne.io import BaseRaw
+from mne_bids import BIDSPath
+
+from braindecode.datasets.base import BaseDataset
+
+from .. import downloader
+from ..bids_eeg_metadata import enrich_from_participants
+from ..logging import logger
+from ..paths import get_default_cache_dir
+
+
+class EEGDashBaseDataset(BaseDataset):
+    """A single EEG recording dataset.
+
+    Represents a single EEG recording, typically hosted on a remote server (like AWS S3)
+    and cached locally upon first access. This class is a subclass of
+    :class:`braindecode.datasets.BaseDataset` and can be used with braindecode's
+    preprocessing and training pipelines.
+
+    Parameters
+    ----------
+    record : dict
+        A fully resolved metadata record for the data to load.
+    cache_dir : str
+        The local directory where the data will be cached.
+    s3_bucket : str, optional
+        The S3 bucket to download data from. If not provided, defaults to the
+        OpenNeuro bucket.
+    **kwargs
+        Additional keyword arguments passed to the
+        :class:`braindecode.datasets.BaseDataset` constructor.
+
+    """
+
+    _AWS_BUCKET = "s3://openneuro.org"
+
+    def __init__(
+        self,
+        record: dict[str, Any],
+        cache_dir: str,
+        s3_bucket: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(None, **kwargs)
+        self.record = record
+        self.cache_dir = Path(cache_dir)
+        self.bids_kwargs = self._get_raw_bids_args()
+
+        if s3_bucket:
+            self.s3_bucket = s3_bucket
+            self.s3_open_neuro = False
+        else:
+            self.s3_bucket = self._AWS_BUCKET
+            self.s3_open_neuro = True
+
+        # Compute a dataset folder name under cache_dir that encodes preprocessing
+        # (e.g., bdf, mini) to avoid overlapping with the original dataset cache.
+        self.dataset_folder = record.get("dataset", "")
+        # TODO: remove this hack when competition is over
+        if s3_bucket:
+            suffixes: list[str] = []
+            bucket_lower = str(s3_bucket).lower()
+            if "bdf" in bucket_lower:
+                suffixes.append("bdf")
+            if "mini" in bucket_lower:
+                suffixes.append("mini")
+            if suffixes:
+                self.dataset_folder = f"{self.dataset_folder}-{'-'.join(suffixes)}"
+
+        # Place files under the dataset-specific folder (with suffix if any)
+        rel = Path(record["bidspath"])  # usually starts with dataset id
+        if rel.parts and rel.parts[0] == record.get("dataset"):
+            rel = Path(self.dataset_folder, *rel.parts[1:])
+        else:
+            rel = Path(self.dataset_folder) / rel
+        self.filecache = self.cache_dir / rel
+        self.bids_root = self.cache_dir / self.dataset_folder
+
+        self.bidspath = BIDSPath(
+            root=self.bids_root,
+            datatype="eeg",
+            suffix="eeg",
+            **self.bids_kwargs,
+        )
+
+        self.s3file = downloader.get_s3path(self.s3_bucket, record["bidspath"])
+        self.bids_dependencies = record["bidsdependencies"]
+        self.bids_dependencies_original = record["bidsdependencies"]
+        # TODO: removing temporary fix for BIDS dependencies path
+        # when the competition is over and dataset is digested properly
+        if not self.s3_open_neuro:
+            self.bids_dependencies = [
+                dep.split("/", 1)[1] for dep in self.bids_dependencies
+            ]
+
+        self._raw = None
+
+    def _get_raw_bids_args(self) -> dict[str, Any]:
+        """Extract BIDS-related arguments from the metadata record."""
+        desired_fields = ["subject", "session", "task", "run"]
+        return {k: self.record[k] for k in desired_fields if self.record[k]}
+
+    def _ensure_raw(self) -> None:
+        """Ensure the raw data file and its dependencies are cached locally."""
+        # TO-DO: remove this once is fixed on the our side
+        # for the competition
+        if not self.s3_open_neuro:
+            self.bidspath = self.bidspath.update(extension=".bdf")
+            self.filecache = self.filecache.with_suffix(".bdf")
+
+        if not os.path.exists(self.filecache):  # not preload
+            if self.bids_dependencies:
+                downloader.download_dependencies(
+                    s3_bucket=self.s3_bucket,
+                    bids_dependencies=self.bids_dependencies,
+                    bids_dependencies_original=self.bids_dependencies_original,
+                    cache_dir=self.cache_dir,
+                    dataset_folder=self.dataset_folder,
+                    record=self.record,
+                    s3_open_neuro=self.s3_open_neuro,
+                )
+            self.filecache = downloader.download_s3_file(
+                self.s3file, self.filecache, self.s3_open_neuro
+            )
+            self.filenames = [self.filecache]
+        if self._raw is None:
+            try:
+                # mne-bids can emit noisy warnings to stderr; keep user logs clean
+                _stderr_buffer = io.StringIO()
+                with redirect_stderr(_stderr_buffer):
+                    self._raw = mne_bids.read_raw_bids(
+                        bids_path=self.bidspath, verbose="ERROR"
+                    )
+                # Enrich Raw.info and description with participants.tsv extras
+                enrich_from_participants(
+                    self.bids_root, self.bidspath, self._raw, self.description
+                )
+
+            except Exception as e:
+                logger.error(
+                    f"Error while reading BIDS file: {self.bidspath}\n"
+                    "This may be due to a missing or corrupted file.\n"
+                    "Please check the file and try again.\n"
+                    "Usually erasing the local cache and re-downloading helps.\n"
+                    f"`rm {self.bidspath}`"
+                )
+                logger.error(f"Exception: {e}")
+                logger.error(traceback.format_exc())
+                raise e
+
+    def __len__(self) -> int:
+        """Return the number of samples in the dataset."""
+        if self._raw is None:
+            if (
+                self.record["ntimes"] is None
+                or self.record["sampling_frequency"] is None
+            ):
+                self._ensure_raw()
+            else:
+                # FIXME: this is a bit strange and should definitely not change as a side effect
+                #  of accessing the data (which it will, since ntimes is the actual length but rounded down)
+                return int(self.record["ntimes"] * self.record["sampling_frequency"])
+        return len(self._raw)
+
+    @property
+    def raw(self) -> BaseRaw:
+        """The MNE Raw object for this recording.
+
+        Accessing this property triggers the download and caching of the data
+        if it has not been accessed before.
+
+        Returns
+        -------
+        mne.io.BaseRaw
+            The loaded MNE Raw object.
+
+        """
+        if self._raw is None:
+            self._ensure_raw()
+        return self._raw
+
+    @raw.setter
+    def raw(self, raw: BaseRaw):
+        self._raw = raw
+
+
+class EEGDashBaseRaw(BaseRaw):
+    """MNE BaseRaw wrapper for automatic S3 data fetching.
+
+    This class extends :class:`mne.io.BaseRaw` to automatically fetch data
+    from an S3 bucket and cache it locally when data is first accessed.
+    It is intended for internal use within the EEGDash ecosystem.
+
+    Parameters
+    ----------
+    input_fname : str
+        The path to the file on the S3 bucket (relative to the bucket root).
+    metadata : dict
+        The metadata record for the recording, containing information like
+        sampling frequency, channel names, etc.
+    preload : bool, default False
+        If True, preload the data into memory.
+    cache_dir : str, optional
+        Local directory for caching data. If None, a default directory is used.
+    bids_dependencies : list of str, default []
+        A list of BIDS metadata files to download alongside the main recording.
+    verbose : str, int, or None, default None
+        The MNE verbosity level.
+
+    See Also
+    --------
+    mne.io.Raw : The base class for Raw objects in MNE.
+
+    """
+
+    _AWS_BUCKET = "s3://openneuro.org"
+
+    def __init__(
+        self,
+        input_fname: str,
+        metadata: dict[str, Any],
+        preload: bool = False,
+        *,
+        cache_dir: str | None = None,
+        bids_dependencies: list[str] | None = None,
+        verbose: Any = None,
+    ):
+        # Create a simple RawArray
+        sfreq = metadata["sfreq"]  # Sampling frequency
+        n_times = metadata["n_times"]
+        ch_names = metadata["ch_names"]
+        ch_types = []
+        for ch in metadata["ch_types"]:
+            chtype = ch.lower()
+            if chtype == "heog" or chtype == "veog":
+                chtype = "eog"
+            ch_types.append(chtype)
+        info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
+
+        self.s3file = downloader.get_s3path(self._AWS_BUCKET, input_fname)
+        self.cache_dir = Path(cache_dir) if cache_dir else get_default_cache_dir()
+        self.filecache = self.cache_dir / input_fname
+        if bids_dependencies is None:
+            bids_dependencies = []
+        self.bids_dependencies = bids_dependencies
+
+        if preload and not os.path.exists(self.filecache):
+            self.filecache = downloader.download_s3_file(
+                self.s3file, self.filecache, self.s3_open_neuro
+            )
+            self.filenames = [self.filecache]
+            preload = self.filecache
+
+        super().__init__(
+            info,
+            preload,
+            last_samps=[n_times - 1],
+            orig_format="single",
+            verbose=verbose,
+        )
+
+    def _read_segment(
+        self, start=0, stop=None, sel=None, data_buffer=None, *, verbose=None
+    ):
+        """Read a segment of data, downloading if necessary."""
+        if not os.path.exists(self.filecache):  # not preload
+            if self.bids_dependencies:  # this is use only to sidecars for now
+                downloader.download_dependencies(
+                    s3_bucket=self._AWS_BUCKET,
+                    bids_dependencies=self.bids_dependencies,
+                    bids_dependencies_original=None,
+                    cache_dir=self.cache_dir,
+                    dataset_folder=self.filecache,
+                    record={},
+                    s3_open_neuro=self.s3_open_neuro,
+                )
+            self.filecache = downloader.download_s3_file(
+                self.s3file, self.filecache, self.s3_open_neuro
+            )
+            self.filenames = [self.filecache]
+        else:  # not preload and file is not cached
+            self.filenames = [self.filecache]
+        return super()._read_segment(start, stop, sel, data_buffer, verbose=verbose)
+
+    def _read_segment_file(self, data, idx, fi, start, stop, cals, mult):
+        """Read a chunk of data from a local file."""
+        _read_segments_file(self, data, idx, fi, start, stop, cals, mult, dtype="<f4")
+
+
+__all__ = ["EEGDashBaseDataset", "EEGDashBaseRaw"]