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 _setup_kkt_and_rhs_structure(self): # First setup the diagonal blocks for ndx, nlp in self._nlps.items(): sub_kkt = BlockMatrix(nbrows=2, nbcols=2) n = nlp.n_primals() + nlp.n_eq_constraints( ) + 2 * nlp.n_ineq_constraints() sub_kkt.set_row_size(0, n) sub_kkt.set_col_size(0, n) if ndx == 0: sub_kkt.set_row_size(1, 0) sub_kkt.set_col_size(1, 0) else: sub_kkt.set_row_size(1, self._num_states) sub_kkt.set_col_size(1, self._num_states) row_1 = BlockMatrix(nbrows=1, nbcols=4) if ndx == 0: row_1.set_row_size(0, 0) else: row_1.set_row_size(0, self._num_states) row_1.set_col_size(0, nlp.n_primals()) row_1.set_col_size(1, nlp.n_ineq_constraints()) row_1.set_col_size(2, nlp.n_eq_constraints()) row_1.set_col_size(3, nlp.n_ineq_constraints()) row_1.set_block(0, 0, self._link_backward_matrices[ndx]) sub_kkt.set_block(1, 0, row_1) sub_kkt.set_block(0, 1, row_1.transpose()) self._kkt.set_block(ndx, ndx, sub_kkt) sub_rhs = BlockVector(2) sub_rhs.set_block(0, np.zeros(n)) if ndx == 0: sub_rhs.set_block(1, np.zeros(0)) else: sub_rhs.set_block(1, np.zeros(self._num_states)) self._rhs.set_block(ndx, sub_rhs) # Setup the border blocks for ndx, nlp in self._nlps.items(): nlp = self._nlps[ndx] block = BlockMatrix(nbrows=2, nbcols=2) sub_block = BlockMatrix(nbrows=self._num_time_blocks, nbcols=4) sub_block.set_col_size(0, nlp.n_primals()) sub_block.set_col_size(1, nlp.n_ineq_constraints()) sub_block.set_col_size(2, nlp.n_eq_constraints()) sub_block.set_col_size(3, nlp.n_ineq_constraints()) for sub_ndx in range(self._num_time_blocks): if sub_ndx == self._num_time_blocks - 1: sub_block.set_row_size(sub_ndx, 0) else: sub_block.set_row_size(sub_ndx, self._num_states) sub_block.set_block(ndx, 0, self._link_forward_matrices[ndx]) block.set_block(0, 0, sub_block) block.set_block( 1, 1, -self._link_backward_coupling_matrices[ndx].transpose()) self._kkt.set_block(self._num_time_blocks, ndx, block) self._kkt.set_block(ndx, self._num_time_blocks, block.transpose()) # Setup the bottom right block block = BlockMatrix(2, 2) rhs_block = BlockVector(2) sub_block = BlockMatrix(1, self._num_time_blocks) sub_rhs_block = BlockVector(self._num_time_blocks) for ndx in range(self._num_time_blocks): sub_block.set_block( 0, ndx, -self._link_forward_coupling_matrices[ndx].transpose()) if ndx == self._num_time_blocks - 1: sub_rhs_block.set_block(ndx, np.zeros(0)) else: sub_rhs_block.set_block(ndx, np.zeros(self._num_states)) rhs_block.set_block(0, sub_rhs_block) rhs_block.set_block(1, np.zeros(self._total_num_coupling_vars)) block.set_block(1, 0, sub_block) block.set_block(0, 1, sub_block.transpose()) self._kkt.set_block(self._num_time_blocks, self._num_time_blocks, block) self._rhs.set_block(self._num_time_blocks, rhs_block)