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

eegdash/features/datasets.py

@@ -0,0 +1,456 @@
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import warnings
+from collections.abc import Callable, Iterable
+from typing import Dict, no_type_check
+
+import numpy as np
+import pandas as pd
+from joblib import Parallel, delayed
+
+from braindecode.datasets.base import (
+    BaseConcatDataset,
+    EEGWindowsDataset,
+    _create_description,
+)
+
+
+class FeaturesDataset(EEGWindowsDataset):
+    """Returns samples from a pandas DataFrame object along with a target.
+
+    Dataset which serves samples from a pandas DataFrame object along with a
+    target. The target is unique for the dataset, and is obtained through the
+    `description` attribute.
+
+    Parameters
+    ----------
+    features : a pandas DataFrame
+        Tabular data.
+    description : dict | pandas.Series | None
+        Holds additional description about the continuous signal / subject.
+    transform : callable | None
+        On-the-fly transform applied to the example before it is returned.
+    """
+
+    def __init__(
+        self,
+        features: pd.DataFrame,
+        metadata: pd.DataFrame | None = None,
+        description: dict | pd.Series | None = None,
+        transform: Callable | None = None,
+        raw_info: Dict | None = None,
+        raw_preproc_kwargs: Dict | None = None,
+        window_kwargs: Dict | None = None,
+        window_preproc_kwargs: Dict | None = None,
+        features_kwargs: Dict | None = None,
+    ):
+        self.features = features
+        self.n_features = features.columns.size
+        self.metadata = metadata
+        self._description = _create_description(description)
+        self.transform = transform
+        self.raw_info = raw_info
+        self.raw_preproc_kwargs = raw_preproc_kwargs
+        self.window_kwargs = window_kwargs
+        self.window_preproc_kwargs = window_preproc_kwargs
+        self.features_kwargs = features_kwargs
+
+        self.crop_inds = metadata.loc[
+            :, ["i_window_in_trial", "i_start_in_trial", "i_stop_in_trial"]
+        ].to_numpy()
+        self.y = metadata.loc[:, "target"].to_list()
+
+    def __getitem__(self, index):
+        crop_inds = self.crop_inds[index].tolist()
+        X = self.features.iloc[index].to_numpy()
+        X = X.copy()
+        X.astype("float32")
+        if self.transform is not None:
+            X = self.transform(X)
+        y = self.y[index]
+        return X, y, crop_inds
+
+    def __len__(self):
+        return len(self.features.index)
+
+
+def _compute_stats(
+    ds: FeaturesDataset,
+    return_count=False,
+    return_mean=False,
+    return_var=False,
+    ddof=1,
+    numeric_only=False,
+):
+    res = []
+    if return_count:
+        res.append(ds.features.count(numeric_only=numeric_only))
+    if return_mean:
+        res.append(ds.features.mean(numeric_only=numeric_only))
+    if return_var:
+        res.append(ds.features.var(ddof=ddof, numeric_only=numeric_only))
+    return tuple(res)
+
+
+def _pooled_var(counts, means, variances, ddof):
+    count = counts.sum(axis=0)
+    mean = np.sum((counts / count) * means, axis=0)
+    var = np.sum(((counts - ddof) / (count - ddof)) * variances, axis=0)
+    var[:] += np.sum((counts / (count - ddof)) * (means**2), axis=0)
+    var[:] -= (count / (count - ddof)) * (mean**2)
+    var[:] = var.clip(min=0)
+    return count, mean, var
+
+
+class FeaturesConcatDataset(BaseConcatDataset):
+    """A base class for concatenated datasets.
+
+    Holds either mne.Raw or mne.Epoch in self.datasets and has
+    a pandas DataFrame with additional description.
+
+    Parameters
+    ----------
+    list_of_ds : list
+        list of BaseDataset, BaseConcatDataset or WindowsDataset
+    target_transform : callable | None
+        Optional function to call on targets before returning them.
+
+    """
+
+    def __init__(
+        self,
+        list_of_ds: list[FeaturesDataset] | None = None,
+        target_transform: Callable | None = None,
+    ):
+        # if we get a list of FeaturesConcatDataset, get all the individual datasets
+        if list_of_ds and isinstance(list_of_ds[0], FeaturesConcatDataset):
+            list_of_ds = [d for ds in list_of_ds for d in ds.datasets]
+        super().__init__(list_of_ds)
+
+        self.target_transform = target_transform
+
+    def split(
+        self,
+        by: str | list[int] | list[list[int]] | dict[str, list[int]],
+    ) -> dict[str, FeaturesConcatDataset]:
+        """Split the dataset based on information listed in its description.
+
+        The format could be based on a DataFrame or based on indices.
+
+        Parameters
+        ----------
+        by : str | list | dict
+            If ``by`` is a string, splitting is performed based on the
+            description DataFrame column with this name.
+            If ``by`` is a (list of) list of integers, the position in the first
+            list corresponds to the split id and the integers to the
+            datapoints of that split.
+            If a dict then each key will be used in the returned
+            splits dict and each value should be a list of int.
+
+        Returns
+        -------
+        splits : dict
+            A dictionary with the name of the split (a string) as key and the
+            dataset as value.
+        """
+        if isinstance(by, str):
+            split_ids = {
+                k: list(v) for k, v in self.description.groupby(by).groups.items()
+            }
+        elif isinstance(by, dict):
+            split_ids = by
+        else:
+            # assume list(int)
+            if not isinstance(by[0], list):
+                by = [by]
+            # assume list(list(int))
+            split_ids = {split_i: split for split_i, split in enumerate(by)}
+
+        return {
+            str(split_name): FeaturesConcatDataset(
+                [self.datasets[ds_ind] for ds_ind in ds_inds],
+                target_transform=self.target_transform,
+            )
+            for split_name, ds_inds in split_ids.items()
+        }
+
+    def get_metadata(self) -> pd.DataFrame:
+        """Concatenate the metadata and description of the wrapped Epochs.
+
+        Returns
+        -------
+        metadata : pd.DataFrame
+            DataFrame containing as many rows as there are windows in the
+            BaseConcatDataset, with the metadata and description information
+            for each window.
+        """
+        if not all([isinstance(ds, FeaturesDataset) for ds in self.datasets]):
+            raise TypeError(
+                "Metadata dataframe can only be computed when all "
+                "datasets are FeaturesDataset."
+            )
+
+        all_dfs = list()
+        for ds in self.datasets:
+            df = ds.metadata
+            for k, v in ds.description.items():
+                df[k] = v
+            all_dfs.append(df)
+
+        return pd.concat(all_dfs)
+
+    def save(self, path: str, overwrite: bool = False, offset: int = 0):
+        """Save datasets to files by creating one subdirectory for each dataset:
+        path/
+            0/
+                0-feat.parquet
+                metadata_df.pkl
+                description.json
+                raw-info.fif (if raw info was saved)
+                raw_preproc_kwargs.json (if raws were preprocessed)
+                window_kwargs.json (if this is a windowed dataset)
+                window_preproc_kwargs.json  (if windows were preprocessed)
+                features_kwargs.json
+            1/
+                1-feat.parquet
+                metadata_df.pkl
+                description.json
+                raw-info.fif (if raw info was saved)
+                raw_preproc_kwargs.json (if raws were preprocessed)
+                window_kwargs.json (if this is a windowed dataset)
+                window_preproc_kwargs.json  (if windows were preprocessed)
+                features_kwargs.json
+
+        Parameters
+        ----------
+        path : str
+            Directory in which subdirectories are created to store
+             -feat.parquet and .json files to.
+        overwrite : bool
+            Whether to delete old subdirectories that will be saved to in this
+            call.
+        offset : int
+            If provided, the integer is added to the id of the dataset in the
+            concat. This is useful in the setting of very large datasets, where
+            one dataset has to be processed and saved at a time to account for
+            its original position.
+        """
+        if len(self.datasets) == 0:
+            raise ValueError("Expect at least one dataset")
+        path_contents = os.listdir(path)
+        n_sub_dirs = len([os.path.isdir(e) for e in path_contents])
+        for i_ds, ds in enumerate(self.datasets):
+            # remove subdirectory from list of untouched files / subdirectories
+            if str(i_ds + offset) in path_contents:
+                path_contents.remove(str(i_ds + offset))
+            # save_dir/i_ds/
+            sub_dir = os.path.join(path, str(i_ds + offset))
+            if os.path.exists(sub_dir):
+                if overwrite:
+                    shutil.rmtree(sub_dir)
+                else:
+                    raise FileExistsError(
+                        f"Subdirectory {sub_dir} already exists. Please select"
+                        f" a different directory, set overwrite=True, or "
+                        f"resolve manually."
+                    )
+            # save_dir/{i_ds+offset}/
+            os.makedirs(sub_dir)
+            # save_dir/{i_ds+offset}/{i_ds+offset}-feat.parquet
+            self._save_features(sub_dir, ds, i_ds, offset)
+            # save_dir/{i_ds+offset}/metadata_df.pkl
+            self._save_metadata(sub_dir, ds)
+            # save_dir/{i_ds+offset}/description.json
+            self._save_description(sub_dir, ds.description)
+            # save_dir/{i_ds+offset}/raw-info.fif
+            self._save_raw_info(sub_dir, ds)
+            # save_dir/{i_ds+offset}/raw_preproc_kwargs.json
+            # save_dir/{i_ds+offset}/window_kwargs.json
+            # save_dir/{i_ds+offset}/window_preproc_kwargs.json
+            # save_dir/{i_ds+offset}/features_kwargs.json
+            self._save_kwargs(sub_dir, ds)
+        if overwrite:
+            # 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(
+                    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 "
+                    f"encounter a mix of differently preprocessed "
+                    f"datasets!",
+                    UserWarning,
+                )
+        # if path contains files or directories that were not touched, raise
+        # warning
+        if path_contents:
+            warnings.warn(
+                f"Chosen directory {path} contains other "
+                f"subdirectories or files {path_contents}."
+            )
+
+    @staticmethod
+    def _save_features(sub_dir, ds, i_ds, offset):
+        parquet_file_name = f"{i_ds + offset}-feat.parquet"
+        parquet_file_path = os.path.join(sub_dir, parquet_file_name)
+        ds.features.to_parquet(parquet_file_path)
+
+    @staticmethod
+    def _save_raw_info(sub_dir, ds):
+        if hasattr(ds, "raw_info"):
+            fif_file_name = "raw-info.fif"
+            fif_file_path = os.path.join(sub_dir, fif_file_name)
+            ds.raw_info.save(fif_file_path)
+
+    @staticmethod
+    def _save_kwargs(sub_dir, ds):
+        for kwargs_name in [
+            "raw_preproc_kwargs",
+            "window_kwargs",
+            "window_preproc_kwargs",
+            "features_kwargs",
+        ]:
+            if hasattr(ds, kwargs_name):
+                kwargs_file_name = ".".join([kwargs_name, "json"])
+                kwargs_file_path = os.path.join(sub_dir, kwargs_file_name)
+                kwargs = getattr(ds, kwargs_name)
+                if kwargs is not None:
+                    with open(kwargs_file_path, "w") as f:
+                        json.dump(kwargs, f)
+
+    def to_dataframe(
+        self, include_metadata=False, include_target=False, include_crop_inds=False
+    ):
+        if include_metadata or (include_target and include_crop_inds):
+            dataframes = [
+                ds.metadata.join(ds.features, how="right", lsuffix="_metadata")
+                for ds in self.datasets
+            ]
+        elif include_target:
+            dataframes = [
+                ds.features.join(ds.metadata["target"], how="left", rsuffix="_metadata")
+                for ds in self.datasets
+            ]
+        elif include_crop_inds:
+            dataframes = [
+                ds.metadata.drop("target", axis="columns").join(
+                    ds.features, how="right", lsuffix="_metadata"
+                )
+                for ds in self.datasets
+            ]
+        else:
+            dataframes = [ds.features for ds in self.datasets]
+        return pd.concat(dataframes, axis=0, ignore_index=True)
+
+    def _numeric_columns(self):
+        return self.datasets[0].features.select_dtypes(include=np.number).columns
+
+    def count(self, numeric_only=False, n_jobs=1):
+        stats = Parallel(n_jobs)(
+            delayed(_compute_stats)(ds, return_count=True, numeric_only=numeric_only)
+            for ds in self.datasets
+        )
+        counts = np.array([s[0] for s in stats])
+        count = counts.sum(axis=0)
+        return pd.Series(count, index=self._numeric_columns())
+
+    def mean(self, numeric_only=False, n_jobs=1):
+        stats = Parallel(n_jobs)(
+            delayed(_compute_stats)(
+                ds, return_count=True, return_mean=True, numeric_only=numeric_only
+            )
+            for ds in self.datasets
+        )
+        counts, means = np.array([s[0] for s in stats]), np.array([s[1] for s in stats])
+        count = counts.sum(axis=0, keepdims=True)
+        mean = np.sum((counts / count) * means, axis=0)
+        return pd.Series(mean, index=self._numeric_columns())
+
+    def var(self, ddof=1, numeric_only=False, n_jobs=1):
+        stats = Parallel(n_jobs)(
+            delayed(_compute_stats)(
+                ds,
+                return_count=True,
+                return_mean=True,
+                return_var=True,
+                ddof=ddof,
+                numeric_only=numeric_only,
+            )
+            for ds in self.datasets
+        )
+        counts, means, variances = (
+            np.array([s[0] for s in stats]),
+            np.array([s[1] for s in stats]),
+            np.array([s[2] for s in stats]),
+        )
+        _, _, var = _pooled_var(counts, means, variances, ddof)
+        return pd.Series(var, index=self._numeric_columns())
+
+    def std(self, ddof=1, numeric_only=False, n_jobs=1):
+        return np.sqrt(self.var(ddof=ddof, numeric_only=numeric_only, n_jobs=n_jobs))
+
+    def zscore(self, ddof=1, numeric_only=False, eps=0, n_jobs=1):
+        stats = Parallel(n_jobs)(
+            delayed(_compute_stats)(
+                ds,
+                return_count=True,
+                return_mean=True,
+                return_var=True,
+                ddof=ddof,
+                numeric_only=numeric_only,
+            )
+            for ds in self.datasets
+        )
+        counts, means, variances = (
+            np.array([s[0] for s in stats]),
+            np.array([s[1] for s in stats]),
+            np.array([s[2] for s in stats]),
+        )
+        _, mean, var = _pooled_var(counts, means, variances, ddof)
+        std = np.sqrt(var) + eps
+        for ds in self.datasets:
+            ds.features = (ds.features - mean) / std
+
+    @staticmethod
+    def _enforce_inplace_operations(func_name, kwargs):
+        if "inplace" in kwargs and kwargs["inplace"] is False:
+            raise ValueError(
+                f"{func_name} only works inplace, please change "
+                + "to inplace=True (default)."
+            )
+        kwargs["inplace"] = True
+
+    def fillna(self, *args, **kwargs):
+        FeaturesConcatDataset._enforce_inplace_operations("fillna", kwargs)
+        for ds in self.datasets:
+            ds.features.fillna(*args, **kwargs)
+
+    def replace(self, *args, **kwargs):
+        FeaturesConcatDataset._enforce_inplace_operations("replace", kwargs)
+        for ds in self.datasets:
+            ds.features.replace(*args, **kwargs)
+
+    def interpolate(self, *args, **kwargs):
+        FeaturesConcatDataset._enforce_inplace_operations("interpolate", kwargs)
+        for ds in self.datasets:
+            ds.features.interpolate(*args, **kwargs)
+
+    def dropna(self, *args, **kwargs):
+        FeaturesConcatDataset._enforce_inplace_operations("dropna", kwargs)
+        for ds in self.datasets:
+            ds.features.dropna(*args, **kwargs)
+
+    def drop(self, *args, **kwargs):
+        FeaturesConcatDataset._enforce_inplace_operations("drop", kwargs)
+        for ds in self.datasets:
+            ds.features.drop(*args, **kwargs)
+
+    def join(self, concat_dataset: FeaturesConcatDataset, **kwargs):
+        assert len(self.datasets) == len(concat_dataset.datasets)
+        for ds1, ds2 in zip(self.datasets, concat_dataset.datasets):
+            assert len(ds1) == len(ds2)
+            ds1.features.join(ds2, **kwargs)

