Generate SPD matrices for EEG-based BCI

generate_spd_mat.py

import numpy as np
from numpy.random import chisquare
from scipy.stats import lognorm
import scipy as sp

import pyriemann

from sklearn.pipeline import Pipeline
from sklearn.decomposition import PCA
from pyriemann.tangentspace import TangentSpace

import matplotlib.pyplot as plt


def generate_samples_snr(n, m, snr=10., dof=3, alpha=1e-6):
    """Generate sample with structured and unstructured noise

    No the best generative model in practice, SNR estimation is
    interesting but sample diversity is difficult to harness with
    this formalism.
    """
    U = 2 * np.random.rand(n, n) - 1
    U = U / np.resize(np.linalg.norm(U, axis=0), (n, n))
    C, D = [], []
    for _ in range(m):
        Dk = np.diag(chisquare(dof, n) / dof
                     * np.array([0.5**i for i in range(1, n+1)]))
        D.append(Dk)
        signal = U @ Dk @ U.T
        V = 2 * np.random.rand(n, n) - 1
        V = V / np.resize(np.linalg.norm(V, axis=0), (n, n))
        E = np.diag(chisquare(dof, n) / dof
                    * np.array([0.5**i for i in range(1, n+1)]))
        struct_noise = V @ E @ V.T
        uncorr_noise = alpha * np.eye(n)
        v = np.trace(signal) / (snr * np.trace(struct_noise + uncorr_noise))
        C.append(signal + v*(struct_noise + uncorr_noise))
    LD = 0
    for d in D:
        LD += logm(d)
    P = expm(LD/m)
    P2 = U @ np.array(D).mean(axis=0) @ U.T
    return P, P2, np.array(C)


def generate_dataset_snr(n, m, snr):
    """Generate dataset with structured and unstructured noise.
    """
    _, P, C = generate_samples_snr(n, m, snr=snr, dof=1, alpha=1e-6)
    return P, C


def generate_reference(n=20):
    """Generate a reference matrix P and its parameters

    Parameters
    ----------
    n: int
       SPD matrix dimension

    Returns
    -------
    P: ndarray, shape (n, n)
       Reference matrix
    D: ndarray, shape (n, )
       Diagonal elements
    U: ndarray, shape (n, n)
       Mixing matrix
    """
    dummyMat = np.random.rand(n, 2 * n)
    U, _, _ = np.linalg.svd(dummyMat, full_matrices=True)
    D = np.random.triangular(1, 2, 5, n)
    P = U @ np.diag(D) @ U.T
    return P, D, U


def generate_robust_samples(n, m, D, U, epsilon):
    """Generate m covariance matrix samples from a reference

    Perturbate diagonal elements from the reference to generate
    a list of m matrices.

    Parameters
    ----------
    n: int
       SPD matrices dimension
    m: int
       number of matrices
    D: ndarray, shape (n,)
       Diagonal elements
    U: ndarray, shape(n, n)
       Mixing matrix
    epsilon: float
       Perturbation to apply on the reference

    Returns
    -------
    C: ndarray, shape (m, n, n)
       SPD matrices
    """
    C = np.array([U @ np.diag(D + np.random.normal(0, epsilon, n)) @ U.T
                  for _ in range(m)])
    for i in range(m):
        while np.any(np.linalg.eigvalsh(C[i]) < 0.):
            C[i] = U @ np.diag(D + np.random.normal(0, epsilon, n)) @ U.T
    return C


def generate_dataset_dispersion(n, m, ep):
    """Generate m covariance matrix samples from a reference

    Perturbate diagonal elements from the reference to generate
    a list of m matrices.

    Parameters
    ----------
    n: int
       SPD matrices dimension
    m: int
       number of matrices
    epsilon: float
       Perturbation to apply on the reference

    Returns
    -------
    P: ndarray, shape (n, n)
       Reference matrix
    C: ndarray, shape (m, n, n)
       SPD matrices
    """
    
    P, D, U = generate_reference(n)
    C = generate_robust_samples(n, m, D, U, ep)
    return P, C

def viz_pca_ts_dataset(obs_covs, groundtruth):
    m = obs_covs.shape[0]
    obs_mean = pyriemann.utils.mean.mean_riemann(obs_covs)
    ts = Pipeline([('mapping', TangentSpace(metric='riemann', tsupdate=False)),
                   ('dim_reduc', PCA(n_components=2))])
    ts.fit(np.concatenate((obs_covs, groundtruth[np.newaxis, ...],
                           obs_mean[np.newaxis, ...])))

    C_ts = ts.transform(np.concatenate((obs_covs, groundtruth[np.newaxis, ...],
                                        obs_mean[np.newaxis, ...])))

    fig, ax = plt.subplots(1, 1, figsize=(9, 9))
    ax.set_title("Tangent space, original data")
    ax.scatter(C_ts[0:m, 0], C_ts[0:m, 1], c="g", alpha=0.3, label=r'$C_k$')
    ax.scatter(C_ts[-2, 0], C_ts[-2, 1], c="k",
               label=r'ground truth mean $\mathcal{G}$', marker='*', s=300)
    ax.scatter(C_ts[-1, 0], C_ts[-1, 1], c="r",
               label=r'observed mean $\hat{\mathcal{G}}$', marker='*', s=300)
    _ = ax.legend()
    plt.show()


if __name__ == "__main__":
    n = 10
    m = 1000
    epsilon = 0.4

    P, C = generate_dataset_dispersion(n, m, epsilon)

    # plot reference matrix
    plt.figure()
    plt.imshow(P)
    plt.xticks([])
    _ = plt.yticks([])

    # Plot 10 first covariance matrices
    plt.figure()
    for i in range(10):
        plt.subplot(2, 5, i+1)
        plt.imshow(C[i])
        plt.xticks([])
        plt.yticks([])
        plt.tight_layout()

    viz_pca_ts_dataset(C, P)