File size: 5,039 Bytes
3bb804c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
# Authors: The MNE-Python contributors.
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

import numpy as np
import scipy.linalg

from ..cov import Covariance, _smart_eigh, compute_whitener


def _handle_restr_mat(C_ref, restr_type, info, rank):
    """Get restricting matrix to C_ref rank-dimensional principal subspace.

    Returns matrix of shape (rank, n_chs) used to restrict or
    restrict+rescale (whiten) covariances matrices.
    """
    if C_ref is None or restr_type is None:
        return None
    if restr_type == "whitening":
        C_ref_cov = Covariance(C_ref, info.ch_names, info["bads"], info["projs"], 0)
        restr_mat = compute_whitener(
            C_ref_cov, info, rank=rank, pca=True, verbose="error"
        )[0]
    elif restr_type == "restricting":
        restr_mat = _get_restr_mat(C_ref, info, rank)
    else:
        raise ValueError(
            "restr_type should either be callable or one of "
            "('whitening', 'restricting')"
        )
    return restr_mat


def _smart_ged(S, R, restr_mat=None, R_func=None):
    """Perform smart generalized eigenvalue decomposition (GED) of S and R.

    If restr_mat is provided S and R will be restricted to the principal subspace
    of a reference matrix with rank r (see _handle_restr_mat), then GED is performed
    on the restricted S and R and then generalized eigenvectors are transformed back
    to the original space. The g-eigenvectors matrix is of shape (n_chs, r).
    If callable R_func is provided the GED will be performed on (S, R_func(S,R))
    """
    if restr_mat is None:
        evals, evecs = scipy.linalg.eigh(S, R)
        return evals, evecs

    S_restr = restr_mat @ S @ restr_mat.T
    R_restr = restr_mat @ R @ restr_mat.T
    if R_func is not None:
        R_restr = R_func([S_restr, R_restr])
    evals, evecs_restr = scipy.linalg.eigh(S_restr, R_restr)
    evecs = restr_mat.T @ evecs_restr

    return evals, evecs


def _is_cov_symm(cov, rtol=1e-7, atol=None):
    if atol is None:
        atol = 1e-7 * np.max(np.abs(cov))
    is_symm = scipy.linalg.issymmetric(cov, rtol=rtol, atol=atol)
    return is_symm


def _get_cov_def(cov, eval_tol=None):
    """Get definiteness of symmetric cov matrix.

    All evals in (-eval_tol, eval_tol) will be considered zero,
    while all evals smaller than -eval_tol will be considered
    negative.
    """
    evals = scipy.linalg.eigvalsh(cov)
    if eval_tol is None:
        eval_tol = 1e-7 * np.max(np.abs(evals))
    if np.all(evals > eval_tol):
        return "pos_def"
    elif np.all(evals >= -eval_tol):
        return "pos_semidef"
    else:
        return "indef"


def _is_cov_pos_semidef(cov, eval_tol=None):
    cov_def = _get_cov_def(cov, eval_tol=eval_tol)
    return cov_def in ("pos_def", "pos_semidef")


def _is_cov_pos_def(cov, eval_tol=None):
    cov_def = _get_cov_def(cov, eval_tol=eval_tol)
    return cov_def == "pos_def"


def _smart_ajd(covs, restr_mat=None, weights=None):
    """Perform smart approximate joint diagonalization.

    If restr_mat is provided all the cov matrices will be restricted to the
    principal subspace of a reference matrix with rank r (see _handle_restr_mat),
    then GED is performed on the restricted S and R and then generalized eigenvectors
    are transformed back to the original space.
    The matrix of generalized eigenvectors is of shape (n_chs, r).
    """
    from .csp import _ajd_pham

    if restr_mat is None:
        are_all_pos_def = all([_is_cov_pos_def(cov) for cov in covs])
        if not are_all_pos_def:
            raise ValueError(
                "If C_ref is not provided by covariance estimator, "
                "all the covs should be positive definite"
            )
        evecs, D = _ajd_pham(covs)
        return evecs

    else:
        are_all_pos_semidef = all([_is_cov_pos_semidef(cov) for cov in covs])
        if not are_all_pos_semidef:
            raise ValueError(
                "All the covs should be positive semi-definite for "
                "approximate joint diagonalization"
            )
        covs = np.array([restr_mat @ cov @ restr_mat.T for cov in covs], float)
        evecs_restr, D = _ajd_pham(covs)
        evecs = _normalize_eigenvectors(evecs_restr.T, covs, weights)
        evecs = restr_mat.T @ evecs
        return evecs


def _get_restr_mat(C, info, rank):
    """Get matrix restricting covariance to rank-dimensional principal subspace of C."""
    _, ref_evecs, mask = _smart_eigh(
        C,
        info,
        rank,
        proj_subspace=True,
        do_compute_rank=False,
        log_ch_type="data",
    )
    restr_mat = ref_evecs[mask]
    return restr_mat


def _normalize_eigenvectors(evecs, covs, sample_weights):
    # Here we apply an euclidean mean. See pyRiemann for other metrics
    mean_cov = np.average(covs, axis=0, weights=sample_weights)

    for ii in range(evecs.shape[1]):
        tmp = np.dot(np.dot(evecs[:, ii].T, mean_cov), evecs[:, ii])
        evecs[:, ii] /= np.sqrt(tmp)
    return evecs