Repository: Pyomo/pyomo

Issue:
Add `AMPLRepn` option to `get_incident_variables`
## Summary

`get_incident_variables` doesn't currently have any option that will greedily remove from subexpressions all the terms of the form: `p*x` where `p == 0`. This means that if such a term appears in a nonlinear subexpression, `x` will be considered "incident to" the expression, despite the fact that, for most purposes, we don't want to consider it.

Options to correct this behavior include using `AMPLRepn` and using `generate_standard_repn(expr, compute_values=True)`. Of these, I think I lean towards `AMPLRepn` as it is probably faster, and we already have an option for `generate_standard_repn(expr, compute_values=False)`. This option is also valuable as it will give the same incidence structure that is used in the nl writer.


Discussion:

Resolving patch:
diff --git a/pyomo/contrib/incidence_analysis/config.py b/pyomo/contrib/incidence_analysis/config.py
--- a/pyomo/contrib/incidence_analysis/config.py
+++ b/pyomo/contrib/incidence_analysis/config.py
@@ -13,6 +13,9 @@
 
 import enum
 from pyomo.common.config import ConfigDict, ConfigValue, InEnum
+from pyomo.common.modeling import NOTSET
+from pyomo.repn.plugins.nl_writer import AMPLRepnVisitor, text_nl_template
+from pyomo.repn.util import FileDeterminism, FileDeterminism_to_SortComponents
 
 
 class IncidenceMethod(enum.Enum):
@@ -24,6 +27,14 @@ class IncidenceMethod(enum.Enum):
     standard_repn = 1
     """Use ``pyomo.repn.standard_repn.generate_standard_repn``"""
 
+    standard_repn_compute_values = 2
+    """Use ``pyomo.repn.standard_repn.generate_standard_repn`` with
+    ``compute_values=True``
+    """
+
+    ampl_repn = 3
+    """Use ``pyomo.repn.plugins.nl_writer.AMPLRepnVisitor``"""
+
 
 _include_fixed = ConfigValue(
     default=False,
@@ -54,6 +65,21 @@ class IncidenceMethod(enum.Enum):
 )
 
 
+def _amplrepnvisitor_validator(visitor):
+    if not isinstance(visitor, AMPLRepnVisitor):
+        raise TypeError(
+            "'visitor' config argument should be an instance of AMPLRepnVisitor"
+        )
+    return visitor
+
+
+_ampl_repn_visitor = ConfigValue(
+    default=None,
+    domain=_amplrepnvisitor_validator,
+    description="Visitor used to generate AMPLRepn of each constraint",
+)
+
+
 IncidenceConfig = ConfigDict()
 """Options for incidence graph generation
 
@@ -63,6 +89,9 @@ class IncidenceMethod(enum.Enum):
   should be included.
 - ``method`` -- Method used to identify incident variables. Must be a value of the
   ``IncidenceMethod`` enum.
+- ``_ampl_repn_visitor`` -- Expression visitor used to generate ``AMPLRepn`` of each
+  constraint. Must be an instance of ``AMPLRepnVisitor``. *This option is constructed
+  automatically when needed and should not be set by users!*
 
 """
 
@@ -74,3 +103,46 @@ class IncidenceMethod(enum.Enum):
 
 
 IncidenceConfig.declare("method", _method)
