def expansion_matrix_xu(self):
Pxu = BlockMatrix(self.nblocks + 1, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
Pxu.set_block(sid, sid, nlp.expansion_matrix_xu())
Pxu[self.nblocks, self.nblocks] = coo_matrix((self.nz, 0))
return Pxu
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())
bm.broadcast_block_sizes()
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()))
with self.assertRaises(Exception) as context:
bm -= serial_bm
def test_iadd(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 getZ(nlp, parm_vars):
# Get the Z matrix to compute reduced hessian
parm_vars_name = [x.name for x in parm_vars]
non_parm_vars = [
x for x in nlp.get_pyomo_variables() if x.name not in parm_vars_name
]
Ji = nlp.extract_submatrix_jacobian(
pyomo_variables=parm_vars,
pyomo_constraints=nlp.get_pyomo_constraints())
Jd = nlp.extract_submatrix_jacobian(
pyomo_variables=non_parm_vars,
pyomo_constraints=nlp.get_pyomo_constraints())
#print("Ji")
#print(Ji.todense())
#print("Jd")
#print(Jd.todense())
Zd = spsolve(Jd.tocsc(), Ji.tocsc())
Z = BlockMatrix(2, 1)
Z[0, 0] = Zd
Z[1, 0] = identity(len(parm_vars))
#print("Z")
#print(Z.todense())
# reorder variables to the order in hessian
zorder = getvarorder(nlp, parm_vars, non_parm_vars)
Zorder = Z.tocsc()[zorder, :].todense()
#print("Zorder")
#print(Zorder)
return Zorder
def test_abs(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)
res = abs(bm)
serial_res = abs(serial_bm)
rows, columns = np.nonzero(bm.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()))
def test_setitem(self):
m = BlockMatrix(2, 2)
m[0, 1] = self.block_m
self.assertFalse(m.is_empty_block(0, 1))
self.assertEqual(m.row_block_sizes()[0], self.block_m.shape[0])
self.assertEqual(m.col_block_sizes()[1], self.block_m.shape[1])
self.assertEqual(m[0, 1].shape, self.block_m.shape)
def test_setitem(self):
m = BlockMatrix(2, 2)
m[0, 1] = self.block_m
self.assertFalse(m.is_empty_block(0, 1))
self.assertEqual(m.row_block_sizes()[0], self.block_m.shape[0])
self.assertEqual(m.col_block_sizes()[1], self.block_m.shape[1])
self.assertEqual(m[0, 1].shape, self.block_m.shape)
def jacobian_c(self, x, out=None, **kwargs):
"""Returns the Jacobian of the equalities evaluated at x
Parameters
----------
x : array_like
Array with values of primal variables.
out : BlockMatrix, optional
Output matrix with the structure of the jacobian already defined.
Returns
-------
BlockMatrix
"""
assert x.size == self.nx, 'Dimension mismatch'
if isinstance(x, BlockVector):
assert x.nblocks == self.nblocks + 1
x_ = x
elif isinstance(x, np.ndarray):
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
else:
raise RuntimeError('Input vector format not recognized')
if out is None:
jac_c = BlockMatrix(2 * self.nblocks, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
jac_c.set_block(sid, sid, nlp.jacobian_c(xi))
# coupling matrices Ai
jac_c[sid + self.nblocks,
sid] = self._AB_coo.get_block(sid, sid)
# coupling matrices Bi
jac_c[sid + self.nblocks, self.nblocks] = -identity(self.nz)
return jac_c
else:
assert isinstance(out, BlockMatrix), 'out must be a BlockMatrix'
assert out.bshape == (2 * self.nblocks,
self.nblocks + 1), "Block shape mismatch"
jac_c = out
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
nlp.jacobian_c(xi, out=jac_c.get_block(sid, sid))
Ai = jac_c[sid + self.nblocks, sid]
assert Ai.shape == self._AB_coo.get_block(sid, sid).shape, \
'Block {} mismatch shape'.format((sid + self.nblocks, sid))
assert Ai.nnz == self._AB_coo.get_block(sid, sid).nnz, \
'Block {} mismatch nnz'.format((sid + self.nblocks, sid))
Bi = jac_c[sid + self.nblocks, self.nblocks]
assert Bi.shape == (self.nz, self.nz), \
'Block {} mismatch shape'.format((sid + self.nblocks, self.nblocks))
assert Bi.nnz == self.nz, \
'Block {} mismatch nnz'.format((sid + self.nblocks, self.nblocks))
return jac_c
def compute_init_lam(nlp, x=None, lam_max=1e3):
if x is None:
x = nlp.init_primals()
else:
assert x.size == nlp.n_primals()
nlp.set_primals(x)
assert nlp.n_ineq_constraints(
) == 0, "only supported for equality constrained nlps for now"
nx = nlp.n_primals()
nc = nlp.n_constraints()
# create Jacobian
jac = nlp.evaluate_jacobian()
# create gradient of objective
df = nlp.evaluate_grad_objective()
# create KKT system
kkt = BlockMatrix(2, 2)
kkt.set_block(0, 0, identity(nx))
kkt.set_block(1, 0, jac)
kkt.set_block(0, 1, jac.transpose())
zeros = np.zeros(nc)
rhs = BlockVector(2)
rhs.set_block(0, -df)
rhs.set_block(1, zeros)
flat_kkt = kkt.tocoo().tocsc()
flat_rhs = rhs.flatten()
sol = spsolve(flat_kkt, flat_rhs)
return sol[nlp.n_primals():nlp.n_primals() + nlp.n_constraints()]
def regularize_hessian(self,
kkt: BlockMatrix,
coef: float,
copy_kkt: bool = True) -> BlockMatrix:
if copy_kkt:
kkt = kkt.copy()
for ndx, nlp in self._nlps.items():
nlp.regularize_hessian(kkt=kkt.get_block(ndx, ndx).get_block(0, 0),
coef=coef,
copy_kkt=False)
return kkt
def coupling_matrix(self):
AB = BlockMatrix(self.nblocks + 1, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
col = self._zid_to_vid[sid]
row = np.arange(self.nz, dtype=np.int)
data = np.ones(self.nz)
AB.set_block(
sid, sid,
csr_matrix((data, (row, col)), shape=(self.nz, nlp.nx)))
AB[self.nblocks, self.nblocks] = -identity(self.nz)
return AB
def create_blocks(self, m: np.ndarray, x: np.ndarray):
m = coo_matrix(m)
r = m * x
bm = BlockMatrix(2, 2)
bm.set_block(0, 0, m.copy())
bm.set_block(1, 1, m.copy())
br = BlockVector(2)
br.set_block(0, r.copy())
br.set_block(1, r.copy())
bx = BlockVector(2)
bx.set_block(0, x.copy())
bx.set_block(1, x.copy())
return bm, bx, br
def test_get_block_row_index(self):
m = BlockMatrix(2, 4)
m.set_block(0, 0, coo_matrix((3, 2)))
m.set_block(0, 1, coo_matrix((3, 4)))
m.set_block(0, 2, coo_matrix((3, 3)))
m.set_block(0, 3, coo_matrix((3, 6)))
m.set_block(1, 3, coo_matrix((5, 6)))
brow = m.get_block_row_index(0)
self.assertEqual(brow, 0)
brow = m.get_block_row_index(6)
self.assertEqual(brow, 1)
def expansion_matrix_xu(self):
Pxu = BlockMatrix(self.nblocks + 1, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
Pxu[sid, sid] = nlp.expansion_matrix_xu()
Pxu[self.nblocks, self.nblocks] = empty_matrix(self.nz, 0)
return Pxu
def jacobian_d(self, x, out=None, **kwargs):
"""Returns the Jacobian of the inequalities evaluated at x
Parameters
----------
x : array_like
Array with values of primal variables.
out : COOMatrix, optional
Output matrix with the structure of the jacobian already defined.
Returns
-------
BlockMatrix
"""
assert x.size == self.nx, "Dimension missmatch"
if out is None:
if isinstance(x, BlockVector):
assert x.nblocks == self.nblocks + 1
jac_d = BlockMatrix(self.nblocks, self.nblocks)
for sid, nlp in enumerate(self._nlps):
xi = x[sid]
jac_d[sid, sid] = nlp.jacobian_d(xi)
return jac_d
elif isinstance(x, np.ndarray):
raise NotImplementedError("ToDo")
else:
raise NotImplementedError("ToDo")
def _create_hessian_structure(self):
# Note: This method requires the complicated vars map to be
# created beforehand
hess_lag = BlockMatrix(self.nblocks + 1, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
xi = nlp.x_init()
yi = nlp.y_init()
hess_lag.set_block(sid, sid, nlp.hessian_lag(xi, yi))
hess_lag[self.nblocks, self.nblocks] = coo_matrix((self.nz, self.nz))
flat_hess = hess_lag.tocoo()
self._irows_hess = flat_hess.row
self._jcols_hess = flat_hess.col
self._nnz_hess_lag = flat_hess.nnz
def test_getitem(self):
m = BlockMatrix(3, 3)
for i in range(3):
for j in range(3):
self.assertIsNone(m[i, j])
m[0, 1] = self.block_m
self.assertEqual(m[0, 1].shape, self.block_m.shape)
def setUp(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))
self.block_m = m
bm = BlockMatrix(2, 2)
bm.name = 'basic_matrix'
bm[0, 0] = m
bm[1, 1] = m
bm[0, 1] = m
self.basic_m = bm
self.composed_m = BlockMatrix(2, 2)
self.composed_m[0, 0] = self.block_m
self.composed_m[1, 1] = self.basic_m
def setUp(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 = COOMatrix((data, (row, col)), shape=(4, 4))
self.block_m = m
bm = BlockMatrix(2, 2)
bm.name = 'basic_matrix'
bm[0, 0] = m
bm[1, 1] = m
bm[0, 1] = m
self.basic_m = bm
self.composed_m = BlockMatrix(2, 2)
self.composed_m[0, 0] = self.block_m
self.composed_m[1, 1] = self.basic_m
def jacobian_d(self, x, out=None, **kwargs):
"""Returns the Jacobian of the inequalities evaluated at x
Parameters
----------
x : array_like
Array with values of primal variables.
out : coo_matrix, optional
Output matrix with the structure of the jacobian already defined.
Returns
-------
BlockMatrix
"""
assert x.size == self.nx, "Dimension mismatch"
if isinstance(x, BlockVector):
assert x.nblocks == self.nblocks + 1
x_ = x
elif isinstance(x, np.ndarray):
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
else:
raise RuntimeError('Input vector format not recognized')
if out is None:
jac_d = BlockMatrix(self.nblocks, self.nblocks)
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
jac_d.set_block(sid, sid, nlp.jacobian_d(xi))
return jac_d
else:
assert isinstance(out, BlockMatrix), 'out must be a BlockMatrix'
assert out.bshape == (self.nblocks,
self.nblocks), 'Block shape mismatch'
jac_d = out
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
nlp.jacobian_d(xi, out=jac_d.get_block(sid, sid))
return jac_d
def main(show_plot=True):
if show_plot:
import matplotlib.pylab as plt
instance = create_problem(0.0, 10.0)
# Discretize model using Orthogonal Collocation
discretizer = pyo.TransformationFactory('dae.collocation')
discretizer.apply_to(instance, nfe=100, ncp=3, scheme='LAGRANGE-RADAU')
discretizer.reduce_collocation_points(instance,
var=instance.u,
ncp=1,
contset=instance.t)
# Interface pyomo model with nlp
nlp = PyomoNLP(instance)
x = nlp.create_new_vector('primals')
x.fill(1.0)
nlp.set_primals(x)
lam = nlp.create_new_vector('duals')
lam.fill(1.0)
nlp.set_duals(lam)
# Evaluate jacobian
jac = nlp.evaluate_jacobian()
if show_plot:
plt.spy(jac)
plt.title('Jacobian of the constraints\n')
plt.show()
# Evaluate hessian of the lagrangian
hess_lag = nlp.evaluate_hessian_lag()
if show_plot:
plt.spy(hess_lag)
plt.title('Hessian of the Lagrangian function\n')
plt.show()
# Build KKT matrix
kkt = BlockMatrix(2, 2)
kkt.set_block(0, 0, hess_lag)
kkt.set_block(1, 0, jac)
kkt.set_block(0, 1, jac.transpose())
if show_plot:
plt.spy(kkt.tocoo())
plt.title('KKT system\n')
plt.show()
def test_dot(self):
A_dense = self.basic_m.toarray()
A_block = self.basic_m
x = np.ones(A_dense.shape[1])
block_x = BlockVector(2)
block_x.set_block(0, np.ones(self.block_m.shape[1]))
block_x.set_block(1, np.ones(self.block_m.shape[1]))
flat_res = A_block.dot(x).flatten()
block_res = A_block.dot(block_x)
self.assertTrue(np.allclose(A_dense.dot(x), flat_res))
self.assertTrue(np.allclose(A_dense.dot(x), block_res.flatten()))
self.assertEqual(block_res.bshape[0], 2)
m = BlockMatrix(2, 2)
sub_m = np.array([[1, 0], [0, 1]])
sub_m = coo_matrix(sub_m)
m.set_block(0, 1, sub_m.copy())
m.set_block(1, 0, sub_m.copy())
x = np.arange(4)
res = m * x
self.assertTrue(np.allclose(res.flatten(), np.array([2, 3, 0, 1])))
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)
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)
res.broadcast_block_sizes()
return res
def test_transpose_with_empty_rows(self):
m = BlockMatrix(2, 2)
m.set_row_size(0, 2)
m.set_row_size(1, 2)
m.set_col_size(0, 2)
m.set_col_size(1, 2)
mt = m.transpose()
self.assertEqual(mt.get_row_size(0), 2)
self.assertEqual(mt.get_row_size(1), 2)
self.assertEqual(mt.get_col_size(0), 2)
self.assertEqual(mt.get_col_size(1), 2)
def test_setitem(self):
m = BlockMatrix(2, 2)
m.set_block(0, 1, self.block_m)
self.assertFalse(m.is_empty_block(0, 1))
self.assertEqual(m._brow_lengths[0], self.block_m.shape[0])
self.assertEqual(m._bcol_lengths[1], self.block_m.shape[1])
self.assertEqual(m.get_block(0, 1).shape, self.block_m.shape)
def test_getitem(self):
m = BlockMatrix(3, 3)
for i in range(3):
for j in range(3):
self.assertIsNone(m.get_block(i, j))
m.set_block(0, 1, self.block_m)
self.assertEqual(m.get_block(0, 1).shape, self.block_m.shape)
def test_mumps_linear_solver(self):
A = np.array([[ 1, 7, 3],
[ 7, 4, -5],
[ 3, -5, 6]], dtype=np.double)
A = coo_matrix(A)
A_lower = tril(A)
x1 = np.arange(3) + 1
b1 = A * x1
x2 = np.array(list(reversed(x1)))
b2 = A * x2
solver = MumpsCentralizedAssembledLinearSolver()
solver.do_symbolic_factorization(A)
solver.do_numeric_factorization(A)
x = solver.do_back_solve(b1)
self.assertTrue(np.allclose(x, x1))
x = solver.do_back_solve(b2)
self.assertTrue(np.allclose(x, x2))
solver = MumpsCentralizedAssembledLinearSolver(sym=2)
x = solver.solve(A_lower, b1)
self.assertTrue(np.allclose(x, x1))
block_A = BlockMatrix(2, 2)
block_A.set_row_size(0, 2)
block_A.set_row_size(1, 1)
block_A.set_col_size(0, 2)
block_A.set_col_size(1, 1)
block_A.copyfrom(A)
block_b1 = BlockVector(2)
block_b1.set_block(0, b1[0:2])
block_b1.set_block(1, b1[2:])
block_b2 = BlockVector(2)
block_b2.set_block(0, b2[0:2])
block_b2.set_block(1, b2[2:])
solver = MumpsCentralizedAssembledLinearSolver(icntl_options={10: -3}, cntl_options={2: 1e-16})
solver.do_symbolic_factorization(block_A)
solver.do_numeric_factorization(block_A)
x = solver.do_back_solve(block_b1)
self.assertTrue(np.allclose(x, x1))
x = solver.do_back_solve(block_b2)
self.assertTrue(np.allclose(x, x2))
self.assertEqual(solver.get_infog(15), 3)
def setUp(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))
self.block_m = m
bm = BlockMatrix(2, 2)
bm.name = 'basic_matrix'
bm.set_block(0, 0, m.copy())
bm.set_block(1, 1, m.copy())
bm.set_block(0, 1, m.copy())
self.basic_m = bm
self.dense = np.zeros((8, 8))
self.dense[0:4, 0:4] = m.toarray()
self.dense[0:4, 4:8] = m.toarray()
self.dense[4:8, 4:8] = m.toarray()
self.composed_m = BlockMatrix(2, 2)
self.composed_m.set_block(0, 0, self.block_m.copy())
self.composed_m.set_block(1, 1, self.basic_m.copy())
def test_abs(self):
row = np.array([0, 3, 1, 2, 3, 0])
col = np.array([0, 0, 1, 2, 3, 3])
data = -1.0 * np.array([2., 1, 3, 4, 5, 1])
m = coo_matrix((data, (row, col)), shape=(4, 4))
self.block_m = m
bm = BlockMatrix(2, 2)
bm.set_block(0, 0, m)
bm.set_block(1, 1, m)
bm.set_block(0, 1, m)
abs_flat = abs(bm.tocoo())
abs_mat = abs(bm)
self.assertIsInstance(abs_mat, BlockMatrix)
self.assertTrue(np.allclose(abs_flat.toarray(), abs_mat.toarray()))
def _create_jacobian_structures(self):
# Note: This method requires the complicated vars map to be
# created beforehand
# build general jacobian
jac_g = BlockMatrix(2 * self.nblocks, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
xi = nlp.x_init()
jac_g[sid, sid] = nlp.jacobian_g(xi)
# coupling matrices Ai
scenario_vids = self._zid_to_vid[sid]
col = np.array([vid for vid in scenario_vids])
row = np.arange(0, self.nz)
data = np.ones(self.nz, dtype=np.double)
jac_g[sid + self.nblocks, sid] = coo_matrix((data, (row, col)),
shape=(self.nz, nlp.nx))
# coupling matrices Bi
jac_g[sid + self.nblocks, self.nblocks] = -identity(self.nz)
self._internal_jacobian_g = jac_g
flat_jac_g = jac_g.tocoo()
self._irows_jac_g = flat_jac_g.row
self._jcols_jac_g = flat_jac_g.col
self._nnz_jac_g = flat_jac_g.nnz
# build jacobian equality constraints
jac_c = BlockMatrix(2 * self.nblocks, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
xi = nlp.x_init()
jac_c[sid, sid] = nlp.jacobian_c(xi)
# coupling matrices Ai
scenario_vids = self._zid_to_vid[sid]
col = np.array([vid for vid in scenario_vids])
row = np.arange(0, self.nz)
data = np.ones(self.nz, dtype=np.double)
jac_c[sid + self.nblocks, sid] = coo_matrix((data, (row, col)),
shape=(self.nz, nlp.nx))
# coupling matrices Bi
jac_c[sid + self.nblocks, self.nblocks] = -identity(self.nz)
self._internal_jacobian_c = jac_c
flat_jac_c = jac_c.tocoo()
self._irows_jac_c = flat_jac_c.row
self._jcols_jac_c = flat_jac_c.col
self._nnz_jac_c = flat_jac_c.nnz
# build jacobian inequality constraints
jac_d = BlockMatrix(self.nblocks, self.nblocks)
for sid, nlp in enumerate(self._nlps):
xi = nlp.x_init()
jac_d[sid, sid] = nlp.jacobian_d(xi)
self._internal_jacobian_d = jac_d
flat_jac_d = jac_d.tocoo()
self._irows_jac_d = flat_jac_d.row
self._jcols_jac_d = flat_jac_d.col
self._nnz_jac_d = flat_jac_d.nnz
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()))
class TestBlockMatrix(unittest.TestCase):
def setUp(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))
self.block_m = m
bm = BlockMatrix(2, 2)
bm.name = 'basic_matrix'
bm[0, 0] = m
bm[1, 1] = m
bm[0, 1] = m
self.basic_m = bm
self.composed_m = BlockMatrix(2, 2)
self.composed_m[0, 0] = self.block_m
self.composed_m[1, 1] = self.basic_m
def test_name(self):
self.assertEqual(self.basic_m.name, 'basic_matrix')
self.basic_m.name = 'hola'
self.assertEqual(self.basic_m.name, 'hola')
def test_bshape(self):
self.assertEqual(self.basic_m.bshape, (2, 2))
def test_shape(self):
shape = (self.block_m.shape[0]*2, self.block_m.shape[1]*2)
self.assertEqual(self.basic_m.shape, shape)
def test_tocoo(self):
block = self.block_m
m = self.basic_m
scipy_mat = bmat([[block, block], [None, block]], format='coo')
dinopy_mat = m.tocoo()
drow = np.sort(dinopy_mat.row)
dcol = np.sort(dinopy_mat.col)
ddata = np.sort(dinopy_mat.data)
srow = np.sort(scipy_mat.row)
scol = np.sort(scipy_mat.col)
sdata = np.sort(scipy_mat.data)
self.assertListEqual(drow.tolist(), srow.tolist())
self.assertListEqual(dcol.tolist(), scol.tolist())
self.assertListEqual(ddata.tolist(), sdata.tolist())
def test_tocsr(self):
block = self.block_m
m = self.basic_m
scipy_mat = bmat([[block, block], [None, block]], format='csr')
dinopy_mat = m.tocsr()
dindices = np.sort(dinopy_mat.indices)
dindptr = np.sort(dinopy_mat.indptr)
ddata = np.sort(dinopy_mat.data)
sindices = np.sort(scipy_mat.indices)
sindptr = np.sort(scipy_mat.indptr)
sdata = np.sort(scipy_mat.data)
self.assertListEqual(dindices.tolist(), sindices.tolist())
self.assertListEqual(dindptr.tolist(), sindptr.tolist())
self.assertListEqual(ddata.tolist(), sdata.tolist())
def test_tocsc(self):
block = self.block_m
m = self.basic_m
scipy_mat = bmat([[block, block], [None, block]], format='csc')
dinopy_mat = m.tocsc()
dindices = np.sort(dinopy_mat.indices)
dindptr = np.sort(dinopy_mat.indptr)
ddata = np.sort(dinopy_mat.data)
sindices = np.sort(scipy_mat.indices)
sindptr = np.sort(scipy_mat.indptr)
sdata = np.sort(scipy_mat.data)
self.assertListEqual(dindices.tolist(), sindices.tolist())
self.assertListEqual(dindptr.tolist(), sindptr.tolist())
self.assertListEqual(ddata.tolist(), sdata.tolist())
def test_multiply(self):
# check scalar multiplication
block = self.block_m
m = self.basic_m * 5.0
scipy_mat = bmat([[block, block], [None, block]], format='coo')
mulscipy_mat = scipy_mat * 5.0
dinopy_mat = m.tocoo()
drow = np.sort(dinopy_mat.row)
dcol = np.sort(dinopy_mat.col)
ddata = np.sort(dinopy_mat.data)
srow = np.sort(mulscipy_mat.row)
scol = np.sort(mulscipy_mat.col)
sdata = np.sort(mulscipy_mat.data)
self.assertListEqual(drow.tolist(), srow.tolist())
self.assertListEqual(dcol.tolist(), scol.tolist())
self.assertListEqual(ddata.tolist(), sdata.tolist())
m = 5.0 * self.basic_m
dinopy_mat = m.tocoo()
drow = np.sort(dinopy_mat.row)
dcol = np.sort(dinopy_mat.col)
ddata = np.sort(dinopy_mat.data)
self.assertListEqual(drow.tolist(), srow.tolist())
self.assertListEqual(dcol.tolist(), scol.tolist())
self.assertListEqual(ddata.tolist(), sdata.tolist())
# check dot product with block vector
block = self.block_m
m = self.basic_m
scipy_mat = bmat([[block, block], [None, block]], format='coo')
x = BlockVector(2)
x[0] = np.ones(block.shape[1], dtype=np.float64)
x[1] = np.ones(block.shape[1], dtype=np.float64)
res_scipy = scipy_mat.dot(x.flatten())
res_dinopy = m * x
res_dinopy_flat = m * x.flatten()
self.assertListEqual(res_dinopy.tolist(), res_scipy.tolist())
self.assertListEqual(res_dinopy_flat.tolist(), res_scipy.tolist())
dense_mat = dinopy_mat.todense()
self.basic_m *= 5.0
self.assertTrue(np.allclose(dense_mat, self.basic_m.todense()))
flat_mat = self.basic_m.tocoo()
result = flat_mat * flat_mat
dense_result = result.toarray()
mat = self.basic_m * self.basic_m.tocoo()
dense_mat = mat.toarray()
self.assertTrue(np.allclose(dense_mat, dense_result))
# not supported block matrix times block matrix for now
#with self.assertRaises(Exception) as context:
# mat = self.basic_m * self.basic_m.tocoo()
def test_getitem(self):
m = BlockMatrix(3, 3)
for i in range(3):
for j in range(3):
self.assertIsNone(m[i, j])
m[0, 1] = self.block_m
self.assertEqual(m[0, 1].shape, self.block_m.shape)
def test_setitem(self):
m = BlockMatrix(2, 2)
m[0, 1] = self.block_m
self.assertFalse(m.is_empty_block(0, 1))
self.assertEqual(m.row_block_sizes()[0], self.block_m.shape[0])
self.assertEqual(m.col_block_sizes()[1], self.block_m.shape[1])
self.assertEqual(m[0, 1].shape, self.block_m.shape)
def test_coo_data(self):
m = self.basic_m.tocoo()
data = self.basic_m.coo_data()
self.assertListEqual(m.data.tolist(), data.tolist())
# ToDo: add tests for block matrices with block matrices in it
# ToDo: add tests for matrices with zeros in the diagonal
# ToDo: add tests for block matrices with coo and csc matrices
def test_nnz(self):
self.assertEqual(self.block_m.nnz*3, self.basic_m.nnz)
def test_block_shapes(self):
shapes = self.basic_m.block_shapes()
for i in range(self.basic_m.bshape[0]):
for j in range(self.basic_m.bshape[1]):
self.assertEqual(shapes[i][j], self.block_m.shape)
def test_dot(self):
A_dense = self.basic_m.todense()
A_block = self.basic_m
x = np.ones(A_dense.shape[1])
block_x = BlockVector(2)
block_x[0] = np.ones(self.block_m.shape[1])
block_x[1] = np.ones(self.block_m.shape[1])
flat_res = A_block.dot(x).flatten()
block_res = A_block.dot(block_x)
self.assertTrue(np.allclose(A_dense.dot(x), flat_res))
self.assertTrue(np.allclose(A_dense.dot(x), block_res.flatten()))
self.assertEqual(block_res.bshape[0], 2)
def test_reset_brow(self):
self.basic_m.reset_brow(0)
for j in range(self.basic_m.bshape[1]):
self.assertIsNone(self.basic_m[0, j])
def test_reset_bcol(self):
self.basic_m.reset_bcol(0)
for j in range(self.basic_m.bshape[0]):
self.assertIsNone(self.basic_m[j, 0])
def test_to_scipy(self):
block = self.block_m
m = self.basic_m
scipy_mat = bmat([[block, block], [None, block]], format='coo')
dinopy_mat = m.tocoo()
drow = np.sort(dinopy_mat.row)
dcol = np.sort(dinopy_mat.col)
ddata = np.sort(dinopy_mat.data)
srow = np.sort(scipy_mat.row)
scol = np.sort(scipy_mat.col)
sdata = np.sort(scipy_mat.data)
self.assertListEqual(drow.tolist(), srow.tolist())
self.assertListEqual(dcol.tolist(), scol.tolist())
self.assertListEqual(ddata.tolist(), sdata.tolist())
def test_has_empty_rows(self):
self.assertFalse(self.basic_m.has_empty_rows())
def test_has_empty_cols(self):
self.assertFalse(self.basic_m.has_empty_cols())
def test_transpose(self):
A_dense = self.basic_m.todense()
A_block = self.basic_m
A_dense_t = A_dense.transpose()
A_block_t = A_block.transpose()
self.assertTrue(np.allclose(A_dense_t, A_block_t.todense()))
A_dense = self.composed_m.todense()
A_block = self.composed_m
A_dense_t = A_dense.transpose()
A_block_t = A_block.transpose()
self.assertTrue(np.allclose(A_dense_t, A_block_t.todense()))
def test_repr(self):
self.assertEqual(len(self.basic_m.__repr__()), 17)
#def test_str(self):
# self.assertEqual(len(self.basic_m.__str__()), 328)
def test_set_item(self):
self.basic_m[1, 0] = None
self.assertIsNone(self.basic_m[1, 0])
self.basic_m[1, 1] = None
self.assertIsNone(self.basic_m[1, 1])
self.assertEqual(self.basic_m._brow_lengths[1], 0)
self.basic_m[1, 1] = self.block_m
self.assertEqual(self.basic_m._brow_lengths[1], self.block_m.shape[1])
def test_add(self):
A_dense = self.basic_m.todense()
A_block = self.basic_m
aa = A_dense + A_dense
mm = A_block + A_block
self.assertTrue(np.allclose(aa, mm.todense()))
mm = A_block.__radd__(A_block)
self.assertTrue(np.allclose(aa, mm.todense()))
def test_sub(self):
A_dense = self.basic_m.todense()
A_block = self.basic_m
aa = A_dense - A_dense
mm = A_block - A_block
self.assertTrue(np.allclose(aa, mm.todense()))
mm = A_block.__rsub__(A_block)
self.assertTrue(np.allclose(aa, mm.todense()))
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()))
