Schaefer-Strimmer shrinkage covariance estimator

doc/whatsnew.rst

@@ -14,6 +14,8 @@ v0.2.8.dev
 
 - Correct spectral estimation in :func:`pyriemann.utils.covariance.cross_spectrum` to obtain equivalence with SciPy
 
+- Add Schaefer-Strimmer covariance estimator in :func:`pyriemann.utils.covariance.covariances`, and an example to compare estimators
+
 v0.2.7 (June 2021)
 ------------------
 

examples/signal/README.txt

@@ -0,0 +1,4 @@
+Preprocessing
+-------------------
+
+Using pyRiemann for signal processing and covariance estimation

examples/signal/plot_covariance_estimation.py

@@ -0,0 +1,208 @@
+"""
+===============================================================================
+Estimate covariance with different time windows
+===============================================================================
+
+Covariance estimators comparison for different EEG signal lengths and their
+impact on classification [1]_.
+"""
+# Author: Sylvain Chevallier
+#
+# License: BSD (3-clause)
+
+from matplotlib import pyplot as plt
+from mne import Epochs, pick_types, events_from_annotations
+from mne.io import concatenate_raws
+from mne.io.edf import read_raw_edf
+from mne.datasets import eegbci
+import numpy as np
+import pandas as pd
+from pyriemann.estimation import Covariances
+from pyriemann.utils.distance import distance
+from pyriemann.classification import MDM
+import seaborn as sns
+from sklearn.model_selection import cross_val_score, StratifiedKFold
+from sklearn.pipeline import make_pipeline
+
+
+rs = np.random.RandomState(42)
+
+###############################################################################
+# Estimating covariance on synthetic data
+# ----------------------------------------
+#
+# Generate synthetic data, sampled from a distribution considered as the
+# groundtruth.
+
+n_trials, n_channels, n_times = 10, 5, 1000
+var = 2.0 + 0.1 * rs.randn(n_trials, n_channels)
+A = 2 * rs.rand(n_channels, n_channels) - 1
+A /= np.linalg.norm(A, axis=1)[:, np.newaxis]
+true_cov = np.empty(shape=(n_trials, n_channels, n_channels))
+X = np.empty(shape=(n_trials, n_channels, n_times))
+for i in range(n_trials):
+    true_cov[i] = A @ np.diag(var[i]) @ A.T
+    X[i] = rs.multivariate_normal(
+        np.array([0.0] * n_channels), true_cov[i], size=n_times
+    ).T
+
+###############################################################################
+# Covariances() object offers several estimators: SCM, Ledoit-Wolf (LWF),
+# Schaefer-Strimmer (SCH), oracle approximating shrunk covariance (OAS),
+# minimum covariance determinant (MCD) and others. We will compare the
+# distance of LWF, OAS and SCH estimators with the groundtruth, while
+# increasing epoch length.
+
+estimators = ["lwf", "oas", "sch"]
+w_len = np.linspace(10, n_times, 20, dtype=int)
+dfd = list()
+for est in estimators:
+    for wl in w_len:
+        cov_est = Covariances(estimator=est).transform(X[:, :, :wl])
+        for k in range(n_trials):
+            dist = distance(cov_est[k], true_cov[k], metric="riemann")
+            dfd.append(dict(estimator=est, wlen=wl, dist=dist))
+dfd = pd.DataFrame(dfd)
+
+###############################################################################
+
+fig, ax = plt.subplots(figsize=(6, 4))
+ax.set(xscale="log")
+sns.lineplot(data=dfd, x="wlen", y="dist", hue="estimator", ax=ax)
+ax.set_title("Distance to groundtruth covariance matrix")
+ax.set_xlabel("Number of samples")
+ax.set_ylabel(r"$\delta(\Sigma, \hat{\Sigma})$")
+plt.tight_layout()
+
+###############################################################################
+# Choice of estimator for motor imagery data
+# -------------------------------------------
+# Loading data from PhysioNet MI dataset, for subject 1.
+
+event_id = dict(hands=2, feet=3)
+subject = 1
+runs = [6, 10, 14]  # motor imagery: hands vs feet
+raw_files = [
+    read_raw_edf(f, preload=True, stim_channel="auto")
+    for f in eegbci.load_data(subject, runs)
+]
+raw = concatenate_raws(raw_files)
+picks = pick_types(raw.info, eeg=True, exclude="bads")
+
+# subsample elecs
+picks = picks[::2]
+# Apply band-pass filter
+raw.filter(7.0, 35.0, method="iir", picks=picks, skip_by_annotation="edge")
+events, _ = events_from_annotations(raw, event_id=dict(T1=2, T2=3))
+event_ids = dict(hands=2, feet=3)
+
+###############################################################################
+# Influence of shrinkage to estimate covariance
+# -----------------------------------------------
+# Sample covariance matrix (SCM) estimation could lead to ill-conditionned
+# matrices depending on the quality and quantity of EEG data available.
+# Matrix condition number is the ratio between the highest and lowest
+# eigenvalues: high values indicates ill-conditionned matrices that are not
+# suitable for classification.
+# A common approach to mitigate this issue is to regularize covariance matrices
+# by shrinkage, like in Ledoit-Wolf, Schaefer-Strimmer or oracle estimators.
+
+estimators = ["lwf", "oas", "sch", "scm"]
+tmin = -0.2
+w_len = np.linspace(0.2, 2, 10)
+n_trials = 45
+dfc = list()
+
+for wl in w_len:
+    epochs = Epochs(
+        raw,
+        events,
+        event_ids,
+        tmin,
+        tmin + wl,
+        picks=picks,
+        preload=True,
+        verbose=False,
+    )
+    for est in estimators:
+        cov = Covariances(estimator=est).transform(epochs.get_data())
+        for k in range(len(cov)):
+            ev, _ = np.linalg.eigh(cov[k, :, :])
+            dfc.append(dict(estimator=est, wlen=wl, cond=ev[-1] / ev[0]))
+dfc = pd.DataFrame(dfc)
+
+###############################################################################
+
+fig, ax = plt.subplots(figsize=(6, 4))
+ax.set(yscale="log")
+sns.lineplot(data=dfc, x="wlen", y="cond", hue="estimator", ax=ax)
+ax.set_title("Condition number of estimated covariance matrices")
+ax.set_xlabel("Epoch length (s)")
+ax.set_ylabel(r"$\lambda_{\max}$/$\lambda_{\min}$")
+plt.tight_layout()
+
+###############################################################################
+# Picking a good estimator for classification
+# -----------------------------------------------
+# The choice of covariance estimator have an impact on classification,
+# especially when the covariances are estimated on short time windows.
+
+estimators = ["lwf", "oas", "sch", "scm"]
+tmin = 0.0
+w_len = np.linspace(0.2, 2.0, 5)
+n_trials, n_splits = 45, 3
+dfa = list()
+sc = "balanced_accuracy"
+
+cv = StratifiedKFold(n_splits=n_splits, shuffle=True)
+for wl in w_len:
+    epochs = Epochs(
+        raw,
+        events,
+        event_ids,
+        tmin,
+        tmin + wl,
+        proj=True,
+        picks=picks,
+        preload=True,
+        baseline=None,
+        verbose=False,
+    )
+    X = epochs.get_data()
+    y = np.array([0 if ev == 2 else 1 for ev in epochs.events[:, -1]])
+    for est in estimators:
+        clf = make_pipeline(Covariances(estimator=est), MDM())
+        try:
+            score = cross_val_score(clf, X, y, cv=cv, scoring=sc)
+            dfa += [dict(estimator=est, wlen=wl, accuracy=sc) for sc in score]
+        except ValueError:
+            print(f"{est}: {wl} is not sufficent to estimate a PSD matrix")
+            dfa += [dict(estimator=est, wlen=wl, accuracy=np.nan)] * n_splits
+dfa = pd.DataFrame(dfa)
+
+###############################################################################
+
+fig, ax = plt.subplots(figsize=(6, 4))
+sns.lineplot(
+    data=dfa,
+    x="wlen",
+    y="accuracy",
+    hue="estimator",
+    style="estimator",
+    ax=ax,
+    ci=None,
+    markers=True,
+    dashes=False,
+)
+ax.set_title("Accuracy for different estimators and epoch lengths")
+ax.set_xlabel("Epoch length (s)")
+ax.set_ylabel("Classification accuracy")
+plt.tight_layout()
+
+###############################################################################
+# References
+# ----------
+# .. [1] S. Chevallier, E. Kalunga, Q. Barthélemy, F. Yger. "Riemannian
+#    classification for SSVEP based BCI: offline versus online implementations"
+#    Brain–Computer Interfaces Handbook: Technological and Theoretical Advances
+#    , 2018.

