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

eegdash/features/decorators.py

@@ -0,0 +1,43 @@
+from typing import List, Type
+from collections.abc import Callable
+
+from .extractors import (
+    FeatureExtractor,
+    UnivariateFeature,
+    BivariateFeature,
+    DirectedBivariateFeature,
+    MultivariateFeature,
+)
+from .extractors import _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,209 @@
+from abc import ABC, abstractmethod
+from typing import Dict
+from collections.abc import Callable
+from functools import partial
+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))
+        ]

eegdash/features/feature_bank/__init__.py

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

eegdash/features/feature_bank/complexity.py

@@ -0,0 +1,97 @@
+import numpy as np
+import numba as nb
+from sklearn.neighbors import KDTree
+
+from ..extractors import FeatureExtractor
+from ..decorators import FeaturePredecessor, univariate_feature
+
+
+__all__ = [
+    "EntropyFeatureExtractor",
+    "complexity_approx_entropy",
+    "complexity_sample_entropy",
+    "complexity_svd_entropy",
+    "complexity_lempel_ziv",
+]
+
+
+@nb.njit(cache=True, fastmath=True)
+def _create_embedding(x, dim, lag):
+    y = np.empty(((x.shape[-1] - dim + 1) // lag, dim))
+    for i in range(0, x.shape[-1] - dim + 1, lag):
+        y[i] = x[i : i + dim]
+    return y
+
+
+def _channel_app_samp_entropy_counts(x, m, r, l):
+    x_emb = _create_embedding(x, m, l)
+    kdtree = KDTree(x_emb, metric="chebyshev")
+    return kdtree.query_radius(x_emb, r, count_only=True)
+
+
+class EntropyFeatureExtractor(FeatureExtractor):
+    def preprocess(self, x, m=2, r=0.2, l=1):
+        rr = r * x.std(axis=-1)
+        counts_m = np.empty((*x.shape[:-1], (x.shape[-1] - m + 1) // l))
+        counts_mp1 = np.empty((*x.shape[:-1], (x.shape[-1] - m) // l))
+        for i in np.ndindex(x.shape[:-1]):
+            counts_m[*i, :] = _channel_app_samp_entropy_counts(x[i], m, rr[i], l)
+            counts_mp1[*i, :] = _channel_app_samp_entropy_counts(x[i], m + 1, rr[i], l)
+        return counts_m, counts_mp1
+
+
+@FeaturePredecessor(EntropyFeatureExtractor)
+@univariate_feature
+def complexity_approx_entropy(counts_m, counts_mp1):
+    phi_m = np.log(counts_m / counts_m.shape[-1]).mean(axis=-1)
+    phi_mp1 = np.log(counts_mp1 / counts_mp1.shape[-1]).mean(axis=-1)
+    return phi_m - phi_mp1
+
+
+@FeaturePredecessor(EntropyFeatureExtractor)
+@univariate_feature
+def complexity_sample_entropy(counts_m, counts_mp1):
+    A = np.sum(counts_mp1 - 1, axis=-1)
+    B = np.sum(counts_m - 1, axis=-1)
+    return -np.log(A / B)
+
+
+@univariate_feature
+def complexity_svd_entropy(x, m=10, tau=1):
+    x_emb = np.empty((*x.shape[:-1], (x.shape[-1] - m + 1) // tau, m))
+    for i in np.ndindex(x.shape[:-1]):
+        x_emb[*i, :, :] = _create_embedding(x[i], m, tau)
+    s = np.linalg.svdvals(x_emb)
+    s /= s.sum(axis=-1, keepdims=True)
+    return -np.sum(s * np.log(s), axis=-1)
+
+
+@univariate_feature
+@nb.njit(cache=True, fastmath=True)
+def complexity_lempel_ziv(x, threshold=None):
+    lzc = np.empty(x.shape[:-1])
+    for i in np.ndindex(x.shape[:-1]):
+        t = np.median(x[i]) if threshold is None else threshold
+        s = x[i] > t
+        n = s.shape[0]
+        j, k, l = 0, 1, 1
+        k_max = 1
+        lzc[i] = 1
+        while True:
+            if s[j + k - 1] == s[l + k - 1]:
+                k += 1
+                if l + k > n:
+                    lzc[i] += 1
+                    break
+            else:
+                k_max = np.maximum(k, k_max)
+                j += 1
+                if j == l:
+                    lzc[i] += 1
+                    l += k_max
+                    if l + 1 > n:
+                        break
+                    j, k, k_max = 0, 1, 1
+                else:
+                    k = 1
+    return lzc

eegdash/features/feature_bank/connectivity.py

@@ -0,0 +1,99 @@
+from itertools import chain
+import numpy as np
+from scipy.signal import csd
+
+from ..extractors import FeatureExtractor, BivariateFeature
+from ..decorators import FeaturePredecessor, bivariate_feature
+
+
+__all__ = [
+    "CoherenceFeatureExtractor",
+    "connectivity_magnitude_square_coherence",
+    "connectivity_imaginary_coherence",
+    "connectivity_lagged_coherence",
+]
+
+
+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
+        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
+        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),
+    },
+):
+    # https://neuroimage.usc.edu/brainstorm/Tutorials/Connectivity
+    coher = c.real**2 + c.imag**2
+    return _avg_over_bands(f, coher, bands)
+
+
+@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),
+    },
+):
+    # https://neuroimage.usc.edu/brainstorm/Tutorials/Connectivity
+    coher = c.imag
+    return _avg_over_bands(f, coher, bands)
+
+
+@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),
+    },
+):
+    # https://neuroimage.usc.edu/brainstorm/Tutorials/Connectivity
+    coher = c.imag / np.sqrt(1 - c.real)
+    return _avg_over_bands(f, coher, bands)

