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_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_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 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 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 test_schur_complement(self):
A = BlockMatrix(4, 4)
A.set_block(0, 0, coo_matrix(np.array([[1, 1],
[0, 1]], dtype=np.double)))
A.set_block(1, 1, coo_matrix(np.array([[1, 0],
[0, 1]], dtype=np.double)))
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, 0]], dtype=np.double)))
A.set_block(3, 0, coo_matrix(np.array([[0, -1],
[0, 0]], dtype=np.double)))
A.set_block(3, 1, coo_matrix(np.array([[-1, 0],
[0, -1]], dtype=np.double)))
A.set_block(3, 2, coo_matrix(np.array([[0, 0],
[-1, 0]], dtype=np.double)))
A_upper = A.copy_structure()
A_upper.set_block(0, 3, A.get_block(3, 0).transpose(copy=True))
A_upper.set_block(1, 3, A.get_block(3, 1).transpose(copy=True))
A_upper.set_block(2, 3, A.get_block(3, 2).transpose(copy=True))
rhs = BlockVector(4)
rhs.set_block(0, np.array([1, 0], dtype=np.double))
rhs.set_block(1, np.array([0, 0], dtype=np.double))
rhs.set_block(2, np.array([0, 1], dtype=np.double))
rhs.set_block(3, np.array([1, 1], dtype=np.double))
x1 = np.linalg.solve((A + A_upper).toarray(), rhs.flatten())
sc_solver = parapint.linalg.SchurComplementLinearSolver(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)
inertia1 = sc_solver.get_inertia()
eig = np.linalg.eigvals((A + A_upper).toarray())
pos = np.count_nonzero(eig > 0)
neg = np.count_nonzero(eig < 0)
zero = np.count_nonzero(eig == 0)
inertia2 = (pos, neg, zero)
self.assertTrue(np.allclose(x1, x2.flatten()))
self.assertEqual(inertia1, inertia2)
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 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 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 hessian_lag(self, x, y, out=None, **kwargs):
"""Return the Hessian of the Lagrangian function evaluated at x and y
Parameters
----------
x : array_like
Array with values of primal variables.
y : array_like
Array with values of dual variables.
out : BlockMatrix
Output matrix with the structure of the hessian already defined. Optional
Returns
-------
BlockMatrix
"""
assert x.size == self.nx, 'Dimension mismatch'
assert y.size == self.ng, 'Dimension mismatch'
eval_f_c = kwargs.pop('eval_f_c', True)
if isinstance(x, BlockVector) and isinstance(y, BlockVector):
assert x.nblocks == self.nblocks + 1
assert y.nblocks == 2 * self.nblocks
x_ = x
y_ = y
elif isinstance(x, np.ndarray) and isinstance(y, BlockVector):
assert y.nblocks == 2 * self.nblocks
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
y_ = y
elif isinstance(x, BlockVector) and isinstance(y, np.ndarray):
assert x.nblocks == self.nblocks + 1
x_ = x
block_y = self.create_vector_y()
block_y.copyfrom(y)
y_ = block_y
elif isinstance(x, np.ndarray) and isinstance(y, np.ndarray):
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
block_y = self.create_vector_y()
block_y.copyfrom(y)
y_ = block_y
else:
raise NotImplementedError('Input vector format not recognized')
if out is None:
hess_lag = BlockMatrix(self.nblocks + 1, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
yi = y_.get_block(sid)
hess_lag.set_block(sid, sid,
nlp.hessian_lag(xi, yi, eval_f_c=eval_f_c))
hess_lag[self.nblocks, self.nblocks] = coo_matrix(
(self.nz, self.nz))
return hess_lag
else:
assert isinstance(out, BlockMatrix), \
'out must be a BlockMatrix'
assert out.bshape == (self.nblocks + 1, self.nblocks + 1), \
'Block shape mismatch'
hess_lag = out
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
yi = y_.get_block(sid)
nlp.hessian_lag(xi,
yi,
out=hess_lag.get_block(sid, sid),
eval_f_c=eval_f_c)
Hz = hess_lag[self.nblocks, self.nblocks]
nb = self.nblocks
assert Hz.shape == (self.nz, self.nz), \
'out must have an {}x{} empty matrix in block {}'.format(nb,
nb,
(nb, nb))
assert Hz.nnz == 0, \
'out must have an empty matrix in block {}'.format((nb, nb))
return hess_lag
class TwoStageStochasticNLP(NLP):
"""
Nonlinear program interface for composite NLP that result from
two-stage stochastic programming problems
"""
def __init__(self, nlps, complicating_vars):
"""
Parameters
----------
model: dictionary with scenarios (scenario names to NLPs)
complicating_vars: dictionary with complicated variables
(scenario names to list of variable indices)
"""
if not isinstance(nlps, dict):
raise RuntimeError("Model must be a dictionary")
if not isinstance(complicating_vars, dict):
raise RuntimeError("complicating_vars must be a dictionary")
if len(complicating_vars) != len(nlps):
raise RuntimeError(
"Each scenario must have a list of complicated variables")
# call parent class to set model
super(TwoStageStochasticNLP, self).__init__(None)
# initialize components
self._initialize_nlp_components(nlps, complicating_vars)
def _initialize_nlp_components(self, *args, **kwargs):
nlps = args[0]
complicating_vars = args[1]
aux_counter = 0
n_z = 0
# check inputs
for k, l in complicating_vars.items():
if k not in nlps:
raise RuntimeError("{} not a scenario name".format(k))
if aux_counter == 0:
n_z = len(l)
else:
if len(l) != n_z:
err_msg = "All scenarios must have the same number of complicated variables"
raise RuntimeError(err_msg)
for val in l:
nlp = nlps[k]
if val > nlp.nx:
raise RuntimeError(
"Variable index cannot be greater than number of vars in NLP"
)
aux_counter += 1
# map of scenario name to indices
self._sname_to_sid = dict()
# map of scenario id to scenario name
self._sid_to_sname = list()
# populate containers
ordered_keys = sorted(nlps.keys())
new_dict = OrderedDict()
self._nlps = list()
for k in ordered_keys:
nlp = nlps[k]
if not isinstance(nlp, NLP):
raise RuntimeError("Scenarios must be NLP objects")
self._sname_to_sid[k] = len(self._sid_to_sname)
self._sid_to_sname.append(k)
self._nlps.append(nlp)
# make model a dictionary of original models (PyomoModels or nl-files)
new_dict[k] = nlp.model
self._model = new_dict
# set number of complicated variables
self._nz = n_z
# define map of complicated variables
# this goes from [scnario_id][zid] - > vid
self._zid_to_vid = list()
for sid, sname in enumerate(self._sid_to_sname):
self._zid_to_vid.append(complicating_vars[sname])
# defines vectors
self._create_vectors()
# define sizes
self._nx = self._init_x.size # this includes nz
self._ng = self._init_y.size
self._nc = sum(nlp.nc for nlp in self._nlps) + self.nz * self.nblocks
self._nd = sum(nlp.nd for nlp in self._nlps)
# define structure of jacobians
self._create_jacobian_structures()
# define structure Hessian
self._create_hessian_structure()
# cache coupling matrices
self._AB_csr = self.coupling_matrix()
self._AB_coo = BlockMatrix(self.nblocks + 1, self.nblocks + 1)
nb = self.nblocks
for i in range(nb):
self._AB_coo.set_block(i, i, self._AB_csr.get_block(i, i).tocoo())
self._AB_coo.set_block(nb, nb, self._AB_csr.get_block(nb, nb))
def _make_unmutable_caches(self):
# no need for caches here
pass
def _create_vectors(self):
# Note: This method requires the complicated vars nz to be defined beforehand
# init values and lower and upper bounds
self._init_x = BlockVector(len(self._nlps) + 1)
self._init_y = BlockVector(len(self._nlps) + self.nblocks)
self._lower_x = BlockVector(len(self._nlps) + 1)
self._upper_x = BlockVector(len(self._nlps) + 1)
self._lower_g = BlockVector(len(self._nlps) + self.nblocks)
self._upper_g = BlockVector(len(self._nlps) + self.nblocks)
ndx = 0
for nlp in self._nlps:
self._init_x.set_block(ndx, nlp.x_init())
self._init_y.set_block(ndx, nlp.y_init())
self._lower_x.set_block(ndx, nlp.xl())
self._upper_x.set_block(ndx, nlp.xu())
self._lower_g.set_block(ndx, nlp.gl())
self._upper_g.set_block(ndx, nlp.gu())
ndx += 1
self._init_x.set_block(ndx, np.zeros(self.nz, dtype=np.double))
self._lower_x.set_block(ndx, np.full(self.nz, -np.inf,
dtype=np.double))
self._upper_x.set_block(ndx, np.full(self.nz, np.inf, dtype=np.double))
for i in range(self.nblocks):
self._init_y.set_block(ndx, np.zeros(self.nz, dtype=np.double))
self._lower_g.set_block(ndx, np.zeros(self.nz, dtype=np.double))
self._upper_g.set_block(ndx, np.zeros(self.nz, dtype=np.double))
ndx += 1
# define x maps and masks
self._lower_x_mask = np.isfinite(self._lower_x)
self._lower_x_map = self._lower_x_mask.nonzero()[0]
self._upper_x_mask = np.isfinite(self._upper_x)
self._upper_x_map = self._upper_x_mask.nonzero()[0]
# define gcd maps and masks
bounds_difference = self._upper_g - self._lower_g
abs_bounds_difference = np.absolute(bounds_difference)
tolerance_equalities = 1e-8
self._c_mask = abs_bounds_difference < tolerance_equalities
self._c_map = self._c_mask.nonzero()[0]
self._d_mask = abs_bounds_difference >= tolerance_equalities
self._d_map = self._d_mask.nonzero()[0]
self._lower_g_mask = np.isfinite(
self._lower_g) * self._d_mask + self._c_mask
self._lower_g_map = self._lower_g_mask.nonzero()[0]
self._upper_g_mask = np.isfinite(
self._upper_g) * self._d_mask + self._c_mask
self._upper_g_map = self._upper_g_mask.nonzero()[0]
self._lower_d_mask = pn.isin(self._d_map, self._lower_g_map)
self._upper_d_mask = pn.isin(self._d_map, self._upper_g_map)
# remove empty vectors at the end of lower and upper d
_lower_d_mask = BlockVector(self.nblocks)
for i in range(self.nblocks):
_lower_d_mask.set_block(i, self._lower_d_mask.get_block(i))
self._lower_d_mask = _lower_d_mask
_upper_d_mask = BlockVector(self.nblocks)
for i in range(self.nblocks):
_upper_d_mask.set_block(i, self._upper_d_mask.get_block(i))
self._upper_d_mask = _upper_d_mask
# define lower and upper d maps
self._lower_d_map = pn.where(self._lower_d_mask)[0]
self._upper_d_map = pn.where(self._upper_d_mask)[0]
# get lower and upper d values
self._lower_d = np.compress(self._d_mask, self._lower_g)
self._upper_d = np.compress(self._d_mask, self._upper_g)
# remove empty vectors at the end of lower and upper d
_lower_d = BlockVector(self.nblocks)
_upper_d = BlockVector(self.nblocks)
for i in range(self.nblocks):
_lower_d.set_block(i, self._lower_d.get_block(i))
_upper_d.set_block(i, self._upper_d.get_block(i))
self._lower_d = _lower_d
self._upper_d = _upper_d
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.set_block(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.set_block(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.set_block(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
# ToDo: decide if we cache _irows and _jcols pointers for composite nlp
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
# ToDo: decide if we cache _irows and _jcols pointers for composite nlp
@property
def nblocks(self):
"""
Returns number of blocks (nlps)
"""
return len(self._nlps)
@property
def nz(self):
"""
Return number of complicated variables
"""
return self._nz
def nlps(self):
"""Creates generator scenario name to nlp """
for sid, name in enumerate(self._sid_to_sname):
yield name, self._nlps[sid]
def create_vector_x(self, subset=None):
"""Returns ndarray of primal variables
Parameters
----------
subset : str, optional
determines size of vector.
`l`: only primal variables with lower bounds
`u`: only primal variables with upper bounds
Returns
-------
BlockVector
"""
if subset is None:
subvectors = [np.zeros(nlp.nx, dtype=np.double) for nlp in self._nlps] + \
[np.zeros(self.nz, dtype=np.double)]
return BlockVector(subvectors)
elif subset == 'l':
vectors = list()
for nlp in self._nlps:
nx_l = len(nlp._lower_x_map)
xl = np.zeros(nx_l, dtype=np.double)
vectors.append(xl)
# complicated variables have no lower bounds
vectors.append(np.zeros(0, dtype=np.double))
return BlockVector(vectors)
elif subset == 'u':
vectors = list()
for nlp in self._nlps:
nx_u = len(nlp._upper_x_map)
xu = np.zeros(nx_u, dtype=np.double)
vectors.append(xu)
# complicated variables have no upper bounds
vectors.append(np.zeros(0, dtype=np.double))
return BlockVector(vectors)
else:
raise RuntimeError('Subset not recognized')
def create_vector_y(self, subset=None):
"""Return ndarray of vector of constraints
Parameters
----------
subset : str, optional
determines size of vector.
`c`: only equality constraints
`d`: only inequality constraints
`dl`: only inequality constraints with lower bound
`du`: only inequality constraints with upper bound
Returns
-------
BlockVector
"""
if subset is None:
return BlockVector(
[np.zeros(nlp.ng, dtype=np.double) for nlp in self._nlps] + [
np.zeros(self.nz, dtype=np.double)
for i in range(self.nblocks)
])
elif subset == 'c':
return BlockVector(
[np.zeros(nlp.nc, dtype=np.double) for nlp in self._nlps] + [
np.zeros(self.nz, dtype=np.double)
for i in range(self.nblocks)
])
elif subset == 'd':
return BlockVector(
[np.zeros(nlp.nd, dtype=np.double) for nlp in self._nlps])
elif subset == 'dl' or subset == 'du':
return BlockVector(
[nlp.create_vector_y(subset=subset) for nlp in self._nlps])
else:
raise RuntimeError('Subset not recognized')
def objective(self, x, **kwargs):
"""Returns value of objective function evaluated at x
Parameters
----------
x : array_like
Array with values of primal variables.
Returns
-------
float
"""
if isinstance(x, BlockVector):
return sum(self._nlps[i].objective(x.get_block(i))
for i in range(self.nblocks))
elif isinstance(x, np.ndarray):
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
return sum(self._nlps[i].objective(x_.get_block(i))
for i in range(self.nblocks))
else:
raise NotImplementedError(
"x must be a numpy array or a BlockVector")
def grad_objective(self, x, out=None, **kwargs):
"""Returns gradient of the objective function evaluated at x
Parameters
----------
x : array_like
Array with values of primal variables.
out : array_like
Output array. Its type is preserved and it
must be of the right shape to hold the output.
Returns
-------
array_like
"""
if out is None:
df = self.create_vector_x()
else:
assert isinstance(
out,
BlockVector), 'Composite NLP takes block vector to evaluate g'
assert out.nblocks == self.nblocks + 1
assert out.size == self.nx
df = out
if isinstance(x, BlockVector):
assert x.size == self.nx
assert x.nblocks == self.nblocks + 1
for i in range(self.nblocks):
self._nlps[i].grad_objective(x.get_block(i),
out=df.get_block(i))
return df
elif isinstance(x, np.ndarray):
assert x.size == self.nx
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
for i in range(self.nblocks):
self._nlps[i].grad_objective(x_.get_block(i),
out=df.get_block(i))
return df
else:
raise NotImplementedError(
"x must be a numpy array or a BlockVector")
def evaluate_g(self, x, out=None, **kwargs):
"""Returns general inequality constraints evaluated at x
Parameters
----------
x : array_like
Array with values of primal variables.
out : array_like
Output array. Its type is preserved and it
must be of the right shape to hold the output.
Returns
-------
array_like
"""
if out is None:
res = self.create_vector_y()
else:
assert isinstance(
out,
BlockVector), 'Composite NLP takes block vector to evaluate g'
assert out.nblocks == 2 * self.nblocks
assert out.size == self.ng
res = out
if isinstance(x, BlockVector):
assert x.size == self.nx
assert x.nblocks == self.nblocks + 1
for sid in range(self.nblocks):
# evaluate gi
self._nlps[sid].evaluate_g(x.get_block(sid),
out=res.get_block(sid))
# evaluate coupling Ax-z
A = self._AB_csr.get_block(sid, sid)
res[sid +
self.nblocks] = A * x.get_block(sid) - x[self.nblocks]
return res
elif isinstance(x, np.ndarray):
assert x.size == self.nx
block_x = self.create_vector_x()
block_x.copyfrom(x) # this is expensive
x_ = block_x
for sid in range(self.nblocks):
self._nlps[sid].evaluate_g(x_.get_block(sid),
out=res.get_block(sid))
# evaluate coupling Ax-z
A = self._AB_csr.get_block(sid, sid)
res[sid +
self.nblocks] = A * x_.get_block(sid) - x_[self.nblocks]
return res
else:
raise NotImplementedError(
"x must be a numpy array or a BlockVector")
def evaluate_c(self, x, out=None, **kwargs):
"""Returns the equality constraints evaluated at x
Parameters
----------
x : array_like
Array with values of primal variables.
out : array_like
Output array. Its type is preserved and it
must be of the right shape to hold the output.
Returns
-------
array_like
"""
evaluated_g = kwargs.pop('evaluated_g', None)
if out is None:
res = self.create_vector_y(subset='c')
else:
assert isinstance(
out,
BlockVector), 'Composite NLP takes block vector to evaluate g'
assert out.nblocks == 2 * self.nblocks
assert out.size == self.nc
res = out
if evaluated_g is not None:
assert isinstance(evaluated_g,
BlockVector), 'evaluated_g must be a BlockVector'
assert evaluated_g.nblocks == 2 * self.nblocks
assert evaluated_g.size == self.ng
g = evaluated_g.compress(self._c_mask)
if out is None:
return g
for bid, blk in enumerate(g):
out.set_block(bid, blk)
return out
if isinstance(x, BlockVector):
assert x.size == self.nx
assert x.nblocks == self.nblocks + 1
for sid in range(self.nblocks):
self._nlps[sid].evaluate_c(x.get_block(sid),
out=res.get_block(sid))
A = self._AB_csr.get_block(sid, sid)
res[sid +
self.nblocks] = A * x.get_block(sid) - x[self.nblocks]
return res
elif isinstance(x, np.ndarray):
assert x.size == self.nx
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
for sid in range(self.nblocks):
self._nlps[sid].evaluate_c(x_.get_block(sid),
out=res.get_block(sid))
A = self._AB_csr.get_block(sid, sid)
res[sid +
self.nblocks] = A * x_.get_block(sid) - x_[self.nblocks]
return res
else:
raise NotImplementedError(
'x must be a numpy array or a BlockVector')
def evaluate_d(self, x, out=None, **kwargs):
"""Returns the inequality constraints evaluated at x
Parameters
----------
x : array_like
Array with values of primal variables.
out : array_like
Output array. Its type is preserved and it
must be of the right shape to hold the output.
Returns
-------
array_like
"""
evaluated_g = kwargs.pop('evaluated_g', None)
if out is None:
res = self.create_vector_y(subset='d')
else:
assert isinstance(
out,
BlockVector), 'Composite NLP takes block vector to evaluate g'
assert out.nblocks == self.nblocks
assert out.size == self.nd
res = out
if evaluated_g is not None:
assert isinstance(evaluated_g,
BlockVector), 'evaluated_g must be a BlockVector'
assert evaluated_g.nblocks == 2 * self.nblocks
assert evaluated_g.size == self.ng
d = evaluated_g.compress(self._d_mask)
if out is None:
return BlockVector(
[d.get_block(j) for j in range(self.nblocks)])
for bid in range(self.nblocks):
out.set_block(bid, d.get_block(bid))
return out
if isinstance(x, BlockVector):
assert x.size == self.nx
assert x.nblocks == self.nblocks + 1
for sid in range(self.nblocks):
self._nlps[sid].evaluate_d(x.get_block(sid),
out=res.get_block(sid))
return res
elif isinstance(x, np.ndarray):
assert x.size == self.nx
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
for sid in range(self.nblocks):
self._nlps[sid].evaluate_d(x_.get_block(sid),
out=res.get_block(sid))
return res
else:
raise NotImplementedError(
"x must be a numpy array or a BlockVector")
def jacobian_g(self, x, out=None, **kwargs):
"""Returns the Jacobian of the general inequalities 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_g = BlockMatrix(2 * self.nblocks, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
jac_g.set_block(sid, sid, nlp.jacobian_g(xi))
# coupling matrices Ai
jac_g[sid + self.nblocks,
sid] = self._AB_coo.get_block(sid, sid)
# coupling matrices Bi
jac_g[sid + self.nblocks, self.nblocks] = -identity(self.nz)
return jac_g
else:
assert isinstance(out, BlockMatrix), 'out must be a BlockMatrix'
assert out.bshape == (2 * self.nblocks,
self.nblocks + 1), "Block shape mismatch"
jac_g = out
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
nlp.jacobian_g(xi, out=jac_g.get_block(sid, sid))
Ai = jac_g[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_g[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_g
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 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 hessian_lag(self, x, y, out=None, **kwargs):
"""Return the Hessian of the Lagrangian function evaluated at x and y
Parameters
----------
x : array_like
Array with values of primal variables.
y : array_like
Array with values of dual variables.
out : BlockMatrix
Output matrix with the structure of the hessian already defined. Optional
Returns
-------
BlockMatrix
"""
assert x.size == self.nx, 'Dimension mismatch'
assert y.size == self.ng, 'Dimension mismatch'
eval_f_c = kwargs.pop('eval_f_c', True)
if isinstance(x, BlockVector) and isinstance(y, BlockVector):
assert x.nblocks == self.nblocks + 1
assert y.nblocks == 2 * self.nblocks
x_ = x
y_ = y
elif isinstance(x, np.ndarray) and isinstance(y, BlockVector):
assert y.nblocks == 2 * self.nblocks
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
y_ = y
elif isinstance(x, BlockVector) and isinstance(y, np.ndarray):
assert x.nblocks == self.nblocks + 1
x_ = x
block_y = self.create_vector_y()
block_y.copyfrom(y)
y_ = block_y
elif isinstance(x, np.ndarray) and isinstance(y, np.ndarray):
block_x = self.create_vector_x()
block_x.copyfrom(x)
x_ = block_x
block_y = self.create_vector_y()
block_y.copyfrom(y)
y_ = block_y
else:
raise NotImplementedError('Input vector format not recognized')
if out is None:
hess_lag = BlockMatrix(self.nblocks + 1, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
yi = y_.get_block(sid)
hess_lag.set_block(sid, sid,
nlp.hessian_lag(xi, yi, eval_f_c=eval_f_c))
hess_lag[self.nblocks, self.nblocks] = coo_matrix(
(self.nz, self.nz))
return hess_lag
else:
assert isinstance(out, BlockMatrix), \
'out must be a BlockMatrix'
assert out.bshape == (self.nblocks + 1, self.nblocks + 1), \
'Block shape mismatch'
hess_lag = out
for sid, nlp in enumerate(self._nlps):
xi = x_.get_block(sid)
yi = y_.get_block(sid)
nlp.hessian_lag(xi,
yi,
out=hess_lag.get_block(sid, sid),
eval_f_c=eval_f_c)
Hz = hess_lag[self.nblocks, self.nblocks]
nb = self.nblocks
assert Hz.shape == (self.nz, self.nz), \
'out must have an {}x{} empty matrix in block {}'.format(nb,
nb,
(nb, nb))
assert Hz.nnz == 0, \
'out must have an empty matrix in block {}'.format((nb, nb))
return hess_lag
def block_id(self, scneario_name):
"""
Returns idx of corresponding nlp for scenario_name
Parameters
----------
scenario_name : str
name of scenario
Returns
-------
int
"""
return self._sname_to_sid[scneario_name]
def block_name(self, bid):
"""
Returns scenario name for given bid index
Parameters
----------
bid : int
index of a given scenario
Returns
-------
int
"""
return self._sid_to_sname[bid]
def get_block(self, scneario_name):
"""
Returns nlp corresponding to scenario_name
Parameters
----------
scenario_name : str
name of scenario
Returns
-------
NLP
"""
bid = self._sname_to_sid[scneario_name]
return self._nlps[bid]
def complicated_vars_ids(self, scenario_name):
return self._zid_to_vid[scenario_name]
# ToDo: order of variables?
# ToDo: order of constraints?
def expansion_matrix_xl(self):
Pxl = BlockMatrix(self.nblocks + 1, self.nblocks + 1)
for sid, nlp in enumerate(self._nlps):
Pxl.set_block(sid, sid, nlp.expansion_matrix_xl())
Pxl[self.nblocks, self.nblocks] = coo_matrix((self.nz, 0))
return Pxl
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 expansion_matrix_dl(self):
Pdl = BlockMatrix(self.nblocks, self.nblocks)
for sid, nlp in enumerate(self._nlps):
Pdl.set_block(sid, sid, nlp.expansion_matrix_dl())
return Pdl
def expansion_matrix_du(self):
Pdu = BlockMatrix(self.nblocks, self.nblocks)
for sid, nlp in enumerate(self._nlps):
Pdu.set_block(sid, sid, nlp.expansion_matrix_du())
return Pdu
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 scenarios_order(self):
return [self._sid_to_sname[i] for i in range(self.nblocks)]
