def test_isin(self):
bv = self.bv
test_bv = BlockVector(2)
a = np.array([1.1, 3.3])
b = np.array([5.5, 7.7])
test_bv.set_block(0, a)
test_bv.set_block(1, b)
res = pn.isin(bv, test_bv)
for bid, blk in enumerate(bv):
self.assertEqual(blk.size, res.get_block(bid).size)
res_flat = np.isin(blk, test_bv.get_block(bid))
self.assertTrue(np.allclose(res.get_block(bid), res_flat))
c = np.concatenate([a, b])
res = pn.isin(bv, c)
for bid, blk in enumerate(bv):
self.assertEqual(blk.size, res.get_block(bid).size)
res_flat = np.isin(blk, c)
self.assertTrue(np.allclose(res.get_block(bid), res_flat))
res = pn.isin(bv, test_bv, invert=True)
for bid, blk in enumerate(bv):
self.assertEqual(blk.size, res.get_block(bid).size)
res_flat = np.isin(blk, test_bv.get_block(bid), invert=True)
self.assertTrue(np.allclose(res.get_block(bid), res_flat))
c = np.concatenate([a, b])
res = pn.isin(bv, c, invert=True)
for bid, blk in enumerate(bv):
self.assertEqual(blk.size, res.get_block(bid).size)
res_flat = np.isin(blk, c, invert=True)
self.assertTrue(np.allclose(res.get_block(bid), res_flat))
def set_primal_dual_kkt_solution(self, sol: BlockVector):
for ndx, nlp in self._nlps.items():
nlp.set_primal_dual_kkt_solution(sol.get_block(ndx).get_block(0))
self._delta_primals.set_block(ndx, nlp.get_delta_primals())
self._delta_slacks.set_block(ndx, nlp.get_delta_slacks())
self._delta_duals_eq.get_block(ndx).set_block(
0, nlp.get_delta_duals_eq())
self._delta_duals_ineq.set_block(ndx, nlp.get_delta_duals_ineq())
self._delta_duals_primals_lb.set_block(
ndx, nlp.get_delta_duals_primals_lb())
self._delta_duals_primals_ub.set_block(
ndx, nlp.get_delta_duals_primals_ub())
self._delta_duals_slacks_lb.set_block(
ndx, nlp.get_delta_duals_slacks_lb())
self._delta_duals_slacks_ub.set_block(
ndx, nlp.get_delta_duals_slacks_ub())
self._delta_duals_eq.get_block(ndx).set_block(
1,
sol.get_block(ndx).get_block(1))
self._delta_duals_eq.get_block(ndx).set_block(
2,
sol.get_block(
self._num_time_blocks).get_block(0).get_block(ndx))
self._delta_primals.set_block(
self._num_time_blocks,
sol.get_block(self._num_time_blocks).get_block(1))
def set_duals_ineq(self, duals_ineq: BlockVector):
"""
Parameters
----------
duals_ineq: BlockVector
The values for the duals of the inequality constraints. This BlockVector has one block for
every time block.
"""
for ndx, nlp in self._nlps.items():
nlp.set_duals_ineq(duals_ineq.get_block(ndx))
self._duals_ineq.set_block(ndx, duals_ineq.get_block(ndx))
def set_primals(self, primals: BlockVector):
"""
Set the values of the primal variables for evaluation (i.e., the evaluate_* methods).
Parameters
----------
primals: BlockVector
The values for each primal variable. This BlockVector should have one block for every time block
and one block for the coupling variables.
"""
for ndx, nlp in self._nlps.items():
nlp.set_primals(primals.get_block(ndx))
self._primals.set_block(ndx, primals.get_block(ndx))
self._primals.set_block(self._num_time_blocks,
primals.get_block(self._num_time_blocks))
def set_duals_eq(self, duals_eq: BlockVector):
"""
Parameters
----------
duals_eq: BlockVector
The values for the duals of the equality constraints, including the coupling constraints.
This BlockVector has one block for every time block. Each block is itself a BlockVector with
3 blocks. The first block contains the duals of the equality constraints in the corresponding time
block. The second block has the duals for the coupling constraints linking the states at the
beginning of the time block to the coupling variables between the time block and the previous
time block. The third block has the duals for the coupling constraints linking the states at the
end of the time block to the coupling variables between the time block and the next time block.
"""
for ndx, nlp in self._nlps.items():
sub_block = duals_eq.get_block(ndx)
nlp.set_duals_eq(sub_block.get_block(0))
self._duals_eq.get_block(ndx).set_block(0, sub_block.get_block(0))
self._duals_eq.get_block(ndx).set_block(1, sub_block.get_block(1))
self._duals_eq.get_block(ndx).set_block(2, sub_block.get_block(2))
def set_duals_slacks_ub(self, duals: BlockVector):
for ndx, nlp in self._nlps.items():
nlp.set_duals_slacks_ub(duals.get_block(ndx))
self._duals_slacks_ub.set_block(ndx, duals.get_block(ndx))
def set_slacks(self, slacks: BlockVector):
for ndx, nlp in self._nlps.items():
nlp.set_slacks(slacks.get_block(ndx))
self._slacks.set_block(ndx, slacks.get_block(ndx))
def setUpClass(cls) -> None:
cls.t0 = 0
cls.delta_t = 1
cls.num_finite_elements = 6
cls.constant_control_duration = 2
cls.time_scale = 1
cls.num_time_blocks = 3
cls.with_bounds = True
cls.with_ineq = False
cls.barrier_parameter = 0.1
cls.interface = Problem(t0=cls.t0,
delta_t=cls.delta_t,
num_finite_elements=cls.num_finite_elements,
constant_control_duration=cls.constant_control_duration,
time_scale=cls.time_scale,
num_time_blocks=cls.num_time_blocks,
with_bounds=cls.with_bounds,
with_ineq=cls.with_ineq)
interface = cls.interface
num_time_blocks = cls.num_time_blocks
primals = BlockVector(num_time_blocks + 1)
duals_eq = BlockVector(num_time_blocks)
duals_ineq = BlockVector(num_time_blocks)
duals_primals_lb = BlockVector(num_time_blocks + 1)
duals_primals_ub = BlockVector(num_time_blocks + 1)
duals_slacks_lb = BlockVector(num_time_blocks)
duals_slacks_ub = BlockVector(num_time_blocks)
ownership_map = _get_ownership_map(num_time_blocks, size)
val_map = pe.ComponentMap()
if ownership_map[0] == rank:
m = interface.pyomo_model(0)
val_map[m.x[0]] = 0
val_map[m.x[1]] = 1
val_map[m.x[2]] = 2
val_map[m.p[0]] = 0.5
val_map[m.cons[1]] = 1
val_map[m.cons[2]] = 2
if ownership_map[1] == rank:
m = interface.pyomo_model(1)
val_map[m.x[2]] = 2
val_map[m.x[3]] = 3
val_map[m.x[4]] = 4
val_map[m.p[2]] = 1
val_map[m.cons[3]] = 3
val_map[m.cons[4]] = 4
if ownership_map[2] == rank:
m = interface.pyomo_model(2)
val_map[m.x[4]] = 4
val_map[m.x[5]] = 5
val_map[m.x[6]] = 6
val_map[m.p[4]] = 1.5
val_map[m.cons[5]] = 5
val_map[m.cons[6]] = 6
for ndx in range(num_time_blocks):
primals.set_block(ndx, np.zeros(4, dtype=np.double))
duals_primals_lb.set_block(ndx, np.zeros(4, dtype=np.double))
duals_primals_ub.set_block(ndx, np.zeros(4, dtype=np.double))
duals_ineq.set_block(ndx, np.zeros(0, dtype=np.double))
duals_slacks_lb.set_block(ndx, np.zeros(0, dtype=np.double))
duals_slacks_ub.set_block(ndx, np.zeros(0, dtype=np.double))
sub_duals_eq = BlockVector(3)
sub_duals_eq.set_block(0, np.zeros(2, dtype=np.double))
if ndx == 0:
sub_duals_eq.set_block(1, np.zeros(0, dtype=np.double))
else:
sub_duals_eq.set_block(1, np.zeros(1, dtype=np.double))
if ndx == num_time_blocks - 1:
sub_duals_eq.set_block(2, np.zeros(0, dtype=np.double))
else:
sub_duals_eq.set_block(2, np.zeros(1, dtype=np.double))
duals_eq.set_block(ndx, sub_duals_eq)
primals.set_block(num_time_blocks, np.zeros(2, dtype=np.double))
duals_primals_lb.set_block(num_time_blocks, np.zeros(2, dtype=np.double))
duals_primals_ub.set_block(num_time_blocks, np.zeros(2, dtype=np.double))
local_block_indices = _distribute_blocks(num_time_blocks, rank, size)
for ndx in local_block_indices:
primals.set_block(ndx, np.array([val_map[i] for i in interface.get_pyomo_variables(ndx)], dtype=np.double))
duals_primals_ub.set_block(ndx, np.array([0, 0, 0, ndx], dtype=np.double))
sub_duals_eq = duals_eq.get_block(ndx)
sub_duals_eq.set_block(0, np.array([val_map[i] for i in interface.get_pyomo_constraints(ndx)], dtype=np.double))
if ndx == 0:
sub_duals_eq.set_block(1, np.zeros(0, dtype=np.double))
else:
sub_duals_eq.set_block(1, np.ones(1, dtype=np.double) * ndx)
if ndx == num_time_blocks - 1:
sub_duals_eq.set_block(2, np.zeros(0, dtype=np.double))
else:
sub_duals_eq.set_block(2, np.ones(1, dtype=np.double) * ndx)
primals_flat = primals.flatten()
res = np.zeros(primals_flat.size, dtype=np.double)
comm.Allreduce(primals_flat, res)
primals.copyfrom(res)
duals_primals_lb_flat = duals_primals_lb.flatten()
res = np.zeros(duals_primals_lb_flat.size, dtype=np.double)
comm.Allreduce(duals_primals_lb_flat, res)
duals_primals_lb.copyfrom(res)
duals_primals_ub_flat = duals_primals_ub.flatten()
res = np.zeros(duals_primals_ub_flat.size, dtype=np.double)
comm.Allreduce(duals_primals_ub_flat, res)
duals_primals_ub.copyfrom(res)
duals_eq_flat = duals_eq.flatten()
res = np.zeros(duals_eq_flat.size, dtype=np.double)
comm.Allreduce(duals_eq_flat, res)
duals_eq.copyfrom(res)
primals.set_block(num_time_blocks, np.array([3, 6], dtype=np.double))
duals_primals_lb.set_block(num_time_blocks, np.zeros(2, dtype=np.double))
duals_primals_ub.set_block(num_time_blocks, np.zeros(2, dtype=np.double))
interface.set_primals(primals)
interface.set_duals_eq(duals_eq)
interface.set_duals_ineq(duals_ineq)
interface.set_duals_slacks_lb(duals_slacks_lb)
interface.set_duals_slacks_ub(duals_slacks_ub)
interface.set_duals_primals_lb(duals_primals_lb)
interface.set_duals_primals_ub(duals_primals_ub)
interface.set_barrier_parameter(cls.barrier_parameter)
def sqp(nlp: NLP,
linear_solver: LinearSolverInterface,
max_iter=100,
tol=1e-8,
output=True):
"""
An example of a simple SQP algoritm for
equality-constrained NLPs.
Parameters
----------
nlp: NLP
A PyNumero NLP
max_iter: int
The maximum number of iterations
tol: float
The convergence tolerance
"""
t0 = time.time()
# setup KKT matrix
kkt = BlockMatrix(2, 2)
rhs = BlockVector(2)
# create and initialize the iteration vector
z = BlockVector(2)
z.set_block(0, nlp.get_primals())
z.set_block(1, nlp.get_duals())
if output:
print(
f"{'Iter':<12}{'Objective':<12}{'Primal Infeasibility':<25}{'Dual Infeasibility':<25}{'Elapsed Time':<15}"
)
# main iteration loop
for _iter in range(max_iter):
nlp.set_primals(z.get_block(0))
nlp.set_duals(z.get_block(1))
grad_lag = (nlp.evaluate_grad_objective() +
nlp.evaluate_jacobian_eq().transpose() * z.get_block(1))
residuals = nlp.evaluate_eq_constraints()
if output:
print(
f"{_iter:<12}{nlp.evaluate_objective():<12.2e}{np.abs(residuals).max():<25.2e}{np.abs(grad_lag).max():<25.2e}{time.time()-t0:<15.2e}"
)
if (np.abs(grad_lag).max() <= tol and np.abs(residuals).max() <= tol):
break
kkt.set_block(0, 0, nlp.evaluate_hessian_lag())
kkt.set_block(1, 0, nlp.evaluate_jacobian_eq())
kkt.set_block(0, 1, nlp.evaluate_jacobian_eq().transpose())
rhs.set_block(0, grad_lag)
rhs.set_block(1, residuals)
delta, res = linear_solver.solve(kkt, -rhs)
assert res.status == LinearSolverStatus.successful
z += delta
