eegdash 0.4.1__py3-none-any.whl → 0.4.1.dev185__py3-none-any.whl

eegdash/dataset/dataset.py

@@ -1,16 +1,519 @@
 from pathlib import Path
+from typing import Any
 
+from docstring_inheritance import NumpyDocstringInheritanceInitMeta
+from mne_bids import find_matching_paths
 from rich.console import Console
 from rich.panel import Panel
 from rich.text import Text
 
-from ..api import EEGDashDataset
-from ..bids_eeg_metadata import build_query_from_kwargs
-from ..const import RELEASE_TO_OPENNEURO_DATASET_MAP, SUBJECT_MINI_RELEASE_MAP
+from braindecode.datasets import BaseConcatDataset
+
+from .. import downloader
+from ..bids_eeg_metadata import (
+    build_query_from_kwargs,
+    merge_participants_fields,
+    normalize_key,
+)
+from ..const import (
+    ALLOWED_QUERY_FIELDS,
+    RELEASE_TO_OPENNEURO_DATASET_MAP,
+    SUBJECT_MINI_RELEASE_MAP,
+)
 from ..logging import logger
+from ..paths import get_default_cache_dir
+from .base import EEGDashBaseDataset
+from .bids_dataset import EEGBIDSDataset
 from .registry import register_openneuro_datasets
 
 
