re-use visitor when iterating over constraints

pyomo/contrib/incidence_analysis/incidence.py

@@ -76,34 +76,38 @@ def _get_incident_via_standard_repn(expr, include_fixed, linear_only):
         return unique_variables
 
 
-def _get_incident_via_ampl_repn(expr, linear_only):
-    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)
-    visitor = AMPLRepnVisitor(
-        text_nl_template,
-        subexpression_cache,
-        subexpression_order,
-        external_functions,
-        var_map,
-        used_named_expressions,
-        symbolic_solver_labels,
-        export_defined_variables,
-        sorter,
-    )
-    AMPLRepn.ActiveVisitor = visitor
-    try:
+def _get_incident_via_ampl_repn(expr, linear_only, visitor=None):
+    if visitor 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)
+        visitor = AMPLRepnVisitor(
+            text_nl_template,
+            subexpression_cache,
+            subexpression_order,
+            external_functions,
+            var_map,
+            used_named_expressions,
+            symbolic_solver_labels,
+            export_defined_variables,
+            sorter,
+        )
+        AMPLRepn.ActiveVisitor = visitor
+        try:
+            repn = visitor.walk_expression((expr, None, 0, 1.0))
+        finally:
+            AMPLRepn.ActiveVisitor = None
+    else:
+        var_map = visitor.var_map
         repn = visitor.walk_expression((expr, None, 0, 1.0))
-    finally:
-        AMPLRepn.ActiveVisitor = None
 
     nonlinear_var_ids = [] if repn.nonlinear is None else repn.nonlinear[1]
     nonlinear_vars = [var_map[v_id] for v_id in nonlinear_var_ids]
@@ -158,6 +162,7 @@ def get_incident_variables(expr, **kwds):
        ['x[1]', 'x[2]']
 
     """
+    visitor = kwds.pop("visitor", None)
     config = IncidenceConfig(kwds)
     method = config.method
     include_fixed = config.include_fixed
@@ -173,7 +178,7 @@ def get_incident_variables(expr, **kwds):
     elif method is IncidenceMethod.standard_repn:
         return _get_incident_via_standard_repn(expr, include_fixed, linear_only)
     elif method is IncidenceMethod.ampl_repn:
-        return _get_incident_via_ampl_repn(expr, linear_only)
+        return _get_incident_via_ampl_repn(expr, linear_only, visitor=visitor)
     else:
         raise ValueError(
             f"Unrecognized value {method} for the method used to identify incident"

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 IncidenceConfig, IncidenceMethod
 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 (
@@ -45,6 +45,8 @@
 )
 from pyomo.contrib.incidence_analysis.incidence import get_incident_variables
 from pyomo.contrib.pynumero.asl import AmplInterface
+from pyomo.repn.plugins.nl_writer import AMPLRepnVisitor, AMPLRepn, text_nl_template
+from pyomo.repn.util import FileDeterminism, FileDeterminism_to_SortComponents
 
 pyomo_nlp, pyomo_nlp_available = attempt_import(
     'pyomo.contrib.pynumero.interfaces.pyomo_nlp'
@@ -99,10 +101,37 @@ def get_bipartite_incidence_graph(variables, constraints, **kwds):
     graph.add_nodes_from(range(M), bipartite=0)
     graph.add_nodes_from(range(M, M + N), bipartite=1)
     var_node_map = ComponentMap((v, M + i) for i, v in enumerate(variables))
-    for i, con in enumerate(constraints):
-        for var in get_incident_variables(con.body, **config):
-            if var in var_node_map:
-                graph.add_edge(i, var_node_map[var])
+
+    if config.method == IncidenceMethod.ampl_repn:
+        subexpression_cache = {}
+        subexpression_order = []
+        external_functions = {}
+        used_named_expressions = set()
+        symbolic_solver_labels = False
+        export_defined_variables = False
+        sorter = FileDeterminism_to_SortComponents(FileDeterminism.ORDERED)
+        visitor = AMPLRepnVisitor(
+            text_nl_template,
+            subexpression_cache,
+            subexpression_order,
+            external_functions,
+            var_map,
+            used_named_expressions,
+            symbolic_solver_labels,
+            export_defined_variables,
+            sorter,
+        )
+    else:
+        visitor = None
+
+    AMPLRepn.ActiveVisitor = visitor
+    try:
+        for i, con in enumerate(constraints):
+            for var in get_incident_variables(con.body, visitor=visitor, **config):
+                if var in var_node_map:
+                    graph.add_edge(i, var_node_map[var])
+    finally:
+        AMPLRepn.ActiveVisitor = None
     return graph