eegdash 0.4.0.dev173498563__py3-none-any.whl → 0.4.1.dev185__py3-none-any.whl

eegdash/bids_eeg_metadata.py

@@ -33,12 +33,30 @@ __all__ = [
 
 
 def build_query_from_kwargs(**kwargs) -> dict[str, Any]:
-    """Build and validate a MongoDB query from user-friendly keyword arguments.
+    """Build and validate a MongoDB query from keyword arguments.
+
+    This function converts user-friendly keyword arguments into a valid
+    MongoDB query dictionary. It handles scalar values as exact matches and
+    list-like values as ``$in`` queries. It also performs validation to
+    reject unsupported fields and empty values.
+
+    Parameters
+    ----------
+    **kwargs
+        Keyword arguments representing query filters. Allowed keys are defined
+        in ``eegdash.const.ALLOWED_QUERY_FIELDS``.
+
+    Returns
+    -------
+    dict
+        A MongoDB query dictionary.
+
+    Raises
+    ------
+    ValueError
+        If an unsupported query field is provided, or if a value is None or
+        an empty string/list.
 
-    Improvements:
-    - Reject None values and empty/whitespace-only strings
-    - For list/tuple/set values: strip strings, drop None/empties, deduplicate, and use `$in`
-    - Preserve scalars as exact matches
     """
     # 1. Validate that all provided keys are allowed for querying
     unknown_fields = set(kwargs.keys()) - ALLOWED_QUERY_FIELDS
@@ -89,24 +107,29 @@ def build_query_from_kwargs(**kwargs) -> dict[str, Any]:
 
 
 def load_eeg_attrs_from_bids_file(bids_dataset, bids_file: str) -> dict[str, Any]:
-    """Build the metadata record for a given BIDS file (single recording) in a BIDS dataset.
+    """Build a metadata record for a BIDS file.
 
-    Attributes are at least the ones defined in data_config attributes (set to None if missing),
-    but are typically a superset, and include, among others, the paths to relevant
-    meta-data files needed to load and interpret the file in question.
+    Extracts metadata attributes from a single BIDS EEG file within a given
+    BIDS dataset. The extracted attributes include BIDS entities, file paths,
+    and technical metadata required for database indexing.
 
     Parameters
     ----------
     bids_dataset : EEGBIDSDataset
         The BIDS dataset object containing the file.
     bids_file : str
-        The path to the BIDS file within the dataset.
+        The path to the BIDS file to process.
 
     Returns
     -------
-    dict:
-        A dictionary representing the metadata record for the given file. This is the
-        same format as the records stored in the database.
+    dict
+        A dictionary of metadata attributes for the file, suitable for
+        insertion into the database.
+
+    Raises
+    ------
+    ValueError
+        If ``bids_file`` is not found in the ``bids_dataset``.
 
     """
     if bids_file not in bids_dataset.files:
@@ -198,11 +221,23 @@ def load_eeg_attrs_from_bids_file(bids_dataset, bids_file: str) -> dict[str, Any
 
 
 def normalize_key(key: str) -> str:
-    """Normalize a metadata key for robust matching.
+    """Normalize a string key for robust matching.
+
+    Converts the key to lowercase, replaces non-alphanumeric characters with
+    underscores, and removes leading/trailing underscores. This allows for
+    tolerant matching of keys that may have different capitalization or
+    separators (e.g., "p-factor" becomes "p_factor").
+
+    Parameters
+    ----------
+    key : str
+        The key to normalize.
+
+    Returns
+    -------
+    str
+        The normalized key.
 
-    Lowercase and replace non-alphanumeric characters with underscores, then strip
-    leading/trailing underscores. This allows tolerant matching such as
-    "p-factor" ≈ "p_factor" ≈ "P Factor".
     """
     return re.sub(r"[^a-z0-9]+", "_", str(key).lower()).strip("_")
 
@@ -212,27 +247,27 @@ def merge_participants_fields(
     participants_row: dict[str, Any] | None,
     description_fields: list[str] | None = None,
 ) -> dict[str, Any]:
-    """Merge participants.tsv fields into a dataset description dictionary.
+    """Merge fields from a participants.tsv row into a description dict.
 
-    - Preserves existing entries in ``description`` (no overwrites).
-    - Fills requested ``description_fields`` first, preserving their original names.
-    - Adds all remaining participants columns generically using normalized keys
-      unless a matching requested field already captured them.
+    Enriches a description dictionary with data from a subject's row in
+    ``participants.tsv``. It avoids overwriting existing keys in the
+    description.
 
     Parameters
     ----------
     description : dict
-        Current description to be enriched in-place and returned.
-    participants_row : dict | None
-        A mapping of participants.tsv columns for the current subject.
-    description_fields : list[str] | None
-        Optional list of requested description fields. When provided, matching is
-        performed by normalized names; the original requested field names are kept.
+        The description dictionary to enrich.
+    participants_row : dict or None
+        A dictionary representing a row from ``participants.tsv``. If None,
+        the original description is returned unchanged.
+    description_fields : list of str, optional
+        A list of specific fields to include in the description. Matching is
+        done using normalized keys.
 
     Returns
     -------
     dict
-        The enriched description (same object as input for convenience).
+        The enriched description dictionary.
 
     """
     if not isinstance(description, dict) or not isinstance(participants_row, dict):
@@ -272,10 +307,26 @@ def participants_row_for_subject(
     subject: str,
     id_columns: tuple[str, ...] = ("participant_id", "participant", "subject"),
 ) -> pd.Series | None:
-    """Load participants.tsv and return the row for a subject.
+    """Load participants.tsv and return the row for a specific subject.
+
+    Searches for a subject's data in the ``participants.tsv`` file within a
+    BIDS dataset. It can identify the subject with or without the "sub-"
+    prefix.
+
+    Parameters
+    ----------
+    bids_root : str or Path
+        The root directory of the BIDS dataset.
+    subject : str
+        The subject identifier (e.g., "01" or "sub-01").
+    id_columns : tuple of str, default ("participant_id", "participant", "subject")
+        A tuple of column names to search for the subject identifier.
+
+    Returns
+    -------
+    pandas.Series or None
+        A pandas Series containing the subject's data if found, otherwise None.
 
-    - Accepts either "01" or "sub-01" as the subject identifier.
-    - Returns a pandas Series for the first matching row, or None if not found.
     """
     try:
         participants_tsv = Path(bids_root) / "participants.tsv"
@@ -311,9 +362,28 @@ def participants_extras_from_tsv(
     id_columns: tuple[str, ...] = ("participant_id", "participant", "subject"),
     na_like: tuple[str, ...] = ("", "n/a", "na", "nan", "unknown", "none"),
 ) -> dict[str, Any]:
-    """Return non-identifier, non-empty participants.tsv fields for a subject.
+    """Extract additional participant information from participants.tsv.
+
+    Retrieves all non-identifier and non-empty fields for a subject from
+    the ``participants.tsv`` file.
+
+    Parameters
+    ----------
+    bids_root : str or Path
+        The root directory of the BIDS dataset.
+    subject : str
+        The subject identifier.
+    id_columns : tuple of str, default ("participant_id", "participant", "subject")
+        Column names to be treated as identifiers and excluded from the
+        output.
+    na_like : tuple of str, default ("", "n/a", "na", "nan", "unknown", "none")
+        Values to be considered as "Not Available" and excluded.
+
+    Returns
+    -------
+    dict
+        A dictionary of extra participant information.
 
-    Uses vectorized pandas operations to drop id columns and NA-like values.
     """
     row = participants_row_for_subject(bids_root, subject, id_columns=id_columns)
     if row is None:
@@ -331,10 +401,21 @@ def attach_participants_extras(
     description: Any,
     extras: dict[str, Any],
 ) -> None:
-    """Attach extras to Raw.info and dataset description without overwriting.
+    """Attach extra participant data to a raw object and its description.
+
+    Updates the ``raw.info['subject_info']`` and the description object
+    (dict or pandas Series) with extra data from ``participants.tsv``.
+    It does not overwrite existing keys.
+
+    Parameters
+    ----------
+    raw : mne.io.Raw
+        The MNE Raw object to be updated.
+    description : dict or pandas.Series
+        The description object to be updated.
+    extras : dict
+        A dictionary of extra participant information to attach.
 
-    - Adds to ``raw.info['subject_info']['participants_extras']``.
-    - Adds to ``description`` if dict or pandas Series (only missing keys).
     """
     if not extras:
         return
@@ -375,9 +456,28 @@ def enrich_from_participants(
     raw: Any,
     description: Any,
 ) -> dict[str, Any]:
-    """Convenience wrapper: read participants.tsv and attach extras for this subject.
+    """Read participants.tsv and attach extra info for the subject.
+
+    This is a convenience function that finds the subject from the
+    ``bidspath``, retrieves extra information from ``participants.tsv``,
+    and attaches it to the raw object and its description.
+
+    Parameters
+    ----------
+    bids_root : str or Path
+        The root directory of the BIDS dataset.
+    bidspath : mne_bids.BIDSPath
+        The BIDSPath object for the current data file.
+    raw : mne.io.Raw
+        The MNE Raw object to be updated.
+    description : dict or pandas.Series
+        The description object to be updated.
+
+    Returns
+    -------
+    dict
+        The dictionary of extras that were attached.
 
-    Returns the extras dictionary for further use if needed.
     """
     subject = getattr(bidspath, "subject", None)
     if not subject:

eegdash/const.py

@@ -28,6 +28,8 @@ ALLOWED_QUERY_FIELDS = {
     "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",
@@ -42,6 +44,8 @@ RELEASE_TO_OPENNEURO_DATASET_MAP = {
     "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": [
@@ -287,6 +291,9 @@ SUBJECT_MINI_RELEASE_MAP = {
         "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"],
@@ -322,3 +329,21 @@ config = {
     ],
     "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

@@ -1,7 +1,8 @@
 """Public API for dataset helpers and dynamically generated datasets."""
 
 from . import dataset as _dataset_mod  # triggers dynamic class registration
-from .dataset import EEGChallengeDataset
+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
@@ -17,6 +18,11 @@ for _name in getattr(_dataset_mod, "__all__", []):
         globals()[_name] = _obj
         _dyn_names.append(_name)
 
-__all__ = ["EEGChallengeDataset", "register_openneuro_datasets"] + _dyn_names
+__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: GNU General Public License
+# 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 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"]