def test_io():
rng = numpy.random.default_rng(seed=8001)
# Clean up in case of a failure
ofile = 'test_covar_io.fits'
if os.path.isfile(ofile):
os.remove(ofile)
# Build a bogus covariance matrix
m = numpy.zeros(10, dtype=float)
c = numpy.diag(numpy.full(10-2, 0.2, dtype=float), k=-2) \
+ numpy.diag(numpy.full(10-1, 0.5, dtype=float), k=-1) \
+ numpy.diag(numpy.full(10, 1.0, dtype=float), k=0) \
+ numpy.diag(numpy.full(10-1, 0.5, dtype=float), k=1) \
+ numpy.diag(numpy.full(10-2, 0.2, dtype=float), k=2)
# Draw samples
s = rng.multivariate_normal(m, c, size=100000)
# Instantiate
covar = Covariance.from_samples(s.T, cov_tol=0.1)
# Write
covar.write(ofile)
# Read
_covar = Covariance.from_fits(ofile)
# Should be the same
assert numpy.allclose(covar.toarray(), _covar.toarray()), 'Bad I/O'
# Clean-up
os.remove(ofile)
def test_mult():
# Build a bogus covariance matrix
c = numpy.diag(numpy.full(10-2, 0.2, dtype=float), k=-2) \
+ numpy.diag(numpy.full(10-1, 0.5, dtype=float), k=-1) \
+ numpy.diag(numpy.full(10, 1.0, dtype=float), k=0) \
+ numpy.diag(numpy.full(10-1, 0.5, dtype=float), k=1) \
+ numpy.diag(numpy.full(10-2, 0.2, dtype=float), k=2)
x = numpy.ones(10, dtype=float)
# Uncorrelated
t = numpy.zeros((3, 10), dtype=float)
t[0, 0] = 1.0
t[1, 3] = 1.0
t[2, 6] = 1.0
y = numpy.dot(t, x)
covar = Covariance.from_matrix_multiplication(t, c)
assert numpy.array_equal(
covar.toarray(), numpy.identity(3)), 'Result should be uncorrelated.'
# Correlated by 0.2
t = numpy.zeros((3, 10), dtype=float)
t[0, 0] = 1.0
t[1, 2] = 1.0
t[2, 4] = 1.0
_c = numpy.diag(numpy.full(3-1, 0.2, dtype=float), k=-1) \
+ numpy.diag(numpy.full(3, 1.0, dtype=float), k=0) \
+ numpy.diag(numpy.full(3-1, 0.2, dtype=float), k=1)
y = numpy.dot(t, x)
covar = Covariance.from_matrix_multiplication(t, c)
assert numpy.array_equal(covar.toarray(),
_c), 'Result should have off-diagonals = 0.2'
# Correlated by 0.5 and 0.2
t = numpy.zeros((3, 10), dtype=float)
t[0, 0] = 1.0
t[1, 1] = 1.0
t[2, 2] = 1.0
_c = numpy.diag(numpy.full(3-2, 0.2, dtype=float), k=-2) \
+ numpy.diag(numpy.full(3-1, 0.5, dtype=float), k=-1) \
+ numpy.diag(numpy.full(3, 1.0, dtype=float), k=0) \
+ numpy.diag(numpy.full(3-1, 0.5, dtype=float), k=1) \
+ numpy.diag(numpy.full(3-2, 0.2, dtype=float), k=2)
y = numpy.dot(t, x)
covar = Covariance.from_matrix_multiplication(t, c)
assert numpy.array_equal(covar.toarray(),
_c), 'Result should have off-diagonals = 0.5,0.2'
def test_samples():
rng = numpy.random.default_rng(seed=8001)
# Build a bogus covariance matrix
m = numpy.zeros(10, dtype=float)
c = numpy.diag(numpy.full(10-2, 0.2, dtype=float), k=-2) \
+ numpy.diag(numpy.full(10-1, 0.5, dtype=float), k=-1) \
+ numpy.diag(numpy.full(10, 1.0, dtype=float), k=0) \
+ numpy.diag(numpy.full(10-1, 0.5, dtype=float), k=1) \
+ numpy.diag(numpy.full(10-2, 0.2, dtype=float), k=2)
# Draw samples
s = rng.multivariate_normal(m, c, size=100000)
# Instantiate
covar = Covariance.from_samples(s.T, cov_tol=0.1)
# Check the values are very nearly the same as the input
assert numpy.all(numpy.absolute(c - covar.toarray()) < 0.02), 'Covariances are too different'
# Check that `find` returns the same indices
# NOTE: For some reason the ordering of numpy.where and
# scipy.sparse.find are different. So I run lexsort before checking
# if the arrays are the same.
coo = numpy.array([ [i,j] for i,j in zip(*numpy.where(c > 0))])
srt = numpy.lexsort((coo[:,0], coo[:,1]))
coo = coo[srt,:]
_coo = numpy.array([ [i,j] for i,j,k in zip(*covar.find())])
srt = numpy.lexsort((_coo[:,0], _coo[:,1]))
_coo = _coo[srt,:]
assert numpy.array_equal(coo, _coo), 'Did not find the same indices'
def test_array():
# Construct the Covariance matrix from a pre-calculated array
c = numpy.diag(numpy.full(10-2, 0.2, dtype=float), k=-2) \
+ numpy.diag(numpy.full(10-1, 0.5, dtype=float), k=-1) \
+ numpy.diag(numpy.full(10, 1.0, dtype=float), k=0) \
+ numpy.diag(numpy.full(10-1, 0.5, dtype=float), k=1) \
+ numpy.diag(numpy.full(10-2, 0.2, dtype=float), k=2)
covar = Covariance.from_array(c)
# Should be the same as the identity matrix.
assert numpy.array_equal(covar.toarray(), c), 'Arrays should be identical'
def test_read_drp():
cfg = MaNGAConfig(7815, 3702)
drpfile = remote_data_file(cfg.file_name)
assert os.path.isfile(drpfile), 'Did not find file'
with fits.open(drpfile) as hdu:
covar = Covariance.from_fits(hdu, ivar_ext=None, covar_ext='GCORREL', impose_triu=True,
correlation=True)
var = numpy.ma.power(hdu['IVAR'].data[hdu['GCORREL'].header['BBINDEX']].T.ravel(),
-1).filled(0.0)
covar = covar.apply_new_variance(var)
covar.revert_correlation()
assert numpy.array_equal(var, numpy.diag(covar.toarray())), 'New variance not applied'
# Read the DRP file
print('reading')
drpf = DRPFits(8138, 1901, 'CUBE', directory_path='.', read=True)
# Get the variance at the wavelength used for the g-band correlation
# matrix
bbindex = int(drpf['GCORREL'].header['BBINDEX'])
var = numpy.ma.power(drpf['IVAR'].data[:, :, bbindex].ravel(),
-1).filled(0.0)
# Use the Covariance class to setup the correlation matrix using the
# DRP GCORREL extension
C = Covariance.from_fits(drpf.hdu,
ivar_ext=None,
covar_ext='GCORREL',
correlation=True,
impose_triu=True,
row_major=True)
C.show()
C = C.apply_new_variance(var)
C.show()
C.revert_correlation()
C.show()
# Recalculate the covariance from scratch
recalc_C = drpf.covariance_matrix(bbindex)
recalc_C.show()
# Calculate the difference between the recalculation at the DRP
# output
diff = numpy.ma.MaskedArray(recalc_C.toarray() - C.toarray())
def test_var():
var = numpy.ones(3, dtype=float)
covar = Covariance.from_variance(var)
assert numpy.array_equal(covar.toarray(), numpy.identity(3)), \
'Result should be an identity matrix'