def test_get_block_vector_for_dot_product_3(self):
rank = comm.Get_rank()
rank_ownership = np.array([[1, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]])
m = MPIBlockMatrix(4, 3, rank_ownership, comm)
sub_m = np.array([[1, 0], [0, 1]])
sub_m = coo_matrix(sub_m)
if rank == 0:
m.set_block(3, rank, sub_m.copy())
elif rank == 1:
m.set_block(0, 0, sub_m.copy())
m.set_block(rank, rank, sub_m.copy())
m.set_block(3, rank, sub_m.copy())
else:
m.set_block(rank, rank, sub_m.copy())
m.set_block(3, rank, sub_m.copy())
rank_ownership = np.array([0, 1, 2])
v = MPIBlockVector(3, rank_ownership, comm)
sub_v = np.ones(2)
v.set_block(rank, sub_v.copy())
res = m._get_block_vector_for_dot_product(v)
self.assertIsNot(res, v)
self.assertTrue(np.array_equal(res.get_block(0), sub_v))
if rank == 0:
self.assertIsNone(res.get_block(1))
self.assertIsNone(res.get_block(2))
elif rank == 1:
self.assertTrue(np.array_equal(res.get_block(1), sub_v))
self.assertIsNone(res.get_block(2))
elif rank == 2:
self.assertTrue(np.array_equal(res.get_block(2), sub_v))
self.assertIsNone(res.get_block(1))
def test_get_block_vector_for_dot_product_4(self):
rank = comm.Get_rank()
rank_ownership = np.array([[-1, 1, 2],
[0, 1, 2],
[0, 1, 2],
[0, 1, 2]])
m = MPIBlockMatrix(4, 3, rank_ownership, comm)
sub_m = np.array([[1, 0],
[0, 1]])
sub_m = coo_matrix(sub_m)
m.set_block(0, 0, sub_m.copy())
if rank == 0:
m.set_block(3, rank, sub_m.copy())
else:
m.set_block(rank, rank, sub_m.copy())
m.set_block(3, rank, sub_m.copy())
rank_ownership = np.array([0, 1, 2])
v = MPIBlockVector(3, rank_ownership, comm)
sub_v = np.ones(2)
v.set_block(rank, sub_v.copy())
res = m._get_block_vector_for_dot_product(v)
self.assertIs(res, v)
def test_dot(self):
mat1 = self.square_mpi_mat
mat2 = self.square_mpi_mat2
serial_mat1 = self.square_serial_mat
serial_mat2 = self.square_serial_mat2
rank = comm.Get_rank()
bv1 = MPIBlockVector(2, [0, 1], comm)
if rank == 0:
bv1.set_block(0, np.arange(4, dtype=np.float64))
if rank == 1:
bv1.set_block(1, np.arange(4, dtype=np.float64) + 4)
bv1.broadcast_block_sizes()
serial_bv1 = BlockVector(2)
serial_bv1.set_block(0, np.arange(4, dtype=np.float64))
serial_bv1.set_block(1, np.arange(4, dtype=np.float64) + 4)
res = mat1.dot(bv1)
serial_res = serial_mat1.dot(serial_bv1)
self.assertIsInstance(res, BlockVector)
self.assertEqual(res.nblocks, serial_res.nblocks)
for bid in range(serial_res.nblocks):
self.assertTrue(
np.allclose(res.get_block(bid), serial_res.get_block(bid)))
def test_matvec_1(self):
rank = comm.Get_rank()
rank_ownership = np.array([[0, -1, -1, 0],
[-1, 1, -1, 1],
[-1, -1, 2, 2],
[0, 1, 2, -1]])
m = MPIBlockMatrix(4, 4, rank_ownership, comm)
sub_m = np.array([[1, 0],
[0, 1]])
sub_m = coo_matrix(sub_m)
m.set_block(rank, rank, sub_m.copy())
m.set_block(rank, 3, sub_m.copy())
m.set_block(3, rank, sub_m.copy())
m.set_block(3, 3, sub_m.copy())
m.broadcast_block_sizes()
rank_ownership = np.array([0, 1, 2, -1])
v = MPIBlockVector(4, rank_ownership, comm)
sub_v = np.ones(2)
v.set_block(rank, sub_v.copy())
v.set_block(3, sub_v.copy())
v.broadcast_block_sizes()
res = m.dot(v)
self.assertIsInstance(res, MPIBlockVector)
self.assertTrue(np.array_equal(res.get_block(rank), sub_v*2))
self.assertTrue(np.array_equal(res.get_block(3), sub_v*4))
self.assertTrue(np.array_equal(res.rank_ownership, np.array([0, 1, 2, -1])))
self.assertFalse(res.has_none)
def main():
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
owners = [0, 1, 2, -1]
x = MPIBlockVector(4, rank_owner=owners, mpi_comm=comm)
owners = np.array([[ 0, -1, -1, 0],
[-1, 1, -1, 1],
[-1, -1, 2, 2]])
a = MPIBlockMatrix(3, 4, rank_ownership=owners, mpi_comm=comm)
np.random.seed(0)
x.set_block(3, np.random.uniform(-10, 10, size=10))
np.random.seed(rank)
x.set_block(rank, np.random.uniform(-10, 10, size=10))
a.set_block(rank, rank, random(10, 10, density=0.1))
a.set_block(rank, 3, random(10, 10, density=0.1))
b = a * x # parallel matrix-vector dot product
# check the answer
local_x = x.make_local_copy().flatten()
local_a = a.to_local_array()
local_b = b.make_local_copy().flatten()
err = np.abs(local_a.dot(local_x) - local_b).max()
if rank == 0:
print('error: ', err)
return err
def _build_mpi_block_vector(self,
extra_block: bool = False) -> MPIBlockVector:
if extra_block:
n = self._num_time_blocks + 1
else:
n = self._num_time_blocks
rank_ownership = -1 * np.ones(n, dtype=np.int64)
for ndx in range(self._num_time_blocks):
rank_ownership[ndx] = self._ownership_map[ndx]
vec = MPIBlockVector(nblocks=n,
rank_owner=rank_ownership,
mpi_comm=self._comm)
return vec
def test_mpi_schur_complement(self):
rank_by_index = list()
for ndx in range(3):
for _rank in range(size):
if (ndx - _rank) % size == 0:
rank_by_index.append(_rank)
rank_by_index.append(-1)
A = MPIBlockMatrix(nbrows=4,
nbcols=4,
rank_ownership=[
rank_by_index, rank_by_index, rank_by_index,
rank_by_index
],
mpi_comm=comm)
if rank_by_index[0] == rank:
A.set_block(
0, 0, coo_matrix(np.array([[1, 1], [0, 1]], dtype=np.double)))
if rank_by_index[1] == rank:
A.set_block(
1, 1, coo_matrix(np.array([[1, 0], [0, 1]], dtype=np.double)))
if rank_by_index[2] == rank:
A.set_block(
2, 2, coo_matrix(np.array([[1, 0], [1, 1]], dtype=np.double)))
A.set_block(3, 3,
coo_matrix(np.array([[0, 0], [0, 1]], dtype=np.double)))
if rank_by_index[0] == rank:
A.set_block(
3, 0, coo_matrix(np.array([[0, -1], [0, 0]], dtype=np.double)))
if rank_by_index[1] == rank:
A.set_block(
3, 1, coo_matrix(np.array([[-1, 0], [0, -1]],
dtype=np.double)))
if rank_by_index[2] == rank:
A.set_block(
3, 2, coo_matrix(np.array([[0, 0], [-1, 0]], dtype=np.double)))
A.broadcast_block_sizes()
local_A = BlockMatrix(4, 4)
local_A.set_block(
0, 0, coo_matrix(np.array([[1, 1], [0, 1]], dtype=np.double)))
local_A.set_block(
1, 1, coo_matrix(np.array([[1, 0], [0, 1]], dtype=np.double)))
local_A.set_block(
2, 2, coo_matrix(np.array([[1, 0], [1, 1]], dtype=np.double)))
local_A.set_block(
3, 3, coo_matrix(np.array([[0, 0], [0, 1]], dtype=np.double)))
local_A.set_block(
3, 0, coo_matrix(np.array([[0, -1], [0, 0]], dtype=np.double)))
local_A.set_block(
3, 1, coo_matrix(np.array([[-1, 0], [0, -1]], dtype=np.double)))
local_A.set_block(
3, 2, coo_matrix(np.array([[0, 0], [-1, 0]], dtype=np.double)))
local_A.set_block(0, 3, local_A.get_block(3, 0).transpose(copy=True))
local_A.set_block(1, 3, local_A.get_block(3, 1).transpose(copy=True))
local_A.set_block(2, 3, local_A.get_block(3, 2).transpose(copy=True))
rhs = MPIBlockVector(nblocks=4,
rank_owner=rank_by_index,
mpi_comm=comm)
if rank_by_index[0] == rank:
rhs.set_block(0, np.array([1, 0], dtype=np.double))
if rank_by_index[1] == rank:
rhs.set_block(1, np.array([0, 0], dtype=np.double))
if rank_by_index[2] == rank:
rhs.set_block(2, np.array([0, 1], dtype=np.double))
rhs.set_block(3, np.array([1, 1], dtype=np.double))
rhs.broadcast_block_sizes()
local_rhs = BlockVector(4)
local_rhs.set_block(0, np.array([1, 0], dtype=np.double))
local_rhs.set_block(1, np.array([0, 0], dtype=np.double))
local_rhs.set_block(2, np.array([0, 1], dtype=np.double))
local_rhs.set_block(3, np.array([1, 1], dtype=np.double))
x1 = np.linalg.solve(local_A.toarray(), local_rhs.flatten())
solver_class = parapint.linalg.MPISchurComplementLinearSolver
sc_solver = solver_class(
subproblem_solvers={
ndx: ScipyInterface(compute_inertia=True)
for ndx in range(3)
},
schur_complement_solver=ScipyInterface(compute_inertia=True))
sc_solver.do_symbolic_factorization(A)
sc_solver.do_numeric_factorization(A)
x2 = sc_solver.do_back_solve(rhs)
self.assertTrue(np.allclose(x1, x2.make_local_copy().flatten()))
inertia1 = sc_solver.get_inertia()
eig = np.linalg.eigvals(local_A.toarray())
pos = np.count_nonzero(eig > 0)
neg = np.count_nonzero(eig < 0)
zero = np.count_nonzero(eig == 0)
inertia2 = (pos, neg, zero)
self.assertEqual(inertia1, inertia2)
sc_solver.do_numeric_factorization(A)
x2 = sc_solver.do_back_solve(rhs)
self.assertTrue(np.allclose(x1, x2.make_local_copy().flatten()))
def test_mul(self):
mat1 = self.square_mpi_mat
mat2 = self.square_mpi_mat2
serial_mat1 = self.square_serial_mat
serial_mat2 = self.square_serial_mat2
rank = comm.Get_rank()
bv1 = MPIBlockVector(2, [0, 1], comm)
if rank == 0:
bv1.set_block(0, np.arange(4, dtype=np.float64))
if rank == 1:
bv1.set_block(1, np.arange(4, dtype=np.float64) + 4)
bv1.broadcast_block_sizes()
serial_bv1 = BlockVector(2)
serial_bv1.set_block(0, np.arange(4, dtype=np.float64))
serial_bv1.set_block(1, np.arange(4, dtype=np.float64) + 4)
res = mat1 * bv1
serial_res = serial_mat1 * serial_bv1
self.assertIsInstance(res, BlockVector)
self.assertEqual(res.nblocks, serial_res.nblocks)
for bid in range(serial_res.nblocks):
self.assertTrue(
np.allclose(res.get_block(bid), serial_res.get_block(bid)))
res = mat2 * bv1
serial_res = serial_mat2 * serial_bv1
self.assertIsInstance(res, BlockVector)
self.assertEqual(res.nblocks, serial_res.nblocks)
for bid in range(serial_res.nblocks):
self.assertTrue(
np.allclose(res.get_block(bid), serial_res.get_block(bid)))
bv1 = MPIBlockVector(2, [0, -1], comm)
if rank == 0:
bv1.set_block(0, np.arange(4, dtype=np.float64))
bv1.set_block(1, np.arange(4, dtype=np.float64) + 4)
bv1.broadcast_block_sizes()
res = mat1 * bv1
serial_res = serial_mat1 * serial_bv1
self.assertIsInstance(res, BlockVector)
self.assertEqual(res.nblocks, serial_res.nblocks)
for bid in range(serial_res.nblocks):
self.assertTrue(
np.allclose(res.get_block(bid), serial_res.get_block(bid)))
res = mat2 * bv1
serial_res = serial_mat2 * serial_bv1
self.assertIsInstance(res, BlockVector)
self.assertEqual(res.nblocks, serial_res.nblocks)
for bid in range(serial_res.nblocks):
self.assertTrue(
np.allclose(res.get_block(bid), serial_res.get_block(bid)))
# rectangular matrix
mat1 = self.rectangular_mpi_mat
serial_mat1 = self.rectangular_serial_mat
bv1 = MPIBlockVector(3, [0, 1, 2], comm)
if rank == 0:
bv1.set_block(0, np.arange(4, dtype=np.float64))
if rank == 1:
bv1.set_block(1, np.arange(4, dtype=np.float64) + 4)
if rank == 2:
bv1.set_block(2, np.arange(2, dtype=np.float64) + 8)
bv1.broadcast_block_sizes()
serial_bv1 = BlockVector(3)
serial_bv1.set_block(0, np.arange(4, dtype=np.float64))
serial_bv1.set_block(1, np.arange(4, dtype=np.float64) + 4)
serial_bv1.set_block(2, np.arange(2, dtype=np.float64) + 8)
# with warnings.catch_warnings():
# warnings.simplefilter("ignore")
res = mat1 * bv1
serial_res = serial_mat1 * serial_bv1
self.assertIsInstance(res, BlockVector)
self.assertEqual(serial_res.nblocks, 2)
self.assertEqual(res.nblocks, 2)
for bid in range(serial_res.nblocks):
self.assertTrue(
np.allclose(res.get_block(bid), serial_res.get_block(bid)))
bv1 = MPIBlockVector(3, [0, 1, 0], comm)
if rank == 0:
bv1.set_block(0, np.arange(4, dtype=np.float64))
bv1.set_block(2, np.arange(2, dtype=np.float64) + 8)
if rank == 1:
bv1.set_block(1, np.arange(4, dtype=np.float64) + 4)
bv1.broadcast_block_sizes()
res = mat1 * bv1
serial_res = serial_mat1 * serial_bv1
self.assertIsInstance(res, BlockVector)
self.assertEqual(res.nblocks, serial_res.nblocks)
for bid in range(serial_res.nblocks):
self.assertTrue(
np.allclose(res.get_block(bid), serial_res.get_block(bid)))
res = mat1 * 3.0
serial_res = serial_mat1 * 3.0
self.assertIsInstance(res, MPIBlockMatrix)
rows, columns = np.nonzero(res.ownership_mask)
for i, j in zip(rows, columns):
if res.get_block(i, j) is not None:
self.assertTrue(
np.allclose(
res.get_block(i, j).toarray(),
serial_res.get_block(i, j).toarray()))
else:
self.assertIsNone(serial_res.get_block(i, j))
res = 3.0 * mat1
serial_res = serial_mat1 * 3.0
self.assertIsInstance(res, MPIBlockMatrix)
rows, columns = np.nonzero(res.ownership_mask)
for i, j in zip(rows, columns):
if res.get_block(i, j) is not None:
self.assertTrue(
np.allclose(
res.get_block(i, j).toarray(),
serial_res.get_block(i, j).toarray()))
else:
self.assertIsNone(serial_res.get_block(i, j))
def main():
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
owners = [2, 0, 1, -1]
x = MPIBlockVector(4, rank_owner=owners, mpi_comm=comm)
x.set_block(owners.index(rank), np.ones(3) * (rank + 1))
x.set_block(3, np.array([1, 2, 3]))
y = MPIBlockVector(4, rank_owner=owners, mpi_comm=comm)
y.set_block(owners.index(rank), np.ones(3) * (rank + 1))
y.set_block(3, np.array([1, 2, 3]))
z1: MPIBlockVector = x + y # add x and y
z2 = x.dot(y) # dot product
z3 = np.abs(x).max() # infinity norm
z1_local = z1.make_local_copy()
if rank == 0:
print(z1_local.flatten())
print(z2)
print(z3)
return z1_local, z2, z3