eegdash/features/feature_bank/csp.py

@@ -0,0 +1,102 @@
+import numpy as np
+import numba as nb
+import scipy
+import scipy.linalg
+
+from ..extractors import FitableFeature
+from ..decorators import multivariate_feature
+
+
+__all__ = [
+    "CommonSpatialPattern",
+]
+
+
+@nb.njit(cache=True, fastmath=True, parallel=True)
+def _update_mean_cov(count, mean, cov, x_count, x_mean, x_cov):
+    alpha2 = x_count / count
+    alpha1 = 1 - alpha2
+    cov[:] = alpha1 * (cov + np.outer(mean, mean))
+    cov[:] += alpha2 * (x_cov + np.outer(x_mean, x_mean))
+    mean[:] = alpha1 * mean + alpha2 * x_mean
+    cov[:] -= np.outer(mean, mean)
+
+
+@multivariate_feature
+class CommonSpatialPattern(FitableFeature):
+    def __init__(self):
+        super().__init__()
+
+    def clear(self):
+        self._labels = None
+        self._counts = np.array([0, 0])
+        self._means = np.array([None, None])
+        self._covs = np.array([None, None])
+        self._mean = None
+        self._eigvals = None
+        self._weights = None
+
+    def _update_labels(self, labels):
+        if self._labels is None:
+            self._labels = labels
+        else:
+            for label in labels:
+                if label not in self._labels:
+                    self._labels = np.append(self._labels, label)
+        assert self._labels.shape[0] < 3
+        return self._labels
+
+    def _update_stats(self, l, x):
+        x_count, x_mean, x_cov = x.shape[0], x.mean(axis=0), np.cov(x.T, ddof=0)
+        if self._counts[l] == 0:
+            self._counts[l] = x_count
+            self._means[l] = x_mean
+            self._covs[l] = x_cov
+        else:
+            self._counts[l] += x_count
+            _update_mean_cov(
+                self._counts[l], self._means[l], self._covs[l], x_count, x_mean, x_cov
+            )
+
+    def partial_fit(self, x, y=None):
+        labels = self._update_labels(np.unique(y))
+        for i, l in enumerate(labels):
+            ind = (y == l).nonzero()[0]
+            if ind.shape[0] > 0:
+                xl = self.transform_input(x[ind])
+                self._update_stats(i, xl)
+
+    @staticmethod
+    def transform_input(x):
+        return x.swapaxes(1, 2).reshape(-1, x.shape[1])
+
+    def fit(self):
+        alphas = self._counts / self._counts.sum()
+        self._mean = np.sum(alphas * self._means)
+        for l in range(len(self._labels)):
+            self._covs[l] *= self._counts[l] / (self._counts[1] - 1)
+        l, w = scipy.linalg.eig(self._covs[0], self._covs[0] + self._covs[1])
+        l = l.real
+        ind = l > 0
+        l, w = l[ind], w[:, ind]
+        ord = np.abs(l - 0.5).argsort()[::-1]
+        self._eigvals = l[ord]
+        self._weights = w[:, ord]
+        super().fit()
+
+    def __call__(self, x, n_select=None, crit_select=None):
+        super().__call__()
+        w = self._weights
+        if n_select:
+            w = w[:, :n_select]
+        if crit_select:
+            sel = 0.5 - np.abs(self._eigvals - 0.5) < crit_select
+            w = w[:, sel]
+        if w.shape[-1] == 0:
+            raise RuntimeError(
+                "CSP weights selection criterion is too strict,"
+                + "all weights were filtered out."
+            )
+        proj = (self.transform_input(x) - self._mean) @ w
+        proj = proj.reshape(x.shape[0], x.shape[2], -1).mean(axis=1)
+        return {f"{i}": proj[:, i] for i in range(proj.shape[-1])}

eegdash/features/feature_bank/dimensionality.py