+class EEGDashDataset(BaseConcatDataset, metaclass=NumpyDocstringInheritanceInitMeta):
+    """Create a new EEGDashDataset from a given query or local BIDS dataset directory
+    and dataset name. An EEGDashDataset is pooled collection of EEGDashBaseDataset
+    instances (individual recordings) and is a subclass of braindecode's BaseConcatDataset.
+
+    Examples
+    --------
+    Basic usage with dataset and subject filtering:
+
+    >>> from eegdash import EEGDashDataset
+    >>> dataset = EEGDashDataset(
+    ...     cache_dir="./data",
+    ...     dataset="ds002718",
+    ...     subject="012"
+    ... )
+    >>> print(f"Number of recordings: {len(dataset)}")
+
+    Filter by multiple subjects and specific task:
+
+    >>> subjects = ["012", "013", "014"]
+    >>> dataset = EEGDashDataset(
+    ...     cache_dir="./data",
+    ...     dataset="ds002718",
+    ...     subject=subjects,
+    ...     task="RestingState"
+    ... )
+
+    Load and inspect EEG data from recordings:
+
+    >>> if len(dataset) > 0:
+    ...     recording = dataset[0]
+    ...     raw = recording.load()
+    ...     print(f"Sampling rate: {raw.info['sfreq']} Hz")
+    ...     print(f"Number of channels: {len(raw.ch_names)}")
+    ...     print(f"Duration: {raw.times[-1]:.1f} seconds")
+
+    Advanced filtering with raw MongoDB queries:
+
+    >>> from eegdash import EEGDashDataset
+    >>> query = {
+    ...     "dataset": "ds002718",
+    ...     "subject": {"$in": ["012", "013"]},
+    ...     "task": "RestingState"
+    ... }
+    >>> dataset = EEGDashDataset(cache_dir="./data", query=query)
+
+    Working with dataset collections and braindecode integration:
+
+    >>> # EEGDashDataset is a braindecode BaseConcatDataset
+    >>> for i, recording in enumerate(dataset):
+    ...     if i >= 2:  # limit output
+    ...         break
+    ...     print(f"Recording {i}: {recording.description}")
+    ...     raw = recording.load()
+    ...     print(f"  Channels: {len(raw.ch_names)}, Duration: {raw.times[-1]:.1f}s")
+
+    Parameters
+    ----------
+    cache_dir : str | Path
+        Directory where data are cached locally.
+    query : dict | None
+        Raw MongoDB query to filter records. If provided, it is merged with
+        keyword filtering arguments (see ``**kwargs``) using logical AND.
+        You must provide at least a ``dataset`` (either in ``query`` or
+        as a keyword argument). Only fields in ``ALLOWED_QUERY_FIELDS`` are
+        considered for filtering.
+    dataset : str
+        Dataset identifier (e.g., ``"ds002718"``). Required if ``query`` does
+        not already specify a dataset.
+    task : str | list[str]
+        Task name(s) to filter by (e.g., ``"RestingState"``).
+    subject : str | list[str]
+        Subject identifier(s) to filter by (e.g., ``"NDARCA153NKE"``).
+    session : str | list[str]
+        Session identifier(s) to filter by (e.g., ``"1"``).
+    run : str | list[str]
+        Run identifier(s) to filter by (e.g., ``"1"``).
+    description_fields : list[str]
+        Fields to extract from each record and include in dataset descriptions
+        (e.g., "subject", "session", "run", "task").
+    s3_bucket : str | None
+        Optional S3 bucket URI (e.g., "s3://mybucket") to use instead of the
+        default OpenNeuro bucket when downloading data files.
+    records : list[dict] | None
+        Pre-fetched metadata records. If provided, the dataset is constructed
+        directly from these records and no MongoDB query is performed.
+    download : bool, default True
+        If False, load from local BIDS files only. Local data are expected
+        under ``cache_dir / dataset``; no DB or S3 access is attempted.
+    n_jobs : int
+        Number of parallel jobs to use where applicable (-1 uses all cores).
+    eeg_dash_instance : EEGDash | None
+        Optional existing EEGDash client to reuse for DB queries. If None,
+        a new client is created on demand, not used in the case of no download.
+    **kwargs : dict
+        Additional keyword arguments serving two purposes:
+
+        - Filtering: any keys present in ``ALLOWED_QUERY_FIELDS`` are treated as
+          query filters (e.g., ``dataset``, ``subject``, ``task``, ...).
+        - Dataset options: remaining keys are forwarded to
+          ``EEGDashBaseDataset``.
+
+    """
+
+    def __init__(
+        self,
+        cache_dir: str | Path,
+        query: dict[str, Any] = None,
+        description_fields: list[str] | None = None,
+        s3_bucket: str | None = None,
+        records: list[dict] | None = None,
+        download: bool = True,
+        n_jobs: int = -1,
+        eeg_dash_instance: Any = None,
+        **kwargs,
+    ):
+        # Parameters that don't need validation
+        _suppress_comp_warning: bool = kwargs.pop("_suppress_comp_warning", False)
+        self.s3_bucket = s3_bucket
+        self.records = records
+        self.download = download
+        self.n_jobs = n_jobs
+        self.eeg_dash_instance = eeg_dash_instance
+
+        if description_fields is None:
+            description_fields = [
+                "subject",
+                "session",
+                "run",
+                "task",
+                "age",
+                "gender",
+                "sex",
+            ]
+
+        self.cache_dir = cache_dir
+        if self.cache_dir == "" or self.cache_dir is None:
+            self.cache_dir = get_default_cache_dir()
+            logger.warning(
+                f"Cache directory is empty, using the eegdash default path: {self.cache_dir}"
+            )
+
+        self.cache_dir = Path(self.cache_dir)
+
+        if not self.cache_dir.exists():
+            logger.warning(
+                f"Cache directory does not exist, creating it: {self.cache_dir}"
+            )
+            self.cache_dir.mkdir(exist_ok=True, parents=True)
+
+        # Separate query kwargs from other kwargs passed to the BaseDataset constructor
+        self.query = query or {}
+        self.query.update(
+            {k: v for k, v in kwargs.items() if k in ALLOWED_QUERY_FIELDS}
+        )
+        base_dataset_kwargs = {k: v for k, v in kwargs.items() if k not in self.query}
+        if "dataset" not in self.query:
+            # If explicit records are provided, infer dataset from records
+            if isinstance(records, list) and records and isinstance(records[0], dict):
+                inferred = records[0].get("dataset")
+                if inferred:
+                    self.query["dataset"] = inferred
+                else:
+                    raise ValueError("You must provide a 'dataset' argument")
+            else:
+                raise ValueError("You must provide a 'dataset' argument")
+
+        # Decide on a dataset subfolder name for cache isolation. If using
+        # challenge/preprocessed buckets (e.g., BDF, mini subsets), append
+        # informative suffixes to avoid overlapping with the original dataset.
+        dataset_folder = self.query["dataset"]
+        if self.s3_bucket:
+            suffixes: list[str] = []
+            bucket_lower = str(self.s3_bucket).lower()
+            if "bdf" in bucket_lower:
+                suffixes.append("bdf")
+            if "mini" in bucket_lower:
+                suffixes.append("mini")
+            if suffixes:
+                dataset_folder = f"{dataset_folder}-{'-'.join(suffixes)}"
+
+        self.data_dir = self.cache_dir / dataset_folder
+
+        if (
+            not _suppress_comp_warning
+            and self.query["dataset"] in RELEASE_TO_OPENNEURO_DATASET_MAP.values()
+        ):
+            message_text = Text.from_markup(
+                "[italic]This notice is only for users who are participating in the [link=https://eeg2025.github.io/]EEG 2025 Competition[/link].[/italic]\n\n"
+                "[bold]EEG 2025 Competition Data Notice![/bold]\n"
+                "You are loading one of the datasets that is used in competition, but via `EEGDashDataset`.\n\n"
+                "[bold red]IMPORTANT[/bold red]: \n"
+                "If you download data from `EEGDashDataset`, it is [u]NOT[/u] identical to the official \n"
+                "competition data, which is accessed via `EEGChallengeDataset`. "
+                "The competition data has been downsampled and filtered.\n\n"
+                "[bold]If you are participating in the competition, \nyou must use the `EEGChallengeDataset` object to ensure consistency.[/bold] \n\n"
+                "If you are not participating in the competition, you can ignore this message."
+            )
+            warning_panel = Panel(
+                message_text,
+                title="[yellow]EEG 2025 Competition Data Notice[/yellow]",
+                subtitle="[cyan]Source: EEGDashDataset[/cyan]",
+                border_style="yellow",
+            )
+
+            try:
+                Console().print(warning_panel)
+            except Exception:
+                logger.warning(str(message_text))
+
+        if records is not None:
+            self.records = records
+            datasets = [
+                EEGDashBaseDataset(
+                    record,
+                    self.cache_dir,
+                    self.s3_bucket,
+                    **base_dataset_kwargs,
+                )
+                for record in self.records
+            ]
+        elif not download:  # only assume local data is complete if not downloading
+            if not self.data_dir.exists():
+                raise ValueError(
+                    f"Offline mode is enabled, but local data_dir {self.data_dir} does not exist."
+                )
+            records = self._find_local_bids_records(self.data_dir, self.query)
+            # Try to enrich from local participants.tsv to restore requested fields
+            try:
+                bids_ds = EEGBIDSDataset(
+                    data_dir=str(self.data_dir), dataset=self.query["dataset"]
+                )  # type: ignore[index]
+            except Exception:
+                bids_ds = None
+
+            datasets = []
+            for record in records:
+                # Start with entity values from filename
+                desc: dict[str, Any] = {
+                    k: record.get(k)
+                    for k in ("subject", "session", "run", "task")
+                    if record.get(k) is not None
+                }
+
+                if bids_ds is not None:
+                    try:
+                        rel_from_dataset = Path(record["bidspath"]).relative_to(
+                            record["dataset"]
+                        )  # type: ignore[index]
+                        local_file = (self.data_dir / rel_from_dataset).as_posix()
+                        part_row = bids_ds.subject_participant_tsv(local_file)
+                        desc = merge_participants_fields(
+                            description=desc,
+                            participants_row=part_row
+                            if isinstance(part_row, dict)
+                            else None,
+                            description_fields=description_fields,
+                        )
+                    except Exception:
+                        pass
+
+                datasets.append(
+                    EEGDashBaseDataset(
+                        record=record,
+                        cache_dir=self.cache_dir,
+                        s3_bucket=self.s3_bucket,
+                        description=desc,
+                        **base_dataset_kwargs,
+                    )
+                )
+        elif self.query:
+            if self.eeg_dash_instance is None:
+                # to avoid circular import
+                from ..api import EEGDash
+
+                self.eeg_dash_instance = EEGDash()
+            datasets = self._find_datasets(
+                query=build_query_from_kwargs(**self.query),
+                description_fields=description_fields,
+                base_dataset_kwargs=base_dataset_kwargs,
+            )
+            # We only need filesystem if we need to access S3
+            self.filesystem = downloader.get_s3_filesystem()
+        else:
+            raise ValueError(
+                "You must provide either 'records', a 'data_dir', or a query/keyword arguments for filtering."
+            )
+
+        super().__init__(datasets)
+
+    def _find_local_bids_records(
+        self, dataset_root: Path, filters: dict[str, Any]
+    ) -> list[dict]:
+        """Discover local BIDS EEG files and build minimal records.
+
+        Enumerates EEG recordings under ``dataset_root`` using
+        ``mne_bids.find_matching_paths`` and applies entity filters to produce
+        records suitable for :class:`EEGDashBaseDataset`. No network access is
+        performed, and files are not read.
+
+        Parameters
+        ----------
+        dataset_root : Path
+            Local dataset directory (e.g., ``/path/to/cache/ds005509``).
+        filters : dict
+            Query filters. Must include ``'dataset'`` and may include BIDS
+            entities like ``'subject'``, ``'session'``, etc.
+
+        Returns
+        -------
+        list of dict
+            A list of records, one for each matched EEG file. Each record
+            contains BIDS entities, paths, and minimal metadata for offline use.
+
+        Notes
+        -----
+        Matching is performed for ``datatypes=['eeg']`` and ``suffixes=['eeg']``.
+        The ``bidspath`` is normalized to ensure it starts with the dataset ID,
+        even for suffixed cache directories.
+
+        """
+        dataset_id = filters["dataset"]
+        arg_map = {
+            "subjects": "subject",
+            "sessions": "session",
+            "tasks": "task",
+            "runs": "run",
+        }
+        matching_args: dict[str, list[str]] = {}
+        for finder_key, entity_key in arg_map.items():
+            entity_val = filters.get(entity_key)
+            if entity_val is None:
+                continue
+            if isinstance(entity_val, (list, tuple, set)):
+                entity_vals = list(entity_val)
+                if not entity_vals:
+                    continue
+                matching_args[finder_key] = entity_vals
+            else:
+                matching_args[finder_key] = [entity_val]
+
+        matched_paths = find_matching_paths(
+            root=str(dataset_root),
+            datatypes=["eeg"],
+            suffixes=["eeg"],
+            ignore_json=True,
+            **matching_args,
+        )
+        records_out: list[dict] = []
+
+        for bids_path in matched_paths:
+            # Build bidspath as dataset_id / relative_path_from_dataset_root (POSIX)
+            rel_from_root = (
+                Path(bids_path.fpath)
+                .resolve()
+                .relative_to(Path(bids_path.root).resolve())
+            )
+            bidspath = f"{dataset_id}/{rel_from_root.as_posix()}"
+
+            rec = {
+                "data_name": f"{dataset_id}_{Path(bids_path.fpath).name}",
+                "dataset": dataset_id,
+                "bidspath": bidspath,
+                "subject": (bids_path.subject or None),
+                "session": (bids_path.session or None),
+                "task": (bids_path.task or None),
+                "run": (bids_path.run or None),
+                # minimal fields to satisfy BaseDataset from eegdash
+                "bidsdependencies": [],  # not needed to just run.
+                "modality": "eeg",
+                # minimal numeric defaults for offline length calculation
+                "sampling_frequency": None,
+                "nchans": None,
+                "ntimes": None,
+            }
+            records_out.append(rec)
+
+        return records_out
+
+    def _find_key_in_nested_dict(self, data: Any, target_key: str) -> Any:
+        """Recursively search for a key in nested dicts/lists.
+
+        Performs a case-insensitive and underscore/hyphen-agnostic search.
+
+        Parameters
+        ----------
+        data : Any
+            The nested data structure (dicts, lists) to search.
+        target_key : str
+            The key to search for.
+
+        Returns
+        -------
+        Any
+            The value of the first matching key, or None if not found.
+
+        """
+        norm_target = normalize_key(target_key)
+        if isinstance(data, dict):
+            for k, v in data.items():
+                if normalize_key(k) == norm_target:
+                    return v
+                res = self._find_key_in_nested_dict(v, target_key)
+                if res is not None:
+                    return res
+        elif isinstance(data, list):
+            for item in data:
+                res = self._find_key_in_nested_dict(item, target_key)
+                if res is not None:
+                    return res
+        return None
+
+    def _find_datasets(
+        self,
+        query: dict[str, Any] | None,
+        description_fields: list[str],
+        base_dataset_kwargs: dict,
+    ) -> list[EEGDashBaseDataset]:
+        """Find and construct datasets from a MongoDB query.
+
+        Queries the database, then creates a list of
+        :class:`EEGDashBaseDataset` objects from the results.
+
+        Parameters
+        ----------
+        query : dict, optional
+            The MongoDB query to execute.
+        description_fields : list of str
+            Fields to extract from each record for the dataset description.
+        base_dataset_kwargs : dict
+            Additional keyword arguments to pass to the
+            :class:`EEGDashBaseDataset` constructor.
+
+        Returns
+        -------
+        list of EEGDashBaseDataset
+            A list of dataset objects matching the query.
+
+        """
+        datasets: list[EEGDashBaseDataset] = []
+        self.records = self.eeg_dash_instance.find(query)
+
+        for record in self.records:
+            description: dict[str, Any] = {}
+            # Requested fields first (normalized matching)
+            for field in description_fields:
+                value = self._find_key_in_nested_dict(record, field)
+                if value is not None:
+                    description[field] = value
+            # Merge all participants.tsv columns generically
+            part = self._find_key_in_nested_dict(record, "participant_tsv")
+            if isinstance(part, dict):
+                description = merge_participants_fields(
+                    description=description,
+                    participants_row=part,
+                    description_fields=description_fields,
+                )
+            datasets.append(
+                EEGDashBaseDataset(
+                    record,
+                    cache_dir=self.cache_dir,
+                    s3_bucket=self.s3_bucket,
+                    description=description,
+                    **base_dataset_kwargs,
+                )
+            )
+        return datasets
+
+    # just to fix the docstring inheritance until we solved it in braindecode.
+    def save(self, path, overwrite=False):
+        """Save the dataset to disk.
+
+        Parameters
+        ----------
+        path : str or Path
+            Destination file path.
+        overwrite : bool, default False
+            If True, overwrite existing file.
+
+        Returns
+        -------
+        None
+
+        """
+        return super().save(path, overwrite=overwrite)
+
+
 class EEGChallengeDataset(EEGDashDataset):
     """A dataset helper for the EEG 2025 Challenge.
 
@@ -186,4 +689,4 @@ registered_classes = register_openneuro_datasets(
 )
 
 
-__all__ = ["EEGChallengeDataset"] + list(registered_classes.keys())
+__all__ = ["EEGDashDataset", "EEGChallengeDataset"] + list(registered_classes.keys())

eegdash/features/datasets.py

@@ -128,7 +128,7 @@ def _compute_stats(
     ddof: int = 1,
     numeric_only: bool = False,
 ) -> tuple:
-    """Compute statistics for a single FeaturesDataset."""
+    """Compute statistics for a single :class:`~eegdash.features.datasets.FeaturesDataset`."""
     res = []
     if return_count:
         res.append(ds.features.count(numeric_only=numeric_only))
@@ -159,9 +159,9 @@ def _pooled_var(
 
 
 class FeaturesConcatDataset(BaseConcatDataset):
-    """A concatenated dataset of `FeaturesDataset` objects.
+    """A concatenated dataset of :class:`~eegdash.features.datasets.FeaturesDataset` objects.
 
-    This class holds a list of :class:`FeaturesDataset` instances and allows
+    This class holds a list of :class:`~eegdash.features.datasets.FeaturesDataset` instances and allows
     them to be treated as a single, larger dataset. It provides methods for
 
     splitting, saving, and performing DataFrame-like operations (e.g., `mean`,
@@ -169,8 +169,8 @@ class FeaturesConcatDataset(BaseConcatDataset):
 
     Parameters
     ----------
-    list_of_ds : list of FeaturesDataset
-        A list of :class:`FeaturesDataset` objects to concatenate.
+    list_of_ds : list of ~eegdash.features.datasets.FeaturesDataset
+        A list of :class:`~eegdash.features.datasets.FeaturesDataset` objects to concatenate.
     target_transform : callable, optional
         A function to apply to the target values before they are returned.
 
@@ -211,9 +211,9 @@ class FeaturesConcatDataset(BaseConcatDataset):
 
         Returns
         -------
-        dict[str, FeaturesConcatDataset]
+        dict[str, ~eegdash.features.datasets.FeaturesConcatDataset]
             A dictionary where keys are split names and values are the new
-            :class:`FeaturesConcatDataset` subsets.
+            :class:`~eegdash.features.datasets.FeaturesConcatDataset` subsets.
 
         """
         if isinstance(by, str):
