Merge branch 'incidence-amplrepn' of https://github.com/robbybp/pyomo …

…into incidence-amplrepn

pyomo/contrib/incidence_analysis/config.py

@@ -65,90 +65,22 @@ class IncidenceMethod(enum.Enum):
 )
 
 
-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):
+def _amplrepnvisitor_validator(visitor):
+    if not isinstance(visitor, AMPLRepnVisitor):
         raise TypeError(
             "'visitor' config argument should be an instance of AMPLRepnVisitor"
         )
-    else:
-        amplvisitor = visitor
-    return amplvisitor
+    return visitor
 
 
 _ampl_repn_visitor = ConfigValue(
-    default=_ReconstructVisitor,
+    default=None,
     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()
+IncidenceConfig = ConfigDict()
 """Options for incidence graph generation
 
 - ``include_fixed`` -- Flag indicating whether fixed variables should be included
@@ -157,8 +89,9 @@ def __call__(
   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``.
+- ``_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!*
 
 """
 
@@ -172,4 +105,44 @@ def __call__(
 IncidenceConfig.declare("method", _method)
 
 
-IncidenceConfig.declare("ampl_repn_visitor", _ampl_repn_visitor)
+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

@@ -19,7 +19,10 @@
 from pyomo.repn.plugins.nl_writer import AMPLRepnVisitor, AMPLRepn, text_nl_template
 from pyomo.repn.util import FileDeterminism, FileDeterminism_to_SortComponents
 from pyomo.util.subsystems import TemporarySubsystemManager
-from pyomo.contrib.incidence_analysis.config import IncidenceMethod, IncidenceConfig
+from pyomo.contrib.incidence_analysis.config import (
+    IncidenceMethod,
+    get_config_from_kwds,
+)
 
 
 #
@@ -148,17 +151,24 @@ 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
+    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:
@@ -174,6 +184,5 @@ def get_incident_variables(expr, **kwds):
     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, IncidenceMethod
+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 (
@@ -64,7 +64,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)
 
 
@@ -95,7 +95,7 @@ def get_bipartite_incidence_graph(variables, constraints, **kwds):
     """
     # Note that this ConfigDict contains the visitor that we will re-use
     # when constructing constraints.
-    config = IncidenceConfig(kwds)
+    config = get_config_from_kwds(**kwds)
     _check_unindexed(variables + constraints)
     N = len(variables)
     M = len(constraints)
@@ -168,7 +168,7 @@ def extract_bipartite_subgraph(graph, nodes0, nodes1):
 
 def _generate_variables_in_constraints(constraints, **kwds):
     # Note: We construct a visitor here
-    config = IncidenceConfig(kwds)
+    config = get_config_from_kwds(**kwds)
     known_vars = ComponentSet()
     for con in constraints:
         for var in get_incident_variables(con.body, **config):
@@ -196,7 +196,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))
@@ -279,7 +279,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