@@ -0,0 +1,108 @@
+import numpy as np
+import numba as nb
+from scipy import special
+
+from ..decorators import univariate_feature
+from .signal import signal_zero_crossings
+
+
+__all__ = [
+    "dimensionality_higuchi_fractal_dim",
+    "dimensionality_petrosian_fractal_dim",
+    "dimensionality_katz_fractal_dim",
+    "dimensionality_hurst_exp",
+    "dimensionality_detrended_fluctuation_analysis",
+]
+
+
+@univariate_feature
+@nb.njit(cache=True, fastmath=True)
+def dimensionality_higuchi_fractal_dim(x, k_max=10, eps=1e-7):
+    N = x.shape[-1]
+    hfd = np.empty(x.shape[:-1])
+    log_k = np.vstack((-np.log(np.arange(1, k_max + 1)), np.ones(k_max))).T
+    L_km = np.empty(k_max)
+    L_k = np.empty(k_max)
+    for i in np.ndindex(x.shape[:-1]):
+        for k in range(1, k_max + 1):
+            for m in range(k):
+                L_km[m] = np.mean(np.abs(np.diff(x[*i, m:], n=k)))
+            L_k[k - 1] = (N - 1) * np.sum(L_km[:k]) / (k**3)
+        L_k = np.maximum(L_k, eps)
+        hfd[i] = np.linalg.lstsq(log_k, np.log(L_k))[0][0]
+    return hfd
+
+
+@univariate_feature
+def dimensionality_petrosian_fractal_dim(x):
+    nd = signal_zero_crossings(np.diff(x, axis=-1))
+    log_n = np.log(x.shape[-1])
+    return log_n / (np.log(nd) + log_n)
+
+
+@univariate_feature
+def dimensionality_katz_fractal_dim(x):
+    dists = np.abs(np.diff(x, axis=-1))
+    L = dists.sum(axis=-1)
+    a = dists.mean(axis=-1)
+    log_n = np.log(L / a)
+    d = np.abs(x[..., 1:] - x[..., 0, None]).max(axis=-1)
+    return log_n / (np.log(d / L) + log_n)
+
+
+@univariate_feature
+@nb.njit(cache=True, fastmath=True)
+def _hurst_exp(x, ns, a, gamma_ratios, log_n):
+    h = np.empty(x.shape[:-1])
+    rs = np.empty((ns.shape[0], x.shape[-1] // ns[0]))
+    log_rs = np.empty(ns.shape[0])
+    for i in np.ndindex(x.shape[:-1]):
+        t0 = 0
+        for j, n in enumerate(ns):
+            for k, t0 in enumerate(range(0, x.shape[-1], n)):
+                xj = x[i][t0 : t0 + n]
+                m = np.mean(xj)
+                y = xj - m
+                z = np.cumsum(y)
+                r = np.ptp(z)
+                s = np.sqrt(np.mean(y**2))
+                if s == 0.0:
+                    rs[j, k] = np.nan
+                else:
+                    rs[j, k] = r / s
+            log_rs[j] = np.log(np.nanmean(rs[j, : x.shape[1] // n]))
+            log_rs[j] -= np.log(np.sum(np.sqrt((n - a[:n]) / a[:n])) * gamma_ratios[j])
+        h[i] = 0.5 + np.linalg.lstsq(log_n, log_rs)[0][0]
+    return h
+
+
+@univariate_feature
+def dimensionality_hurst_exp(x):
+    ns = np.unique(np.power(2, np.arange(2, np.log2(x.shape[-1]) - 1)).astype(int))
+    idx = ns > 340
+    gamma_ratios = np.empty(ns.shape[0])
+    gamma_ratios[idx] = 1 / np.sqrt(ns[idx] / 2)
+    gamma_ratios[~idx] = special.gamma((ns[~idx] - 1) / 2) / special.gamma(ns[~idx] / 2)
+    gamma_ratios /= np.sqrt(np.pi)
+    log_n = np.vstack((np.log(ns), np.ones(ns.shape[0]))).T
+    a = np.arange(1, ns[-1], dtype=float)
+    return _hurst_exp(x, ns, a, gamma_ratios, log_n)
+
+
+@univariate_feature
+@nb.njit(cache=True, fastmath=True)
+def dimensionality_detrended_fluctuation_analysis(x):
+    ns = np.unique(np.floor(np.power(2, np.arange(2, np.log2(x.shape[-1]) - 1))))
+    a = np.vstack((np.arange(ns[-1]), np.ones(int(ns[-1])))).T
+    log_n = np.vstack((np.log(ns), np.ones(ns.shape[0]))).T
+    Fn = np.empty(ns.shape[0])
+    alpha = np.empty(x.shape[:-1])
+    for i in np.ndindex(x.shape[:-1]):
+        X = np.cumsum(x[i] - np.mean(x[i]))
+        for j, n in enumerate(ns):
+            n = int(n)
+            Z = np.reshape(X[: n * (X.shape[0] // n)], (n, X.shape[0] // n))
+            Fni2 = np.linalg.lstsq(a[:n], Z)[1] / n
+            Fn[j] = np.sqrt(np.mean(Fni2))
+        alpha[i] = np.linalg.lstsq(log_n, np.log(Fn))[0][0]
+    return alpha