def test_multiply(self):
# test scalar multiplication
m = self.basic_m * 5.0
dense_m = m.todense()
b00 = self.block00.tofullcoo()
b11 = self.block11.tofullcoo()
b10 = self.block10
scipy_m = bmat([[b00, b10.transpose()], [b10, b11]], format='coo')
dense_scipy_m = scipy_m.todense() * 5.0
self.assertTrue(np.allclose(dense_scipy_m, dense_m, atol=1e-3))
m = 5.0 * self.basic_m
dense_m = m.todense()
self.assertTrue(np.allclose(dense_scipy_m, dense_m, atol=1e-3))
# test matrix vector product
m = self.basic_m
x = BlockVector(m.bshape[1])
for i in range(m.bshape[1]):
x[i] = np.ones(m.col_block_sizes()[i], dtype=np.float64)
dinopy_res = m * x
scipy_res = scipy_m * x.flatten()
self.assertListEqual(dinopy_res.tolist(), scipy_res.tolist())
dinopy_res = m * x.flatten()
scipy_res = scipy_m * x.flatten()
self.assertListEqual(dinopy_res.tolist(), scipy_res.tolist())
def test_create_vector_x(self):
x_ = BlockVector(self.n_scenarios + 1)
nz = len(self.complicated_vars_ids)
nx_i = (self.G.shape[0] + nz)
for i in range(self.n_scenarios):
x_[i] = np.zeros(nx_i)
x_[self.n_scenarios] = np.zeros(nz)
self.assertEqual(x_.shape, self.nlp.create_vector_x().shape)
self.assertEqual(x_.nblocks, self.nlp.create_vector_x().nblocks)
self.assertTrue(
np.allclose(x_.block_sizes(),
self.nlp.create_vector_x().block_sizes()))
self.assertListEqual(list(x_.flatten()),
list(self.nlp.create_vector_x().flatten()))
# check for subset
for s in ['l', 'u']:
xs = self.nlp.create_vector_x(subset=s)
xs_ = BlockVector(self.n_scenarios + 1)
for i in range(self.n_scenarios):
xs_[i] = np.zeros(1)
xs_[self.n_scenarios] = np.zeros(0)
self.assertEqual(xs_.shape, xs.shape)
self.assertEqual(xs_.nblocks, xs.nblocks)
self.assertTrue(np.allclose(xs_.block_sizes(), xs.block_sizes()))
self.assertListEqual(list(xs_.flatten()), list(xs.flatten()))
def test_multiply(self):
# test scalar multiplication
m = self.basic_m * 5.0
dense_m = m.todense()
b00 = self.block00.tocoo()
b11 = self.block11.tocoo()
b10 = self.block10
scipy_m = bmat([[b00, b10.transpose()], [b10, b11]], format='coo')
dense_scipy_m = scipy_m.todense() * 5.0
self.assertTrue(np.allclose(dense_scipy_m, dense_m, atol=1e-3))
m = 5.0 * self.basic_m
dense_m = m.todense()
self.assertTrue(np.allclose(dense_scipy_m, dense_m, atol=1e-3))
# test matrix vector product
m = self.basic_m
x = BlockVector(m.bshape[1])
for i in range(m.bshape[1]):
x[i] = np.ones(m.col_block_sizes()[i], dtype=np.float64)
dinopy_res = m * x
scipy_res = scipy_m * x.flatten()
self.assertListEqual(dinopy_res.tolist(), scipy_res.tolist())
dinopy_res = m * x.flatten()
scipy_res = scipy_m * x.flatten()
self.assertListEqual(dinopy_res.tolist(), scipy_res.tolist())
self.basic_m *= 5.0
self.assertTrue(np.allclose(self.basic_m.todense(), dense_m, atol=1e-3))
def test_create_vector_y(self):
nz = len(self.complicated_vars_ids)
ng_i = self.A.shape[0] + nz
y_ = BlockVector(2 * self.n_scenarios)
for i in range(self.n_scenarios):
y_[i] = np.zeros(ng_i)
y_[self.n_scenarios + i] = np.zeros(nz)
y = self.nlp.create_vector_y()
self.assertEqual(y_.shape, y.shape)
self.assertEqual(y_.nblocks, y.nblocks)
self.assertTrue(np.allclose(y_.block_sizes(), y.block_sizes()))
self.assertListEqual(list(y_.flatten()), list(y.flatten()))
# check for equalities
ys_ = BlockVector(2 * self.n_scenarios)
for i in range(self.n_scenarios):
ys_[i] = np.zeros(ng_i)
ys_[self.n_scenarios + i] = np.zeros(nz)
ys = self.nlp.create_vector_y(subset='c')
self.assertEqual(ys_.shape, ys.shape)
self.assertEqual(ys_.nblocks, ys.nblocks)
self.assertTrue(np.allclose(ys_.block_sizes(), ys.block_sizes()))
self.assertListEqual(list(ys_.flatten()), list(ys.flatten()))
# check for inequalities
ys_ = BlockVector(self.n_scenarios)
for i in range(self.n_scenarios):
ys_[i] = np.zeros(0)
ys = self.nlp.create_vector_y(subset='d')
self.assertEqual(ys_.shape, ys.shape)
self.assertEqual(ys_.nblocks, ys.nblocks)
self.assertTrue(np.allclose(ys_.block_sizes(), ys.block_sizes()))
self.assertListEqual(list(ys_.flatten()), list(ys.flatten()))
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))
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))
def compute_init_lam(nlp, x=None, lam_max=1e3):
if x is None:
x = nlp.x_init
else:
assert x.size == nlp.nx
assert nlp.nd == 0, "only supported for equality constrained nlps for now"
nx = nlp.nx
nc = nlp.nc
# create Jacobian
jac_c = nlp.jacobian_g(x)
# create gradient of objective
df = nlp.grad_objective(x)
# create KKT system
kkt = BlockSymMatrix(2)
kkt[0, 0] = identity(nx)
kkt[1, 0] = jac_c
zeros = np.zeros(nc)
rhs = BlockVector([-df, zeros])
flat_kkt = kkt.tocoo().tocsc()
flat_rhs = rhs.flatten()
sol = spsolve(flat_kkt, flat_rhs)
return sol[nlp.nx:nlp.nx + nlp.ng]
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 = BlockSymMatrix(2)
kkt[0, 0] = identity(nx)
kkt[1, 0] = jac
zeros = np.zeros(nc)
rhs = BlockVector([-df, 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 test_get_block_vector_for_dot_product_5(self):
rank = comm.Get_rank()
rank_ownership = np.array([[1, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]])
m = MPIBlockMatrix(4, 3, rank_ownership, comm)
sub_m = np.array([[1, 0], [0, 1]])
sub_m = coo_matrix(sub_m)
if rank == 0:
m.set_block(3, rank, sub_m.copy())
elif rank == 1:
m.set_block(0, 0, sub_m.copy())
m.set_block(rank, rank, sub_m.copy())
m.set_block(3, rank, sub_m.copy())
else:
m.set_block(rank, rank, sub_m.copy())
m.set_block(3, rank, sub_m.copy())
v = BlockVector(3)
sub_v = np.ones(2)
for ndx in range(3):
v.set_block(ndx, sub_v.copy())
res = m._get_block_vector_for_dot_product(v)
self.assertIs(res, v)
v_flat = v.flatten()
res = m._get_block_vector_for_dot_product(v_flat)
self.assertIsInstance(res, BlockVector)
for ndx in range(3):
block = res.get_block(ndx)
self.assertTrue(np.array_equal(block, sub_v))
def test_flatten(self):
v = BlockVector(2)
a = np.arange(5)
b = np.arange(9)
c = np.concatenate([a, b])
v[0] = a
v[1] = b
self.assertListEqual(v.flatten().tolist(), c.tolist())
def test_xu(self):
xu = BlockVector(self.n_scenarios + 1)
nz = len(self.complicated_vars_ids)
nx_i = (self.G.shape[0] + nz)
for i in range(self.n_scenarios):
xu[i] = np.array([np.inf] * nx_i)
xu[i][0] = 100.0
xu[self.n_scenarios] = np.array([np.inf] * nz)
self.assertIsInstance(self.nlp.xu(), BlockVector)
xu_flat = xu.flatten()
self.assertEqual(xu.nblocks, self.nlp.xu().nblocks)
self.assertTrue(
np.allclose(xu.block_sizes(),
self.nlp.xu().block_sizes()))
self.assertListEqual(list(xu_flat), list(self.nlp.xu().flatten()))
def test_x_init(self):
x_init = BlockVector(self.n_scenarios + 1)
nz = len(self.complicated_vars_ids)
nx_i = (self.G.shape[0] + nz)
for i in range(self.n_scenarios):
x_init[i] = np.zeros(nx_i)
x_init[i][0] = 1.0
x_init[self.n_scenarios] = np.zeros(nz)
self.assertIsInstance(self.nlp.x_init(), BlockVector)
x_init_flat = x_init.flatten()
self.assertEqual(x_init.nblocks, self.nlp.x_init().nblocks)
self.assertTrue(
np.allclose(x_init.block_sizes(),
self.nlp.x_init().block_sizes()))
self.assertListEqual(list(x_init_flat),
list(self.nlp.x_init().flatten()))
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 compute_init_lam(nlp, x=None, lam_max=1e3):
if x is None:
x = nlp.x_init
else:
assert x.size == nlp.nx
assert nlp.nd == 0, "only supported for equality constrained nlps for now"
nx = nlp.nx
nc = nlp.nc
# create Jacobian
jac_c = nlp.jacobian_g(x)
# create gradient of objective
df = nlp.grad_objective(x)
diag_ones = np.ones(nx)
irows = np.arange(nx)
jcols = np.arange(nx)
eye = COOSymMatrix((diag_ones, (irows, jcols)), shape=(nx, nx))
# create KKT system
kkt = BlockSymMatrix(2)
kkt[0, 0] = eye
kkt[1, 0] = jac_c
zeros = np.zeros(nc)
rhs = BlockVector([-df, zeros])
flat_kkt = kkt.tofullmatrix().tocsc()
flat_rhs = rhs.flatten()
sol = spsolve(flat_kkt, flat_rhs)
return sol[nlp.nx: nlp.nx + nlp.ng]
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 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)
