Rewrite core.relax_integer_vars transformation

pyomo/core/base/__init__.py

@@ -38,7 +38,7 @@
 from pyomo.core.base.component import name, Component, ModelComponentFactory
 from pyomo.core.base.componentuid import ComponentUID
 from pyomo.core.base.config import PyomoOptions
-from pyomo.core.base.enums import SortComponents, TraversalStrategy
+from pyomo.core.base.enums import SortComponents, TraversalStrategy, VarCollector
 from pyomo.core.base.label import (
     CuidLabeler,
     CounterLabeler,

pyomo/core/base/enums.py

@@ -11,6 +11,7 @@
 
 import enum
 import sys
+from pyomo.common import enums
 
 if sys.version_info[:2] >= (3, 11):
     strictEnum = {'boundary': enum.STRICT}
@@ -93,3 +94,8 @@ def sort_names(flag):
     @staticmethod
     def sort_indices(flag):
         return SortComponents.SORTED_INDICES in SortComponents(flag)
+
+
+class VarCollector(enums.IntEnum):
+    FromVarComponents = 1
+    FromExpressions = 2

pyomo/core/plugins/transform/discrete_vars.py

@@ -14,7 +14,23 @@
 logger = logging.getLogger('pyomo.core')
 
 from pyomo.common import deprecated
-from pyomo.core.base import Transformation, TransformationFactory, Var, Suffix, Reals
+from pyomo.common.config import ConfigDict, ConfigValue, In, IsInstance
+from pyomo.common.deprecation import deprecation_warning
+from pyomo.core.base import (
+    Transformation,
+    TransformationFactory,
+    Var,
+    Suffix,
+    Reals,
+    Block,
+    ReverseTransformationToken,
+    VarCollector,
+    Constraint,
+    Objective,
+)
+from pyomo.core.util import target_list
+from pyomo.gdp import Disjunct
+from pyomo.util.vars_from_expressions import get_vars_from_components
 
 
 #
@@ -25,56 +41,181 @@
     'core.relax_integer_vars', doc="Relax integer variables to continuous counterparts"
 )
 class RelaxIntegerVars(Transformation):
+    CONFIG = ConfigDict('core.relax_integer_vars')
+    CONFIG.declare(
+        'targets',
+        ConfigValue(
+            default=None,
+            domain=target_list,
+            description="target or list of targets that will be relaxed",
+            doc="""
+            This specifies the list of components to relax. If None (default), the
+            entire model is transformed. Note that if the transformation is done
+            out of place, the list of targets should be attached to the model before
+            it is cloned, and the list will specify the targets on the cloned
+            instance.""",
+        ),
+    )
+    CONFIG.declare(
+        'reverse',
+        ConfigValue(
+            default=None,
+            domain=IsInstance(ReverseTransformationToken),
+            description="The token returned by a (forward) call to this "
+            "transformation, if you wish to reverse the transformation.",
+            doc="""
+            This argument should be the reverse transformation token
+            returned by a previous call to this transformation to transform
+            fixed disjunctive state in the given model.
+            If this argument is specified, this call to the transformation
+            will reverse what the transformation did in the call that returned
+            the token. Note that if there are intermediate changes to the model
+            in between the forward and the backward calls to the transformation,
+            the behavior could be unexpected.
+            """,
+        ),
+    )
+    CONFIG.declare(
+        'var_collector',
+        ConfigValue(
+            default=VarCollector.FromVarComponents,
+            domain=In(VarCollector),
+            description="The method for collection the Vars to relax. If "
+            "VarCollector.FromVarComponents (default), any Var component on "
+            "the active tree will be relaxed.",
+            doc="""
+            This specifies the method for collecting the Var components to relax.
+            The default, VarCollector.FromVarComponents, assumes that all relevant
+            Vars are on the active tree. If this is true, then this is the most
+            performant option. However, in more complex cases where some Vars may not
+            be in the active tree (e.g. some are on deactivated Blocks or come from
+            other models), specify VarCollector.FromExpressions to relax all Vars that
+            appear in expressions in the active tree.
+            """,
+        ),
+    )
+    CONFIG.declare(
+        'transform_deactivated_blocks',
+        ConfigValue(
+            default=True,
+            description="[DEPRECATED]: Whether or not to search for Var components to "
+            "relax on deactivated Blocks. True by default",
+        ),
+    )
+    CONFIG.declare(
+        'undo',
+        ConfigValue(
+            default=False,
+            domain=bool,
+            description="[DEPRECATED]: Please use the 'reverse' argument to undo "
+            "the transformation.",
+        ),
+    )
+
     def __init__(self):
-        super(RelaxIntegerVars, self).__init__()
+        super().__init__()
 
     def _apply_to(self, model, **kwds):