+
+
+IncidenceConfig.declare("_ampl_repn_visitor", _ampl_repn_visitor)
+
+
+def get_config_from_kwds(**kwds):
+    """Get an instance of IncidenceConfig from provided keyword arguments.
+
+    If the ``method`` argument is ``IncidenceMethod.ampl_repn`` and no
+    ``AMPLRepnVisitor`` has been provided, a new ``AMPLRepnVisitor`` is
+    constructed. This function should generally be used by callers such
+    as ``IncidenceGraphInterface`` to ensure that a visitor is created then
+    re-used when calling ``get_incident_variables`` in a loop.
+
+    """
+    if (
+        kwds.get("method", None) is IncidenceMethod.ampl_repn
+        and kwds.get("_ampl_repn_visitor", None) is None
+    ):
+        subexpression_cache = {}
+        subexpression_order = []
+        external_functions = {}
+        var_map = {}
+        used_named_expressions = set()
+        symbolic_solver_labels = False
+        # TODO: Explore potential performance benefit of exporting defined variables.
+        # This likely only shows up if we can preserve the subexpression cache across
+        # multiple constraint expressions.
+        export_defined_variables = False
+        sorter = FileDeterminism_to_SortComponents(FileDeterminism.ORDERED)
+        amplvisitor = AMPLRepnVisitor(
+            text_nl_template,
+            subexpression_cache,
+            subexpression_order,
+            external_functions,
+            var_map,
+            used_named_expressions,
+            symbolic_solver_labels,
+            export_defined_variables,
+            sorter,
+        )
+        kwds["_ampl_repn_visitor"] = amplvisitor
+    return IncidenceConfig(kwds)
diff --git a/pyomo/contrib/incidence_analysis/incidence.py b/pyomo/contrib/incidence_analysis/incidence.py
--- a/pyomo/contrib/incidence_analysis/incidence.py
+++ b/pyomo/contrib/incidence_analysis/incidence.py
@@ -17,7 +17,11 @@
 from pyomo.core.expr.numvalue import value as pyo_value
 from pyomo.repn import generate_standard_repn
 from pyomo.util.subsystems import TemporarySubsystemManager
-from pyomo.contrib.incidence_analysis.config import IncidenceMethod, IncidenceConfig
+from pyomo.repn.plugins.nl_writer import AMPLRepn
+from pyomo.contrib.incidence_analysis.config import (
+    IncidenceMethod,
+    get_config_from_kwds,
+)
 
 
 #
@@ -29,7 +33,9 @@ def _get_incident_via_identify_variables(expr, include_fixed):
     return list(identify_variables(expr, include_fixed=include_fixed))
 
 