pyriemann/utils/covariance.py

@@ -30,6 +30,52 @@ def _mcd(X):
     return C
 
 
+def _sch(X):
+    """Schaefer-Strimmer covariance estimator
+
+    Shrinkage estimator using method from [1]_:
+    .. math::
+            \hat{\Sigma} = (1 - \gamma)\Sigma_{scm} + \gamma T
+
+    where :math:`T` is the diagonal target matrix:
+    .. math::
+            T_{i,j} = \{ \Sigma_{scm}^{ii} \text{if} i = j, 0 \text{otherwise} \}
+    Note that the optimal :math:`\gamma` is estimated by the authors' method.
+
+    :param X: Signal matrix, (n_channels, n_times)
+
+    :returns: Schaefer-Strimmer shrinkage covariance matrix, (n_channels, n_channels)
+
+    Notes
+    -----
+    .. versionadded:: 0.2.8.dev
+
+    References
+    ----------
+    .. [1] Schafer, J., and K. Strimmer. 2005. A shrinkage approach to
+        large-scale covariance estimation and implications for functional
+        genomics. Statist. Appl. Genet. Mol. Biol. 4:32.
+    """  # noqa
+    n_times = X.shape[1]
+    X_c = (X.T - X.T.mean(axis=0)).T
+    C_scm = 1. / n_times * X_c @ X_c.T
+
+    # Compute optimal gamma, the weigthing between SCM and srinkage estimator
+    R = (n_times / ((n_times - 1.) * np.outer(X.std(axis=1), X.std(axis=1))))
+    R *= C_scm
+    var_R = (X_c ** 2) @ (X_c ** 2).T - 2 * C_scm * (X_c @ X_c.T)
+    var_R += n_times * C_scm ** 2
+    Xvar = np.outer(X.var(axis=1), X.var(axis=1))
+    var_R *= n_times / ((n_times - 1) ** 3 * Xvar)
+    R -= np.diag(np.diag(R))
+    var_R -= np.diag(np.diag(var_R))
+    gamma = max(0, min(1, var_R.sum() / (R ** 2).sum()))
+
+    sigma = (1. - gamma) * (n_times / (n_times - 1.)) * C_scm
+    shrinkage = gamma * (n_times / (n_times - 1.)) * np.diag(np.diag(C_scm))
+    return sigma + shrinkage
+
+
 def _check_est(est):
     """Check if a given estimator is valid"""
 
