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

eegdash/features/datasets.py

@@ -1,16 +1,19 @@
 from __future__ import annotations
-import os
+
 import json
+import os
 import shutil
 import warnings
-from typing import Dict, no_type_check
 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 (
-    EEGWindowsDataset,
     BaseConcatDataset,
+    EEGWindowsDataset,
     _create_description,
 )
 

eegdash/features/decorators.py

@@ -1,14 +1,14 @@
-from typing import List, Type
 from collections.abc import Callable
+from typing import List, Type
 
 from .extractors import (
-    FeatureExtractor,
-    UnivariateFeature,
     BivariateFeature,
     DirectedBivariateFeature,
+    FeatureExtractor,
     MultivariateFeature,
+    UnivariateFeature,
+    _get_underlying_func,
 )
-from .extractors import _get_underlying_func
 
 
 class FeaturePredecessor:

eegdash/features/extractors.py

@@ -1,7 +1,8 @@
 from abc import ABC, abstractmethod
-from typing import Dict
 from collections.abc import Callable
 from functools import partial
+from typing import Dict
+
 import numpy as np
 from numba.core.dispatcher import Dispatcher
 

eegdash/features/feature_bank/__init__.py

@@ -1,6 +1,6 @@
-from .signal import *
-from .spectral import *
 from .complexity import *
-from .dimensionality import *
 from .connectivity import *
 from .csp import *
+from .dimensionality import *
+from .signal import *
+from .spectral import *

eegdash/features/feature_bank/complexity.py

@@ -1,10 +1,9 @@
-import numpy as np
 import numba as nb
+import numpy as np
 from sklearn.neighbors import KDTree
 
-from ..extractors import FeatureExtractor
 from ..decorators import FeaturePredecessor, univariate_feature
