import numpy as np
import h5py
import math
def reconstruct_symmetric_matrix(matrix_triu):
"""
Reconstructs a symmetric matrix from a flat array containing only the
upper triangle of the matrix.
Inputs:
matrix_triu (np.ndarray): Flat array with size n*(n+1)/2, where the
output symmetric matrix will have shape (n,n).
Returns:
Symmetric matrix with shape (n,n).
"""
# Solve n(n+1)/2 = len(matrix_triu) for n. The output symmetric matrix
# will have shape (n,n).
n = (math.isqrt(1+8*len(matrix_triu)) - 1) // 2
# Create empty square matrix filled with zeros
S = np.zeros((n,n), dtype=matrix_triu.dtype)
# Fill in the upper triangle of the matrix
idx_triu = np.triu_indices(n)
S[idx_triu] = matrix_triu
# Symmetrize the matrix by adding in its transpose
S += S.T
# Divide the diagonal by 2, as it was added twice
idx_diag = np.diag_indices(n)
S[idx_diag] *= 0.5
return S
def main():
# Location of the catalog containing cov and icov:
catalog_h5_loc = "stellar_covariances_catalog_00.h5"
n_stars = 5 # Only load this many stars, for purposes of this example
with h5py.File(catalog_h5_loc, 'r') as f:
# Extract cov and icov arrays. Only the upper triangles are stored.
cov_triu = f['stellar_params_cov_triu'][:n_stars]
icov_triu = f['stellar_params_icov_triu'][:n_stars]
# Reconstruct the cov and icov matrices from the upper triangles
cov = np.stack([reconstruct_symmetric_matrix(x) for x in cov_triu], axis=0)
icov = np.stack([reconstruct_symmetric_matrix(x) for x in icov_triu], axis=0)
# Print the results
with np.printoptions(precision=3, suppress=True,
threshold=5, sign='+'):
print("Reconstructed covariance matrices:")
for c in cov:
print(c)
print('')
print("Reconstructed inverse-covariance matrices:")
for ic in icov:
print(ic)
print('')
# Covariance times inverse covariance should be the identity matrix
cov_icov = np.einsum('ijk, ikl-> ijl', cov, icov)
print("Cov multiplied by iCov")
for ident in cov_icov:
print(ident)
print('')
if __name__=='__main__':
main()