Fix and test issues with loading DR coefficients

Author: shermanjasonaf

pyomo/contrib/pyros/tests/test_grcs.py

@@ -2186,6 +2186,50 @@ def test_pyros_intersection_aux_vars(self):
         self.assertAlmostEqual(m.x1.value, 1)
         self.assertAlmostEqual(m.x2.value, 1)
 
+    @unittest.skipUnless(ipopt_available, "IPOPT is not available.")
+    def test_pyros_dr_interface_ordering(self):
+        """
+        Test PyROS DR interface returns coefficients in expected order.
+        """
+        m = ConcreteModel()
+        m.q1 = Param(initialize=0.5, mutable=True)
+        m.q2 = Param(initialize=0.5, mutable=True)
+        m.x1 = Var(bounds=[1, 1])
+        m.x2 = Var(bounds=[m.q1, m.q1])
+        x3_bound_expr = 3 + m.q1 * m.q2
+        m.x3 = Var(bounds=[x3_bound_expr, x3_bound_expr])
+        m.obj = Objective(expr=m.x1 + m.x2 + m.x3 + m.q2)
+        res = SolverFactory("pyros").solve(
+            model=m,
+            first_stage_variables=[],
+            second_stage_variables=[m.x1, m.x2, m.x3],
+            uncertain_params=[m.q1, m.q2],
+            uncertainty_set=BoxSet([[0, 1]] * 2),
+            local_solver="ipopt",
+            global_solver="ipopt",
+            decision_rule_order=2,
+        )
+        self.assertEqual(
+            res.pyros_termination_condition,
+            pyrosTerminationCondition.robust_feasible,
+        )
+        self.assertEqual(m.x1.value, 1)
+        self.assertAlmostEqual(m.x2.value, 0.5)  # nominal realization
+        self.assertAlmostEqual(m.x3.value, 3.25)
+        # nominal objective
+        self.assertAlmostEqual(res.final_objective_value, 5.25)
+        # test DR coefficients: easily inferred from the
+        # variable bound expressions
+        np.testing.assert_allclose(res.decision_rule_coeffs["static"], [1, 0, 3])
+        np.testing.assert_allclose(
+            res.decision_rule_coeffs["affine"],
+            [[0, 0], [1, 0], [0, 0]],
+        )
+        np.testing.assert_allclose(
+            res.decision_rule_coeffs["quadratic"],
+            [[[0, 0], [0, 0]]] * 2 + [[[0, 0.5], [0.5, 0]]],
+        )
+
 
 @unittest.skipUnless(ipopt_available, "IPOPT not available.")
 class TestPyROSSeparationPriorityOrder(unittest.TestCase):
@@ -2540,6 +2584,10 @@ def test_two_stage_discrete_set_rank2_affine_dr(self):
         self.assertAlmostEqual(res.final_objective_value, 2, places=4)
         self.assertAlmostEqual(m.x.value, 4, places=4)
         self.assertAlmostEqual(m.z.value, -2, places=4)
+        # optimal DR can be computed analytically
+        np.testing.assert_allclose(res.decision_rule_coeffs["static"], [-10/3])
+        np.testing.assert_allclose(res.decision_rule_coeffs["affine"], [[2/3]])
+        self.assertIsNone(res.decision_rule_coeffs["quadratic"])
 
     def test_two_stage_discrete_set_fullrank_affine_dr(self):
         m = self.build_two_stage_model()
@@ -2568,6 +2616,10 @@ def test_two_stage_discrete_set_fullrank_affine_dr(self):
         # variables must be 0
         self.assertAlmostEqual(m.x.value, 0, places=4)
         self.assertAlmostEqual(m.z.value, 0, places=4)
+        # optimal DR can be calculated analytically
+        np.testing.assert_allclose(res.decision_rule_coeffs["static"], [0])
+        np.testing.assert_allclose(res.decision_rule_coeffs["affine"], [[0]])
+        self.assertIsNone(res.decision_rule_coeffs["quadratic"])
 
 
 @unittest.skipUnless(ipopt_available, "IPOPT not available.")

pyomo/contrib/pyros/util.py

@@ -3251,7 +3251,7 @@ def load_final_solution(model_data, master_soln, original_user_var_partitioning)
     for orig_var, master_blk_var in zip(original_model_vars, master_soln_vars):
         orig_var.set_value(master_blk_var.value, skip_validation=True)
 
-    dr_coeffs_tuple = assemble_final_dr_coefficients(soln_master_blk)
+    dr_coeffs_tuple = assemble_final_dr_coefficients(soln_master_blk, config)
     dr_coeffs_dict = dict(
         static=dr_coeffs_tuple[0],
         affine=dr_coeffs_tuple[1] if config.decision_rule_order >= 1 else None,
@@ -3265,7 +3265,7 @@ def load_final_solution(model_data, master_soln, original_user_var_partitioning)
     return dr_coeffs_dict, worst_case_param_realization
 
 
-def assemble_final_dr_coefficients(working_model):
+def assemble_final_dr_coefficients(working_model, config):
     """
     Assemble final DR coefficients into matrices.
     """
@@ -3286,11 +3286,12 @@ def assemble_final_dr_coefficients(working_model):
     param_to_idx_map = ComponentMap(
         (param, idx) for idx, param in enumerate(uncertain_params_set)
     )
-    for ss_idx, ss_var in enumerate(second_stage_vars_set):
-        if ss_var not in effective_second_stage_vars_set:
-            static_coeffs[ss_idx] = ss_var.value
+    for ss_idx, ss_var in enumerate(config.second_stage_variables):
+        working_ss_var = working_model.user_model.find_component(ss_var)
+        if working_ss_var not in effective_second_stage_vars_set:
+            static_coeffs[ss_idx] = working_ss_var.value
         else:
-            indexed_dr_var = working_model.eff_ss_var_to_dr_var_map[ss_var]
+            indexed_dr_var = working_model.eff_ss_var_to_dr_var_map[working_ss_var]
             for dr_vdata in indexed_dr_var.values():
                 dr_monomial_params = working_model.dr_var_to_param_combo_map[dr_vdata]
                 if not dr_monomial_params: