eegdash 0.0.8__py3-none-any.whl → 0.1.0__py3-none-any.whl

eegdash/features/serialization.py

@@ -0,0 +1,87 @@
+"""
+Convenience functions for storing and loading of features datasets.
+
+see also: https://github.com/braindecode/braindecode//blob/master/braindecode/datautil/serialization.py#L165-L229
+"""
+
+import json
+from pathlib import Path
+
+import pandas as pd
+from joblib import Parallel, delayed
+from mne.io import read_info
+
+from braindecode.datautil.serialization import _load_kwargs_json
+
+from .datasets import (
+    FeaturesConcatDataset,
+    FeaturesDataset,
+)
+
+
+def load_features_concat_dataset(path, ids_to_load=None, n_jobs=1):
+    """Load a stored FeaturesConcatDataset of FeaturesDatasets from files.
+
+    Parameters
+    ----------
+    path: str | pathlib.Path
+        Path to the directory of the .fif / -epo.fif and .json files.
+    ids_to_load: list of int | None
+        Ids of specific files to load.
+    n_jobs: int
+        Number of jobs to be used to read files in parallel.
+
+    Returns
+    -------
+    concat_dataset: FeaturesConcatDataset of FeaturesDatasets
+    """
+    # Make sure we always work with a pathlib.Path
+    path = Path(path)
+
+    # else we have a dataset saved in the new way with subdirectories in path
+    # for every dataset with description.json and -feat.parquet,
+    # target_name.json, raw_preproc_kwargs.json, window_kwargs.json,
+    # window_preproc_kwargs.json, features_kwargs.json
+    if ids_to_load is None:
+        ids_to_load = [p.name for p in path.iterdir()]
+        ids_to_load = sorted(ids_to_load, key=lambda i: int(i))
+    ids_to_load = [str(i) for i in ids_to_load]
+
+    datasets = Parallel(n_jobs)(delayed(_load_parallel)(path, i) for i in ids_to_load)
+    return FeaturesConcatDataset(datasets)
+
+
+def _load_parallel(path, i):
+    sub_dir = path / i
+
+    parquet_name_pattern = "{}-feat.parquet"
+    parquet_file_name = parquet_name_pattern.format(i)
+    parquet_file_path = sub_dir / parquet_file_name
+
+    features = pd.read_parquet(parquet_file_path)
+
+    description_file_path = sub_dir / "description.json"
+    description = pd.read_json(description_file_path, typ="series")
+
+    raw_info_file_path = sub_dir / "raw-info.fif"
+    raw_info = None
+    if raw_info_file_path.exists():
+        raw_info = read_info(raw_info_file_path)
+
+    raw_preproc_kwargs = _load_kwargs_json("raw_preproc_kwargs", sub_dir)
+    window_kwargs = _load_kwargs_json("window_kwargs", sub_dir)
+    window_preproc_kwargs = _load_kwargs_json("window_preproc_kwargs", sub_dir)
+    features_kwargs = _load_kwargs_json("features_kwargs", sub_dir)
+    metadata = pd.read_pickle(path / i / "metadata_df.pkl")
+
+    dataset = FeaturesDataset(
+        features,
+        metadata=metadata,
+        description=description,
+        raw_info=raw_info,
+        raw_preproc_kwargs=raw_preproc_kwargs,
+        window_kwargs=window_kwargs,
+        window_preproc_kwargs=window_preproc_kwargs,
+        features_kwargs=features_kwargs,
+    )
+    return dataset

eegdash/features/utils.py

@@ -0,0 +1,116 @@
+import copy
+from collections.abc import Callable
+from typing import Dict, List
+
+import numpy as np
+import pandas as pd
+from joblib import Parallel, delayed
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from braindecode.datasets.base import (
+    BaseConcatDataset,
+    EEGWindowsDataset,
+    WindowsDataset,
+)
+
+from .datasets import FeaturesConcatDataset, FeaturesDataset
+from .extractors import FeatureExtractor
+
+
+def _extract_features_from_windowsdataset(
+    win_ds: EEGWindowsDataset | WindowsDataset,
+    feature_extractor: FeatureExtractor,
+    batch_size: int = 512,
+):
+    metadata = win_ds.metadata
+    if not win_ds.targets_from == "metadata":
+        metadata = copy.deepcopy(metadata)
+        metadata["orig_index"] = metadata.index
+        metadata.set_index(
+            ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"],
+            drop=False,
+            inplace=True,
+        )
+    win_dl = DataLoader(win_ds, batch_size=batch_size, shuffle=False, drop_last=False)
+    features_dict = dict()
+    ch_names = win_ds.raw.ch_names
+    for X, y, crop_inds in win_dl:
+        X = X.numpy()
+        if hasattr(y, "tolist"):
+            y = y.tolist()
+        win_dict = dict()
+        win_dict.update(
+            feature_extractor(X, _batch_size=X.shape[0], _ch_names=ch_names)
+        )
+        if not win_ds.targets_from == "metadata":
+            metadata.loc[crop_inds, "target"] = y
+        for k, v in win_dict.items():
+            if k not in features_dict:
+                features_dict[k] = []
+            features_dict[k].extend(v)
+    features_df = pd.DataFrame(features_dict)
+    if not win_ds.targets_from == "metadata":
+        metadata.set_index("orig_index", drop=False, inplace=True)
+        metadata.reset_index(drop=True, inplace=True)
+        metadata.drop("orig_index", axis=1, inplace=True)
+
+    # FUTURE: truly support WindowsDataset objects
+    return FeaturesDataset(
+        features_df,
+        metadata=metadata,
+        description=win_ds.description,
+        raw_info=win_ds.raw.info,
+        raw_preproc_kwargs=win_ds.raw_preproc_kwargs,
+        window_kwargs=win_ds.window_kwargs,
+        features_kwargs=feature_extractor.features_kwargs,
+    )
+
+
+def extract_features(
+    concat_dataset: BaseConcatDataset,
+    features: FeatureExtractor | Dict[str, Callable] | List[Callable],
+    *,
+    batch_size: int = 512,
+    n_jobs: int = 1,
+):
+    if isinstance(features, list):
+        features = dict(enumerate(features))
+    if not isinstance(features, FeatureExtractor):
+        features = FeatureExtractor(features)
+    feature_ds_list = list(
+        tqdm(
+            Parallel(n_jobs=n_jobs, return_as="generator")(
+                delayed(_extract_features_from_windowsdataset)(
+                    win_ds, features, batch_size
+                )
+                for win_ds in concat_dataset.datasets
+            ),
+            total=len(concat_dataset.datasets),
+            desc="Extracting features",
+        )
+    )
+    return FeaturesConcatDataset(feature_ds_list)
+
+
+def fit_feature_extractors(
+    concat_dataset: BaseConcatDataset,
+    features: FeatureExtractor | Dict[str, Callable] | List[Callable],
+    batch_size: int = 8192,
+):
+    if isinstance(features, list):
+        features = dict(enumerate(features))
+    if not isinstance(features, FeatureExtractor):
+        features = FeatureExtractor(features)
+    if not features._is_fitable:
+        return features
+    features.clear()
+    concat_dl = DataLoader(
+        concat_dataset, batch_size=batch_size, shuffle=False, drop_last=False
+    )
+    for X, y, _ in tqdm(
+        concat_dl, total=len(concat_dl), desc="Fitting feature extractors"
+    ):
+        features.partial_fit(X.numpy(), y=np.array(y))
+    features.fit()
+    return features