eegdash/features/decorators.py

@@ -0,0 +1,43 @@
+from collections.abc import Callable
+from typing import List, Type
+
+from .extractors import (
+    BivariateFeature,
+    DirectedBivariateFeature,
+    FeatureExtractor,
+    MultivariateFeature,
+    UnivariateFeature,
+    _get_underlying_func,
+)
+
+
+class FeaturePredecessor:
+    def __init__(self, *parent_extractor_type: List[Type]):
+        parent_cls = parent_extractor_type
+        if not parent_cls:
+            parent_cls = [FeatureExtractor]
+        for p_cls in parent_cls:
+            assert issubclass(p_cls, FeatureExtractor)
+        self.parent_extractor_type = parent_cls
+
+    def __call__(self, func: Callable):
+        f = _get_underlying_func(func)
+        f.parent_extractor_type = self.parent_extractor_type
+        return func
+
+
+class FeatureKind:
+    def __init__(self, feature_kind: MultivariateFeature):
+        self.feature_kind = feature_kind
+
+    def __call__(self, func):
+        f = _get_underlying_func(func)
+        f.feature_kind = self.feature_kind
+        return func
+
+
+# Syntax sugar
+univariate_feature = FeatureKind(UnivariateFeature())
+bivariate_feature = FeatureKind(BivariateFeature())
+directed_bivariate_feature = FeatureKind(DirectedBivariateFeature())
+multivariate_feature = FeatureKind(MultivariateFeature())

