Enable nonadjustable multistage DRs

Author: shermanjasonaf

pyomo/contrib/pyros/config.py

@@ -154,6 +154,35 @@ def uncertain_param_data_validator(uncertain_obj):
             )
 
 
+class MultistageInputDataStandardizer(object):
+    """
+    Standardizer for modeling objects representing multi-stage
+    variables or uncertain parameters. Standard form is a
+    list of lists of `cdatatype` objects.
+    """
+    def __init__(self, ctype, cdatatype):
+        self.ctype = ctype
+        self.cdatatype = cdatatype
+
+    def __call__(self, obj):
+        if isinstance(obj, self.ctype):
+            return [list(obj.values())]
+        if isinstance(obj, self.cdatatype):
+            return [[obj]]
+        if isinstance(obj, (list, tuple)):
+            base_std = InputDataStandardizer(self.ctype, self.cdatatype)
+            if all(isinstance(itm, (list, tuple)) for itm in obj):
+                ans = list()
+                for item in obj:
+                    ans.append(base_std(item))
+            else:
+                ans = [base_std(obj)]
+        else:
+            raise TypeError("Not supported")
+
+        return ans
+
+
 class InputDataStandardizer:
     """
     Domain validator for an object that is castable to
@@ -544,6 +573,26 @@ def pyros_config():
         ),
     )
 
+    # for multi-stage
+    CONFIG.declare("nested_second_stage_variables", ConfigValue(
+        default=[],
+        domain=MultistageInputDataStandardizer(Var, VarData),
+        description=(
+            "A two-dimensional list for extending the second-stage "
+            "variables to multi-stage variables. Each list specifies "
+            "the variables adjusted in each stage."
+        )
+    ))
+    CONFIG.declare("nested_uncertain_params", ConfigValue(
+        default=[],
+        domain=MultistageInputDataStandardizer(Param, ParamData),
+        description=(
+            "A two-dimensional list for extending the uncertain parameters "
+            "variables to multi-stage context. Each list specifies "
+            "the parameters realized by the corresponding stage."
+        )
+    ))
+
     # ================================================
     # === Required User Inputs
     # ================================================

pyomo/contrib/pyros/pyros.py

@@ -424,20 +424,45 @@ def solve(
             config, user_var_partitioning = self._resolve_and_validate_pyros_args(
                 model, **kwds
             )
+
+            # for multi-stage: confirm there's a match
+            config.nested_second_stage_variables = second_stage_variables
+            config.nested_uncertain_params = uncertain_params
+            matching_stages = (
+                len(config.nested_uncertain_params)
+                == len(config.nested_second_stage_variables)
+            )
+            if config.second_stage_variables and not matching_stages:
+                raise ValueError(
+                    "Number of stages mismatch: "
+                    f"there are {len(config.nested_second_stage_variables)} "
+                    "stages of second-stage variables, but "
+                    f"{len(config.nested_uncertain_params)} "
+                    "stages of uncertain parameters."
+                )
+
             self._log_config_user_values(
                 logger=config.progress_logger,
                 config=config,
                 exclude_options=(
                     self._DEFAULT_CONFIG_USER_OPTIONS
                     + ["nominal_uncertain_param_vals"]
                     * (not nominal_param_vals_in_kwds)
+                    # for multi-stage
+                    + ["nested_second_stage_variables"]
+                    + ["nested_uncertain_params"]
                 ),
                 level=logging.INFO,
             )
             self._log_config(
                 logger=config.progress_logger,
                 config=config,
-                exclude_options=None,
+                exclude_options=(
+                    self._DEFAULT_CONFIG_USER_OPTIONS
+                    # for multi-stage
+                    + ["nested_second_stage_variables"]
+                    + ["nested_uncertain_params"]
+                ),
                 level=logging.DEBUG,
             )
             model_data.config = config

pyomo/contrib/pyros/util.py

@@ -1933,6 +1933,11 @@ def setup_working_model(model_data, user_var_partitioning):
     orig_temp_util_block.user_var_partitioning = VariablePartitioning(
         **user_var_partitioning._asdict()
     )
+    # for multi-stage
+    orig_temp_util_block.nested_second_stage_variables = (
+        config.nested_second_stage_variables
+    )
+    orig_temp_util_block.nested_uncertain_params = config.nested_uncertain_params
 
     # now set up working model
     model_data.working_model = working_model = ConcreteModel()
@@ -1961,6 +1966,14 @@ def setup_working_model(model_data, user_var_partitioning):
         **working_temp_util_block.user_var_partitioning._asdict()
     )
 
+    # for multi-stage
+    working_model.nested_second_stage_variables = (
+        working_temp_util_block.nested_second_stage_variables
+    )
+    working_model.nested_uncertain_params = (
+        working_temp_util_block.nested_uncertain_params
+    )
+
     # we are done with the util blocks
     delattr(original_model, temp_util_block_attr_name)
     delattr(working_model.user_model, temp_util_block_attr_name)
@@ -2017,6 +2030,22 @@ def setup_working_model(model_data, user_var_partitioning):
         else:
             working_model.original_active_inequality_cons.append(con)
 
+    # for multi-stage: map adjustable vars to stages
+    # and also uncertain params (assuming all were originally Params)
+    working_model.adj_var_to_stage_map = ComponentMap()
+    working_model.param_to_stage_map = ComponentMap()
+    ss_vars_param_zip = zip(
+        working_model.nested_second_stage_variables,
+        working_model.nested_uncertain_params,
+    )
+    for stage_idx, (varlist, paramlist) in enumerate(ss_vars_param_zip):
+        working_model.adj_var_to_stage_map.update(
+            (var, stage_idx) for var in varlist
+        )
+        working_model.param_to_stage_map.update(
+            (param, stage_idx) for param in paramlist
+        )
+
 
 def standardize_inequality_constraints(model_data):
     """
