def test_get_block_vector_for_dot_product_5(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())
v = BlockVector(3)
sub_v = np.ones(2)
for ndx in range(3):
v.set_block(ndx, sub_v.copy())
res = m._get_block_vector_for_dot_product(v)
self.assertIs(res, v)
v_flat = v.flatten()
res = m._get_block_vector_for_dot_product(v_flat)
self.assertIsInstance(res, BlockVector)
for ndx in range(3):
block = res.get_block(ndx)
self.assertTrue(np.array_equal(block, sub_v))
def test_get_block_vector_for_dot_product_2(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([-1, 1, 2])
v = MPIBlockVector(3, rank_ownership, comm)
sub_v = np.ones(2)
v.set_block(0, sub_v.copy())
if rank != 0:
v.set_block(rank, sub_v.copy())
res = m._get_block_vector_for_dot_product(v)
self.assertIs(res, v)
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_1(self):
rank = comm.Get_rank()
rank_ownership = np.array([[0, 1, 2],
[1, 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(rank, rank, sub_m.copy())
m.set_block(3, rank, sub_m.copy())
m.broadcast_block_sizes()
rank_ownership = np.array([0, 1, 2])
v = MPIBlockVector(3, rank_ownership, comm)
sub_v = np.ones(2)
v.set_block(rank, sub_v)
v.broadcast_block_sizes()
res = m._get_block_vector_for_dot_product(v)
self.assertIs(res, v)