@@ -40,6 +86,7 @@ def _check_est(est):
         'lwf': _lwf,
         'oas': _oas,
         'mcd': _mcd,
+        'sch': _sch,
         'corr': np.corrcoef
     }
 
@@ -65,7 +112,7 @@ def covariances(X, estimator='cov'):
     X : ndarray, shape (n_trials, n_channels, n_times)
         ndarray of trials.
 
-    estimator : {'cov', 'scm', 'lwf', 'oas', 'mcd', 'corr'} (default: 'scm')
+    estimator : {'cov', 'scm', 'lwf', 'oas', 'mcd', 'sch', 'corr'} (default: 'scm')
         covariance matrix estimator:
 
         * 'cov' for numpy based covariance matrix,
@@ -78,6 +125,8 @@ def covariances(X, estimator='cov'):
           https://scikit-learn.org/stable/modules/generated/sklearn.covariance.OAS.html
         * 'mcd' for minimum covariance determinant matrix,
           https://scikit-learn.org/stable/modules/generated/sklearn.covariance.MinCovDet.html
+        * 'sch' for Schaefer-Strimmer covariance,
+          http://doi.org/10.2202/1544-6115.1175,
         * 'corr' for correlation coefficient matrix,
           https://numpy.org/doc/stable/reference/generated/numpy.corrcoef.html
 