eegdash/features/extractors.py

@@ -0,0 +1,210 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from functools import partial
+from typing import Dict
+
+import numpy as np
+from numba.core.dispatcher import Dispatcher
+
+
+def _get_underlying_func(func):
+    f = func
+    if isinstance(f, partial):
+        f = f.func
+    if isinstance(f, Dispatcher):
+        f = f.py_func
+    return f
+
+
+class FitableFeature(ABC):
+    def __init__(self):
+        self._is_fitted = False
+        self.clear()
+
+    @abstractmethod
+    def clear(self):
+        pass
+
+    @abstractmethod
+    def partial_fit(self, *x, y=None):
+        pass
+
+    def fit(self):
+        self._is_fitted = True
+
+    def __call__(self, *args, **kwargs):
+        if not self._is_fitted:
+            raise RuntimeError(
+                f"{self.__class__} cannot be called, it has to be fitted first."
+            )
+
+
+class FeatureExtractor(FitableFeature):
+    def __init__(
+        self, feature_extractors: Dict[str, Callable], **preprocess_kwargs: Dict
+    ):
+        self.feature_extractors_dict = self._validate_execution_tree(feature_extractors)
+        self._is_fitable = self._check_is_fitable(feature_extractors)
+        super().__init__()
+
+        # bypassing FeaturePredecessor to avoid circular import
+        if not hasattr(self, "parent_extractor_type"):
+            self.parent_extractor_type = [FeatureExtractor]
+
+        self.preprocess_kwargs = preprocess_kwargs
+        if self.preprocess_kwargs is None:
+            self.preprocess_kwargs = dict()
+        self.features_kwargs = {
+            "preprocess_kwargs": preprocess_kwargs,
+        }
+        for fn, fe in feature_extractors.items():
+            if isinstance(fe, FeatureExtractor):
+                self.features_kwargs[fn] = fe.features_kwargs
+            if isinstance(fe, partial):
+                self.features_kwargs[fn] = fe.keywords
+
+    def _validate_execution_tree(self, feature_extractors):
+        for fname, f in feature_extractors.items():
+            f = _get_underlying_func(f)
+            pe_type = getattr(f, "parent_extractor_type", [FeatureExtractor])
+            assert type(self) in pe_type
+        return feature_extractors
+
+    def _check_is_fitable(self, feature_extractors):
+        is_fitable = False
+        for fname, f in feature_extractors.items():
+            if isinstance(f, FeatureExtractor):
+                is_fitable = f._is_fitable
+            else:
+                f = _get_underlying_func(f)
+                if isinstance(f, FitableFeature):
+                    is_fitable = True
+            if is_fitable:
+                break
+        return is_fitable
+
+    def preprocess(self, *x, **kwargs):
+        return (*x,)
+
+    def feature_channel_names(self, ch_names):
+        return [""]
+
+    def __call__(self, *x, _batch_size=None, _ch_names=None):
+        assert _batch_size is not None
+        assert _ch_names is not None
+        if self._is_fitable:
+            super().__call__()
+        results_dict = dict()
+        z = self.preprocess(*x, **self.preprocess_kwargs)
+        for fname, f in self.feature_extractors_dict.items():
+            if isinstance(f, FeatureExtractor):
+                r = f(*z, _batch_size=_batch_size, _ch_names=_ch_names)
+            else:
+                r = f(*z)
+            f = _get_underlying_func(f)
+            if hasattr(f, "feature_kind"):
+                r = f.feature_kind(r, _ch_names=_ch_names)
+            if not isinstance(fname, str) or not fname:
+                if isinstance(f, FeatureExtractor) or not hasattr(f, "__name__"):
+                    fname = ""
+                else:
+                    fname = f.__name__
+            if isinstance(r, dict):
+                if fname:
+                    fname += "_"
+                for k, v in r.items():
+                    self._add_feature_to_dict(results_dict, fname + k, v, _batch_size)
+            else:
+                self._add_feature_to_dict(results_dict, fname, r, _batch_size)
+        return results_dict
+
+    def _add_feature_to_dict(self, results_dict, name, value, batch_size):
+        if not isinstance(value, np.ndarray):
+            results_dict[name] = value
+        else:
+            assert value.shape[0] == batch_size
+            results_dict[name] = value
+
+    def clear(self):
+        if not self._is_fitable:
+            return
+        for fname, f in self.feature_extractors_dict.items():
+            f = _get_underlying_func(f)
+            if isinstance(f, FitableFeature):
+                f.clear()
+
+    def partial_fit(self, *x, y=None):
+        if not self._is_fitable:
+            return
+        z = self.preprocess(*x, **self.preprocess_kwargs)
+        for fname, f in self.feature_extractors_dict.items():
+            f = _get_underlying_func(f)
+            if isinstance(f, FitableFeature):
+                f.partial_fit(*z, y=y)
+
+    def fit(self):
+        if not self._is_fitable:
+            return
+        for fname, f in self.feature_extractors_dict.items():
+            f = _get_underlying_func(f)
+            if isinstance(f, FitableFeature):
+                f.fit()
+        super().fit()
+
+
+class MultivariateFeature:
+    def __call__(self, x, _ch_names=None):
+        assert _ch_names is not None
+        f_channels = self.feature_channel_names(_ch_names)
+        if isinstance(x, dict):
+            r = dict()
+            for k, v in x.items():
+                r.update(self._array_to_dict(v, f_channels, k))
+            return r
+        return self._array_to_dict(x, f_channels)
+
+    @staticmethod
+    def _array_to_dict(x, f_channels, name=""):
+        assert isinstance(x, np.ndarray)
+        if len(f_channels) == 0:
+            assert x.ndim == 1
+            if name:
+                return {name: x}
+            return x
+        assert x.shape[1] == len(f_channels)
+        x = x.swapaxes(0, 1)
+        names = [f"{name}_{ch}" for ch in f_channels] if name else f_channels
+        return dict(zip(names, x))
+
+    def feature_channel_names(self, ch_names):
+        return []
+
+
+class UnivariateFeature(MultivariateFeature):
+    def feature_channel_names(self, ch_names):
+        return ch_names
+
+
+class BivariateFeature(MultivariateFeature):
+    def __init__(self, *args, channel_pair_format="{}<>{}"):
+        super().__init__(*args)
+        self.channel_pair_format = channel_pair_format
+
+    @staticmethod
+    def get_pair_iterators(n):
+        return np.triu_indices(n, 1)
+
+    def feature_channel_names(self, ch_names):
+        return [
+            self.channel_pair_format.format(ch_names[i], ch_names[j])
+            for i, j in zip(*self.get_pair_iterators(len(ch_names)))
+        ]
+
+
+class DirectedBivariateFeature(BivariateFeature):
+    @staticmethod
+    def get_pair_iterators(n):
+        return [
+            np.append(a, b)
+            for a, b in zip(np.tril_indices(n, -1), np.triu_indices(n, 1))
+        ]