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

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

eegdash/features/feature_bank/signal.py

@@ -0,0 +1,103 @@
+import numbers
+import numpy as np
+from scipy import stats
+
+from ..decorators import univariate_feature
+
+
+__all__ = [
+    "signal_mean",
+    "signal_variance",
+    "signal_skewness",
+    "signal_kurtosis",
+    "signal_std",
+    "signal_root_mean_square",
+    "signal_peak_to_peak",
+    "signal_quantile",
+    "signal_zero_crossings",
+    "signal_line_length",
+    "signal_hjorth_activity",
+    "signal_hjorth_mobility",
+    "signal_hjorth_complexity",
+    "signal_decorrelation_time",
+]
+
+
+@univariate_feature
+def signal_mean(x):
+    return x.mean(axis=-1)
+
+
+@univariate_feature
+def signal_variance(x, **kwargs):
+    return x.var(axis=-1, **kwargs)
+
+
+@univariate_feature
+def signal_std(x, **kwargs):
+    return x.std(axis=-1, **kwargs)
+
+
+@univariate_feature
+def signal_skewness(x, **kwargs):
+    return stats.skew(x, axis=x.ndim - 1, **kwargs)
+
+
+@univariate_feature
+def signal_kurtosis(x, **kwargs):
+    return stats.kurtosis(x, axis=x.ndim - 1, **kwargs)
+
+
+@univariate_feature
+def signal_root_mean_square(x):
+    return np.sqrt(np.power(x, 2).mean(axis=-1))
+
+
+@univariate_feature
+def signal_peak_to_peak(x, **kwargs):
+    return np.ptp(x, axis=-1, **kwargs)
+
+
+@univariate_feature
+def signal_quantile(x, q: numbers.Number = 0.5, **kwargs):
+    return np.quantile(x, q=q, axis=-1, **kwargs)
+
+
+@univariate_feature
+def signal_line_length(x):
+    return np.abs(np.diff(x, axis=-1)).mean(axis=-1)
+
+
+@univariate_feature
+def signal_zero_crossings(x, threshold=1e-15):
+    zero_ind = np.logical_and(x > -threshold, x < threshold)
+    zero_cross = np.diff(zero_ind, axis=-1).astype(int).sum(axis=-1)
+    y = x.copy()
+    y[zero_ind] = 0
+    zero_cross += np.sum(np.signbit(y[..., :-1]) != np.signbit(y[..., 1:]), axis=-1)
+    return zero_cross
+
+
+@univariate_feature
+def signal_hjorth_mobility(x):
+    return np.diff(x, axis=-1).std(axis=-1) / x.std(axis=-1)
+
+
+@univariate_feature
+def signal_hjorth_complexity(x):
+    return np.diff(x, 2, axis=-1).std(axis=-1) / x.std(axis=-1)
+
+
+@univariate_feature
+def signal_decorrelation_time(x, fs=1):
+    f = np.fft.fft(x - x.mean(axis=-1, keepdims=True), axis=-1)
+    ac = np.fft.ifft(f.real**2 + f.imag**2, axis=-1)[..., : x.shape[-1] // 2]
+    dct = np.empty(x.shape[:-1])
+    for i in np.ndindex(x.shape[:-1]):
+        dct[i] = np.searchsorted(ac[i] <= 0, True)
+    return dct / fs
+
+
+# =================================  Aliases  =================================
+
+signal_hjorth_activity = signal_variance

eegdash/features/feature_bank/spectral.py

@@ -0,0 +1,134 @@
+import numpy as np
+import numba as nb
+from scipy.signal import welch
+
+from ..extractors import FeatureExtractor
+from ..decorators import FeaturePredecessor, univariate_feature
+
+
+__all__ = [
+    "SpectralFeatureExtractor",
+    "NormalizedSpectralFeatureExtractor",
+    "DBSpectralFeatureExtractor",
+    "spectral_root_total_power",
+    "spectral_moment",
+    "spectral_entropy",
+    "spectral_edge",
+    "spectral_slope",
+    "spectral_bands_power",
+    "spectral_hjorth_activity",
+    "spectral_hjorth_mobility",
+    "spectral_hjorth_complexity",
+]
+
+
+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
+        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]
+        return f, p
+
+
+@FeaturePredecessor(SpectralFeatureExtractor)
+class NormalizedSpectralFeatureExtractor(FeatureExtractor):
+    def preprocess(self, *x):
+        return (*x[:-1], x[-1] / x[-1].sum(axis=-1, keepdims=True))
+
+
+@FeaturePredecessor(SpectralFeatureExtractor)
+class DBSpectralFeatureExtractor(FeatureExtractor):
+    def preprocess(self, *x, eps=1e-15):
+        return (*x[:-1], 10 * np.log10(x[-1] + eps))
+
+
+@FeaturePredecessor(SpectralFeatureExtractor)
+@univariate_feature
+def spectral_root_total_power(f, p):
+    return np.sqrt(p.sum(axis=-1))
+
+
+@FeaturePredecessor(NormalizedSpectralFeatureExtractor)
+@univariate_feature
+def spectral_moment(f, p):
+    return np.sum(f * p, axis=-1)
+
+
+@FeaturePredecessor(SpectralFeatureExtractor)
+@univariate_feature
+def spectral_hjorth_activity(f, p):
+    return np.sum(p, axis=-1)
+
+
+@FeaturePredecessor(NormalizedSpectralFeatureExtractor)
+@univariate_feature
+def spectral_hjorth_mobility(f, p):
+    return np.sqrt(np.sum(np.power(f, 2) * p, axis=-1))
+
+
+@FeaturePredecessor(NormalizedSpectralFeatureExtractor)
+@univariate_feature
+def spectral_hjorth_complexity(f, p):
+    return np.sqrt(np.sum(np.power(f, 4) * p, axis=-1))
+
+
+@FeaturePredecessor(NormalizedSpectralFeatureExtractor)
+@univariate_feature
+def spectral_entropy(f, p):
+    idx = p > 0
+    plogp = np.zeros_like(p)
+    plogp[idx] = p[idx] * np.log(p[idx])
+    return -np.sum(plogp, axis=-1)
+
+
+@FeaturePredecessor(NormalizedSpectralFeatureExtractor)
+@univariate_feature
+@nb.njit(cache=True, fastmath=True)
+def spectral_edge(f, p, edge=0.9):
+    se = np.empty(p.shape[:-1])
+    for i in np.ndindex(p.shape[:-1]):
+        se[i] = f[np.searchsorted(np.cumsum(p[i]), edge)]
+    return se
+
+
+@FeaturePredecessor(DBSpectralFeatureExtractor)
+@univariate_feature
+def spectral_slope(f, p):
+    log_f = np.vstack((np.log(f), np.ones(f.shape[0]))).T
+    r = np.linalg.lstsq(log_f, p.reshape(-1, p.shape[-1]).T)[0]
+    r = r.reshape(2, *p.shape[:-1])
+    return {"exp": r[0], "int": r[1]}
+
+
+@FeaturePredecessor(
+    SpectralFeatureExtractor,
+    NormalizedSpectralFeatureExtractor,
+    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