move AMPLRepnVisitor construction into ConfigValue validation

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, AMPLRepn, text_nl_template
+from pyomo.repn.util import FileDeterminism, FileDeterminism_to_SortComponents
 
 
 class IncidenceMethod(enum.Enum):
@@ -62,7 +65,92 @@ class IncidenceMethod(enum.Enum):
 )
 
 
-IncidenceConfig = ConfigDict()
+class _ReconstructVisitor:
+    pass
+
+
+def _amplrepnvisitor_validator(visitor=_ReconstructVisitor):
+    # This checks for and returns a valid AMPLRepnVisitor, but I don't want
+    # to construct this if we're not using IncidenceMethod.ampl_repn.
+    # It is not necessarily the end of the world if we construct this, however,
+    # as the code should still work.
+    if visitor is _ReconstructVisitor:
+        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,
+        )
+    elif not isinstance(visitor, AMPLRepnVisitor):
+        raise TypeError(
+            "'visitor' config argument should be an instance of AMPLRepnVisitor"
+        )
+    else:
+        amplvisitor = visitor
+    return amplvisitor
+
+
+_ampl_repn_visitor = ConfigValue(
+    default=_ReconstructVisitor,
+    domain=_amplrepnvisitor_validator,
+    description="Visitor used to generate AMPLRepn of each constraint",
+)
+
+
+class _IncidenceConfigDict(ConfigDict):
+
+    def __call__(
+        self,
+        value=NOTSET,
+        default=NOTSET,
+        domain=NOTSET,
+        description=NOTSET,
+        doc=NOTSET,
+        visibility=NOTSET,
+        implicit=NOTSET,
+        implicit_domain=NOTSET,
+        preserve_implicit=False,
+    ):
+        init_value = value
+        new = super().__call__(
+            value=value,
+            default=default,
+            domain=domain,
+            description=description,
+            doc=doc,
+            visibility=visibility,
+            implicit=implicit,
+            implicit_domain=implicit_domain,
+            preserve_implicit=preserve_implicit,
+        )
+
+        if (
+            new.method == IncidenceMethod.ampl_repn
+            and "ampl_repn_visitor" not in init_value
+        ):
+            new.ampl_repn_visitor = _ReconstructVisitor
+
+        return new
+
+
+
+IncidenceConfig = _IncidenceConfigDict()
 """Options for incidence graph generation
 
 - ``include_fixed`` -- Flag indicating whether fixed variables should be included
@@ -71,6 +159,8 @@ 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``.
 
 """
 
@@ -82,3 +172,6 @@ class IncidenceMethod(enum.Enum):
 
 
 IncidenceConfig.declare("method", _method)
+
+
+IncidenceConfig.declare("ampl_repn_visitor", _ampl_repn_visitor)

pyomo/contrib/incidence_analysis/incidence.py

@@ -80,38 +80,14 @@ def _get_incident_via_standard_repn(
         return unique_variables
 
 
-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
+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()
@@ -172,11 +148,11 @@ 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
     linear_only = config.linear_only
+    amplrepnvisitor = config.ampl_repn_visitor
     if linear_only and method is IncidenceMethod.identify_variables:
         raise RuntimeError(
             "linear_only=True is not supported when using identify_variables"
@@ -194,7 +170,7 @@ def get_incident_variables(expr, **kwds):
             expr, include_fixed, linear_only, compute_values=True
         )
     elif method is IncidenceMethod.ampl_repn:
-        return _get_incident_via_ampl_repn(expr, linear_only, visitor=visitor)
+        return _get_incident_via_ampl_repn(expr, linear_only, amplrepnvisitor)
     else:
         raise ValueError(
             f"Unrecognized value {method} for the method used to identify incident"

pyomo/contrib/incidence_analysis/interface.py

@@ -93,6 +93,8 @@ def get_bipartite_incidence_graph(variables, constraints, **kwds):
     ``networkx.Graph``
 
     """
+    # Note that this ConfigDict contains the visitor that we will re-use
+    # when constructing constraints.
     config = IncidenceConfig(kwds)
     _check_unindexed(variables + constraints)
     N = len(variables)
@@ -101,38 +103,10 @@ 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))
-
-    if config.method == IncidenceMethod.ampl_repn:
-        subexpression_cache = {}
-        subexpression_order = []
-        external_functions = {}
-        var_map = {}
-        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
+    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])
     return graph
 
 
@@ -193,46 +167,14 @@ def extract_bipartite_subgraph(graph, nodes0, nodes1):
 
 
 def _generate_variables_in_constraints(constraints, **kwds):
+    # Note: We construct a visitor here
     config = IncidenceConfig(kwds)
-
-    if config.method == IncidenceMethod.ampl_repn:
-        subexpression_cache = {}
-        subexpression_order = []
-        external_functions = {}
-        var_map = {}
-        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:
-        known_vars = ComponentSet()
-        for con in constraints:
-            for var in get_incident_variables(con.body, visitor=visitor, **config):
-                if var not in known_vars:
-                    known_vars.add(var)
-                    yield var
-    finally:
-        # NOTE: I believe this is only guaranteed to be called when the
-        # generator is garbage collected. This could lead to some nasty
-        # bug where ActiveVisitor is set for longer than we intend.
-        # TODO: Convert this into a function. (or yield from variables
-        # after this try/finally.
-        AMPLRepn.ActiveVisitor = None
+    known_vars = ComponentSet()
+    for con in constraints:
+        for var in get_incident_variables(con.body, **config):
+            if var not in known_vars:
+                known_vars.add(var)
+                yield var
 
 
 def get_structural_incidence_matrix(variables, constraints, **kwds):
@@ -329,7 +271,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):