Add ampl_repn option to Incidence Analysis

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)

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."
         )

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