#!/usr/bin/env python3
# 3rd
import numpy as np
# todo: docstring
[docs]def ensure_array(x):
if not isinstance(x, np.ndarray):
return np.array(x)
else:
return x
[docs]def cov2err(cov):
""" Convert covariance matrix (or array of covariance matrices of equal
shape) to error array (or array thereof).
Args:
cov: [n x ] nbins x nbins array
Returns
[n x ] nbins array
"""
cov = ensure_array(cov)
if cov.ndim == 2:
return np.sqrt(cov.diagonal())
elif cov.ndim == 3:
return np.sqrt(cov.diagonal(axis1=1, axis2=2))
else:
raise ValueError("Wrong dimensions.")
[docs]def cov2corr(cov):
""" Convert covariance matrix (or array of covariance matrices of equal
shape) to correlation matrix (or array thereof).
Args:
cov: [n x ] nbins x nbins array
Returns
[n x ] nbins x nbins array
"""
cov = ensure_array(cov)
err = cov2err(cov)
if cov.ndim == 2:
return cov / np.einsum("i,j->ij", err, err)
elif cov.ndim == 3:
return cov / np.einsum("ki,kj->kij", err, err)
else:
raise ValueError("Wrong dimensions")
[docs]def corr2cov(corr, err):
""" Convert correlation matrix (or array of covariance matrices of equal
shape) together with error array (or array thereof) to covariance
matrix (or array thereof).
Args:
corr: [n x ] nbins x nbins array
err: [n x ] nbins array
Returns
[n x ] nbins x nbins array
"""
corr = ensure_array(corr)
err = ensure_array(err)
if corr.ndim == 2:
return np.einsum("ij,i,j->ij", corr, err, err)
elif corr.ndim == 3:
return np.einsum("kij,ki,kj->kij", corr, err, err)
else:
raise ValueError("Wrong dimensions")
[docs]def rel2abs_cov(cov, data):
""" Convert relative covariance matrix to absolute covariance matrix
Args:
cov: n x nbins x nbins array
data: n x nbins array
Returns:
n x nbins x nbins array
"""
cov = ensure_array(cov)
data = ensure_array(data)
assert cov.ndim == data.ndim + 1
if data.ndim == 1:
return np.einsum("ij,i,j->ij", cov, data, data)
elif data.ndim == 2:
return np.einsum("kij,ki,kj->kij", cov, data, data)
else:
raise ValueError("Wrong dimensions")
[docs]def abs2rel_cov(cov, data):
""" Convert covariance matrix to relative covariance matrix
Args:
cov: n x nbins x nbins array
data: n x nbins array
Returns:
n x nbins x nbins array
"""
cov = ensure_array(cov)
data = ensure_array(data)
assert cov.ndim == data.ndim + 1
if data.ndim == 1:
nbins = len(data)
return cov / data.reshape((nbins, 1)) / data.reshape((1, nbins))
elif data.ndim == 2:
n, nbins = data.shape
return cov / data.reshape((n, nbins, 1)) / data.reshape((n, 1, nbins))
else:
raise ValueError("Wrong dimensions")