Merge pull request #3186 from Robbybp/scc-performance-2

Improve performance of `solve_strongly_connected_components` for models with named expressions

Author: mrmundt

Diff for: pyomo/contrib/incidence_analysis/scc_solver.py

@@ -18,13 +18,14 @@
     IncidenceGraphInterface,
     _generate_variables_in_constraints,
 )
+from pyomo.contrib.incidence_analysis.config import IncidenceMethod
 
 
 _log = logging.getLogger(__name__)
 
 
 def generate_strongly_connected_components(
-    constraints, variables=None, include_fixed=False
+    constraints, variables=None, include_fixed=False, igraph=None
 ):
     """Yield in order ``_BlockData`` that each contain the variables and
     constraints of a single diagonal block in a block lower triangularization
@@ -41,9 +42,12 @@ def generate_strongly_connected_components(
     variables: List of Pyomo variable data objects
         Variables that may participate in strongly connected components.
         If not provided, all variables in the constraints will be used.
-    include_fixed: Bool
+    include_fixed: Bool, optional
         Indicates whether fixed variables will be included when
         identifying variables in constraints.
+    igraph: IncidenceGraphInterface, optional
+        Incidence graph containing (at least) the provided constraints
+        and variables.
 
     Yields
     ------
@@ -55,11 +59,17 @@ def generate_strongly_connected_components(
     """
     if variables is None:
         variables = list(
-            _generate_variables_in_constraints(constraints, include_fixed=include_fixed)
+            _generate_variables_in_constraints(
+                constraints,
+                include_fixed=include_fixed,
+                method=IncidenceMethod.ampl_repn,
+            )
         )
 
     assert len(variables) == len(constraints)
-    igraph = IncidenceGraphInterface()
+    if igraph is None:
+        igraph = IncidenceGraphInterface()
+
     var_blocks, con_blocks = igraph.block_triangularize(
         variables=variables, constraints=constraints
     )
