def test_idiv(self):
row = np.array([0, 3, 1, 2, 3, 0])
col = np.array([0, 0, 1, 2, 3, 3])
data = np.array([2., 1, 3, 4, 5, 1])
m = coo_matrix((data, (row, col)), shape=(4, 4))
rank = comm.Get_rank()
# create mpi matrix
rank_ownership = [[0, -1], [-1, 1]]
bm = MPIBlockMatrix(2, 2, rank_ownership, comm)
if rank == 0:
bm.set_block(0, 0, m)
if rank == 1:
bm.set_block(1, 1, m)
bm.broadcast_block_sizes()
serial_bm = BlockMatrix(2, 2)
serial_bm.set_block(0, 0, m)
serial_bm.set_block(1, 1, m)
bm /= 2.0
serial_bm /= 2.0
rows, columns = np.nonzero(bm.ownership_mask)
for i, j in zip(rows, columns):
if bm.get_block(i, j) is not None:
self.assertTrue(
np.allclose(
bm.get_block(i, j).toarray(),
serial_bm.get_block(i, j).toarray()))
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())
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())
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 test_isub(self):
row = np.array([0, 3, 1, 2, 3, 0])
col = np.array([0, 0, 1, 2, 3, 3])
data = np.array([2., 1, 3, 4, 5, 1])
m = coo_matrix((data, (row, col)), shape=(4, 4))
rank = comm.Get_rank()
# create mpi matrix
rank_ownership = [[0, -1], [-1, 1]]
bm = MPIBlockMatrix(2, 2, rank_ownership, comm)
if rank == 0:
bm.set_block(0, 0, m.copy())
if rank == 1:
bm.set_block(1, 1, m.copy())
serial_bm = BlockMatrix(2, 2)
serial_bm.set_block(0, 0, m.copy())
serial_bm.set_block(1, 1, m.copy())
bm -= bm
serial_bm -= serial_bm
rows, columns = np.nonzero(bm.ownership_mask)
for i, j in zip(rows, columns):
if bm.get_block(i, j) is not None:
self.assertTrue(
np.allclose(
bm.get_block(i, j).toarray(),
serial_bm.get_block(i, j).toarray()))
bm -= serial_bm
serial_bm -= serial_bm
self._compare_mpi_and_serial_block_matrices(bm, serial_bm)
def _build_mpi_block_matrix(self,
extra_row: bool = False,
extra_col: bool = False) -> MPIBlockMatrix:
if extra_row:
nrows = self._num_time_blocks + 1
else:
nrows = self._num_time_blocks
if extra_col:
ncols = self._num_time_blocks + 1
else:
ncols = self._num_time_blocks
rank_ownership = -1 * np.ones((nrows, ncols), dtype=np.int64)
for ndx in range(self._num_time_blocks):
rank_ownership[ndx, ndx] = self._ownership_map[ndx]
if extra_row:
rank_ownership[self._num_time_blocks,
ndx] = self._ownership_map[ndx]
if extra_col:
rank_ownership[
ndx, self._num_time_blocks] = self._ownership_map[ndx]
mat = MPIBlockMatrix(nbrows=nrows,
nbcols=ncols,
rank_ownership=rank_ownership,
mpi_comm=self._comm)
return mat
def test_setitem(self):
row = np.array([0, 3, 1, 2, 3, 0])
col = np.array([0, 0, 1, 2, 3, 3])
data = np.array([2., 1, 3, 4, 5, 1])
m = coo_matrix((data, (row, col)), shape=(4, 4))
rank = comm.Get_rank()
# create mpi matrix
rank_ownership = [[0, -1], [-1, 1]]
bm = MPIBlockMatrix(2, 2, rank_ownership, comm)
bm.set_block(0, 1, m)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
self.assertTrue((m == bm.get_block(0, 1)).toarray().all())
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_compression_matrix(compression_mask):
"""
Return a sparse matrix CM of ones such that
compressed_vector = CM*full_vector based on the
compression mask
Parameters
----------
compression_mask: np.ndarray or pyomo.contrib.pynumero.sparse.block_vector.BlockVector
Returns
-------
cm: coo_matrix or BlockMatrix
The compression matrix
"""
if isinstance(compression_mask, BlockVector):
n = compression_mask.nblocks
res = BlockMatrix(nbrows=n, nbcols=n)
for ndx, block in enumerate(compression_mask):
sub_matrix = build_compression_matrix(block)
res.set_block(ndx, ndx, sub_matrix)
return res
elif type(compression_mask) is np.ndarray:
cols = compression_mask.nonzero()[0]
nnz = len(cols)
rows = np.arange(nnz, dtype=np.int)
data = np.ones(nnz)
return coo_matrix((data, (rows, cols)),
shape=(nnz, len(compression_mask)))
elif isinstance(compression_mask, mpi_block_vector.MPIBlockVector):
from pyomo.contrib.pynumero.sparse.mpi_block_matrix import MPIBlockMatrix
n = compression_mask.nblocks
rank_ownership = np.ones((n, n), dtype=np.int64) * -1
for i in range(n):
rank_ownership[i, i] = compression_mask.rank_ownership[i]
res = MPIBlockMatrix(nbrows=n,
nbcols=n,
rank_ownership=rank_ownership,
mpi_comm=compression_mask.mpi_comm,
assert_correct_owners=False)
for ndx in compression_mask.owned_blocks:
block = compression_mask.get_block(ndx)
sub_matrix = build_compression_matrix(block)
res.set_block(ndx, ndx, sub_matrix)
return res
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)
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 setUpClass(cls):
# test problem 1
row = np.array([0, 3, 1, 2, 3, 0])
col = np.array([0, 0, 1, 2, 3, 3])
data = np.array([2., 1, 3, 4, 5, 1])
m = coo_matrix((data, (row, col)), shape=(4, 4))
rank = comm.Get_rank()
# create mpi matrix
rank_ownership = [[0, -1], [-1, 1]]
bm = MPIBlockMatrix(2, 2, rank_ownership, comm)
if rank == 0:
bm.set_block(0, 0, m)
if rank == 1:
bm.set_block(1, 1, m)
# create serial matrix image
serial_bm = BlockMatrix(2, 2)
serial_bm.set_block(0, 0, m)
serial_bm.set_block(1, 1, m)
cls.square_serial_mat = serial_bm
bm.broadcast_block_sizes()
cls.square_mpi_mat = bm
# create mpi matrix
rank_ownership = [[0, -1], [-1, 1]]
bm = MPIBlockMatrix(2, 2, rank_ownership, comm)
if rank == 0:
bm.set_block(0, 0, m)
if rank == 1:
bm.set_block(1, 1, m)
cls.square_mpi_mat_no_broadcast = bm
# create matrix with shared blocks
rank_ownership = [[0, -1], [-1, 1]]
bm = MPIBlockMatrix(2, 2, rank_ownership, comm)
if rank == 0:
bm.set_block(0, 0, m)
if rank == 1:
bm.set_block(1, 1, m)
bm.set_block(0, 1, m)
bm.broadcast_block_sizes()
cls.square_mpi_mat2 = bm
# create serial matrix image
serial_bm = BlockMatrix(2, 2)
serial_bm.set_block(0, 0, m)
serial_bm.set_block(1, 1, m)
serial_bm.set_block(0, 1, m)
cls.square_serial_mat2 = serial_bm
row = np.array([0, 1, 2, 3])
col = np.array([0, 1, 0, 1])
data = np.array([1., 1., 1., 1.])
m2 = coo_matrix((data, (row, col)), shape=(4, 2))
rank_ownership = [[0, -1, 0], [-1, 1, -1]]
bm = MPIBlockMatrix(2, 3, rank_ownership, comm)
if rank == 0:
bm.set_block(0, 0, m)
bm.set_block(0, 2, m2)
if rank == 1:
bm.set_block(1, 1, m)
bm.broadcast_block_sizes()
cls.rectangular_mpi_mat = bm
bm = BlockMatrix(2, 3)
bm.set_block(0, 0, m)
bm.set_block(0, 2, m2)
bm.set_block(1, 1, m)
cls.rectangular_serial_mat = bm
def test_reset_bcol(self):
row = np.array([0, 3, 1, 2, 3, 0])
col = np.array([0, 0, 1, 2, 3, 3])
data = np.array([2., 1, 3, 4, 5, 1])
m = coo_matrix((data, (row, col)), shape=(4, 4))
rank = comm.Get_rank()
# create mpi matrix
rank_ownership = [[0, -1], [-1, 1]]
bm = MPIBlockMatrix(2, 2, rank_ownership, comm)
if rank == 0:
bm.set_block(0, 0, m)
if rank == 1:
bm.set_block(1, 1, m)
bm.broadcast_block_sizes()
serial_bm = BlockMatrix(2, 2)
serial_bm.set_block(0, 0, m)
serial_bm.set_block(1, 1, m)
self.assertTrue(
np.allclose(serial_bm.row_block_sizes(), bm.row_block_sizes()))
bm.reset_bcol(0)
serial_bm.reset_bcol(0)
self.assertTrue(
np.allclose(serial_bm.col_block_sizes(), bm.col_block_sizes()))
bm.reset_bcol(1)
serial_bm.reset_bcol(1)
self.assertTrue(
np.allclose(serial_bm.col_block_sizes(), bm.col_block_sizes()))
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_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_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)