@@ -91,20 +140,24 @@ def covariances(X, estimator='cov'):
     .. [1] https://scikit-learn.org/stable/modules/covariance.html
     """  # noqa
     est = _check_est(estimator)
-    Nt, Ne, Ns = X.shape
-    covmats = np.zeros((Nt, Ne, Ne))
-    for i in range(Nt):
+    n_trials, n_channels, n_times = X.shape
+    covmats = np.zeros((n_trials, n_channels, n_channels))
+    for i in range(n_trials):
         covmats[i, :, :] = est(X[i, :, :])
     return covmats
 
 
 def covariances_EP(X, P, estimator='cov'):
     """Special form covariance matrix."""
     est = _check_est(estimator)
-    Nt, Ne, Ns = X.shape
-    Np, Ns = P.shape
-    covmats = np.zeros((Nt, Ne + Np, Ne + Np))
-    for i in range(Nt):
+    n_trials, n_channels, n_times = X.shape
+    n_proto, n_times_P = P.shape
+    if n_times_P != n_times:
+        raise ValueError(
+            f"X and P do not have the same n_times: {n_times} and {n_times_P}"
+        )
+    covmats = np.zeros((n_trials, n_channels + n_proto, n_channels + n_proto))
+    for i in range(n_trials):
         covmats[i, :, :] = est(np.concatenate((P, X[i, :, :]), axis=0))
     return covmats
 
@@ -117,10 +170,10 @@ def eegtocov(sig, window=128, overlapp=0.5, padding=True, estimator='cov'):
         padd = np.zeros((int(window / 2), sig.shape[1]))
         sig = np.concatenate((padd, sig, padd), axis=0)
 
-    Ns, Ne = sig.shape
+    n_times, n_channels = sig.shape
     jump = int(window * overlapp)
     ix = 0
-    while (ix + window < Ns):
+    while (ix + window < n_times):
         X.append(est(sig[ix:ix + window, :].T))
         ix = ix + jump
 

tests/test_tangentspace.py

@@ -1,6 +1,7 @@
 """Test tangent space functions."""
 import numpy as np
 from numpy.testing import assert_array_almost_equal
+import pytest
 
 from pyriemann.tangentspace import TangentSpace, FGDA
 
@@ -37,6 +38,19 @@ def test_TangentSpace_transform():
     ts.transform(covset)
 
 
+@pytest.mark.parametrize('shape', [(10, 9), (10, 9, 8), (10), (12, 8, 8)])
+def test_TangentSpace_transform_dim(shape):
+    """Test transform input shape, could be TS vector or covmat"""
+    n_trials, n_channels = 10, 3
+    covset = generate_cov(n_trials, n_channels)
+    ts = TangentSpace(metric='riemann')
+    ts.fit(covset)
+
+    X = np.zeros(shape=shape)
+    with pytest.raises(ValueError):
+        ts.transform(X)
+
+
 def test_TangentSpace_transform_without_fit():
     """Test transform of Tangent Space without fit."""
     covset = generate_cov(10, 3)