-        options = kwds.pop('options', {})
-        if kwds.get('undo', options.get('undo', False)):
+        if not model.ctype in (Block, Disjunct):
+            raise ValueError(
+                "Transformation called on %s of type %s. 'model' "
+                "must be a ConcreteModel or Block." % (model.name, model.ctype)
+            )
+        config = self.CONFIG(kwds.pop('options', {}))
+        config.set_value(kwds)
+
+        if config.undo:
+            deprecation_warning(
+                "The 'undo' argument is deprecated. Please use the 'reverse' "
+                "argument to undo the transformation.",
+                version='6.9.3.dev0',
+            )
             for v, d in model._relaxed_integer_vars[None].values():
                 bounds = v.bounds
                 v.domain = d
                 v.setlb(bounds[0])
                 v.setub(bounds[1])
             model.del_component("_relaxed_integer_vars")
             return
-        # True by default, you can specify False if you want
-        descend = kwds.get(
-            'transform_deactivated_blocks',
-            options.get('transform_deactivated_blocks', True),
-        )
-        active = None if descend else True
 
-        # Relax the model
-        relaxed_vars = {}
-        _base_model_vars = model.component_data_objects(
-            Var, active=active, descend_into=True
-        )
-        for var in _base_model_vars:
+        targets = (model,) if config.targets is None else config.targets
+
+        if config.reverse is None:
+            reverse_dict = {}
+            # Relax the model
+            reverse_token = ReverseTransformationToken(
+                self.__class__, model, targets, reverse_dict
+            )
+        else:
+            # reverse the transformation
+            reverse_token = config.reverse
+            reverse_token.check_token_valid(self.__class__, model, targets)
+            reverse_dict = reverse_token.reverse_dict
+            for v, d in reverse_dict.values():
+                lb, ub = v.bounds
+                v.domain = d
+                v.setlb(lb)
+                v.setub(ub)
+            return
+
+        ### [ESJ 4/29/25]: This can go away when we remove 'undo'
+        model._relaxed_integer_vars = Suffix(direction=Suffix.LOCAL)
+        model._relaxed_integer_vars[None] = reverse_dict
+        ###
+
+        for t in targets:
+            if isinstance(t, Block):
+                blocks = t.values() if t.is_indexed() else (t,)
+                for block in blocks:
+                    self._relax_block(block, config, reverse_dict)
+            elif t.ctype is Var:
+                self._relax_var(t, reverse_dict)
+            else:
+                raise ValueError(
+                    "Target '%s' was not a Block or Var. It was of type "
+                    "'%s' and cannot be transformed." % (t.name, type(t))
+                )
+
+        return reverse_token
+
+    def _relax_block(self, block, config, reverse_dict):
+        self._relax_vars_from_block(block, config, reverse_dict)
+
+        for b in block.component_data_objects(Block, active=None, descend_into=True):
+            if not b.active:
+                if config.transform_deactivated_blocks:
+                    deprecation_warning(
+                        "The `transform_deactivated_blocks` arguments is deprecated. "
+                        "Either specify deactivated Blocks as targets to activate them "
+                        "if transforming them is the desired behavior.",
+                        version='6.9.3.dev0',
+                    )
+                else:
+                    continue
+            self._relax_vars_from_block(b, config, reverse_dict)
+
+    def _relax_vars_from_block(self, block, config, reverse_dict):
+        if config.var_collector is VarCollector.FromVarComponents:
+            model_vars = block.component_data_objects(Var, descend_into=False)
+        else:
+            model_vars = get_vars_from_components(
+                block, ctype=(Constraint, Objective), descend_into=False
+            )
+        for var in model_vars:
+            if id(var) not in reverse_dict:
+                self._relax_var(var, reverse_dict)
+
+    def _relax_var(self, v, reverse_dict):
+        var_datas = v.values() if v.is_indexed() else (v,)
+        for var in var_datas:
             if not var.is_integer():
                 continue
-            # Note: some indexed components can only have their
-            # domain set on the parent component (the individual
-            # indices cannot be set independently)
-            _c = var.parent_component()
-            try:
-                lb, ub = var.bounds
-                _domain = var.domain
-                var.domain = Reals
-                var.setlb(lb)
-                var.setub(ub)
-                relaxed_vars[id(var)] = (var, _domain)
-            except:
-                if id(_c) in relaxed_vars:
-                    continue
-                _domain = _c.domain
-                lb, ub = _c.bounds
-                _c.domain = Reals
-                _c.setlb(lb)
-                _c.setub(ub)
-                relaxed_vars[id(_c)] = (_c, _domain)
-        model._relaxed_integer_vars = Suffix(direction=Suffix.LOCAL)
-        model._relaxed_integer_vars[None] = relaxed_vars
+            lb, ub = var.bounds
+            _domain = var.domain
+            var.domain = Reals
+            var.setlb(lb)
+            var.setub(ub)
+            reverse_dict[id(var)] = (var, _domain)
 
 
 @TransformationFactory.register(