@@ -73,7 +83,7 @@ def generate_strongly_connected_components(
 
 
 def solve_strongly_connected_components(
-    block, solver=None, solve_kwds=None, calc_var_kwds=None
+    block, *, solver=None, solve_kwds=None, use_calc_var=True, calc_var_kwds=None
 ):
     """Solve a square system of variables and equality constraints by
     solving strongly connected components individually.
@@ -98,6 +108,9 @@ def solve_strongly_connected_components(
         a solve method.
     solve_kwds: Dictionary
         Keyword arguments for the solver's solve method
+    use_calc_var: Bool
+        Whether to use ``calculate_variable_from_constraint`` for one-by-one
+        square system solves
     calc_var_kwds: Dictionary
         Keyword arguments for calculate_variable_from_constraint
 
@@ -112,23 +125,28 @@ def solve_strongly_connected_components(
         calc_var_kwds = {}
 
     igraph = IncidenceGraphInterface(
-        block, active=True, include_fixed=False, include_inequality=False
+        block,
+        active=True,
+        include_fixed=False,
+        include_inequality=False,
+        method=IncidenceMethod.ampl_repn,
     )
     constraints = igraph.constraints
     variables = igraph.variables
 
     res_list = []
     log_blocks = _log.isEnabledFor(logging.DEBUG)
-    for scc, inputs in generate_strongly_connected_components(constraints, variables):
-        with TemporarySubsystemManager(to_fix=inputs):
+    for scc, inputs in generate_strongly_connected_components(
+        constraints, variables, igraph=igraph
+    ):
+        with TemporarySubsystemManager(to_fix=inputs, remove_bounds_on_fix=True):
             N = len(scc.vars)
-            if N == 1:
+            if N == 1 and use_calc_var:
                 if log_blocks:
                     _log.debug(f"Solving 1x1 block: {scc.cons[0].name}.")
                 results = calculate_variable_from_constraint(
                     scc.vars[0], scc.cons[0], **calc_var_kwds
                 )
-                res_list.append(results)
             else:
                 if solver is None:
                     var_names = [var.name for var in scc.vars.values()][:10]
@@ -142,5 +160,5 @@ def solve_strongly_connected_components(
                 if log_blocks:
                     _log.debug(f"Solving {N}x{N} block.")
                 results = solver.solve(scc, **solve_kwds)
-                res_list.append(results)
+            res_list.append(results)
     return res_list

Diff for: pyomo/util/subsystems.py

@@ -17,18 +17,35 @@
 
 from pyomo.core.base.constraint import Constraint
 from pyomo.core.base.expression import Expression
+from pyomo.core.base.objective import Objective
 from pyomo.core.base.external import ExternalFunction
 from pyomo.core.expr.visitor import StreamBasedExpressionVisitor
 from pyomo.core.expr.numeric_expr import ExternalFunctionExpression
-from pyomo.core.expr.numvalue import native_types
+from pyomo.core.expr.numvalue import native_types, NumericValue
 
 
 class _ExternalFunctionVisitor(StreamBasedExpressionVisitor):
+    def __init__(self, descend_into_named_expressions=True):
+        super().__init__()
+        self._descend_into_named_expressions = descend_into_named_expressions
+        self.named_expressions = []
+
     def initializeWalker(self, expr):
         self._functions = []
         self._seen = set()
         return True, None
 
+    def beforeChild(self, parent, child, index):
+        if child.__class__ in native_types:
+            return False, None
+        elif (
+            not self._descend_into_named_expressions
+            and child.is_named_expression_type()
+        ):
+            self.named_expressions.append(child)
+            return False, None
+        return True, None
+
     def exitNode(self, node, data):
         if type(node) is ExternalFunctionExpression:
             if id(node) not in self._seen:
@@ -38,26 +55,35 @@ def exitNode(self, node, data):
     def finalizeResult(self, result):
         return self._functions
 
-    def enterNode(self, node):
-        pass
-
-    def acceptChildResult(self, node, data, child_result, child_idx):
-        pass
-
-    def acceptChildResult(self, node, data, child_result, child_idx):
-        if child_result.__class__ in native_types:
-            return False, None
-        return child_result.is_expression_type(), None
-
 
 def identify_external_functions(expr):
     yield from _ExternalFunctionVisitor().walk_expression(expr)
 
 
 def add_local_external_functions(block):
     ef_exprs = []
-    for comp in block.component_data_objects((Constraint, Expression), active=True):
-        ef_exprs.extend(identify_external_functions(comp.expr))
+    named_expressions = []
+    visitor = _ExternalFunctionVisitor(descend_into_named_expressions=False)
+    for comp in block.component_data_objects(
+        (Constraint, Expression, Objective), active=True
+    ):
+        ef_exprs.extend(visitor.walk_expression(comp.expr))
+    named_expr_set = ComponentSet(visitor.named_expressions)
+    # List of unique named expressions
+    named_expressions = list(named_expr_set)
+    while named_expressions:
+        expr = named_expressions.pop()
+        # Clear named expression cache so we don't re-check named expressions
+        # we've seen before.
+        visitor.named_expressions.clear()
+        ef_exprs.extend(visitor.walk_expression(expr))
+        # Only add to the stack named expressions that we have
+        # not encountered yet.
+        for local_expr in visitor.named_expressions:
+            if local_expr not in named_expr_set:
+                named_expressions.append(local_expr)
+                named_expr_set.add(local_expr)
+
     unique_functions = []
     fcn_set = set()
     for expr in ef_exprs:
@@ -148,7 +174,14 @@ class TemporarySubsystemManager(object):
 
     """
 
-    def __init__(self, to_fix=None, to_deactivate=None, to_reset=None, to_unfix=None):
+    def __init__(
+        self,
+        to_fix=None,
+        to_deactivate=None,
+        to_reset=None,
+        to_unfix=None,
+        remove_bounds_on_fix=False,
+    ):
         """
         Arguments
         ---------
@@ -168,6 +201,8 @@ def __init__(self, to_fix=None, to_deactivate=None, to_reset=None, to_unfix=None
             List of var data objects to be temporarily unfixed. These are
             restored to their original status on exit from this object's
             context manager.
+        remove_bounds_on_fix: Bool
+            Whether bounds should be removed temporarily for fixed variables
 
         """
         if to_fix is None:
@@ -194,6 +229,8 @@ def __init__(self, to_fix=None, to_deactivate=None, to_reset=None, to_unfix=None
         self._con_was_active = None
         self._comp_original_value = None
         self._var_was_unfixed = None
+        self._remove_bounds_on_fix = remove_bounds_on_fix
+        self._fixed_var_bounds = None
 
     def __enter__(self):
         to_fix = self._vars_to_fix
@@ -203,8 +240,13 @@ def __enter__(self):
         self._var_was_fixed = [(var, var.fixed) for var in to_fix + to_unfix]
         self._con_was_active = [(con, con.active) for con in to_deactivate]
         self._comp_original_value = [(comp, comp.value) for comp in to_set]
+        self._fixed_var_bounds = [(var.lb, var.ub) for var in to_fix]
 
         for var in self._vars_to_fix:
+            if self._remove_bounds_on_fix:
+                # TODO: Potentially override var.domain as well?
+                var.setlb(None)
+                var.setub(None)
             var.fix()
 
         for con in self._cons_to_deactivate:
@@ -223,6 +265,11 @@ def __exit__(self, ex_type, ex_val, ex_bt):
                 var.fix()
             else:
                 var.unfix()
+        if self._remove_bounds_on_fix:
+            for var, (lb, ub) in zip(self._vars_to_fix, self._fixed_var_bounds):
+                var.setlb(lb)
+                var.setub(ub)
+
         for con, was_active in self._con_was_active:
             if was_active:
                 con.activate()

Diff for: pyomo/util/tests/test_subsystems.py

@@ -292,17 +292,29 @@ def test_generate_dont_fix_inputs_with_fixed_var(self):
         self.assertFalse(m.v3.fixed)
         self.assertTrue(m.v4.fixed)
 
-    def _make_model_with_external_functions(self):
+    def _make_model_with_external_functions(self, named_expressions=False):
         m = pyo.ConcreteModel()
         gsl = find_GSL()
         m.bessel = pyo.ExternalFunction(library=gsl, function="gsl_sf_bessel_J0")
         m.fermi = pyo.ExternalFunction(library=gsl, function="gsl_sf_fermi_dirac_m1")
         m.v1 = pyo.Var(initialize=1.0)
         m.v2 = pyo.Var(initialize=2.0)
         m.v3 = pyo.Var(initialize=3.0)
+        if named_expressions:
+            m.subexpr = pyo.Expression(pyo.PositiveIntegers)
+            m.subexpr[1] = 2 * m.fermi(m.v1)
+            m.subexpr[2] = m.bessel(m.v1) - m.bessel(m.v2)
+            m.subexpr[3] = m.subexpr[2] + m.v3**2
+            subexpr1 = m.subexpr[1]
+            subexpr2 = m.subexpr[2]
+            subexpr3 = m.subexpr[3]
+        else:
+            subexpr1 = 2 * m.fermi(m.v1)
+            subexpr2 = m.bessel(m.v1) - m.bessel(m.v2)
+            subexpr3 = subexpr2 + m.v3**2
         m.con1 = pyo.Constraint(expr=m.v1 == 0.5)
-        m.con2 = pyo.Constraint(expr=2 * m.fermi(m.v1) + m.v2**2 - m.v3 == 1.0)
-        m.con3 = pyo.Constraint(expr=m.bessel(m.v1) - m.bessel(m.v2) + m.v3**2 == 2.0)
+        m.con2 = pyo.Constraint(expr=subexpr1 + m.v2**2 - m.v3 == 1.0)
+        m.con3 = pyo.Constraint(expr=subexpr3 == 2.0)
         return m
 
     @unittest.skipUnless(find_GSL(), "Could not find the AMPL GSL library")
@@ -329,6 +341,15 @@ def test_identify_external_functions(self):
         pred_fcn_data = {(gsl, "gsl_sf_bessel_J0"), (gsl, "gsl_sf_fermi_dirac_m1")}
         self.assertEqual(fcn_data, pred_fcn_data)
 
+    @unittest.skipUnless(find_GSL(), "Could not find the AMPL GSL library")
+    def test_local_external_functions_with_named_expressions(self):
+        m = self._make_model_with_external_functions(named_expressions=True)
+        variables = list(m.component_data_objects(pyo.Var))
+        constraints = list(m.component_data_objects(pyo.Constraint, active=True))
+        b = create_subsystem_block(constraints, variables)
+        self.assertTrue(isinstance(b._gsl_sf_bessel_J0, pyo.ExternalFunction))
+        self.assertTrue(isinstance(b._gsl_sf_fermi_dirac_m1, pyo.ExternalFunction))
+
     def _solve_ef_model_with_ipopt(self):
         m = self._make_model_with_external_functions()
         ipopt = pyo.SolverFactory("ipopt")
@@ -362,6 +383,33 @@ def test_with_external_function(self):
         self.assertAlmostEqual(m.v2.value, m_full.v2.value)
         self.assertAlmostEqual(m.v3.value, m_full.v3.value)
 
+    @unittest.skipUnless(find_GSL(), "Could not find the AMPL GSL library")
+    @unittest.skipUnless(
+        pyo.SolverFactory("ipopt").available(), "ipopt is not available"
+    )
+    def test_with_external_function_in_named_expression(self):
+        m = self._make_model_with_external_functions(named_expressions=True)
+        subsystem = ([m.con2, m.con3], [m.v2, m.v3])
+
+        m.v1.set_value(0.5)
+        block = create_subsystem_block(*subsystem)
+        ipopt = pyo.SolverFactory("ipopt")
+        with TemporarySubsystemManager(to_fix=list(block.input_vars.values())):
+            ipopt.solve(block)
+
+        # Correct values obtained by solving with Ipopt directly
+        # in another script.
+        self.assertEqual(m.v1.value, 0.5)
+        self.assertFalse(m.v1.fixed)
+        self.assertAlmostEqual(m.v2.value, 1.04816, delta=1e-5)
+        self.assertAlmostEqual(m.v3.value, 1.34356, delta=1e-5)
+
+        # Result obtained by solving the full system
+        m_full = self._solve_ef_model_with_ipopt()
+        self.assertAlmostEqual(m.v1.value, m_full.v1.value)
+        self.assertAlmostEqual(m.v2.value, m_full.v2.value)
+        self.assertAlmostEqual(m.v3.value, m_full.v3.value)
+
     @unittest.skipUnless(find_GSL(), "Could not find the AMPL GSL library")
     def test_external_function_with_potential_name_collision(self):
         m = self._make_model_with_external_functions()