-
+from ..extractors import FeatureExtractor
 
 __all__ = [
     "EntropyFeatureExtractor",

eegdash/features/feature_bank/connectivity.py

@@ -1,10 +1,11 @@
 from itertools import chain
+
 import numpy as np
 from scipy.signal import csd
 
-from ..extractors import FeatureExtractor, BivariateFeature
 from ..decorators import FeaturePredecessor, bivariate_feature
-
+from ..extractors import BivariateFeature, FeatureExtractor
+from . import utils
 
 __all__ = [
     "CoherenceFeatureExtractor",
@@ -18,82 +19,41 @@ class CoherenceFeatureExtractor(FeatureExtractor):
     def preprocess(self, x, **kwargs):
         f_min = kwargs.pop("f_min") if "f_min" in kwargs else None
         f_max = kwargs.pop("f_max") if "f_max" in kwargs else None
+        assert "fs" in kwargs and "nperseg" in kwargs
         kwargs["axis"] = -1
         n = x.shape[1]
         idx_x, idx_y = BivariateFeature.get_pair_iterators(n)
         ix, iy = list(chain(range(n), idx_x)), list(chain(range(n), idx_y))
         f, s = csd(x[:, ix], x[:, iy], **kwargs)
-        if f_min is not None or f_max is not None:
-            f_min_idx = f > f_min if f_min is not None else True
-            f_max_idx = f < f_max if f_max is not None else True
-            idx = np.logical_and(f_min_idx, f_max_idx)
-            f = f[idx]
-            s = s[..., idx]
-        sx, sxy = np.split(s, [n], axis=1)
-        sxx, syy = sx[:, idx_x].real, sx[:, idx_y].real
+        f_min, f_max = utils.get_valid_freq_band(
+            kwargs["fs"], x.shape[-1], f_min, f_max
+        )
+        f, s = utils.slice_freq_band(f, s, f_min=f_min, f_max=f_max)
+        p, sxy = np.split(s, [n], axis=1)
+        sxx, syy = p[:, idx_x].real, p[:, idx_y].real
         c = sxy / np.sqrt(sxx * syy)
         return f, c
 
 
-def _avg_over_bands(f, x, bands):
-    bands_avg = dict()
-    for k, v in bands.items():
-        assert isinstance(k, str)
-        assert isinstance(v, tuple)
-        assert len(v) == 2
-        assert v[0] < v[1]
-        mask = np.logical_and(f > v[0], f < v[1])
-        avg = x[..., mask].mean(axis=-1)
-        bands_avg[k] = avg
-    return bands_avg
-
-
 @FeaturePredecessor(CoherenceFeatureExtractor)
 @bivariate_feature
-def connectivity_magnitude_square_coherence(
-    f,
-    c,
-    bands={
-        "delta": (1, 4.5),
-        "theta": (4.5, 8),
-        "alpha": (8, 12),
-        "beta": (12, 30),
-    },
-):
+def connectivity_magnitude_square_coherence(f, c, bands=utils.DEFAULT_FREQ_BANDS):
     # https://neuroimage.usc.edu/brainstorm/Tutorials/Connectivity
     coher = c.real**2 + c.imag**2
-    return _avg_over_bands(f, coher, bands)
+    return utils.reduce_freq_bands(f, coher, bands, np.mean)
 
 
 @FeaturePredecessor(CoherenceFeatureExtractor)
 @bivariate_feature
-def connectivity_imaginary_coherence(
-    f,
-    c,
-    bands={
-        "delta": (1, 4.5),
-        "theta": (4.5, 8),
-        "alpha": (8, 12),
-        "beta": (12, 30),
-    },
-):
+def connectivity_imaginary_coherence(f, c, bands=utils.DEFAULT_FREQ_BANDS):
     # https://neuroimage.usc.edu/brainstorm/Tutorials/Connectivity
     coher = c.imag
-    return _avg_over_bands(f, coher, bands)
+    return utils.reduce_freq_bands(f, coher, bands, np.mean)
 
 
 @FeaturePredecessor(CoherenceFeatureExtractor)
 @bivariate_feature
-def connectivity_lagged_coherence(
-    f,
-    c,
-    bands={
-        "delta": (1, 4.5),
-        "theta": (4.5, 8),
-        "alpha": (8, 12),
-        "beta": (12, 30),
-    },
-):
+def connectivity_lagged_coherence(f, c, bands=utils.DEFAULT_FREQ_BANDS):
     # https://neuroimage.usc.edu/brainstorm/Tutorials/Connectivity
     coher = c.imag / np.sqrt(1 - c.real)
-    return _avg_over_bands(f, coher, bands)
+    return utils.reduce_freq_bands(f, coher, bands, np.mean)

eegdash/features/feature_bank/csp.py

@@ -1,11 +1,10 @@
-import numpy as np
 import numba as nb
+import numpy as np
 import scipy
 import scipy.linalg
 
-from ..extractors import FitableFeature
 from ..decorators import multivariate_feature
-
+from ..extractors import FitableFeature
 
 __all__ = [
     "CommonSpatialPattern",

eegdash/features/feature_bank/dimensionality.py

@@ -1,11 +1,10 @@
-import numpy as np
 import numba as nb
+import numpy as np
 from scipy import special
 
 from ..decorators import univariate_feature
 from .signal import signal_zero_crossings
 
-
 __all__ = [
     "dimensionality_higuchi_fractal_dim",
     "dimensionality_petrosian_fractal_dim",

eegdash/features/feature_bank/signal.py

@@ -1,10 +1,10 @@
 import numbers
+
 import numpy as np
 from scipy import stats
 
 from ..decorators import univariate_feature
 
-
 __all__ = [
     "signal_mean",
     "signal_variance",

eegdash/features/feature_bank/spectral.py

@@ -1,10 +1,10 @@
-import numpy as np
 import numba as nb
+import numpy as np
 from scipy.signal import welch
 
-from ..extractors import FeatureExtractor
 from ..decorators import FeaturePredecessor, univariate_feature
-
+from ..extractors import FeatureExtractor
+from . import utils
 
 __all__ = [
     "SpectralFeatureExtractor",
@@ -26,14 +26,13 @@ class SpectralFeatureExtractor(FeatureExtractor):
     def preprocess(self, x, **kwargs):
         f_min = kwargs.pop("f_min") if "f_min" in kwargs else None
         f_max = kwargs.pop("f_max") if "f_max" in kwargs else None
+        assert "fs" in kwargs
         kwargs["axis"] = -1
         f, p = welch(x, **kwargs)
-        if f_min is not None or f_max is not None:
-            f_min_idx = f > f_min if f_min is not None else True
-            f_max_idx = f < f_max if f_max is not None else True
-            idx = np.logical_and(f_min_idx, f_max_idx)
-            f = f[idx]
-            p = p[..., idx]
+        f_min, f_max = utils.get_valid_freq_band(
+            kwargs["fs"], x.shape[-1], f_min, f_max
+        )
+        f, p = utils.slice_freq_band(f, p, f_min=f_min, f_max=f_max)
         return f, p
 
 
@@ -113,22 +112,5 @@ def spectral_slope(f, p):
     DBSpectralFeatureExtractor,
 )
 @univariate_feature
-def spectral_bands_power(
-    f,
-    p,
-    bands={
-        "delta": (1, 4.5),
-        "theta": (4.5, 8),
-        "alpha": (8, 12),
-        "beta": (12, 30),
-    },
-):
-    bands_power = dict()
-    for k, v in bands.items():
-        assert isinstance(k, str)
-        assert isinstance(v, tuple)
-        assert len(v) == 2
-        mask = np.logical_and(f > v[0], f < v[1])
-        power = p[..., mask].sum(axis=-1)
-        bands_power[k] = power
-    return bands_power
+def spectral_bands_power(f, p, bands=utils.DEFAULT_FREQ_BANDS):
+    return utils.reduce_freq_bands(f, p, bands, np.sum)

eegdash/features/feature_bank/utils.py

@@ -0,0 +1,48 @@
+import numpy as np
+
+DEFAULT_FREQ_BANDS = {
+    "delta": (1, 4.5),
+    "theta": (4.5, 8),
+    "alpha": (8, 12),
+    "beta": (12, 30),
+}
+
+
+def get_valid_freq_band(fs, n, f_min=None, f_max=None):
+    f0 = 2 * fs / n
+    f1 = fs / 2
+    if f_min is None:
+        f_min = f0
+    else:
+        assert f_min >= f0
+    if f_max is None:
+        f_max = f1
+    else:
+        assert f_max <= f1
+    return f_min, f_max
+
+
+def slice_freq_band(f, *x, f_min=None, f_max=None):
+    if f_min is None and f_max is None:
+        return f, *x
+    else:
+        f_min_idx = f >= f_min if f_min is not None else True
+        f_max_idx = f <= f_max if f_max is not None else True
+        idx = np.logical_and(f_min_idx, f_max_idx)
+        f = f[idx]
+        xl = [*x]
+        for i, xi in enumerate(xl):
+            xl[i] = xi[..., idx]
+        return f, *xl
+
+
+def reduce_freq_bands(f, x, bands, reduce_func=np.sum):
+    x_bands = dict()
+    for k, lims in bands.items():
+        assert isinstance(k, str)
+        assert len(lims) == 2 and lims[0] <= lims[1]
+        assert lims[0] >= f[0] and lims[1] <= f[-1]
+        mask = np.logical_and(f >= lims[0], f < lims[1])
+        xf = x[..., mask]
+        x_bands[k] = reduce_func(xf, axis=-1)
+    return x_bands

eegdash/features/serialization.py

@@ -9,13 +9,13 @@ 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 (
-    FeaturesDataset,
     FeaturesConcatDataset,
+    FeaturesDataset,
 )
 
 

eegdash/features/utils.py

@@ -1,18 +1,20 @@
-from typing import Dict, List
-from collections.abc import Callable
 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 tqdm import tqdm
 from torch.utils.data import DataLoader
+from tqdm import tqdm
+
 from braindecode.datasets.base import (
+    BaseConcatDataset,
     EEGWindowsDataset,
     WindowsDataset,
-    BaseConcatDataset,
 )
 
-from .datasets import FeaturesDataset, FeaturesConcatDataset
+from .datasets import FeaturesConcatDataset, FeaturesDataset
 from .extractors import FeatureExtractor
 
 
@@ -53,7 +55,7 @@ def _extract_features_from_windowsdataset(
         metadata.reset_index(drop=True, inplace=True)
         metadata.drop("orig_index", axis=1, inplace=True)
 
-    # FUTURE: truely support WindowsDataset objects
+    # FUTURE: truly support WindowsDataset objects
     return FeaturesDataset(
         features_df,
         metadata=metadata,