-def _get_incident_via_standard_repn(expr, include_fixed, linear_only):
+def _get_incident_via_standard_repn(
+    expr, include_fixed, linear_only, compute_values=False
+):
     if include_fixed:
         to_unfix = [
             var for var in identify_variables(expr, include_fixed=True) if var.fixed
@@ -39,7 +45,9 @@ def _get_incident_via_standard_repn(expr, include_fixed, linear_only):
         context = nullcontext()
 
     with context:
-        repn = generate_standard_repn(expr, compute_values=False, quadratic=False)
+        repn = generate_standard_repn(
+            expr, compute_values=compute_values, quadratic=False
+        )
 
     linear_vars = []
     # Check coefficients to make sure we don't include linear variables with
@@ -74,6 +82,36 @@ def _get_incident_via_standard_repn(expr, include_fixed, linear_only):
         return unique_variables
 
 
+def _get_incident_via_ampl_repn(expr, linear_only, visitor):
+    var_map = visitor.var_map
+    orig_activevisitor = AMPLRepn.ActiveVisitor
+    AMPLRepn.ActiveVisitor = visitor
+    try:
+        repn = visitor.walk_expression((expr, None, 0, 1.0))
+    finally:
+        AMPLRepn.ActiveVisitor = orig_activevisitor
+
+    nonlinear_var_ids = [] if repn.nonlinear is None else repn.nonlinear[1]
+    nonlinear_var_id_set = set()
+    unique_nonlinear_var_ids = []
+    for v_id in nonlinear_var_ids:
+        if v_id not in nonlinear_var_id_set:
+            nonlinear_var_id_set.add(v_id)
+            unique_nonlinear_var_ids.append(v_id)
+
+    nonlinear_vars = [var_map[v_id] for v_id in unique_nonlinear_var_ids]
+    linear_only_vars = [
+        var_map[v_id]
+        for v_id, coef in repn.linear.items()
+        if coef != 0.0 and v_id not in nonlinear_var_id_set
+    ]
+    if linear_only:
+        return linear_only_vars
+    else:
+        variables = linear_only_vars + nonlinear_vars
+        return variables
+
+
 def get_incident_variables(expr, **kwds):
     """Get variables that participate in an expression
 
@@ -112,21 +150,38 @@ def get_incident_variables(expr, **kwds):
        ['x[1]', 'x[2]']
 
     """
-    config = IncidenceConfig(kwds)
+    config = get_config_from_kwds(**kwds)
     method = config.method
     include_fixed = config.include_fixed
     linear_only = config.linear_only
+    amplrepnvisitor = config._ampl_repn_visitor
+
+    # Check compatibility of arguments
     if linear_only and method is IncidenceMethod.identify_variables:
         raise RuntimeError(
             "linear_only=True is not supported when using identify_variables"
         )
+    if include_fixed and method is IncidenceMethod.ampl_repn:
+        raise RuntimeError("include_fixed=True is not supported when using ampl_repn")
+    if method is IncidenceMethod.ampl_repn and amplrepnvisitor is None:
+        # Developer error, this should never happen!
+        raise RuntimeError("_ampl_repn_visitor must be provided when using ampl_repn")
+
+    # Dispatch to correct method
     if method is IncidenceMethod.identify_variables:
         return _get_incident_via_identify_variables(expr, include_fixed)
     elif method is IncidenceMethod.standard_repn:
-        return _get_incident_via_standard_repn(expr, include_fixed, linear_only)
+        return _get_incident_via_standard_repn(
+            expr, include_fixed, linear_only, compute_values=False
+        )
+    elif method is IncidenceMethod.standard_repn_compute_values:
+        return _get_incident_via_standard_repn(
+            expr, include_fixed, linear_only, compute_values=True
+        )
+    elif method is IncidenceMethod.ampl_repn:
+        return _get_incident_via_ampl_repn(expr, linear_only, amplrepnvisitor)
     else:
         raise ValueError(
             f"Unrecognized value {method} for the method used to identify incident"
-            f" variables. Valid options are {IncidenceMethod.identify_variables}"
-            f" and {IncidenceMethod.standard_repn}."
+            f" variables. See the IncidenceMethod enum for valid methods."
         )
diff --git a/pyomo/contrib/incidence_analysis/interface.py b/pyomo/contrib/incidence_analysis/interface.py
--- a/pyomo/contrib/incidence_analysis/interface.py
+++ b/pyomo/contrib/incidence_analysis/interface.py
@@ -29,7 +29,7 @@
     plotly,
 )
 from pyomo.common.deprecation import deprecated
