def test_mask(self):
msk_arr = np.random.choice([0, 1], size=(1, 128))
msk_arr = comm.bcast(msk_arr, root=0)
if not mpirank:
dat_arr = np.random.rand(2, 128)
else:
dat_arr = np.random.rand(1, 128)
cov_arr = np.random.rand(
mpi_arrange(128)[1] - mpi_arrange(128)[0], 128)
# mask by methods
dat_msk = mask_obs(dat_arr, msk_arr)
cov_msk = mask_cov(cov_arr, msk_arr)
# mask manually
test_dat = dat_arr * msk_arr
test_dat = test_dat[test_dat != 0]
dat_msk = dat_msk[dat_msk != 0]
assert np.allclose(test_dat, dat_msk)
#
cov_mat = np.vstack(comm.allgather(cov_arr))
cov_mat = cov_mat * msk_arr
cov_mat = np.transpose(cov_mat)
cov_mat = cov_mat * msk_arr
cov_mat = np.transpose(cov_mat)
cov_mat = cov_mat[cov_mat != 0]
test_cov = np.vstack(comm.allgather(cov_msk))
test_cov = test_cov[test_cov != 0]
assert np.allclose(test_cov, test_cov)
def test_slogdet_odd(self):
cols = 32
rows = mpi_arrange(cols)[1] - mpi_arrange(cols)[0]
arr = np.random.rand(rows, cols)
sign, logdet = mpi_slogdet(arr)
full_arr = np.vstack(comm.allgather(arr))
test_sign, test_logdet = np.linalg.slogdet(full_arr)
assert sign == test_sign
assert np.allclose(logdet, test_logdet)
def test_lu_solve_odd(self):
cols = 32
rows = mpi_arrange(cols)[1] - mpi_arrange(cols)[0]
arr = np.random.rand(rows, cols)
full_arr = np.vstack(comm.allgather(arr))
brr = np.random.rand(1, cols)
comm.Bcast(brr, root=0)
xrr = mpi_lu_solve(arr, brr)
test_xrr = (np.linalg.solve(full_arr, brr.T)).T
assert np.allclose(xrr[0], test_xrr[0])
def test_trans(self):
if not mpirank:
arr = np.random.rand(2, 128)
else:
arr = np.random.rand(1, 128)
test_arr = mpi_trans(arr)
full_arr = np.transpose(np.vstack(comm.allgather(arr)))
local_begin, local_end = mpi_arrange(full_arr.shape[0])
part_arr = full_arr[local_begin:local_end]
assert np.allclose(part_arr, test_arr)
def test_mpi_local(self):
if not mpirank:
arr_a = np.random.rand(32, 128)
else:
arr_a = None
test_a = mpi_local(arr_a)
arr_a = comm.bcast(arr_a, root=0)
local_a_begin, local_a_end = mpi_arrange(arr_a.shape[0])
part_a = arr_a[local_a_begin:local_a_end, :]
part_a = part_a.reshape(1, -1)
test_a = test_a.reshape(1, -1)
assert np.allclose(part_a[0], test_a[0])
def test_mult(self):
if not mpirank:
arr_a = np.random.rand(2, 128)
else:
arr_a = np.random.rand(1, 128)
arr_b = mpi_trans(arr_a)
test_c = mpi_mult(arr_a, arr_b)
# make comparison
full_a = np.vstack(comm.allgather(arr_a))
full_b = np.vstack(comm.allgather(arr_b))
full_c = np.dot(full_a, full_b)
local_begin, local_end = mpi_arrange(full_c.shape[0])
part_c = (full_c[local_begin:local_end]).reshape(1, -1)
test_c = test_c.reshape(1, -1)
assert np.allclose(part_c[0], test_c[0])