@@ -252,7 +252,8 @@ class FeaturesConcatDataset(BaseConcatDataset):
         Raises
         ------
         TypeError
-            If any of the contained datasets is not a :class:`FeaturesDataset`.
+            If any of the contained datasets is not a
+            :class:`~eegdash.features.datasets.FeaturesDataset`.
 
         """
         if not all([isinstance(ds, FeaturesDataset) for ds in self.datasets]):

eegdash/features/decorators.py

@@ -71,7 +71,7 @@ class FeatureKind:
 
     Parameters
     ----------
-    feature_kind : MultivariateFeature
+    feature_kind : ~eegdash.features.extractors.MultivariateFeature
         An instance of a feature kind class, such as
         :class:`~eegdash.features.extractors.UnivariateFeature` or
         :class:`~eegdash.features.extractors.BivariateFeature`.
@@ -104,7 +104,7 @@ class FeatureKind:
 univariate_feature = FeatureKind(UnivariateFeature())
 """Decorator to mark a feature as univariate.
 
-This is a convenience instance of :class:`FeatureKind` pre-configured for
+This is a convenience instance of :class:`~eegdash.features.decorators.FeatureKind` pre-configured for
 univariate features.
 """
 
@@ -139,6 +139,6 @@ def bivariate_feature(func: Callable, directed: bool = False) -> Callable:
 multivariate_feature = FeatureKind(MultivariateFeature())
 """Decorator to mark a feature as multivariate.
 
-This is a convenience instance of :class:`FeatureKind` pre-configured for
+This is a convenience instance of :class:`~eegdash.features.decorators.FeatureKind` pre-configured for
 multivariate features, which operate on all channels simultaneously.
 """