import numpy as np import autograd.numpy as anp import math import pyerrors as pe import pytest np.random.seed(0) def get_real_matrix(dimension): base_matrix = np.empty((dimension, dimension), dtype=object) for (n, m), entry in np.ndenumerate(base_matrix): exponent_real = np.random.normal(0, 1) exponent_imag = np.random.normal(0, 1) base_matrix[n, m] = pe.Obs([np.random.normal(1.0, 0.1, 100)], ['t']) return base_matrix def get_complex_matrix(dimension): base_matrix = np.empty((dimension, dimension), dtype=object) for (n, m), entry in np.ndenumerate(base_matrix): exponent_real = np.random.normal(0, 1) exponent_imag = np.random.normal(0, 1) base_matrix[n, m] = pe.CObs(pe.Obs([np.random.normal(1.0, 0.1, 100)], ['t']), pe.Obs([np.random.normal(1.0, 0.1, 100)], ['t'])) return base_matrix def test_matmul(): for dim in [4, 6]: for const in [1, pe.cov_Obs(1.0, 0.002, 'cov')]: my_list = [] length = 100 + np.random.randint(200) for i in range(dim ** 2): my_list.append(pe.Obs([np.random.rand(length), np.random.rand(length + 1)], ['t1', 't2'])) my_array = const * np.array(my_list).reshape((dim, dim)) tt = pe.linalg.matmul(my_array, my_array) - my_array @ my_array for t, e in np.ndenumerate(tt): assert e.is_zero(), t my_list = [] length = 100 + np.random.randint(200) for i in range(dim ** 2): my_list.append(pe.CObs(pe.Obs([np.random.rand(length), np.random.rand(length + 1)], ['t1', 't2']), pe.Obs([np.random.rand(length), np.random.rand(length + 1)], ['t1', 't2']))) my_array = np.array(my_list).reshape((dim, dim)) * const tt = pe.linalg.matmul(my_array, my_array) - my_array @ my_array for t, e in np.ndenumerate(tt): assert e.is_zero(), t def test_jack_matmul(): tt = get_real_matrix(8) check1 = pe.linalg.jack_matmul(tt, tt) - pe.linalg.matmul(tt, tt) [o.gamma_method() for o in check1.ravel()] assert np.all([o.is_zero_within_error(0.1) for o in check1.ravel()]) assert np.all([o.dvalue < 0.001 for o in check1.ravel()]) trace1 = np.trace(check1) trace1.gamma_method() assert trace1.dvalue < 0.001 tr = np.random.rand(8, 8) check2 = pe.linalg.jack_matmul(tt, tr) - pe.linalg.matmul(tt, tr) [o.gamma_method() for o in check2.ravel()] assert np.all([o.is_zero_within_error(0.1) for o in check2.ravel()]) assert np.all([o.dvalue < 0.001 for o in check2.ravel()]) trace2 = np.trace(check2) trace2.gamma_method() assert trace2.dvalue < 0.001 tt2 = get_complex_matrix(8) check3 = pe.linalg.jack_matmul(tt2, tt2) - pe.linalg.matmul(tt2, tt2) [o.gamma_method() for o in check3.ravel()] assert np.all([o.real.is_zero_within_error(0.1) for o in check3.ravel()]) assert np.all([o.imag.is_zero_within_error(0.1) for o in check3.ravel()]) assert np.all([o.real.dvalue < 0.001 for o in check3.ravel()]) assert np.all([o.imag.dvalue < 0.001 for o in check3.ravel()]) trace3 = np.trace(check3) trace3.gamma_method() assert trace3.real.dvalue < 0.001 assert trace3.imag.dvalue < 0.001 tr2 = np.random.rand(8, 8) + 1j * np.random.rand(8, 8) check4 = pe.linalg.jack_matmul(tt2, tr2) - pe.linalg.matmul(tt2, tr2) [o.gamma_method() for o in check4.ravel()] assert np.all([o.real.is_zero_within_error(0.1) for o in check4.ravel()]) assert np.all([o.imag.is_zero_within_error(0.1) for o in check4.ravel()]) assert np.all([o.real.dvalue < 0.001 for o in check4.ravel()]) assert np.all([o.imag.dvalue < 0.001 for o in check4.ravel()]) trace4 = np.trace(check4) trace4.gamma_method() assert trace4.real.dvalue < 0.001 assert trace4.imag.dvalue < 0.001 def test_einsum(): def _perform_real_check(arr): [o.gamma_method() for o in arr] assert np.all([o.is_zero_within_error(0.001) for o in arr]) assert np.all([o.dvalue < 0.001 for o in arr]) def _perform_complex_check(arr): [o.gamma_method() for o in arr] assert np.all([o.real.is_zero_within_error(0.001) for o in arr]) assert np.all([o.real.dvalue < 0.001 for o in arr]) assert np.all([o.imag.is_zero_within_error(0.001) for o in arr]) assert np.all([o.imag.dvalue < 0.001 for o in arr]) tt = [get_real_matrix(4), get_real_matrix(3)] q = np.tensordot(tt[0], tt[1], 0) c1 = tt[1] @ q c2 = pe.linalg.einsum('ij,abjd->abid', tt[1], q) check1 = c1 - c2 _perform_real_check(check1.ravel()) check2 = np.trace(tt[0]) - pe.linalg.einsum('ii', tt[0]) _perform_real_check([check2]) check3 = np.trace(tt[1]) - pe.linalg.einsum('ii', tt[1]) _perform_real_check([check3]) tt = [get_real_matrix(4), np.random.random((3, 3))] q = np.tensordot(tt[0], tt[1], 0) c1 = tt[1] @ q c2 = pe.linalg.einsum('ij,abjd->abid', tt[1], q) check1 = c1 - c2 _perform_real_check(check1.ravel()) tt = [get_complex_matrix(4), get_complex_matrix(3)] q = np.tensordot(tt[0], tt[1], 0) c1 = tt[1] @ q c2 = pe.linalg.einsum('ij,abjd->abid', tt[1], q) check1 = c1 - c2 _perform_complex_check(check1.ravel()) check2 = np.trace(tt[0]) - pe.linalg.einsum('ii', tt[0]) _perform_complex_check([check2]) check3 = np.trace(tt[1]) - pe.linalg.einsum('ii', tt[1]) _perform_complex_check([check3]) tt = [get_complex_matrix(4), np.random.random((3, 3))] q = np.tensordot(tt[0], tt[1], 0) c1 = tt[1] @ q c2 = pe.linalg.einsum('ij,abjd->abid', tt[1], q) check1 = c1 - c2 _perform_complex_check(check1.ravel()) def test_multi_dot(): for dim in [4, 6]: my_list = [] length = 1000 + np.random.randint(200) for i in range(dim ** 2): my_list.append(pe.Obs([np.random.rand(length), np.random.rand(length + 1)], ['t1', 't2'])) my_array = pe.cov_Obs(1.0, 0.002, 'cov') * np.array(my_list).reshape((dim, dim)) tt = pe.linalg.matmul(my_array, my_array, my_array, my_array) - my_array @ my_array @ my_array @ my_array for t, e in np.ndenumerate(tt): assert e.is_zero(), t my_list = [] length = 1000 + np.random.randint(200) for i in range(dim ** 2): my_list.append(pe.CObs(pe.Obs([np.random.rand(length), np.random.rand(length + 1)], ['t1', 't2']), pe.Obs([np.random.rand(length), np.random.rand(length + 1)], ['t1', 't2']))) my_array = np.array(my_list).reshape((dim, dim)) * pe.cov_Obs(1.0, 0.002, 'cov') tt = pe.linalg.matmul(my_array, my_array, my_array, my_array) - my_array @ my_array @ my_array @ my_array for t, e in np.ndenumerate(tt): assert e.is_zero(), t def test_jack_multi_dot(): for dim in [2, 4, 8]: my_array = get_real_matrix(dim) tt = pe.linalg.jack_matmul(my_array, my_array, my_array) - pe.linalg.matmul(my_array, my_array, my_array) for t, e in np.ndenumerate(tt): e.gamma_method() assert e.is_zero_within_error(0.01) assert e.is_zero(atol=1e-1), t assert np.isclose(e.value, 0.0) def test_matmul_irregular_histories(): dim = 2 length = 500 standard_array = [] for i in range(dim ** 2): standard_array.append(pe.Obs([np.random.normal(1.1, 0.2, length)], ['ens1'])) standard_matrix = np.array(standard_array).reshape((dim, dim)) * pe.cov_Obs(1.0, 0.002, 'cov') * pe.pseudo_Obs(0.1, 0.002, 'qr') for idl in [range(1, 501, 2), range(250, 273), [2, 8, 19, 20, 78]]: irregular_array = [] for i in range(dim ** 2): irregular_array.append(pe.Obs([np.random.normal(1.1, 0.2, len(idl))], ['ens1'], idl=[idl])) irregular_matrix = np.array(irregular_array).reshape((dim, dim)) * pe.cov_Obs([1.0, 1.0], [[0.001,0.0001], [0.0001, 0.002]], 'norm')[0] t1 = standard_matrix @ irregular_matrix t2 = pe.linalg.matmul(standard_matrix, irregular_matrix) assert np.all([o.is_zero() for o in (t1 - t2).ravel()]) assert np.all([o.is_merged for o in t1.ravel()]) assert np.all([o.is_merged for o in t2.ravel()]) def test_irregular_matrix_inverse(): dim = 3 length = 500 for idl in [range(8, 508, 10), range(250, 273), [2, 8, 19, 20, 78, 99, 828, 10548979]]: irregular_array = [] for i in range(dim ** 2): irregular_array.append(pe.Obs([np.random.normal(1.1, 0.2, len(idl)), np.random.normal(0.25, 0.1, 10)], ['ens1', 'ens2'], idl=[idl, range(1, 11)])) irregular_matrix = np.array(irregular_array).reshape((dim, dim)) * pe.cov_Obs(1.0, 0.002, 'cov') * pe.pseudo_Obs(1.0, 0.002, 'ens2|r23') invertible_irregular_matrix = np.identity(dim) + irregular_matrix @ irregular_matrix.T inverse = pe.linalg.inv(invertible_irregular_matrix) assert np.allclose(np.linalg.inv(np.vectorize(lambda x: x.value)(invertible_irregular_matrix)) - np.vectorize(lambda x: x.value)(inverse), 0.0) check1 = pe.linalg.matmul(invertible_irregular_matrix, inverse) assert np.all([o.is_zero() for o in (check1 - np.identity(dim)).ravel()]) check2 = invertible_irregular_matrix @ inverse assert np.all([o.is_zero() for o in (check2 - np.identity(dim)).ravel()]) def test_matrix_inverse(): content = [] for t in range(9): exponent = np.random.normal(3, 5) content.append(pe.pseudo_Obs(2 + 10 ** exponent, 10 ** (exponent - 1), 't')) content.append(1.0) # Add 1.0 as a float matrix = np.diag(content) inverse_matrix = pe.linalg.inv(matrix) assert all([o.is_zero() for o in np.diag(matrix) * np.diag(inverse_matrix) - 1]) def test_complex_matrix_inverse(): dimension = 4 base_matrix = np.empty((dimension, dimension), dtype=object) matrix = np.empty((dimension, dimension), dtype=complex) for (n, m), entry in np.ndenumerate(base_matrix): exponent_real = np.random.normal(2, 3) exponent_imag = np.random.normal(2, 3) base_matrix[n, m] = pe.CObs(pe.pseudo_Obs(2 + 10 ** exponent_real, 10 ** (exponent_real - 1), 't'), pe.pseudo_Obs(2 + 10 ** exponent_imag, 10 ** (exponent_imag - 1), 't')) # Construct invertible matrix obs_matrix = np.identity(dimension) + base_matrix @ base_matrix.T for (n, m), entry in np.ndenumerate(obs_matrix): matrix[n, m] = entry.real.value + 1j * entry.imag.value inverse_matrix = np.linalg.inv(matrix) inverse_obs_matrix = pe.linalg.inv(obs_matrix) for (n, m), entry in np.ndenumerate(inverse_matrix): assert np.isclose(inverse_matrix[n, m].real, inverse_obs_matrix[n, m].real.value) assert np.isclose(inverse_matrix[n, m].imag, inverse_obs_matrix[n, m].imag.value) def test_matrix_functions(): dim = 4 matrix = [] for i in range(dim): row = [] for j in range(dim): row.append(pe.pseudo_Obs(np.random.rand(), 0.2 + 0.1 * np.random.rand(), 'e1')) matrix.append(row) matrix = np.array(matrix) @ np.identity(dim) # Check inverse of matrix inv = pe.linalg.inv(matrix) check_inv = matrix @ inv for (i, j), entry in np.ndenumerate(check_inv): entry.gamma_method() if(i == j): assert math.isclose(entry.value, 1.0, abs_tol=1e-9), 'value ' + str(i) + ',' + str(j) + ' ' + str(entry.value) else: assert math.isclose(entry.value, 0.0, abs_tol=1e-9), 'value ' + str(i) + ',' + str(j) + ' ' + str(entry.value) assert math.isclose(entry.dvalue, 0.0, abs_tol=1e-9), 'dvalue ' + str(i) + ',' + str(j) + ' ' + str(entry.dvalue) # Check Cholesky decomposition sym = np.dot(matrix, matrix.T) cholesky = pe.linalg.cholesky(sym) check = cholesky @ cholesky.T for (i, j), entry in np.ndenumerate(check): diff = entry - sym[i, j] assert diff.is_zero() # Check eigh e, v = pe.linalg.eigh(sym) for i in range(dim): tmp = sym @ v[:, i] - v[:, i] * e[i] for j in range(dim): assert tmp[j].is_zero() # Check svd u, v, vh = pe.linalg.svd(sym) diff = sym - u @ np.diag(v) @ vh for (i, j), entry in np.ndenumerate(diff): assert entry.is_zero() def test_complex_matrix_operations(): dimension = 4 base_matrix = np.empty((dimension, dimension), dtype=object) for (n, m), entry in np.ndenumerate(base_matrix): exponent_real = np.random.normal(3, 5) exponent_imag = np.random.normal(3, 5) base_matrix[n, m] = pe.CObs(pe.pseudo_Obs(2 + 10 ** exponent_real, 10 ** (exponent_real - 1), 't'), pe.pseudo_Obs(2 + 10 ** exponent_imag, 10 ** (exponent_imag - 1), 't')) for other in [2, 2.3, (1 - 0.1j), (0 + 2.1j)]: ta = base_matrix * other tb = other * base_matrix diff = ta - tb for (i, j), entry in np.ndenumerate(diff): assert entry.is_zero() ta = base_matrix + other tb = other + base_matrix diff = ta - tb for (i, j), entry in np.ndenumerate(diff): assert entry.is_zero() ta = base_matrix - other tb = other - base_matrix diff = ta + tb for (i, j), entry in np.ndenumerate(diff): assert entry.is_zero() ta = base_matrix / other tb = other / base_matrix diff = ta * tb - 1 for (i, j), entry in np.ndenumerate(diff): assert entry.is_zero()