Skip to content

Commit

Permalink
Substitute ParameterizedQuadaraticRepnVisitor for standard repn
Browse files Browse the repository at this point in the history
  • Loading branch information
shermanjasonaf committed Aug 20, 2024
1 parent 39fdfde commit 80c2dd6
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 42 deletions.
12 changes: 6 additions & 6 deletions pyomo/contrib/pyros/tests/test_preprocessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
75 changes: 39 additions & 36 deletions pyomo/contrib/pyros/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 "
Expand Down Expand Up @@ -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 "
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 80c2dd6

Please sign in to comment.