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
12import pyomo.common.unittest as unittest
13from pyomo.common.collections import ComponentSet, ComponentMap
14from pyomo.core.expr.visitor import identify_variables
15import pyomo.environ as pyo
16
17from pyomo.contrib.pynumero.dependencies import (
18    numpy as np, numpy_available, scipy, scipy_available
19)
20
21if not (numpy_available and scipy_available):
22    raise unittest.SkipTest("Pynumero needs scipy and numpy to run NLP tests")
23
24from pyomo.common.dependencies.scipy import sparse as sps
25
26from pyomo.contrib.pynumero.asl import AmplInterface
27if not AmplInterface.available():
28    raise unittest.SkipTest(
29        "Pynumero needs the ASL extension to run cyipopt tests")
30
31from pyomo.contrib.pynumero.algorithms.solvers.cyipopt_solver import (
32    cyipopt_available,
33)
34from pyomo.contrib.pynumero.interfaces.external_pyomo_model import (
35    ExternalPyomoModel,
36    get_hessian_of_constraint,
37)
38from pyomo.contrib.pynumero.interfaces.external_grey_box import (
39    ExternalGreyBoxBlock,
40    ScalarExternalGreyBoxBlock,
41    IndexedExternalGreyBoxBlock,
42)
43from pyomo.contrib.pynumero.interfaces.pyomo_grey_box_nlp import (
44    PyomoNLPWithGreyBoxBlocks,
45)
46from pyomo.contrib.pynumero.interfaces.tests.external_grey_box_models import (
47    PressureDropTwoOutputsWithHessian,
48)
49
50if not pyo.SolverFactory("ipopt").available():
51    raise unittest.SkipTest(
52        "Need IPOPT to run ExternalPyomoModel tests"
53        )
54
55
56def _make_external_model():
57    m = pyo.ConcreteModel()
58    m.a = pyo.Var()
59    m.b = pyo.Var()
60    m.r = pyo.Var()
61    m.x = pyo.Var()
62    m.y = pyo.Var()
63    m.x_out = pyo.Var()
64    m.y_out = pyo.Var()
65    m.c_out_1 = pyo.Constraint(expr=m.x_out - m.x == 0)
66    m.c_out_2 = pyo.Constraint(expr=m.y_out - m.y == 0)
67    m.c_ex_1 = pyo.Constraint(expr=
68            m.x**3 - 2*m.y == m.a**2 + m.b**3 - m.r**3 - 2
69            )
70    m.c_ex_2 = pyo.Constraint(expr=
71            m.x + m.y**3 == m.a**3 + 2*m.b**2 + m.r**2 + 1
72            )
73    return m
74
75
76def _add_linking_constraints(m):
77    m.a = pyo.Var()
78    m.b = pyo.Var()
79    m.r = pyo.Var()
80
81    n_inputs = 3
82
83    def linking_constraint_rule(m, i):
84        if i == 0:
85            return m.a == m.ex_block.inputs["input_0"]
86        elif i == 1:
87            return m.b == m.ex_block.inputs["input_1"]
88        elif i == 2:
89            return m.r == m.ex_block.inputs["input_2"]
90
91    m.linking_constraint = pyo.Constraint(range(n_inputs),
92            rule=linking_constraint_rule)
93
94
95def _add_nonlinear_linking_constraints(m):
96    # Nonlinear linking constraints are unusual. They are useful
97    # here to test correct combination of Hessians from multiple
98    # sources, however.
99    m.a = pyo.Var()
100    m.b = pyo.Var()
101    m.r = pyo.Var()
102
103    n_inputs = 3
104
105    def linking_constraint_rule(m, i):
106        if i == 0:
107            return m.a**2 - 0.5*m.ex_block.inputs["input_0"]**2 == 0
108        elif i == 1:
109            return m.b**2 - 0.5*m.ex_block.inputs["input_1"]**2 == 0
110        elif i == 2:
111            return m.r**2 - 0.5*m.ex_block.inputs["input_2"]**2 == 0
112
113    m.linking_constraint = pyo.Constraint(range(n_inputs),
114            rule=linking_constraint_rule)
115
116
117def make_dynamic_model():
118    m = pyo.ConcreteModel()
119    m.time = pyo.Set(initialize=[0, 1, 2])
120    m = pyo.ConcreteModel()
121
122    m.time = pyo.Set(initialize=[0, 1, 2])
123    t0 = m.time.first()
124
125    m.h = pyo.Var(m.time, initialize=1.0)
126    m.dhdt = pyo.Var(m.time, initialize=1.0)
127    m.flow_in = pyo.Var(m.time, bounds=(0, None), initialize=1.0)
128    m.flow_out = pyo.Var(m.time, initialize=1.0)
129
130    m.flow_coef = pyo.Param(initialize=2.0, mutable=True)
131
132    def h_diff_eqn_rule(m, t):
133        return m.dhdt[t] - (m.flow_in[t] - m.flow_out[t]) == 0
134    m.h_diff_eqn = pyo.Constraint(m.time, rule=h_diff_eqn_rule)
135
136    def dhdt_disc_eqn_rule(m, t):
137        if t == m.time.first():
138            return pyo.Constraint.Skip
139        else:
140            t_prev = m.time.prev(t)
141            delta_t = (t - t_prev)
142            return m.dhdt[t] - delta_t*(m.h[t] - m.h[t_prev]) == 0
143    m.dhdt_disc_eqn = pyo.Constraint(m.time, rule=dhdt_disc_eqn_rule)
144
145    def flow_out_eqn(m, t):
146        return m.flow_out[t] == m.flow_coef*m.h[t]**0.5
147    m.flow_out_eqn = pyo.Constraint(m.time, rule=flow_out_eqn)
148
149    return m
150
151
152class TestExternalGreyBoxBlock(unittest.TestCase):
153
154    def test_construct_scalar(self):
155        m = pyo.ConcreteModel()
156        m.ex_block = ExternalGreyBoxBlock(concrete=True)
157        block = m.ex_block
158        self.assertIs(type(block), ScalarExternalGreyBoxBlock)
159
160        m_ex = _make_external_model()
161        input_vars = [m_ex.a, m_ex.b, m_ex.r, m_ex.x_out, m_ex.y_out]
162        external_vars = [m_ex.x, m_ex.y]
163        residual_cons = [m_ex.c_out_1, m_ex.c_out_2]
164        external_cons = [m_ex.c_ex_1, m_ex.c_ex_2]
165        ex_model = ExternalPyomoModel(
166                input_vars,
167                external_vars,
168                residual_cons,
169                external_cons,
170                )
171        block.set_external_model(ex_model)
172
173        self.assertEqual(len(block.inputs), len(input_vars))
174        self.assertEqual(len(block.outputs), 0)
175        self.assertEqual(len(block._equality_constraint_names), 2)
176
177    def test_construct_indexed(self):
178        block = ExternalGreyBoxBlock([0, 1, 2], concrete=True)
179        self.assertIs(type(block), IndexedExternalGreyBoxBlock)
180
181        m_ex = _make_external_model()
182        input_vars = [m_ex.a, m_ex.b, m_ex.r, m_ex.x_out, m_ex.y_out]
183        external_vars = [m_ex.x, m_ex.y]
184        residual_cons = [m_ex.c_out_1, m_ex.c_out_2]
185        external_cons = [m_ex.c_ex_1, m_ex.c_ex_2]
186        ex_model = ExternalPyomoModel(
187                input_vars,
188                external_vars,
189                residual_cons,
190                external_cons,
191                )
192
193        for i in block:
194            b = block[i]
195            b.set_external_model(ex_model)
196            self.assertEqual(len(b.inputs), len(input_vars))
197            self.assertEqual(len(b.outputs), 0)
198            self.assertEqual(len(b._equality_constraint_names), 2)
199
200    @unittest.skipUnless(cyipopt_available, "cyipopt is not available")
201    def test_solve_square(self):
202        m = pyo.ConcreteModel()
203        m.ex_block = ExternalGreyBoxBlock(concrete=True)
204        block = m.ex_block
205
206        m_ex = _make_external_model()
207        input_vars = [m_ex.a, m_ex.b, m_ex.r, m_ex.x_out, m_ex.y_out]
208        external_vars = [m_ex.x, m_ex.y]
209        residual_cons = [m_ex.c_out_1, m_ex.c_out_2]
210        external_cons = [m_ex.c_ex_1, m_ex.c_ex_2]
211        ex_model = ExternalPyomoModel(
212                input_vars,
213                external_vars,
214                residual_cons,
215                external_cons,
216                )
217        block.set_external_model(ex_model)
218
219        _add_linking_constraints(m)
220
221        m.a.fix(1)
222        m.b.fix(2)
223        m.r.fix(3)
224
225        m.obj = pyo.Objective(expr=0)
226
227        solver = pyo.SolverFactory("cyipopt")
228        solver.solve(m)
229
230        self.assertFalse(m_ex.a.fixed)
231        self.assertFalse(m_ex.b.fixed)
232        self.assertFalse(m_ex.r.fixed)
233
234        m_ex.a.fix(1)
235        m_ex.b.fix(2)
236        m_ex.r.fix(3)
237        ipopt = pyo.SolverFactory("ipopt")
238        ipopt.solve(m_ex)
239
240        x = m.ex_block.inputs["input_3"]
241        y = m.ex_block.inputs["input_4"]
242        self.assertAlmostEqual(m_ex.x.value, x.value, delta=1e-8)
243        self.assertAlmostEqual(m_ex.y.value, y.value, delta=1e-8)
244
245    @unittest.skipUnless(cyipopt_available, "cyipopt is not available")
246    def test_optimize(self):
247        m = pyo.ConcreteModel()
248        m.ex_block = ExternalGreyBoxBlock(concrete=True)
249        block = m.ex_block
250
251        m_ex = _make_external_model()
252        input_vars = [m_ex.a, m_ex.b, m_ex.r, m_ex.x_out, m_ex.y_out]
253        external_vars = [m_ex.x, m_ex.y]
254        residual_cons = [m_ex.c_out_1, m_ex.c_out_2]
255        external_cons = [m_ex.c_ex_1, m_ex.c_ex_2]
256        ex_model = ExternalPyomoModel(
257                input_vars,
258                external_vars,
259                residual_cons,
260                external_cons,
261                )
262        block.set_external_model(ex_model)
263
264        a = m.ex_block.inputs["input_0"]
265        b = m.ex_block.inputs["input_1"]
266        r = m.ex_block.inputs["input_2"]
267        x = m.ex_block.inputs["input_3"]
268        y = m.ex_block.inputs["input_4"]
269        m.obj = pyo.Objective(expr=
270                (x-2.0)**2 + (y-2.0)**2 + (a-2.0)**2 + (b-2.0)**2 + (r-2.0)**2
271                )
272
273        # Solve with external model embedded
274        solver = pyo.SolverFactory("cyipopt")
275        solver.solve(m)
276
277        m_ex.obj = pyo.Objective(expr=
278                (m_ex.x-2.0)**2 + (m_ex.y-2.0)**2 + (m_ex.a-2.0)**2 +
279                (m_ex.b-2.0)**2 + (m_ex.r-2.0)**2
280                )
281        m_ex.a.set_value(0.0)
282        m_ex.b.set_value(0.0)
283        m_ex.r.set_value(0.0)
284        m_ex.y.set_value(0.0)
285        m_ex.x.set_value(0.0)
286
287        # Solve external model, now with same objective function
288        ipopt = pyo.SolverFactory("ipopt")
289        ipopt.solve(m_ex)
290
291        # Make sure full space and reduced space solves give same
292        # answers.
293        self.assertAlmostEqual(m_ex.a.value, a.value, delta=1e-8)
294        self.assertAlmostEqual(m_ex.b.value, b.value, delta=1e-8)
295        self.assertAlmostEqual(m_ex.r.value, r.value, delta=1e-8)
296        self.assertAlmostEqual(m_ex.x.value, x.value, delta=1e-8)
297        self.assertAlmostEqual(m_ex.y.value, y.value, delta=1e-8)
298
299    def test_construct_dynamic(self):
300        m = make_dynamic_model()
301        time = m.time
302        t0 = m.time.first()
303
304        inputs = [m.h, m.dhdt, m.flow_in]
305        ext_vars = [m.flow_out]
306        residuals = [m.h_diff_eqn]
307        ext_cons = [m.flow_out_eqn]
308
309        external_model_dict = {
310                t: ExternalPyomoModel(
311                    [var[t] for var in inputs],
312                    [var[t] for var in ext_vars],
313                    [con[t] for con in residuals],
314                    [con[t] for con in ext_cons],
315                    )
316                for t in time
317                }
318
319        reduced_space = pyo.Block(concrete=True)
320        reduced_space.external_block = ExternalGreyBoxBlock(
321                time,
322                external_model=external_model_dict,
323                )
324        block = reduced_space.external_block
325        block[t0].deactivate()
326        self.assertIs(type(block), IndexedExternalGreyBoxBlock)
327
328        for t in time:
329            b = block[t]
330            self.assertEqual(len(b.inputs), len(inputs))
331            self.assertEqual(len(b.outputs), 0)
332            self.assertEqual(len(b._equality_constraint_names), len(residuals))
333
334        reduced_space.diff_var = pyo.Reference(m.h)
335        reduced_space.deriv_var = pyo.Reference(m.dhdt)
336        reduced_space.input_var = pyo.Reference(m.flow_in)
337        reduced_space.disc_eqn = pyo.Reference(m.dhdt_disc_eqn)
338
339        pyomo_vars = list(reduced_space.component_data_objects(pyo.Var))
340        pyomo_cons = list(reduced_space.component_data_objects(pyo.Constraint))
341        # NOTE: Variables in the EGBB are not found by component_data_objects
342        self.assertEqual(len(pyomo_vars), len(inputs)*len(time))
343        # "Constraints" defined by the EGBB are not found either, although
344        # this is expected.
345        self.assertEqual(len(pyomo_cons), len(time)-1)
346
347        reduced_space._obj = pyo.Objective(expr=0)
348
349        # This is required to avoid a failure in the implicit function
350        # evaluation when "initializing" (?) the PNLPwGBB.
351        # Why exactly is function evaluation necessary for this
352        # initialization again?
353        block[:].inputs[:].set_value(1.0)
354
355        # This is necessary for these variables to appear in the PNLPwGBB.
356        # Otherwise they don't appear in any "real" constraints of the
357        # PyomoNLP.
358        reduced_space.const_input_eqn = pyo.Constraint(
359            expr=reduced_space.input_var[2] - reduced_space.input_var[1] == 0
360            )
361
362        nlp = PyomoNLPWithGreyBoxBlocks(reduced_space)
363        self.assertEqual(
364                nlp.n_primals(),
365                # EGBB "inputs", dhdt, and flow_in exist for t != t0.
366                # h exists for all time.
367                (2+len(inputs))*(len(time)-1) + len(time),
368                )
369        self.assertEqual(
370                nlp.n_constraints(),
371                # EGBB equality constraints and disc_eqn exist for t != t0.
372                # const_input_eqn is a single constraint
373                (len(residuals)+1)*(len(time)-1) + 1,
374                )
375
376    @unittest.skipUnless(cyipopt_available, "cyipopt is not available")
377    def test_solve_square_dynamic(self):
378        # Create the "external model"
379        m = make_dynamic_model()
380        time = m.time
381        t0 = m.time.first()
382        m.h[t0].fix(1.2)
383        m.flow_in.fix(1.5)
384
385        # Create the block that will hold the reduced space model.
386        reduced_space = pyo.Block(concrete=True)
387        reduced_space.diff_var = pyo.Reference(m.h)
388        reduced_space.deriv_var = pyo.Reference(m.dhdt)
389        reduced_space.input_var = pyo.Reference(m.flow_in)
390        reduced_space.disc_eq = pyo.Reference(m.dhdt_disc_eqn)
391
392        reduced_space.external_block = ExternalGreyBoxBlock(time)
393        block = reduced_space.external_block
394        block[t0].deactivate()
395        for t in time:
396            # TODO: skipping time.first() necessary?
397            if t != t0:
398                input_vars = [m.h[t], m.dhdt[t]]
399                external_vars = [m.flow_out[t]]
400                residual_cons = [m.h_diff_eqn[t]]
401                external_cons = [m.flow_out_eqn[t]]
402                external_model = ExternalPyomoModel(
403                        input_vars,
404                        external_vars,
405                        residual_cons,
406                        external_cons,
407                        )
408                block[t].set_external_model(external_model)
409
410        n_inputs = len(input_vars)
411        def linking_constraint_rule(m, i, t):
412            if t == t0:
413                return pyo.Constraint.Skip
414            if i == 0:
415                return m.diff_var[t] == m.external_block[t].inputs["input_0"]
416            elif i == 1:
417                return m.deriv_var[t] == m.external_block[t].inputs["input_1"]
418
419        reduced_space.linking_constraint = pyo.Constraint(
420                range(n_inputs),
421                time,
422                rule=linking_constraint_rule,
423                )
424        # Initialize new variables
425        for t in time:
426            if t != t0:
427                block[t].inputs["input_0"].set_value(m.h[t].value)
428                block[t].inputs["input_1"].set_value(m.dhdt[t].value)
429
430        reduced_space._obj = pyo.Objective(expr=0)
431
432        solver = pyo.SolverFactory("cyipopt")
433        results = solver.solve(reduced_space, tee=True)
434
435        # Full space square model was solved in a separate script
436        # to obtain these values.
437        h_target = [1.2, 0.852923, 0.690725]
438        dhdt_target = [-0.690890, -0.347077, -0.162198]
439        flow_out_target = [2.190980, 1.847077, 1.662198]
440        for t in time:
441            if t == t0:
442                continue
443            values = [m.h[t].value, m.dhdt[t].value, m.flow_out[t].value]
444            target_values = [h_target[t], dhdt_target[t], flow_out_target[t]]
445            self.assertStructuredAlmostEqual(values, target_values, delta=1e-5)
446
447    @unittest.skipUnless(cyipopt_available, "cyipopt is not available")
448    def test_optimize_dynamic(self):
449        # Create the "external model"
450        m = make_dynamic_model()
451        time = m.time
452        t0 = m.time.first()
453        m.h[t0].fix(1.2)
454        m.flow_in[t0].fix(1.5)
455
456        m.obj = pyo.Objective(expr=sum(
457            (m.h[t] - 2.0)**2 for t in m.time if t != t0
458            ))
459
460        # Create the block that will hold the reduced space model.
461        reduced_space = pyo.Block(concrete=True)
462        reduced_space.diff_var = pyo.Reference(m.h)
463        reduced_space.deriv_var = pyo.Reference(m.dhdt)
464        reduced_space.input_var = pyo.Reference(m.flow_in)
465        reduced_space.disc_eq = pyo.Reference(m.dhdt_disc_eqn)
466        reduced_space.objective = pyo.Reference(m.obj)
467
468        reduced_space.external_block = ExternalGreyBoxBlock(time)
469        block = reduced_space.external_block
470        block[t0].deactivate()
471        for t in time:
472            # TODO: skipping time.first() necessary?
473            if t != t0:
474                input_vars = [m.h[t], m.dhdt[t], m.flow_in[t]]
475                external_vars = [m.flow_out[t]]
476                residual_cons = [m.h_diff_eqn[t]]
477                external_cons = [m.flow_out_eqn[t]]
478                external_model = ExternalPyomoModel(
479                        input_vars,
480                        external_vars,
481                        residual_cons,
482                        external_cons,
483                        )
484                block[t].set_external_model(external_model)
485
486        n_inputs = len(input_vars)
487        def linking_constraint_rule(m, i, t):
488            if t == t0:
489                return pyo.Constraint.Skip
490            if i == 0:
491                return m.diff_var[t] == m.external_block[t].inputs["input_0"]
492            elif i == 1:
493                return m.deriv_var[t] == m.external_block[t].inputs["input_1"]
494            elif i == 2:
495                return m.input_var[t] == m.external_block[t].inputs["input_2"]
496
497        reduced_space.linking_constraint = pyo.Constraint(
498                range(n_inputs),
499                time,
500                rule=linking_constraint_rule,
501                )
502        # Initialize new variables
503        for t in time:
504            if t != t0:
505                block[t].inputs["input_0"].set_value(m.h[t].value)
506                block[t].inputs["input_1"].set_value(m.dhdt[t].value)
507                block[t].inputs["input_2"].set_value(m.flow_in[t].value)
508
509        solver = pyo.SolverFactory("cyipopt")
510        results = solver.solve(reduced_space)
511
512        # These values were obtained by solving this problem in the full
513        # space in a separate script.
514        h_target = [1.2, 2.0, 2.0]
515        dhdt_target = [-0.690890, 0.80, 0.0]
516        flow_in_target = [1.5, 3.628427, 2.828427]
517        flow_out_target = [2.190890, 2.828427, 2.828427]
518        for t in time:
519            if t == t0:
520                continue
521            values = [m.h[t].value, m.dhdt[t].value,
522                    m.flow_out[t].value, m.flow_in[t].value]
523            target_values = [h_target[t], dhdt_target[t],
524                    flow_out_target[t], flow_in_target[t]]
525            self.assertStructuredAlmostEqual(values, target_values, delta=1e-5)
526
527    @unittest.skipUnless(cyipopt_available, "cyipopt is not available")
528    def test_optimize_dynamic_references(self):
529        """
530        When when pre-existing variables are attached to the EGBB
531        as references, linking constraints are no longer necessary.
532        """
533        # Create the "external model"
534        m = make_dynamic_model()
535        time = m.time
536        t0 = m.time.first()
537        m.h[t0].fix(1.2)
538        m.flow_in[t0].fix(1.5)
539
540        m.obj = pyo.Objective(expr=sum(
541            (m.h[t] - 2.0)**2 for t in m.time if t != t0
542            ))
543
544        # Create the block that will hold the reduced space model.
545        reduced_space = pyo.Block(concrete=True)
546        reduced_space.diff_var = pyo.Reference(m.h)
547        reduced_space.deriv_var = pyo.Reference(m.dhdt)
548        reduced_space.input_var = pyo.Reference(m.flow_in)
549        reduced_space.disc_eq = pyo.Reference(m.dhdt_disc_eqn)
550        reduced_space.objective = pyo.Reference(m.obj)
551
552        reduced_space.external_block = ExternalGreyBoxBlock(time)
553        block = reduced_space.external_block
554        block[t0].deactivate()
555        for t in time:
556            # TODO: is skipping time.first() necessary?
557            if t != t0:
558                input_vars = [m.h[t], m.dhdt[t], m.flow_in[t]]
559                external_vars = [m.flow_out[t]]
560                residual_cons = [m.h_diff_eqn[t]]
561                external_cons = [m.flow_out_eqn[t]]
562                external_model = ExternalPyomoModel(
563                        input_vars,
564                        external_vars,
565                        residual_cons,
566                        external_cons,
567                        )
568                block[t].set_external_model(external_model, inputs=input_vars)
569
570        solver = pyo.SolverFactory("cyipopt")
571        results = solver.solve(reduced_space)
572
573        # These values were obtained by solving this problem in the full
574        # space in a separate script.
575        h_target = [1.2, 2.0, 2.0]
576        dhdt_target = [-0.690890, 0.80, 0.0]
577        flow_in_target = [1.5, 3.628427, 2.828427]
578        flow_out_target = [2.190890, 2.828427, 2.828427]
579        for t in time:
580            if t == t0:
581                continue
582            values = [m.h[t].value, m.dhdt[t].value,
583                    m.flow_out[t].value, m.flow_in[t].value]
584            target_values = [h_target[t], dhdt_target[t],
585                    flow_out_target[t], flow_in_target[t]]
586            self.assertStructuredAlmostEqual(values, target_values, delta=1e-5)
587
588
589class TestPyomoNLPWithGreyBoxBLocks(unittest.TestCase):
590
591    def test_set_and_evaluate(self):
592        m = pyo.ConcreteModel()
593        m.ex_block = ExternalGreyBoxBlock(concrete=True)
594        block = m.ex_block
595
596        m_ex = _make_external_model()
597        input_vars = [m_ex.a, m_ex.b, m_ex.r, m_ex.x_out, m_ex.y_out]
598        external_vars = [m_ex.x, m_ex.y]
599        residual_cons = [m_ex.c_out_1, m_ex.c_out_2]
600        external_cons = [m_ex.c_ex_1, m_ex.c_ex_2]
601        ex_model = ExternalPyomoModel(
602                input_vars,
603                external_vars,
604                residual_cons,
605                external_cons,
606                )
607        block.set_external_model(ex_model)
608
609        a = m.ex_block.inputs["input_0"]
610        b = m.ex_block.inputs["input_1"]
611        r = m.ex_block.inputs["input_2"]
612        x = m.ex_block.inputs["input_3"]
613        y = m.ex_block.inputs["input_4"]
614        m.obj = pyo.Objective(expr=
615                (x-2.0)**2 + (y-2.0)**2 + (a-2.0)**2 + (b-2.0)**2 + (r-2.0)**2
616                )
617
618        _add_linking_constraints(m)
619
620        nlp = PyomoNLPWithGreyBoxBlocks(m)
621
622        # Set primals in model, get primals in nlp
623        # set/get duals
624        # evaluate constraints
625        # evaluate Jacobian
626        # evaluate Hessian
627        self.assertEqual(nlp.n_primals(), 8)
628
629        # PyomoNLPWithGreyBoxBlocks sorts variables by name
630        primals_names = [
631                "a",
632                "b",
633                "ex_block.inputs[input_0]",
634                "ex_block.inputs[input_1]",
635                "ex_block.inputs[input_2]",
636                "ex_block.inputs[input_3]",
637                "ex_block.inputs[input_4]",
638                "r",
639                ]
640        self.assertEqual(nlp.primals_names(), primals_names)
641        np.testing.assert_equal(np.zeros(8), nlp.get_primals())
642
643        primals = np.array([0, 1, 2, 3, 4, 5, 6, 7])
644        nlp.set_primals(primals)
645        np.testing.assert_equal(primals, nlp.get_primals())
646        nlp.load_state_into_pyomo()
647
648        for name, val in zip(primals_names, primals):
649            var = m.find_component(name)
650            self.assertEqual(var.value, val)
651
652        constraint_names = [
653                "linking_constraint[0]",
654                "linking_constraint[1]",
655                "linking_constraint[2]",
656                "ex_block.residual_0",
657                "ex_block.residual_1",
658                ]
659        self.assertEqual(constraint_names, nlp.constraint_names())
660        residuals = np.array([
661                -2.0,
662                -2.0,
663                3.0,
664                # These values were obtained by solving the same system
665                # with Ipopt in another script. It may be better to do
666                # the solve in this test in case the system changes.
667                5.0-(-3.03051522),
668                6.0-3.583839997,
669                ])
670        np.testing.assert_allclose(residuals, nlp.evaluate_constraints(),
671                rtol=1e-8)
672
673        duals = np.array([1, 2, 3, 4, 5])
674        nlp.set_duals(duals)
675
676        self.assertEqual(ex_model.residual_con_multipliers, [4, 5])
677        np.testing.assert_equal(nlp.get_duals(), duals)
678
679    def test_jacobian(self):
680        m = pyo.ConcreteModel()
681        m.ex_block = ExternalGreyBoxBlock(concrete=True)
682        block = m.ex_block
683
684        m_ex = _make_external_model()
685        input_vars = [m_ex.a, m_ex.b, m_ex.r, m_ex.x_out, m_ex.y_out]
686        external_vars = [m_ex.x, m_ex.y]
687        residual_cons = [m_ex.c_out_1, m_ex.c_out_2]
688        external_cons = [m_ex.c_ex_1, m_ex.c_ex_2]
689        ex_model = ExternalPyomoModel(
690                input_vars,
691                external_vars,
692                residual_cons,
693                external_cons,
694                )
695        block.set_external_model(ex_model)
696
697        a = m.ex_block.inputs["input_0"]
698        b = m.ex_block.inputs["input_1"]
699        r = m.ex_block.inputs["input_2"]
700        x = m.ex_block.inputs["input_3"]
701        y = m.ex_block.inputs["input_4"]
702        m.obj = pyo.Objective(expr=
703                (x-2.0)**2 + (y-2.0)**2 + (a-2.0)**2 + (b-2.0)**2 + (r-2.0)**2
704                )
705
706        _add_linking_constraints(m)
707
708        nlp = PyomoNLPWithGreyBoxBlocks(m)
709        primals = np.array([0, 1, 2, 3, 4, 5, 6, 7])
710        nlp.set_primals(primals)
711        jac = nlp.evaluate_jacobian()
712
713        # Variable and constraint orders (verified by previous test):
714        # Rows:
715        # [
716        #     "linking_constraint[0]",
717        #     "linking_constraint[1]",
718        #     "linking_constraint[2]",
719        #     "ex_block.residual_0",
720        #     "ex_block.residual_1",
721        # ]
722        # Cols:
723        # [
724        #     "a",
725        #     "b",
726        #     "ex_block.inputs[input_0]",
727        #     "ex_block.inputs[input_1]",
728        #     "ex_block.inputs[input_2]",
729        #     "ex_block.inputs[input_3]",
730        #     "ex_block.inputs[input_4]",
731        #     "r",
732        # ]
733        row = [
734                0, 0,
735                1, 1,
736                2, 2,
737                3, 3, 3, 3, 3,
738                4, 4, 4, 4, 4,
739                ]
740        col = [
741                0, 2,
742                1, 3,
743                7, 4,
744                2, 3, 4, 5, 6,
745                2, 3, 4, 5, 6,
746                ]
747        data = [
748                1, -1,
749                1, -1,
750                1, -1,
751                -0.16747094, -1.00068434, 1.72383729, 1, 0,
752                -0.30708535, -0.28546127, -0.25235924, 0, 1,
753                ]
754        self.assertEqual(len(row), len(jac.row))
755        rcd_dict = dict(((i, j), val) for i, j, val in zip(row, col, data))
756        for i, j, val in zip(jac.row, jac.col, jac.data):
757            self.assertIn((i, j), rcd_dict)
758            self.assertAlmostEqual(rcd_dict[i, j], val, delta=1e-8)
759
760    def test_hessian_1(self):
761        # Test with duals equal vector of ones.
762        m = pyo.ConcreteModel()
763        m.ex_block = ExternalGreyBoxBlock(concrete=True)
764        block = m.ex_block
765
766        m_ex = _make_external_model()
767        input_vars = [m_ex.a, m_ex.b, m_ex.r, m_ex.x_out, m_ex.y_out]
768        external_vars = [m_ex.x, m_ex.y]
769        residual_cons = [m_ex.c_out_1, m_ex.c_out_2]
770        external_cons = [m_ex.c_ex_1, m_ex.c_ex_2]
771        ex_model = ExternalPyomoModel(
772                input_vars,
773                external_vars,
774                residual_cons,
775                external_cons,
776                )
777        block.set_external_model(ex_model)
778
779        a = m.ex_block.inputs["input_0"]
780        b = m.ex_block.inputs["input_1"]
781        r = m.ex_block.inputs["input_2"]
782        x = m.ex_block.inputs["input_3"]
783        y = m.ex_block.inputs["input_4"]
784        m.obj = pyo.Objective(expr=
785                (x-2.0)**2 + (y-2.0)**2 + (a-2.0)**2 + (b-2.0)**2 + (r-2.0)**2
786                )
787
788        _add_nonlinear_linking_constraints(m)
789
790        nlp = PyomoNLPWithGreyBoxBlocks(m)
791        primals = np.array([0, 1, 2, 3, 4, 5, 6, 7])
792        duals = np.array([1, 1, 1, 1, 1])
793        nlp.set_primals(primals)
794        nlp.set_duals(duals)
795        hess = nlp.evaluate_hessian_lag()
796
797        # Variable order (verified by a previous test):
798        # [
799        #     "a",
800        #     "b",
801        #     "ex_block.inputs[input_0]",
802        #     "ex_block.inputs[input_1]",
803        #     "ex_block.inputs[input_2]",
804        #     "ex_block.inputs[input_3]",
805        #     "ex_block.inputs[input_4]",
806        #     "r",
807        # ]
808        row = [0, 1, 7]
809        col = [0, 1, 7]
810        # Data entries are influenced by multiplier values.
811        # Here these are just ones.
812        data = [2.0, 2.0, 2.0]
813        # ^ These variables only appear in linking constraints
814        rcd_dict = dict(((i, j), val) for i, j, val in zip(row, col, data))
815
816        # These are the coordinates of the Hessian corresponding to
817        # external variables with true nonzeros. The coordinates have
818        # terms due to objective, linking constraints, and external
819        # constraints. Values were extracted from the external model
820        # while writing this test, which is just meant to verify
821        # that the different Hessians combined properly.
822        ex_block_nonzeros = {
823                (2, 2): 2.0 + (-1.0) + (-0.10967928) + (-0.25595929),
824                (2, 3): (-0.10684633) + (0.05169308),
825                (3, 2): (-0.10684633) + (0.05169308),
826                (2, 4): (0.19329898) + (0.03823075),
827                (4, 2): (0.19329898) + (0.03823075),
828                (3, 3): 2.0 + (-1.0) + (-1.31592135) + (-0.0241836),
829                (3, 4): (1.13920361) + (0.01063667),
830                (4, 3): (1.13920361) + (0.01063667),
831                (4, 4): 2.0 + (-1.0) + (-1.0891866) + (0.01190218),
832                (5, 5): 2.0,
833                (6, 6): 2.0,
834                }
835        rcd_dict.update(ex_block_nonzeros)
836
837        # Because "external Hessians" are computed by factorizing matrices,
838        # we have dense blocks in the Hessian for now.
839        ex_block_coords = [2, 3, 4, 5, 6]
840        for i, j in itertools.product(ex_block_coords, ex_block_coords):
841            row.append(i)
842            col.append(j)
843            if (i, j) not in rcd_dict:
844                rcd_dict[i, j] = 0.0
845
846        self.assertEqual(len(row), len(hess.row))
847        for i, j, val in zip(hess.row, hess.col, hess.data):
848            self.assertIn((i, j), rcd_dict)
849            self.assertAlmostEqual(rcd_dict[i, j], val, delta=1e-8)
850
851    def test_hessian_2(self):
852        # Test with duals different than vector of ones
853        m = pyo.ConcreteModel()
854        m.ex_block = ExternalGreyBoxBlock(concrete=True)
855        block = m.ex_block
856
857        m_ex = _make_external_model()
858        input_vars = [m_ex.a, m_ex.b, m_ex.r, m_ex.x_out, m_ex.y_out]
859        external_vars = [m_ex.x, m_ex.y]
860        residual_cons = [m_ex.c_out_1, m_ex.c_out_2]
861        external_cons = [m_ex.c_ex_1, m_ex.c_ex_2]
862        ex_model = ExternalPyomoModel(
863                input_vars,
864                external_vars,
865                residual_cons,
866                external_cons,
867                )
868        block.set_external_model(ex_model)
869
870        a = m.ex_block.inputs["input_0"]
871        b = m.ex_block.inputs["input_1"]
872        r = m.ex_block.inputs["input_2"]
873        x = m.ex_block.inputs["input_3"]
874        y = m.ex_block.inputs["input_4"]
875        m.obj = pyo.Objective(expr=
876                (x-2.0)**2 + (y-2.0)**2 + (a-2.0)**2 + (b-2.0)**2 + (r-2.0)**2
877                )
878
879        _add_nonlinear_linking_constraints(m)
880
881        nlp = PyomoNLPWithGreyBoxBlocks(m)
882        primals = np.array([0, 1, 2, 3, 4, 5, 6, 7])
883        duals = np.array([4.4, -3.3, 2.2, -1.1, 0.0])
884        nlp.set_primals(primals)
885        nlp.set_duals(duals)
886        hess = nlp.evaluate_hessian_lag()
887
888        # Variable order (verified by a previous test):
889        # [
890        #     "a",
891        #     "b",
892        #     "ex_block.inputs[input_0]",
893        #     "ex_block.inputs[input_1]",
894        #     "ex_block.inputs[input_2]",
895        #     "ex_block.inputs[input_3]",
896        #     "ex_block.inputs[input_4]",
897        #     "r",
898        # ]
899        row = [0, 1, 7]
900        col = [0, 1, 7]
901        # Data entries are influenced by multiplier values.
902        data = [4.4*2.0, -3.3*2.0, 2.2*2.0]
903        # ^ These variables only appear in linking constraints
904        rcd_dict = dict(((i, j), val) for i, j, val in zip(row, col, data))
905
906        # These are the coordinates of the Hessian corresponding to
907        # external variables with true nonzeros. The coordinates have
908        # terms due to objective, linking constraints, and external
909        # constraints. Values were extracted from the external model
910        # while writing this test, which is just meant to verify
911        # that the different Hessians combined properly.
912        ex_block_nonzeros = {
913                (2, 2): 2.0 + 4.4*(-1.0) + -1.1*(-0.10967928),
914                (2, 3): -1.1*(-0.10684633),
915                (3, 2): -1.1*(-0.10684633),
916                (2, 4): -1.1*(0.19329898),
917                (4, 2): -1.1*(0.19329898),
918                (3, 3): 2.0 + (-3.3)*(-1.0) + -1.1*(-1.31592135),
919                (3, 4): -1.1*(1.13920361),
920                (4, 3): -1.1*(1.13920361),
921                (4, 4): 2.0 + 2.2*(-1.0) + -1.1*(-1.0891866),
922                (5, 5): 2.0,
923                (6, 6): 2.0,
924                }
925        rcd_dict.update(ex_block_nonzeros)
926
927        # Because "external Hessians" are computed by factorizing matrices,
928        # we have dense blocks in the Hessian for now.
929        ex_block_coords = [2, 3, 4, 5, 6]
930        for i, j in itertools.product(ex_block_coords, ex_block_coords):
931            row.append(i)
932            col.append(j)
933            if (i, j) not in rcd_dict:
934                rcd_dict[i, j] = 0.0
935
936        self.assertEqual(len(row), len(hess.row))
937        for i, j, val in zip(hess.row, hess.col, hess.data):
938            self.assertIn((i, j), rcd_dict)
939            self.assertAlmostEqual(rcd_dict[i, j], val, delta=1e-8)
940
941
942class TestExternalGreyBoxBlockWithReferences(unittest.TestCase):
943    """
944    Tests for ExternalGreyBoxBlock with existing variables used
945    as inputs and outputs
946    """
947
948    def _create_pressure_drop_model(self):
949        """
950        Create a Pyomo model with pure ExternalGreyBoxModel embedded.
951        """
952        m = pyo.ConcreteModel()
953
954        # Create variables that the external block will use
955        m.Pin = pyo.Var()
956        m.c = pyo.Var()
957        m.F = pyo.Var()
958        m.P2 = pyo.Var()
959        m.Pout = pyo.Var()
960
961        # Create some random constraints and objective. These variables
962        # need to appear somewhere other than the external block.
963        m.Pin_con = pyo.Constraint(expr=m.Pin == 5.0)
964        m.c_con = pyo.Constraint(expr=m.c == 1.0)
965        m.F_con = pyo.Constraint(expr=m.F == 10.0)
966        m.P2_con = pyo.Constraint(expr=m.P2 <= 5.0)
967        m.obj = pyo.Objective(expr=(m.Pout - 3.0)**2)
968
969        cons = [m.c_con, m.F_con, m.Pin_con, m.P2_con]
970        inputs = [m.Pin, m.c, m.F]
971        outputs = [m.P2, m.Pout]
972
973        # This is "model 3" from the external_grey_box_models.py file.
974        ex_model = PressureDropTwoOutputsWithHessian()
975        m.egb = ExternalGreyBoxBlock()
976        m.egb.set_external_model(
977                ex_model,
978                inputs=inputs,
979                outputs=outputs,
980                )
981        return m
982
983    def test_pressure_drop_model(self):
984        m = self._create_pressure_drop_model()
985
986        cons = [m.c_con, m.F_con, m.Pin_con, m.P2_con]
987        inputs = [m.Pin, m.c, m.F]
988        outputs = [m.P2, m.Pout]
989
990        pyomo_variables = list(m.component_data_objects(pyo.Var))
991        pyomo_constraints = list(m.component_data_objects(pyo.Constraint))
992
993        # The references to inputs and outputs are not picked up twice,
994        # as EGBB does not have ctype Block
995        self.assertEqual(len(pyomo_variables), len(inputs)+len(outputs))
996        self.assertEqual(len(pyomo_constraints), len(cons))
997
998        # Test the inputs and outputs attributes on egb
999        self.assertIs(m.egb.inputs.ctype, pyo.Var)
1000        self.assertIs(m.egb.outputs.ctype, pyo.Var)
1001
1002        self.assertEqual(len(m.egb.inputs), len(inputs))
1003        self.assertEqual(len(m.egb.outputs), len(outputs))
1004
1005        for i in range(len(inputs)):
1006            self.assertIs(inputs[i], m.egb.inputs[i])
1007        for i in range(len(outputs)):
1008            self.assertIs(outputs[i], m.egb.outputs[i])
1009
1010    def test_pressure_drop_model_nlp(self):
1011        m = self._create_pressure_drop_model()
1012
1013        cons = [m.c_con, m.F_con, m.Pin_con, m.P2_con]
1014        inputs = [m.Pin, m.c, m.F]
1015        outputs = [m.P2, m.Pout]
1016
1017        nlp = PyomoNLPWithGreyBoxBlocks(m)
1018
1019        n_primals = len(inputs) + len(outputs)
1020        n_eq_con = len(cons) + len(outputs)
1021        self.assertEqual(nlp.n_primals(), n_primals)
1022        self.assertEqual(nlp.n_constraints(), n_eq_con)
1023
1024        constraint_names = [
1025                "c_con",
1026                "F_con",
1027                "Pin_con",
1028                "P2_con",
1029                "egb.output_constraints[P2]",
1030                "egb.output_constraints[Pout]",
1031                ]
1032        primals = inputs + outputs
1033        nlp_constraints = nlp.constraint_names()
1034        nlp_vars = nlp.primals_names()
1035
1036        con_idx_map = {}
1037        for name in constraint_names:
1038            # Quadratic scan to get constraint indices is not ideal.
1039            # Could this map be created while PyNLPwGBB is being constructed?
1040            con_idx_map[name] = nlp_constraints.index(name)
1041
1042        var_idx_map = ComponentMap()
1043        for var in primals:
1044            name = var.name
1045            var_idx_map[var] = nlp_vars.index(name)
1046
1047        incident_vars = {con.name: list(identify_variables(con.expr))
1048                for con in cons}
1049        incident_vars["egb.output_constraints[P2]"] = inputs + [outputs[0]]
1050        incident_vars["egb.output_constraints[Pout]"] = inputs + [outputs[1]]
1051
1052        expected_nonzeros = set()
1053        for con, varlist in incident_vars.items():
1054            i = con_idx_map[con]
1055            for var in varlist:
1056                j = var_idx_map[var]
1057                expected_nonzeros.add((i, j))
1058
1059        self.assertEqual(len(expected_nonzeros), nlp.nnz_jacobian())
1060
1061        jac = nlp.evaluate_jacobian()
1062        for i, j in zip(jac.row, jac.col):
1063            self.assertIn((i, j), expected_nonzeros)
1064
1065
1066if __name__ == '__main__':
1067    unittest.main()
1068