eegdash 0.3.8__py3-none-any.whl → 0.3.9.dev129__py3-none-any.whl

eegdash/downloader.py

@@ -0,0 +1,176 @@
+import re
+import tempfile
+from pathlib import Path
+from typing import Any
+from urllib.parse import urlsplit
+
+import mne
+import numpy as np
+import s3fs
+import xarray as xr
+from fsspec.callbacks import TqdmCallback
+
+
+def get_s3_filesystem():
+    """Returns an S3FileSystem object."""
+    return s3fs.S3FileSystem(anon=True, client_kwargs={"region_name": "us-east-2"})
+
+
+def get_s3path(s3_bucket: str, filepath: str) -> str:
+    """Helper to form an AWS S3 URI for the given relative filepath."""
+    return f"{s3_bucket}/{filepath}"
+
+
+def download_s3_file(s3_path: str, local_path: Path, s3_open_neuro: bool):
+    """Download function that gets the raw EEG data from S3."""
+    filesystem = get_s3_filesystem()
+    if not s3_open_neuro:
+        s3_path = re.sub(r"(^|/)ds\d{6}/", r"\1", s3_path, count=1)
+        # TODO: remove this hack when competition is over
+        if s3_path.endswith(".set"):
+            s3_path = s3_path[:-4] + ".bdf"
+            local_path = local_path.with_suffix(".bdf")
+
+    local_path.parent.mkdir(parents=True, exist_ok=True)
+    _filesystem_get(filesystem=filesystem, s3path=s3_path, filepath=local_path)
+
+    return local_path
+
+
+def download_dependencies(
+    s3_bucket: str,
+    bids_dependencies: list[str],
+    bids_dependencies_original: list[str],
+    cache_dir: Path,
+    dataset_folder: Path,
+    record: dict[str, Any],
+    s3_open_neuro: bool,
+):
+    """Download all BIDS dependency files from S3 and cache them locally."""
+    filesystem = get_s3_filesystem()
+    for i, dep in enumerate(bids_dependencies):
+        if not s3_open_neuro:
+            if dep.endswith(".set"):
+                dep = dep[:-4] + ".bdf"
+
+        s3path = get_s3path(s3_bucket, dep)
+        if not s3_open_neuro:
+            dep = bids_dependencies_original[i]
+
+        dep_path = Path(dep)
+        if dep_path.parts and dep_path.parts[0] == record.get("dataset"):
+            dep_local = Path(dataset_folder, *dep_path.parts[1:])
+        else:
+            dep_local = Path(dataset_folder) / dep_path
+        filepath = cache_dir / dep_local
+        if not s3_open_neuro:
+            if filepath.suffix == ".set":
+                filepath = filepath.with_suffix(".bdf")
+
+        if not filepath.exists():
+            filepath.parent.mkdir(parents=True, exist_ok=True)
+            _filesystem_get(filesystem=filesystem, s3path=s3path, filepath=filepath)
+
+
+def _filesystem_get(filesystem: s3fs.S3FileSystem, s3path: str, filepath: Path):
+    """Helper to download a file from S3 with a progress bar."""
+    info = filesystem.info(s3path)
+    size = info.get("size") or info.get("Size")
+
+    callback = TqdmCallback(
+        size=size,
+        tqdm_kwargs=dict(
+            desc=f"Downloading {Path(s3path).name}",
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+            dynamic_ncols=True,
+            leave=True,
+            mininterval=0.2,
+            smoothing=0.1,
+            miniters=1,
+            bar_format="{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} "
+            "[{elapsed}<{remaining}, {rate_fmt}]",
+        ),
+    )
+    filesystem.get(s3path, str(filepath), callback=callback)
+    return filepath
+
+
+def load_eeg_from_s3(s3path: str):
+    """Load EEG data from an S3 URI into an ``xarray.DataArray``.
+
+    Preserves the original filename, downloads sidecar files when applicable
+    (e.g., ``.fdt`` for EEGLAB, ``.vmrk``/``.eeg`` for BrainVision), and uses
+    MNE's direct readers.
+
+    Parameters
+    ----------
+    s3path : str
+        An S3 URI (should start with "s3://").
+
+    Returns
+    -------
+    xr.DataArray
+        EEG data with dimensions ``("channel", "time")``.
+
+    Raises
+    ------
+    ValueError
+        If the file extension is unsupported.
+
+    """
+    filesystem = get_s3_filesystem()
+    # choose a temp dir so sidecars can be colocated
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Derive local filenames from the S3 key to keep base name consistent
+        s3_key = urlsplit(s3path).path  # e.g., "/dsXXXX/sub-.../..._eeg.set"
+        basename = Path(s3_key).name
+        ext = Path(basename).suffix.lower()
+        local_main = Path(tmpdir) / basename
+
+        # Download main file
+        with (
+            filesystem.open(s3path, mode="rb") as fsrc,
+            open(local_main, "wb") as fdst,
+        ):
+            fdst.write(fsrc.read())
+
+        # Determine and fetch any required sidecars
+        sidecars: list[str] = []
+        if ext == ".set":  # EEGLAB
+            sidecars = [".fdt"]
+        elif ext == ".vhdr":  # BrainVision
+            sidecars = [".vmrk", ".eeg", ".dat", ".raw"]
+
+        for sc_ext in sidecars:
+            sc_key = s3_key[: -len(ext)] + sc_ext
+            sc_uri = f"s3://{urlsplit(s3path).netloc}{sc_key}"
+            try:
+                # If sidecar exists, download next to the main file
+                info = filesystem.info(sc_uri)
+                if info:
+                    sc_local = Path(tmpdir) / Path(sc_key).name
+                    with (
+                        filesystem.open(sc_uri, mode="rb") as fsrc,
+                        open(sc_local, "wb") as fdst,
+                    ):
+                        fdst.write(fsrc.read())
+            except Exception:
+                # Sidecar not present; skip silently
+                pass
+
+        # Read using appropriate MNE reader
+        raw = mne.io.read_raw(str(local_main), preload=True, verbose=False)
+
+        data = raw.get_data()
+        fs = raw.info["sfreq"]
+        max_time = data.shape[1] / fs
+        time_steps = np.linspace(0, max_time, data.shape[1]).squeeze()
+        channel_names = raw.ch_names
+
+        return xr.DataArray(
+            data=data,
+            dims=["channel", "time"],
+            coords={"time": time_steps, "channel": channel_names},
+        )