-from pyomo.contrib.incidence_analysis.config import IncidenceConfig
+from pyomo.contrib.incidence_analysis.config import get_config_from_kwds
 from pyomo.contrib.incidence_analysis.matching import maximum_matching
 from pyomo.contrib.incidence_analysis.connected import get_independent_submatrices
 from pyomo.contrib.incidence_analysis.triangularize import (
@@ -62,7 +62,7 @@ def _check_unindexed(complist):
 
 
 def get_incidence_graph(variables, constraints, **kwds):
-    config = IncidenceConfig(kwds)
+    config = get_config_from_kwds(**kwds)
     return get_bipartite_incidence_graph(variables, constraints, **config)
 
 
@@ -91,7 +91,9 @@ def get_bipartite_incidence_graph(variables, constraints, **kwds):
     ``networkx.Graph``
 
     """
-    config = IncidenceConfig(kwds)
+    # Note that this ConfigDict contains the visitor that we will re-use
+    # when constructing constraints.
+    config = get_config_from_kwds(**kwds)
     _check_unindexed(variables + constraints)
     N = len(variables)
     M = len(constraints)
@@ -163,7 +165,8 @@ def extract_bipartite_subgraph(graph, nodes0, nodes1):
 
 
 def _generate_variables_in_constraints(constraints, **kwds):
-    config = IncidenceConfig(kwds)
+    # Note: We construct a visitor here
+    config = get_config_from_kwds(**kwds)
     known_vars = ComponentSet()
     for con in constraints:
         for var in get_incident_variables(con.body, **config):
@@ -191,7 +194,7 @@ def get_structural_incidence_matrix(variables, constraints, **kwds):
         Entries are 1.0.
 
     """
-    config = IncidenceConfig(kwds)
+    config = get_config_from_kwds(**kwds)
     _check_unindexed(variables + constraints)
     N, M = len(variables), len(constraints)
     var_idx_map = ComponentMap((v, i) for i, v in enumerate(variables))
@@ -266,7 +269,6 @@ class IncidenceGraphInterface(object):
         ``evaluate_jacobian_eq`` method instead of ``evaluate_jacobian``
         rather than checking constraint expression types.
 
-
     """
 
     def __init__(self, model=None, active=True, include_inequality=True, **kwds):
@@ -275,7 +277,7 @@ def __init__(self, model=None, active=True, include_inequality=True, **kwds):
         # to cache the incidence graph for fast analysis later on.
         # WARNING: This cache will become invalid if the user alters their
         # model.
-        self._config = IncidenceConfig(kwds)
+        self._config = get_config_from_kwds(**kwds)
         if model is None:
             self._incidence_graph = None
             self._variables = None


Test patch:
diff --git a/pyomo/contrib/incidence_analysis/tests/test_incidence.py b/pyomo/contrib/incidence_analysis/tests/test_incidence.py
--- a/pyomo/contrib/incidence_analysis/tests/test_incidence.py
+++ b/pyomo/contrib/incidence_analysis/tests/test_incidence.py
@@ -56,44 +56,56 @@ def test_basic_incidence(self):
 
     def test_incidence_with_fixed_variable(self):
         m = pyo.ConcreteModel()
-        m.x = pyo.Var([1, 2, 3])
+        m.x = pyo.Var([1, 2, 3], initialize=1.0)
         expr = m.x[1] + m.x[1] * m.x[2] + m.x[1] * pyo.exp(m.x[3])
         m.x[2].fix()
         variables = self._get_incident_variables(expr)
         var_set = ComponentSet(variables)
         self.assertEqual(var_set, ComponentSet([m.x[1], m.x[3]]))
 
-    def test_incidence_with_mutable_parameter(self):
+    def test_incidence_with_named_expression(self):
         m = pyo.ConcreteModel()
         m.x = pyo.Var([1, 2, 3])
-        m.p = pyo.Param(mutable=True, initialize=None)
-        expr = m.x[1] + m.p * m.x[1] * m.x[2] + m.x[1] * pyo.exp(m.x[3])
+        m.subexpr = pyo.Expression(pyo.Integers)
+        m.subexpr[1] = m.x[1] * pyo.exp(m.x[3])
+        expr = m.x[1] + m.x[1] * m.x[2] + m.subexpr[1]
         variables = self._get_incident_variables(expr)
         self.assertEqual(ComponentSet(variables), ComponentSet(m.x[:]))
 
 
-class TestIncidenceStandardRepn(unittest.TestCase, _TestIncidence):
-    def _get_incident_variables(self, expr, **kwds):
-        method = IncidenceMethod.standard_repn
-        return get_incident_variables(expr, method=method, **kwds)
+class _TestIncidenceLinearOnly(object):
+    """Tests for methods that support linear_only"""
 
-    def test_assumed_standard_repn_behavior(self):
+    def _get_incident_variables(self, expr):
+        raise NotImplementedError(
+            "_TestIncidenceLinearOnly should not be used directly"
+        )
+
+    def test_linear_only(self):
         m = pyo.ConcreteModel()
-        m.x = pyo.Var([1, 2])
-        m.p = pyo.Param(initialize=0.0)
+        m.x = pyo.Var([1, 2, 3])
 
-        # We rely on variables with constant coefficients of zero not appearing
-        # in the standard repn (as opposed to appearing with explicit
-        # coefficients of zero).
-        expr = m.x[1] + 0 * m.x[2]
-        repn = generate_standard_repn(expr)
-        self.assertEqual(len(repn.linear_vars), 1)
-        self.assertIs(repn.linear_vars[0], m.x[1])
+        expr = 2 * m.x[1] + 4 * m.x[2] * m.x[1] - m.x[1] * pyo.exp(m.x[3])
+        variables = self._get_incident_variables(expr, linear_only=True)
+        self.assertEqual(len(variables), 0)
 
-        expr = m.p * m.x[1] + m.x[2]
-        repn = generate_standard_repn(expr)
-        self.assertEqual(len(repn.linear_vars), 1)
-        self.assertIs(repn.linear_vars[0], m.x[2])
+        expr = 2 * m.x[1] + 2 * m.x[2] * m.x[3] + 3 * m.x[2]
+        variables = self._get_incident_variables(expr, linear_only=True)
+        self.assertEqual(ComponentSet(variables), ComponentSet([m.x[1]]))
+
+        m.x[3].fix(2.5)
+        expr = 2 * m.x[1] + 2 * m.x[2] * m.x[3] + 3 * m.x[2]
+        variables = self._get_incident_variables(expr, linear_only=True)
+        self.assertEqual(ComponentSet(variables), ComponentSet([m.x[1], m.x[2]]))
+
+
+class _TestIncidenceLinearCancellation(object):
+    """Tests for methods that perform linear cancellation"""
+
+    def _get_incident_variables(self, expr):
+        raise NotImplementedError(
+            "_TestIncidenceLinearCancellation should not be used directly"
+        )
 
     def test_zero_coef(self):
         m = pyo.ConcreteModel()
@@ -113,23 +125,6 @@ def test_variable_minus_itself(self):
         var_set = ComponentSet(variables)
         self.assertEqual(var_set, ComponentSet([m.x[2], m.x[3]]))
 
-    def test_linear_only(self):
-        m = pyo.ConcreteModel()
-        m.x = pyo.Var([1, 2, 3])
-
-        expr = 2 * m.x[1] + 4 * m.x[2] * m.x[1] - m.x[1] * pyo.exp(m.x[3])
-        variables = self._get_incident_variables(expr, linear_only=True)
-        self.assertEqual(len(variables), 0)
-
-        expr = 2 * m.x[1] + 2 * m.x[2] * m.x[3] + 3 * m.x[2]
-        variables = self._get_incident_variables(expr, linear_only=True)
-        self.assertEqual(ComponentSet(variables), ComponentSet([m.x[1]]))
-
-        m.x[3].fix(2.5)
-        expr = 2 * m.x[1] + 2 * m.x[2] * m.x[3] + 3 * m.x[2]
-        variables = self._get_incident_variables(expr, linear_only=True)
-        self.assertEqual(ComponentSet(variables), ComponentSet([m.x[1], m.x[2]]))
-
     def test_fixed_zero_linear_coefficient(self):
         m = pyo.ConcreteModel()
         m.x = pyo.Var([1, 2, 3])
@@ -148,6 +143,9 @@ def test_fixed_zero_linear_coefficient(self):
         variables = self._get_incident_variables(expr)
         self.assertEqual(ComponentSet(variables), ComponentSet([m.x[1], m.x[2]]))
 
+    # NOTE: This test assumes that all methods that support linear cancellation
+    # accept a linear_only argument. If this changes, this test will need to be
+    # moved.
     def test_fixed_zero_coefficient_linear_only(self):
         m = pyo.ConcreteModel()
         m.x = pyo.Var([1, 2, 3])
@@ -159,6 +157,35 @@ def test_fixed_zero_coefficient_linear_only(self):
         self.assertEqual(len(variables), 1)
         self.assertIs(variables[0], m.x[3])
 
+
+class TestIncidenceStandardRepn(
+    unittest.TestCase,
+    _TestIncidence,
+    _TestIncidenceLinearOnly,
+    _TestIncidenceLinearCancellation,
+):
+    def _get_incident_variables(self, expr, **kwds):
+        method = IncidenceMethod.standard_repn
+        return get_incident_variables(expr, method=method, **kwds)
+
+    def test_assumed_standard_repn_behavior(self):
+        m = pyo.ConcreteModel()
+        m.x = pyo.Var([1, 2])
+        m.p = pyo.Param(initialize=0.0)
+
+        # We rely on variables with constant coefficients of zero not appearing
+        # in the standard repn (as opposed to appearing with explicit
+        # coefficients of zero).
+        expr = m.x[1] + 0 * m.x[2]
+        repn = generate_standard_repn(expr)
+        self.assertEqual(len(repn.linear_vars), 1)
+        self.assertIs(repn.linear_vars[0], m.x[1])
+
+        expr = m.p * m.x[1] + m.x[2]
+        repn = generate_standard_repn(expr)
+        self.assertEqual(len(repn.linear_vars), 1)
+        self.assertIs(repn.linear_vars[0], m.x[2])
+
     def test_fixed_none_linear_coefficient(self):
         m = pyo.ConcreteModel()
         m.x = pyo.Var([1, 2, 3])
@@ -168,6 +195,14 @@ def test_fixed_none_linear_coefficient(self):
         variables = self._get_incident_variables(expr)
         self.assertEqual(ComponentSet(variables), ComponentSet([m.x[1], m.x[2]]))
 
+    def test_incidence_with_mutable_parameter(self):
+        m = pyo.ConcreteModel()
+        m.x = pyo.Var([1, 2, 3])
+        m.p = pyo.Param(mutable=True, initialize=None)
+        expr = m.x[1] + m.p * m.x[1] * m.x[2] + m.x[1] * pyo.exp(m.x[3])
+        variables = self._get_incident_variables(expr)
+        self.assertEqual(ComponentSet(variables), ComponentSet(m.x[:]))
+
 
 class TestIncidenceIdentifyVariables(unittest.TestCase, _TestIncidence):
     def _get_incident_variables(self, expr, **kwds):
@@ -192,6 +227,36 @@ def test_variable_minus_itself(self):
         var_set = ComponentSet(variables)
         self.assertEqual(var_set, ComponentSet(m.x[:]))
 
+    def test_incidence_with_mutable_parameter(self):
+        m = pyo.ConcreteModel()
+        m.x = pyo.Var([1, 2, 3])
+        m.p = pyo.Param(mutable=True, initialize=None)
+        expr = m.x[1] + m.p * m.x[1] * m.x[2] + m.x[1] * pyo.exp(m.x[3])
+        variables = self._get_incident_variables(expr)
+        self.assertEqual(ComponentSet(variables), ComponentSet(m.x[:]))
+
+
+class TestIncidenceAmplRepn(
+    unittest.TestCase,
+    _TestIncidence,
+    _TestIncidenceLinearOnly,
+    _TestIncidenceLinearCancellation,
+):
+    def _get_incident_variables(self, expr, **kwds):
+        method = IncidenceMethod.ampl_repn
+        return get_incident_variables(expr, method=method, **kwds)
+
+
+class TestIncidenceStandardRepnComputeValues(
+    unittest.TestCase,
+    _TestIncidence,
+    _TestIncidenceLinearOnly,
+    _TestIncidenceLinearCancellation,
+):
+    def _get_incident_variables(self, expr, **kwds):
+        method = IncidenceMethod.standard_repn_compute_values
+        return get_incident_variables(expr, method=method, **kwds)
+
 
 if __name__ == "__main__":
     unittest.main()