#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Functions for estimating the precision matrix based on the covariance of
either the residuals (temporal based precision matrix) or of the measurements
(instance based precision matrix)
"""
from collections.abc import Iterable
import numpy as np
from rsatoolbox.data import average_dataset_by
from rsatoolbox.util.data_utils import get_unique_inverse
from scipy.linalg import sqrtm
def _check_demean(matrix):
"""
checks that an input has 2 or 3 dimensions and subtracts the mean.
returns a 2D matrix for covariance/precision computation and the
degrees of freedom
Args:
matrix (np.ndarray):
n_conditions x n_channels
Returns:
numpy.ndarray:
demeaned matrix
"""
assert isinstance(matrix, np.ndarray), "input must be ndarray"
if matrix.ndim in [1, 2]:
matrix = matrix - np.mean(matrix, axis=0, keepdims=True)
dof = matrix.shape[0] - 1
elif matrix.ndim == 3:
matrix -= np.mean(matrix, axis=2, keepdims=True)
dof = (matrix.shape[0] - 1) * matrix.shape[2]
matrix = matrix.transpose(0, 2, 1).reshape(
matrix.shape[0] * matrix.shape[2], matrix.shape[1])
else:
raise ValueError('Matrix for covariance estimation has wrong # of dimensions!')
return matrix, dof
def _estimate_covariance(matrix, dof, method):
""" calls the right covariance estimation function based on the ""method" argument
Args:
matrix (np.ndarray):
n_conditions x n_channels
dof (int):
degrees of freedom
method (string):
which estimator to use
Returns:
numpy.ndarray, numpy.ndarray:
cov_mat: n_channels x n_channels sample covariance matrix
"""
matrix, dof_nat = _check_demean(matrix)
if dof is None:
dof = dof_nat
# calculate sample covariance matrix s
if method == 'shrinkage_eye':
cov_mat = _covariance_eye(matrix, dof)
elif method == 'shrinkage_diag':
cov_mat = _covariance_diag(matrix, dof)
elif method == 'diag':
cov_mat = _variance(matrix, dof)
elif method == 'full':
cov_mat = _covariance_full(matrix, dof)
return cov_mat
def _variance(matrix, dof):
"""
returns the vector of variances per measurement channel.
The formula used here implies that the mean was already removed.
Args:
matrix (np.ndarray):
n_conditions x n_channels
Returns:
numpy.ndarray:
variance vector
"""
return np.diag(np.einsum('ij, ij-> j', matrix, matrix) / dof)
def _covariance_full(matrix, dof):
"""
computes the sample covariance matrix from a 2d-array.
matrix should be demeaned before!
Args:
matrix (np.ndarray):
n_conditions x n_channels
Returns:
numpy.ndarray, numpy.ndarray:
s_mean: n_channels x n_channels sample covariance matrix
"""
return np.einsum('ij, ik-> jk', matrix, matrix, optimize=True) / dof
def _covariance_eye(matrix, dof):
"""
computes the sample covariance matrix from a 2d-array.
matrix should be demeaned before!
Computes an optimal shrinkage estimate of a sample covariance matrix
as described by the following publication:
Ledoit and Wolfe (2004): "A well-conditioned
estimator for large-dimensional covariance matrices"
Args:
matrix (np.ndarray):
n_conditions x n_channels
Returns:
numpy.ndarray, numpy.ndarray:
s_mean: n_channels x n_channels sample covariance matrix
xt_x:
Einstein summation form of the matrix product
of the 2d-array with itself
"""
s_sum = np.zeros((matrix.shape[1], matrix.shape[1]))
s2_sum = np.zeros((matrix.shape[1], matrix.shape[1]))
for m_line in matrix:
xt_x = np.outer(m_line, m_line)
s_sum += xt_x
s2_sum += xt_x ** 2
s = s_sum / matrix.shape[0]
b2 = np.sum(s2_sum / matrix.shape[0] - s * s) / matrix.shape[0]
# calculate the scalar estimators to find the optimal shrinkage:
# m, d^2, b^2 as in Ledoit & Wolfe paper
m = np.sum(np.diag(s)) / s.shape[0]
d2 = np.sum((s - m * np.eye(s.shape[0])) ** 2)
b2 = min(d2, b2)
# shrink covariance matrix
s_shrink = b2 / d2 * m * np.eye(s.shape[0]) \
+ (d2 - b2) / d2 * s
# correction for degrees of freedom
s_shrink = s_shrink * matrix.shape[0] / dof
return s_shrink
def _covariance_diag(matrix, dof):
"""
computes the sample covariance matrix from a 2d-array.
matrix should be demeaned before!
Computes an optimal shrinkage estimate of a sample covariance matrix
as described by the following publication:
Schäfer, J., & Strimmer, K. (2005). "A Shrinkage Approach to Large-Scale
Covariance Matrix Estimation and Implications for Functional Genomics.""
Args:
matrix (np.ndarray):
n_conditions x n_channels
Returns:
numpy.ndarray, numpy.ndarray:
s_mean: n_channels x n_channels sample covariance matrix
xt_x:
Einstein summation form of the matrix product
of the 2d-array with itself
"""
s_sum = np.zeros((matrix.shape[1], matrix.shape[1]))
s2_sum = np.zeros((matrix.shape[1], matrix.shape[1]))
for m_line in matrix:
xt_x = np.outer(m_line, m_line)
s_sum += xt_x
s2_sum += xt_x ** 2
s = s_sum / dof
var = np.diag(s)
std = np.sqrt(var)
s_mean = s_sum / np.expand_dims(std, 0) / np.expand_dims(std, 1) / (matrix.shape[0] - 1)
s2_mean = s2_sum / np.expand_dims(var, 0) / np.expand_dims(var, 1) / (matrix.shape[0] - 1)
var_hat = matrix.shape[0] / dof ** 2 \
* (s2_mean - s_mean ** 2)
mask = ~np.eye(s.shape[0], dtype=bool)
lamb = np.sum(var_hat[mask]) / np.sum(s_mean[mask] ** 2)
lamb = max(min(lamb, 1), 0)
scaling = np.eye(s.shape[0]) + (1 - lamb) * mask
s_shrink = s * scaling
return s_shrink
[docs]
def cov_from_residuals(residuals, dof=None, method='shrinkage_diag'):
"""
Estimates a covariance matrix from measurements. Allows for shrinkage estimates.
Use 'method' to choose which estimation method is used.
Args:
residuals(numpy.ndarray or list of these): n_residuals x n_channels
matrix of residuals
dof(int or list of int): degrees of freedom for covariance estimation
defaults to n_res - 1, should be corrected for the number
of regressors in a GLM if applicable.
method(str): which estimate to use:
'diag': provides a diagonal matrix, i.e. univariate noise normalizer
'full': computes the sample covariance without shrinkage
'shrinkage_eye': shrinks the data covariance towards a multiple of the identity.
'shrinkage_diag': shrinks the covariance matrix towards the diagonal covariance matrix.
Returns:
numpy.ndarray (or list): sigma_p: covariance matrix over channels
"""
if not isinstance(residuals, np.ndarray):
cov_mat = []
for i, residual in enumerate(residuals):
if dof is None:
cov_mat.append(cov_from_residuals(
residual, method=method))
elif isinstance(dof, Iterable):
cov_mat.append(cov_from_residuals(
residuals, method=method, dof=dof[i]))
else:
cov_mat.append(cov_from_residuals(
residual, method=method, dof=dof))
else:
cov_mat = _estimate_covariance(residuals, dof, method)
return cov_mat
[docs]
def prec_from_residuals(residuals, dof=None, method='shrinkage_diag'):
"""
Estimates the covariance matrix from residuals and finds its multiplicative
inverse (= the precision matrix)
Use 'method' to choose which estimation method is used.
Args:
residuals(numpy.ndarray or list of these): n_residuals x n_channels
matrix of residuals
dof(int or list of int): degrees of freedom for covariance estimation
defaults to n_res - 1, should be corrected for the number
of regressors in a GLM if applicable.
method(str): which estimate to use:
'diag': provides a diagonal matrix, i.e. univariate noise normalizer
'full': computes the sample covariance without shrinkage
'shrinkage_eye': shrinks the data covariance towards a multiple of the identity.
'shrinkage_diag': shrinks the covariance matrix towards the diagonal covariance matrix.
Returns:
numpy.ndarray (or list): sigma_p: precision matrix over channels
"""
cov = cov_from_residuals(residuals=residuals, dof=dof, method=method)
if not isinstance(cov, np.ndarray):
prec = [None] * len(cov)
for i, cov_i in enumerate(cov):
prec[i] = np.linalg.inv(cov_i)
elif len(cov.shape) > 2:
prec = np.zeros(cov.shape)
for i, cov_i in enumerate(cov):
prec[i] = np.linalg.inv(cov_i)
else:
prec = np.linalg.inv(cov)
return prec
[docs]
def cov_from_measurements(dataset, obs_desc, dof=None, method='shrinkage_diag'):
"""
Estimates a covariance matrix from measurements. Allows for shrinkage estimates.
Use 'method' to choose which estimation method is used.
Args:
dataset(data.Dataset):
rsatoolbox Dataset object
dof(int or list of int): degrees of freedom for covariance estimation
defaults to n_res - 1, should be corrected for the number
of regressors in a GLM if applicable.
method(str): which estimate to use:
'diag': provides a diagonal matrix, i.e. univariate noise normalizer
'full': computes the sample covariance without shrinkage
'shrinkage_eye': shrinks the data covariance towards a multiple of the identity.
'shrinkage_diag': shrinks the covariance matrix towards the diagonal covariance matrix.
Returns:
numpy.ndarray (or list): sigma_p: covariance matrix over channels
"""
if isinstance(dataset, Iterable):
cov_mat = []
for i, dat in enumerate(dataset):
if dof is None:
cov_mat.append(cov_from_unbalanced(
dat, obs_desc=obs_desc, method=method))
elif isinstance(dof, Iterable):
cov_mat.append(cov_from_unbalanced(
dat, obs_desc=obs_desc, method=method, dof=dof[i]))
else:
cov_mat.append(cov_from_unbalanced(
dat, obs_desc=obs_desc, method=method, dof=dof))
else:
assert "Dataset" in str(type(dataset)), "Provided object is not a dataset"
assert obs_desc in dataset.obs_descriptors.keys(), \
"obs_desc not contained in the dataset's obs_descriptors"
tensor, _ = dataset.get_measurements_tensor(obs_desc)
# calculate sample covariance matrix s
cov_mat = _estimate_covariance(tensor, dof, method)
return cov_mat
[docs]
def prec_from_measurements(dataset, obs_desc, dof=None, method='shrinkage_diag'):
"""
Estimates the covariance matrix from measurements and finds its multiplicative
inverse (= the precision matrix)
Use 'method' to choose which estimation method is used.
Args:
residuals(numpy.ndarray or list of these): n_residuals x n_channels
matrix of residuals
dof(int or list of int): degrees of freedom for covariance estimation
defaults to n_res - 1, should be corrected for the number
of regressors in a GLM if applicable.
method(str): which estimate to use:
'diag': provides a diagonal matrix, i.e. univariate noise normalizer
'full': computes the sample covariance without shrinkage
'shrinkage_eye': shrinks the data covariance towards a multiple of the identity.
'shrinkage_diag': shrinks the covariance matrix towards the diagonal covariance matrix.
Returns:
numpy.ndarray (or list): sigma_p: precision matrix over channels
"""
cov = cov_from_measurements(dataset, obs_desc, dof=dof, method=method)
if not isinstance(cov, np.ndarray):
prec = [None] * len(cov)
for i, cov_i in enumerate(cov):
prec[i] = np.linalg.inv(cov_i)
elif len(cov.shape) > 2:
prec = np.zeros(cov.shape)
for i, cov_i in enumerate(cov):
prec[i] = np.linalg.inv(cov_i)
else:
prec = np.linalg.inv(cov)
return prec
[docs]
def cov_from_unbalanced(dataset, obs_desc, dof=None, method='shrinkage_diag'):
"""
Estimates a covariance matrix from an unbalanced dataset, i.e. from a
dataset that contains different numbers of samples for different
stimuli.
Args:
dataset(data.Dataset):
rsatoolbox Dataset object
dof(int or list of int): degrees of freedom for covariance estimation
defaults to n_measurements - n_stimuli, should be corrected
if this is not the case
method(str): which estimate to use:
'diag': provides a diagonal matrix, i.e. univariate noise normalizer
'full': computes the sample covariance without shrinkage
'shrinkage_eye': shrinks the data covariance towards a multiple of the identity.
'shrinkage_diag': shrinks the covariance matrix towards the diagonal covariance matrix.
Returns:
numpy.ndarray (or list): sigma_p: covariance matrix over channels
"""
if isinstance(dataset, Iterable):
cov_mat = []
for i, dat in enumerate(dataset):
if dof is None:
cov_mat.append(cov_from_unbalanced(
dat, obs_desc=obs_desc, method=method))
elif isinstance(dof, Iterable):
cov_mat.append(cov_from_unbalanced(
dat, obs_desc=obs_desc, method=method, dof=dof[i]))
else:
cov_mat.append(cov_from_unbalanced(
dat, obs_desc=obs_desc, method=method, dof=dof))
else:
assert "Dataset" in str(type(dataset)), "Provided object is not a dataset"
assert obs_desc in dataset.obs_descriptors.keys(), \
"obs_desc not contained in the dataset's obs_descriptors"
matrix = dataset.measurements.copy()
means, values, _ = average_dataset_by(dataset, obs_desc)
values, inverse = get_unique_inverse(dataset.obs_descriptors[obs_desc])
matrix -= means[inverse]
# calculate sample covariance matrix s
if dof is None:
dof = matrix.shape[0] - len(values)
cov_mat = _estimate_covariance(matrix, dof, method)
return cov_mat
[docs]
def prec_from_unbalanced(dataset, obs_desc, dof=None, method='shrinkage_diag'):
"""
Estimates the covariance matrix from measurements and finds its multiplicative
inverse (= the precision matrix)
Use 'method' to choose which estimation method is used.
Args:
residuals(numpy.ndarray or list of these): n_residuals x n_channels
matrix of residuals
dof(int or list of int): degrees of freedom for covariance estimation
defaults to n_res - 1, should be corrected for the number
of regressors in a GLM if applicable.
method(str): which estimate to use:
'diag': provides a diagonal matrix, i.e. univariate noise normalizer
'full': computes the sample covariance without shrinkage
'shrinkage_eye': shrinks the data covariance towards a multiple of the identity.
'shrinkage_diag': shrinks the covariance matrix towards the diagonal covariance matrix.
Returns:
numpy.ndarray (or list): sigma_p: precision matrix over channels
"""
cov = cov_from_unbalanced(dataset, obs_desc, dof=dof, method=method)
if not isinstance(cov, np.ndarray):
prec = [None] * len(cov)
for i, cov_i in enumerate(cov):
prec[i] = np.linalg.inv(cov_i)
elif len(cov.shape) > 2:
prec = np.zeros(cov.shape)
for i, cov_i in enumerate(cov):
prec[i] = np.linalg.inv(cov_i)
else:
prec = np.linalg.inv(cov)
return prec
[docs]
def sigmak_from_measurements(dataset, obs_descriptor, cv_descriptor, noise=None):
"""
Estimates sigma_k, the matrix encoding the noise variance/covariance among the k conditions when two
conditions are measured in the same partition (e.g., due to shared fMRI noise
from the sluggishness of the HRF when two conditions are adjacent in time). If a noise matrix is provided,
prewhitening is performed on the data before computing sigma_k (make sure to do this if using
Mahalanobis or crossnobis distance). Assumes that sigma_k is constant across partitions, implementing
equation 36 from Diedrichsen et al. (2016), "On the distribution of cross-validated Mahalanobis distances."
Args:
dataset(data.Dataset):
rsatoolbox Dataset object
obs_descriptor (String):
descriptor defining experimental conditions
cv_descriptor (String):
descriptor defining crossvalidation folds/partitions
noise (numpy.ndarray):
dataset.n_channel x dataset.n_channel
precision matrix for noise between channels
default: identity matrix, i.e. euclidean distance
Returns:
numpy.ndarray: sigma_k: noise covariance matrix over conditions
n_conditions x n_conditions
"""
n_channels = dataset.n_channel
if noise is None:
noise = np.eye(n_channels)
else:
if noise.shape != (n_channels, n_channels):
raise ValueError("noise must have shape n_channel x n_channel")
noise_sqrt = sqrtm(noise) # take matrix square root to get whitening matrix
dataset_whitened = dataset.copy()
dataset_whitened.measurements = dataset.measurements @ noise_sqrt
# Compute mean patterns per condition
U_mean, conds, _ = average_dataset_by(dataset_whitened, obs_descriptor)
n_cond = len(conds)
cv_folds = np.unique(np.array(dataset_whitened.obs_descriptors[cv_descriptor]))
pair_counts = np.zeros((n_cond, n_cond)) # tally how many partitions each condition pair occurs in
# Compute sigma_k per partition, then average
sigma_ks = []
for fold in cv_folds:
U_fold = np.zeros(U_mean.shape) * np.nan # fold activations
dataset_fold = dataset_whitened.subset_obs(cv_descriptor, fold)
dataset_fold, fold_conds, _ = average_dataset_by(dataset_fold, obs_descriptor)
# Get indices mapping from subsetted conditions to full set, fill out U_fold
inds = [np.where(conds == c)[0][0] for c in fold_conds]
U_fold[inds, :] = dataset_fold
sigma_k_fold = ((U_fold - U_mean) @ ((U_fold - U_mean).T))
sigma_ks.append(sigma_k_fold)
# Increment pair counts for all condition pairs present in this fold
pair_counts += np.isfinite(sigma_k_fold).astype(int)
# Finally add all sigma_ks and divide by pair counts to get average
sigma_k = np.nansum(sigma_ks, axis=0) / (pair_counts - 1) / n_channels
# Any pairs that never occurred together should be set to zero
sigma_k[np.isnan(sigma_k)] = 0.0
return sigma_k