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 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()))