eegdash/features/datasets.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import json
 import os
 import shutil
-import warnings
 from collections.abc import Callable
 from typing import Dict, List
 
@@ -17,6 +16,8 @@ from braindecode.datasets.base import (
     _create_description,
 )
 
+from ..logging import logger
+
 
 class FeaturesDataset(EEGWindowsDataset):
     """Returns samples from a pandas DataFrame object along with a target.
@@ -283,7 +284,7 @@ class FeaturesConcatDataset(BaseConcatDataset):
             # the following will be True for all datasets preprocessed and
             # stored in parallel with braindecode.preprocessing.preprocess
             if i_ds + 1 + offset < n_sub_dirs:
-                warnings.warn(
+                logger.warning(
                     f"The number of saved datasets ({i_ds + 1 + offset}) "
                     f"does not match the number of existing "
                     f"subdirectories ({n_sub_dirs}). You may now "
@@ -294,7 +295,7 @@ class FeaturesConcatDataset(BaseConcatDataset):
         # if path contains files or directories that were not touched, raise
         # warning
         if path_contents:
-            warnings.warn(
+            logger.warning(
                 f"Chosen directory {path} contains other "
                 f"subdirectories or files {path_contents}."
             )

eegdash/hbn/preprocessing.py

@@ -1,11 +1,9 @@
-import logging
-
 import mne
 import numpy as np
 
 from braindecode.preprocessing import Preprocessor
 
-logger = logging.getLogger("eegdash")
+from ..logging import logger
 
 
 class hbn_ec_ec_reannotation(Preprocessor):

eegdash/hbn/windows.py

@@ -7,8 +7,6 @@ from mne_bids import get_bids_path_from_fname
 
 from braindecode.datasets.base import BaseConcatDataset
 
-logger = logging.getLogger("eegdash")
-
 
 def build_trial_table(events_df: pd.DataFrame) -> pd.DataFrame:
     """One row per contrast trial with stimulus/response metrics."""

eegdash/logging.py

@@ -0,0 +1,23 @@
+import logging
+
+from rich.logging import RichHandler
+
+# Get the root logger
+root_logger = logging.getLogger()
+
+# --- This is the key part ---
+# 1. Remove any handlers that may have been added by default
+root_logger.handlers = []
+
+# 2. Add your RichHandler
+root_logger.addHandler(RichHandler(rich_tracebacks=True, markup=True))
+# ---------------------------
+
+# 3. Set the level for the root logger
+root_logger.setLevel(logging.INFO)
+
+# Now, get your package-specific logger. It will inherit the
+# configuration from the root logger we just set up.
+logger = logging.getLogger("eegdash")
+
+logger.setLevel(logging.INFO)

{eegdash-0.3.8.dist-info → eegdash-0.3.9.dev129.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: eegdash
-Version: 0.3.8
+Version: 0.3.9.dev129
 Summary: EEG data for machine learning
 Author-email: Young Truong <dt.young112@gmail.com>, Arnaud Delorme <adelorme@gmail.com>, Aviv Dotan <avivd220@gmail.com>, Oren Shriki <oren70@gmail.com>, Bruno Aristimunha <b.aristimunha@gmail.com>
 License-Expression: GPL-3.0-only
@@ -25,7 +25,7 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: braindecode>=1.0
-Requires-Dist: mne_bids>=0.16.0
+Requires-Dist: mne_bids>=0.17.0
 Requires-Dist: numba
 Requires-Dist: numpy
 Requires-Dist: pandas
@@ -41,6 +41,7 @@ Requires-Dist: pymatreader
 Requires-Dist: eeglabio
 Requires-Dist: tabulate
 Requires-Dist: docstring_inheritance
+Requires-Dist: rich
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
@@ -63,6 +64,7 @@ Requires-Dist: memory_profiler; extra == "docs"
 Requires-Dist: ipython; extra == "docs"
 Requires-Dist: lightgbm; extra == "docs"
 Requires-Dist: plotly; extra == "docs"
+Requires-Dist: nbformat; extra == "docs"
 Provides-Extra: all
 Requires-Dist: eegdash[docs]; extra == "all"
 Requires-Dist: eegdash[dev]; extra == "all"

{eegdash-0.3.8.dist-info → eegdash-0.3.9.dev129.dist-info}/RECORD

@@ -1,17 +1,19 @@
-eegdash/__init__.py,sha256=i7v1mHVxMkRTcAuf2bt0COJ12jmgUmFF5aZBNvMY154,277
-eegdash/api.py,sha256=fzZ69wC8Rn0yP3RGQR08k7tQwNxFz-B8UYvEUJ_9hfg,40585
-eegdash/bids_eeg_metadata.py,sha256=LZrGPGVdnGUbZlD4M_aAW4kEItzwTTeZFicH-jyqDyc,9712
+eegdash/__init__.py,sha256=WLMYpZuTIJ5DnzhwE40o4aqt8iwVpxdmOn9ezY7hw60,284
+eegdash/api.py,sha256=Eh-CxLOhDthmVeAubezXPu0eSOhRuYBN23iGIV75EVM,37675
+eegdash/bids_eeg_metadata.py,sha256=XUkQp2M8zQ_wH5JC8lQiVR0TWssOjSdrdEWkgwGjiZ8,13699
 eegdash/const.py,sha256=qdFBEL9kIrsj9CdxbXhBkR61R3CrTGSaj5Iq0YOACIs,7313
-eegdash/data_utils.py,sha256=DZ-B03VleA9-mOUzGXcS4N18dVC2uFkFGXMFsKK8nUc,34166
+eegdash/data_utils.py,sha256=dePEXcJefo3gmD534bb576p8v9jErDNLDHQoDGHqH-g,26006
+eegdash/downloader.py,sha256=B-8u0c39F4inV-v_WgYZrUKAGPQOdfqC2RX1qNrRQYM,5808
+eegdash/logging.py,sha256=SZdB7WLT5b2okecWpvLx4UWUxg3DiA11Z5d9lhYdDyc,616
 eegdash/mongodb.py,sha256=GD3WgA253oFgpzOHrYaj4P1mRjNtDMT5Oj4kVvHswjI,2006
 eegdash/paths.py,sha256=246xkectTxDAYcREs1Qma_F1Y-oSmLlb0hn0F2Za5Ss,866
 eegdash/utils.py,sha256=7TfQ9D0LrAJ7FgnSXEvWgeHWK2QqaqS-_WcWXD86ObQ,408
 eegdash/dataset/__init__.py,sha256=Qmzki5G8GaFlzTb10e4SmC3WkKuJyo1Ckii15tCEHAo,157
-eegdash/dataset/dataset.py,sha256=5ku9I-JTC8DNult_m1Dgem1OQnSPazp_Xdktq4s_2s4,6829
+eegdash/dataset/dataset.py,sha256=DbQX_ajPExWE-6DmqCHWEDyJV-6eZI-dY7gi8y125BQ,7070
 eegdash/dataset/dataset_summary.csv,sha256=XF0vdHz77DFyVLTaET8lL5gQQ4r-q1xAfSDWH5GTPLA,23655
-eegdash/dataset/registry.py,sha256=4L8KHlCrY7LT2SDAjKSiEuwklZVKbP-KZceoMDM3zO4,4332
+eegdash/dataset/registry.py,sha256=genOqAuf9cQBnHhPqRwfLP7S1XsnkLot6sLyJozPtf4,4150
 eegdash/features/__init__.py,sha256=BXNhjvL4_SSFAY1lcP9nyGpkbJNtoOMH4AHlF6OyABo,4078
-eegdash/features/datasets.py,sha256=kU1DO70ArSIy-LF1hHD2NN4iT-kJrI0mVpSkyV_OSeI,18301
+eegdash/features/datasets.py,sha256=eV4d86EU4fu1yoIMdPQnot6YZDRGG4qE9h77lk7iVhU,18317
 eegdash/features/decorators.py,sha256=v0qaJz_dcX703p1fvFYbAIXmwK3d8naYGlq7fRVKn_w,1313
 eegdash/features/extractors.py,sha256=H7h6tP3dKoRcjDJpWWAo0ppmokCq5QlhqMcehYwYV9s,6845
 eegdash/features/inspect.py,sha256=PmbWhx5H_WqpnorUpWONUSkUtaIHkZblRa_Xyk7Szyc,1569
@@ -26,10 +28,10 @@ eegdash/features/feature_bank/signal.py,sha256=3Tb8z9gX7iZipxQJ9DSyy30JfdmW58kgv
 eegdash/features/feature_bank/spectral.py,sha256=bNB7skusePs1gX7NOU6yRlw_Gr4UOCkO_ylkCgybzug,3319
 eegdash/features/feature_bank/utils.py,sha256=DGh-Q7-XFIittP7iBBxvsJaZrlVvuY5mw-G7q6C-PCI,1237
 eegdash/hbn/__init__.py,sha256=U8mK64napnKU746C5DOwkX7W7sg3iW5kb_cVv2pfFq0,394
-eegdash/hbn/preprocessing.py,sha256=7S_TTRKPKEk47tTnh2D6WExBt4cctAMxUxGDjJqq5lU,2221
-eegdash/hbn/windows.py,sha256=DU_QruLOHQOttZbXCgtO4mjKaG3E5STWjMQ0_s-g0gw,9929
-eegdash-0.3.8.dist-info/licenses/LICENSE,sha256=asisR-xupy_NrQBFXnx6yqXeZcYWLvbAaiETl25iXT0,931
-eegdash-0.3.8.dist-info/METADATA,sha256=BGXm7_rvJvGlDTbL-Zifh_-CeuIqdFyT4evdSeEkTMY,10294
-eegdash-0.3.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-eegdash-0.3.8.dist-info/top_level.txt,sha256=zavO69HQ6MyZM0aQMR2zUS6TAFc7bnN5GEpDpOpFZzU,8
-eegdash-0.3.8.dist-info/RECORD,,
+eegdash/hbn/preprocessing.py,sha256=zMkDFXQxWLn-Iy0cHRk-JL5c479HpclgIAGm-xuKPFg,2196
+eegdash/hbn/windows.py,sha256=Obv4L2hP8ACancUawkMYbiusfvhAw-mG_vqyDDMwOJ8,9890
+eegdash-0.3.9.dev129.dist-info/licenses/LICENSE,sha256=asisR-xupy_NrQBFXnx6yqXeZcYWLvbAaiETl25iXT0,931
+eegdash-0.3.9.dev129.dist-info/METADATA,sha256=_G43SOhR91fq4TW8pEem03XOxPDerz57mH0CejyMtvc,10362
+eegdash-0.3.9.dev129.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+eegdash-0.3.9.dev129.dist-info/top_level.txt,sha256=zavO69HQ6MyZM0aQMR2zUS6TAFc7bnN5GEpDpOpFZzU,8
+eegdash-0.3.9.dev129.dist-info/RECORD,,