@@ -3140,6 +3169,7 @@ def add_decision_rule_constraints(model_data):
     indexed_dr_var_list = model_data.working_model.first_stage.decision_rule_vars
     uncertain_params = model_data.working_model.effective_uncertain_params
     degree = config.decision_rule_order
+    param_to_stage_map = model_data.working_model.param_to_stage_map
 
     model_data.working_model.second_stage.decision_rule_eqns = decision_rule_eqns = (
         Constraint(range(len(effective_second_stage_vars)))
@@ -3162,6 +3192,9 @@ def add_decision_rule_constraints(model_data):
         ComponentMap()
     )
 
+    # for multi-stage
+    model_data.working_model.dr_var_to_stage_map = dr_var_to_stage_map = ComponentMap()
+
     # set up uncertain parameter combinations for
     # construction of the monomials of the DR expressions
     monomial_param_combos = []
@@ -3192,6 +3225,10 @@ def add_decision_rule_constraints(model_data):
             # associated monomial with respect to the uncertain params
             dr_var_to_exponent_map[dr_var] = len(param_combo)
             dr_var_to_param_combo_map[dr_var] = param_combo
+            # for multi-stage
+            dr_var_to_stage_map[dr_var] = max(
+                param_to_stage_map[param] for param in param_combo
+            ) if param_combo else 0
 
         # declare constraint on model
         decision_rule_eqns[idx] = dr_expression - eff_ss_var == 0
@@ -3213,14 +3250,25 @@ def enforce_dr_degree(working_blk, config, degree):
     degree : int
         Degree of the DR polynomials that is to be enforced.
     """
-    for indexed_dr_var in working_blk.first_stage.decision_rule_vars:
+    # account for multi-stage setup
+    dr_var_to_stage_map = working_blk.dr_var_to_stage_map
+    adj_var_to_stage_map = working_blk.adj_var_to_stage_map
+    for eff_ss_var, indexed_dr_var in working_blk.eff_ss_var_to_dr_var_map.items():
         for dr_var in indexed_dr_var.values():
             dr_var_degree = working_blk.dr_var_to_exponent_map[dr_var]
-            if dr_var_degree > degree:
+            dr_var_stage = dr_var_to_stage_map[dr_var]
+            ss_var_stage = adj_var_to_stage_map[eff_ss_var]
+            if (dr_var_degree > degree) or (dr_var_stage > ss_var_stage):
                 dr_var.fix(0)
             else:
                 dr_var.unfix()
 
+            # config.progress_logger.debug(
+            #     f"{dr_var.name}, SS var {eff_ss_var.name}, "
+            #     f"{dr_var_degree=}, {ss_var_stage=}, {dr_var_stage=}, "
+            #     f"{dr_var.fixed=}, DR > SS stage {dr_var_stage > ss_var_stage}"
+            # )
+
 
 def load_final_solution(model_data, master_soln, original_user_var_partitioning):
     """