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)