diff --git a/pyomo/contrib/pyros/tests/test_preprocessor.py b/pyomo/contrib/pyros/tests/test_preprocessor.py index 5268ebc9dfa..144c3644862 100644 --- a/pyomo/contrib/pyros/tests/test_preprocessor.py +++ b/pyomo/contrib/pyros/tests/test_preprocessor.py @@ -1950,17 +1950,17 @@ def test_coefficient_matching_correct_constraints_added(self): assertExpressionsEqual( self, first_stage_eq_cons["coeff_matching_eq_con_coeff_1"].expr, - 2.5 + m.x1 + (-5) * (m.x1 * m.x2) + m.x1**3 == 0, + m.x1 ** 3 + 0.5 + 5 * m.x1 * m.x2 * (-1) + (-1) * (m.x1 + 2) * (-1) == 0, ) assertExpressionsEqual( self, first_stage_eq_cons["coeff_matching_eq_con_coeff_2"].expr, - (-1) + m.x2 == 0, + m.x2 - 1 == 0, ) assertExpressionsEqual( self, first_stage_eq_cons["coeff_matching_eq_con_2_coeff_1"].expr, - (-1) + m.x2 == 0, + m.x2 - 1 == 0, ) def test_reformulate_nonlinear_state_var_independent_eq_con(self): @@ -2050,7 +2050,7 @@ def test_reformulate_nonlinear_state_var_independent_eq_con(self): assertExpressionsEqual( self, wm.first_stage.equality_cons["coeff_matching_eq_con_2_coeff_1"].expr, - (-1) + m.x1 == 0, + m.x1 - 1 == 0, ) # separation priorities were also updated @@ -2578,7 +2578,7 @@ def test_preprocessor_coefficient_matching( assertExpressionsEqual( self, fs_eqs["coeff_matching_var_z5_uncertain_eq_bound_con_coeff_1"].expr, - -1 + fs.decision_rule_vars[1][1] == 0, + fs.decision_rule_vars[1][1] - 1 == 0, ) assertExpressionsEqual( self, @@ -2612,7 +2612,7 @@ def test_preprocessor_coefficient_matching( assertExpressionsEqual( self, fs_eqs["coeff_matching_var_z5_uncertain_eq_bound_con_coeff_1"].expr, - -1 + fs.decision_rule_vars[1][1] == 0, + fs.decision_rule_vars[1][1] - 1 == 0, ) assertExpressionsEqual( self, diff --git a/pyomo/contrib/pyros/util.py b/pyomo/contrib/pyros/util.py index 7bc81a97033..58595ff3e35 100644 --- a/pyomo/contrib/pyros/util.py +++ b/pyomo/contrib/pyros/util.py @@ -52,7 +52,7 @@ ) from pyomo.core.util import prod from pyomo.opt import SolverFactory, TerminationCondition as tc -from pyomo.repn.standard_repn import generate_standard_repn +from pyomo.repn.parameterized_quadratic import ParameterizedQuadraticRepnVisitor import pyomo.repn.plugins.nl_writer as pyomo_nl_writer import pyomo.repn.ampl as pyomo_ampl_repn from pyomo.util.vars_from_expressions import get_vars_from_components @@ -989,6 +989,25 @@ def preprocess(self, user_var_partitioning): return preprocess_model_data(self, user_var_partitioning) +def setup_quadratic_expression_visitor( + wrt, + subexpression_cache=None, + var_map=None, + var_order=None, + sorter=None, +): + """Setup a parameterized quadratic expression walker.""" + visitor = ParameterizedQuadraticRepnVisitor( + subexpression_cache={} if subexpression_cache is None else subexpression_cache, + var_map={} if var_map is None else var_map, + var_order={} if var_order is None else var_order, + sorter=sorter, + wrt=wrt, + ) + visitor.expand_nonlinear_products = True + return visitor + + class BoundType: """ Indicator for whether a bound on a variable/constraint @@ -1299,23 +1318,16 @@ def get_effective_var_partitioning(model_data): # tolerance. if len(adj_vars_in_con) == 1: adj_var_in_con = next(iter(adj_vars_in_con)) - ccon_expr_repn = generate_standard_repn( - expr=ccon.body - ccon.upper, quadratic=False, compute_values=True + visitor = setup_quadratic_expression_visitor(wrt=[]) + ccon_expr_repn = visitor.walk_expression(expr=ccon.body - ccon.upper) + adj_var_appears_linearly = ( + adj_var_in_con + not in ComponentSet(identify_variables(ccon_expr_repn.nonlinear)) + and id(adj_var_in_con) in ComponentSet(ccon_expr_repn.linear) ) - adj_var_appears_linearly = adj_var_in_con not in ComponentSet( - ccon_expr_repn.nonlinear_vars - ) and adj_var_in_con in ComponentSet(ccon_expr_repn.linear_vars) if adj_var_appears_linearly: - # get coefficient by summation just in case - # standard repn does not simplify completely - var_linear_coeff = sum( - lcoeff - for lvar, lcoeff in zip( - ccon_expr_repn.linear_vars, ccon_expr_repn.linear_coefs - ) - if lvar is adj_var_in_con - ) - if abs(var_linear_coeff) > PRETRIANGULAR_VAR_COEFF_TOL: + adj_var_linear_coeff = ccon_expr_repn.linear[id(adj_var_in_con)] + if abs(adj_var_linear_coeff) > PRETRIANGULAR_VAR_COEFF_TOL: new_pretriangular_con_var_map[ccon] = adj_var_in_con config.progress_logger.debug( f" The variable {adj_var_in_con.name!r} is " @@ -2197,18 +2209,10 @@ def reformulate_state_var_independent_eq_cons(model_data): # uncertain parameters only. thus, only the proxy # variables for the uncertain parameters are unfixed # during the analysis - for var in originally_unfixed_vars: - var.fix() - expr_repn = generate_standard_repn( - expr=con_expr_after_all_substitutions, compute_values=False - ) - - # ensure state of every variable remains unchanged - # when done - for var in originally_unfixed_vars: - var.unfix() + visitor = setup_quadratic_expression_visitor(wrt=originally_unfixed_vars) + expr_repn = visitor.walk_expression(con_expr_after_all_substitutions) - if expr_repn.nonlinear_expr is not None: + if expr_repn.nonlinear is not None: config.progress_logger.debug( f"Equality constraint {con.name!r} " "is state-variable independent, but cannot be written " @@ -2241,26 +2245,25 @@ def reformulate_state_var_independent_eq_cons(model_data): else: polynomial_repn_coeffs = ( [expr_repn.constant] - + list(expr_repn.linear_coefs) - + list(expr_repn.quadratic_coefs) + + list(expr_repn.linear.values()) + + ( + [] if expr_repn.quadratic is None + else list(expr_repn.quadratic.values()) + ) ) for coeff_idx, coeff_expr in enumerate(polynomial_repn_coeffs): - simplified_coeff_expr = generate_standard_repn( - expr=coeff_expr, compute_values=True - ).to_expression() - # for robust satisfaction of the original equality # constraint, all polynomial coefficients must be # equal to zero. so for each coefficient, # we either check for trivial robust # feasibility/infeasibility, or add a constraint # restricting the coefficient expression to value 0 - if isinstance(simplified_coeff_expr, tuple(native_types)): + if isinstance(coeff_expr, tuple(native_types)): # coefficient is a constant; # check value to determine # trivial feasibility/infeasibility robust_infeasible = not math.isclose( - a=simplified_coeff_expr, + a=coeff_expr, b=0, rel_tol=COEFF_MATCH_REL_TOL, abs_tol=COEFF_MATCH_ABS_TOL, @@ -2288,7 +2291,7 @@ def reformulate_state_var_independent_eq_cons(model_data): # and DR variables. add matching constraint new_con_name = f"coeff_matching_{con_idx}_coeff_{coeff_idx}" working_model.first_stage.equality_cons[new_con_name] = ( - simplified_coeff_expr == 0 + coeff_expr == 0 ) new_con = working_model.first_stage.equality_cons[new_con_name] coefficient_matching_cons.append(new_con)