1#  ___________________________________________________________________________
2#
3#  Pyomo: Python Optimization Modeling Objects
4#  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
5#  Under the terms of Contract DE-NA0003525 with National Technology and
6#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
7#  rights in this software.
8#  This software is distributed under the 3-clause BSD License.
9#  ___________________________________________________________________________
10
11import itertools
12from pyomo.environ import SolverFactory
13from pyomo.core.base.var import Var
14from pyomo.core.base.constraint import Constraint
15from pyomo.core.base.objective import Objective
16from pyomo.core.expr.visitor import identify_variables
17from pyomo.common.collections import ComponentSet
18from pyomo.util.subsystems import (
19        create_subsystem_block,
20        TemporarySubsystemManager,
21        )
22from pyomo.contrib.pynumero.interfaces.pyomo_nlp import PyomoNLP
23from pyomo.contrib.pynumero.interfaces.external_grey_box import (
24        ExternalGreyBoxModel,
25        )
26import numpy as np
27import scipy.sparse as sps
28
29
30def _dense_to_full_sparse(matrix):
31    """
32    Used to convert a dense matrix (2d NumPy array) to SciPy sparse matrix
33    with explicit coordinates for every entry, including zeros. This is
34    used because _ExternalGreyBoxAsNLP methods rely on receiving sparse
35    matrices where sparsity structure does not change between calls.
36    This is difficult to achieve for matrices obtained via the implicit
37    function theorem unless an entry is returned for every coordinate
38    of the matrix.
39
40    Note that this does not mean that the Hessian of the entire NLP will
41    be dense, only that the block corresponding to this external model
42    will be dense.
43    """
44    # TODO: Allow methods to hard-code Jacobian/Hessian sparsity structure
45    # in the case it is known a priori.
46    # TODO: Decompose matrices to infer maximum fill-in sparsity structure.
47    nrow, ncol = matrix.shape
48    row = []
49    col = []
50    data = []
51    for i, j in itertools.product(range(nrow), range(ncol)):
52        row.append(i)
53        col.append(j)
54        data.append(matrix[i,j])
55    row = np.array(row)
56    col = np.array(col)
57    data = np.array(data)
58    return sps.coo_matrix((data, (row, col)), shape=(nrow, ncol))
59
60
61def get_hessian_of_constraint(constraint, wrt1=None, wrt2=None, nlp=None):
62    constraints = [constraint]
63    if wrt1 is None and wrt2 is None:
64        variables = list(identify_variables(constraint.expr, include_fixed=False))
65        wrt1 = variables
66        wrt2 = variables
67    elif wrt1 is not None and wrt2 is not None:
68        variables = wrt1 + wrt2
69    elif wrt1 is not None: # but wrt2 is None
70        wrt2 = wrt1
71        variables = wrt1
72    else:
73        # wrt2 is not None and wrt1 is None
74        wrt1 = wrt2
75        variables = wrt1
76
77    if nlp is None:
78        block = create_subsystem_block(constraints, variables=variables)
79        # Could fix input_vars so I don't evaluate the Hessian with respect
80        # to variables I don't care about...
81
82        # HUGE HACK: Variables not included in a constraint are not written
83        # to the nl file, so we cannot take the derivative with respect to
84        # them, even though we know this derivative is zero. To work around,
85        # we make sure all variables appear on the block in the form of a
86        # dummy constraint. Then we can take derivatives of any constraint
87        # with respect to them. Conveniently, the extract_submatrix_
88        # call deals with extracting the variables and constraint we care
89        # about, in the proper order.
90        block._dummy_var = Var()
91        block._dummy_con = Constraint(expr=sum(variables) == block._dummy_var)
92        block._obj = Objective(expr=0.0)
93        nlp = PyomoNLP(block)
94
95    saved_duals = nlp.get_duals()
96    saved_obj_factor = nlp.get_obj_factor()
97    temp_duals = np.zeros(len(saved_duals))
98
99    # NOTE: This makes some assumption about how the Lagrangian is constructed.
100    # TODO: Define the convention we assume and convert if necessary.
101    idx = nlp.get_constraint_indices(constraints)[0]
102    temp_duals[idx] = 1.0
103    nlp.set_duals(temp_duals)
104    nlp.set_obj_factor(0.0)
105
106    # NOTE: The returned matrix preserves explicit zeros. I.e. it contains
107    # coordinates for every entry that could possibly be nonzero.
108    submatrix = nlp.extract_submatrix_hessian_lag(wrt1, wrt2)
109
110    nlp.set_obj_factor(saved_obj_factor)
111    nlp.set_duals(saved_duals)
112    return submatrix
113
114
115class ExternalPyomoModel(ExternalGreyBoxModel):
116    """
117    This is an ExternalGreyBoxModel used to create an exteral model
118    from existing Pyomo components. Given a system of variables and
119    equations partitioned into "input" and "external" variables and
120    "residual" and "external" equations, this class computes the
121    residual of the "residual equations," as well as their Jacobian
122    and Hessian, as a function of only the inputs.
123
124    Pyomo components:
125        f(x, y) == 0 # "Residual equations"
126        g(x, y) == 0 # "External equations", dim(g) == dim(y)
127
128    Effective constraint seen by this "external model":
129        F(x) == f(x, y(x)) == 0
130        where y(x) solves g(x, y) == 0
131
132    """
133
134    def __init__(self,
135            input_vars,
136            external_vars,
137            residual_cons,
138            external_cons,
139            solver=None,
140            ):
141        if solver is None:
142            solver = SolverFactory("ipopt")
143        self._solver = solver
144
145        # We only need this block to construct the NLP, which wouldn't
146        # be necessary if we could compute Hessians of Pyomo constraints.
147        self._block = create_subsystem_block(
148                residual_cons+external_cons,
149                input_vars+external_vars,
150                )
151        self._block._obj = Objective(expr=0.0)
152        self._nlp = PyomoNLP(self._block)
153
154        assert len(external_vars) == len(external_cons)
155
156        self.input_vars = input_vars
157        self.external_vars = external_vars
158        self.residual_cons = residual_cons
159        self.external_cons = external_cons
160
161        self.residual_con_multipliers = [None for _ in residual_cons]
162
163    def n_inputs(self):
164        return len(self.input_vars)
165
166    def n_equality_constraints(self):
167        return len(self.residual_cons)
168
169    # I would like to try to get by without using the following "name" methods.
170    def input_names(self):
171        return ["input_%i" % i for i in range(self.n_inputs())]
172    def equality_constraint_names(self):
173        return ["residual_%i" % i for i in range(self.n_equality_constraints())]
174
175    def set_input_values(self, input_values):
176        solver = self._solver
177        external_cons = self.external_cons
178        external_vars = self.external_vars
179        input_vars = self.input_vars
180
181        for var, val in zip(input_vars, input_values):
182            var.set_value(val)
183
184        _temp = create_subsystem_block(external_cons, variables=external_vars)
185        possible_input_vars = ComponentSet(input_vars)
186        #for var in _temp.input_vars.values():
187        #    # TODO: Is this check necessary?
188        #    assert var in possible_input_vars
189
190        with TemporarySubsystemManager(to_fix=list(_temp.input_vars.values())):
191            solver.solve(_temp)
192
193        # Should we create the NLP from the original block or the temp block?
194        # Need to create it from the original block because temp block won't
195        # have residual constraints, whose derivatives are necessary.
196        self._nlp = PyomoNLP(self._block)
197
198    def set_equality_constraint_multipliers(self, eq_con_multipliers):
199        for i, val in enumerate(eq_con_multipliers):
200            self.residual_con_multipliers[i] = val
201
202    def evaluate_equality_constraints(self):
203        return self._nlp.extract_subvector_constraints(self.residual_cons)
204
205    def evaluate_jacobian_equality_constraints(self):
206        nlp = self._nlp
207        x = self.input_vars
208        y = self.external_vars
209        f = self.residual_cons
210        g = self.external_cons
211        jfx = nlp.extract_submatrix_jacobian(x, f)
212        jfy = nlp.extract_submatrix_jacobian(y, f)
213        jgx = nlp.extract_submatrix_jacobian(x, g)
214        jgy = nlp.extract_submatrix_jacobian(y, g)
215
216        nf = len(f)
217        nx = len(x)
218        n_entries = nf*nx
219
220        # TODO: Does it make sense to cast dydx to a sparse matrix?
221        # My intuition is that it does only if jgy is "decomposable"
222        # in the strongly connected component sense, which is probably
223        # not usually the case.
224        dydx = -1 * sps.linalg.splu(jgy.tocsc()).solve(jgx.toarray())
225        # NOTE: PyNumero block matrices require this to be a sparse matrix
226        # that contains coordinates for every entry that could possibly
227        # be nonzero. Here, this is all of the entries.
228        dfdx = jfx + jfy.dot(dydx)
229
230        return _dense_to_full_sparse(dfdx)
231
232    def evaluate_jacobian_external_variables(self):
233        nlp = self._nlp
234        x = self.input_vars
235        y = self.external_vars
236        g = self.external_cons
237        jgx = nlp.extract_submatrix_jacobian(x, g)
238        jgy = nlp.extract_submatrix_jacobian(y, g)
239        jgy_csc = jgy.tocsc()
240        dydx = -1 * sps.linalg.splu(jgy_csc).solve(jgx.toarray())
241        return dydx
242
243    def evaluate_hessian_external_variables(self):
244        nlp = self._nlp
245        x = self.input_vars
246        y = self.external_vars
247        g = self.external_cons
248        jgx = nlp.extract_submatrix_jacobian(x, g)
249        jgy = nlp.extract_submatrix_jacobian(y, g)
250        jgy_csc = jgy.tocsc()
251        jgy_fact = sps.linalg.splu(jgy_csc)
252        dydx = -1 * jgy_fact.solve(jgx.toarray())
253
254        ny = len(y)
255        nx = len(x)
256
257        hgxx = np.array([
258            get_hessian_of_constraint(con, x, nlp=nlp).toarray() for con in g
259            ])
260        hgxy = np.array([
261            get_hessian_of_constraint(con, x, y, nlp=nlp).toarray() for con in g
262            ])
263        hgyy = np.array([
264            get_hessian_of_constraint(con, y, nlp=nlp).toarray() for con in g
265            ])
266
267        # This term is sparse, but we do not exploit it.
268        term1 = hgxx
269
270        # This is what we want.
271        # prod[i,j,k] = sum(hgxy[i,:,j] * dydx[:,k])
272        prod = hgxy.dot(dydx)
273        # Swap the second and third axes of the tensor
274        term2 = prod + prod.transpose((0, 2, 1))
275        # The term2 tensor could have some sparsity worth exploiting.
276
277        # matrix.dot(tensor) is not what we want, so we reverse the order of the
278        # product. Exploit symmetry of hgyy to only perform one transpose.
279        term3 = hgyy.dot(dydx).transpose((0, 2, 1)).dot(dydx)
280
281        rhs = term1 + term2 + term3
282
283        rhs.shape = (ny, nx*nx)
284        sol = jgy_fact.solve(rhs)
285        sol.shape = (ny, nx, nx)
286        d2ydx2 = -sol
287
288        return d2ydx2
289
290    def evaluate_hessians_of_residuals(self):
291        """
292        This method computes the Hessian matrix of each equality
293        constraint individually, rather than the sum of Hessians
294        times multipliers.
295        """
296        nlp = self._nlp
297        x = self.input_vars
298        y = self.external_vars
299        f = self.residual_cons
300        g = self.external_cons
301        jfx = nlp.extract_submatrix_jacobian(x, f)
302        jfy = nlp.extract_submatrix_jacobian(y, f)
303
304        dydx = self.evaluate_jacobian_external_variables()
305
306        ny = len(y)
307        nf = len(f)
308        nx = len(x)
309
310        hfxx = np.array([
311            get_hessian_of_constraint(con, x, nlp=nlp).toarray() for con in f
312            ])
313        hfxy = np.array([
314            get_hessian_of_constraint(con, x, y, nlp=nlp).toarray() for con in f
315            ])
316        hfyy = np.array([
317            get_hessian_of_constraint(con, y, nlp=nlp).toarray() for con in f
318            ])
319
320        d2ydx2 = self.evaluate_hessian_external_variables()
321
322        term1 = hfxx
323        prod = hfxy.dot(dydx)
324        term2 = prod + prod.transpose((0, 2, 1))
325        term3 = hfyy.dot(dydx).transpose((0, 2, 1)).dot(dydx)
326
327        d2ydx2.shape = (ny, nx*nx)
328        term4 = jfy.dot(d2ydx2)
329        term4.shape = (nf, nx, nx)
330
331        d2fdx2 = term1 + term2 + term3 + term4
332        return d2fdx2
333
334    def evaluate_hessian_equality_constraints(self):
335        """
336        This method actually evaluates the sum of Hessians times
337        multipliers, i.e. the term in the Hessian of the Lagrangian
338        due to these equality constraints.
339        """
340        d2fdx2 = self.evaluate_hessians_of_residuals()
341        multipliers = self.residual_con_multipliers
342
343        sum_ = sum(mult*matrix for mult, matrix in zip(multipliers, d2fdx2))
344        # Return a sparse matrix with every entry accounted for because it
345        # is difficult to determine rigorously which coordinates
346        # _could possibly_ be nonzero.
347        sparse = _dense_to_full_sparse(sum_)
348        return sps.tril(sparse)
349