GDP: Use private_data for all mappings between original and transformed components

pyomo/gdp/plugins/bigm.py

@@ -13,6 +13,7 @@
 
 import logging
 
+from pyomo.common.autoslots import AutoSlots
 from pyomo.common.collections import ComponentMap
 from pyomo.common.config import ConfigDict, ConfigValue
 from pyomo.common.gc_manager import PauseGC
@@ -58,6 +59,26 @@
 logger = logging.getLogger('pyomo.gdp.bigm')
 
 
+class _BigMData(AutoSlots.Mixin):
+    __slots__ = ('bigm_src',)
+
+    def __init__(self):
+        # we will keep a map of constraints (hashable, ha!) to a tuple to
+        # indicate what their M value is and where it came from, of the form:
+        # ((lower_value, lower_source, lower_key), (upper_value, upper_source,
+        # upper_key)), where the first tuple is the information for the lower M,
+        # the second tuple is the info for the upper M, source is the Suffix or
+        # argument dictionary and None if the value was calculated, and key is
+        # the key in the Suffix or argument dictionary, and None if it was
+        # calculated. (Note that it is possible the lower or upper is
+        # user-specified and the other is not, hence the need to store
+        # information for both.)
+        self.bigm_src = {}
+
+
+Block.register_private_data_initializer(_BigMData)
+
+
 @TransformationFactory.register(
     'gdp.bigm', doc="Relax disjunctive model using big-M terms."
 )
@@ -94,15 +115,8 @@ class BigM_Transformation(GDP_to_MIP_Transformation, _BigM_MixIn):
     name beginning "_pyomo_gdp_bigm_reformulation".  That Block will
     contain an indexed Block named "relaxedDisjuncts", which will hold
     the relaxed disjuncts.  This block is indexed by an integer
-    indicating the order in which the disjuncts were relaxed.
-    Each block has a dictionary "_constraintMap":
-
-        'srcConstraints': ComponentMap(<transformed constraint>:
-                                       <src constraint>)
-        'transformedConstraints': ComponentMap(<src constraint>:
-                                               <transformed constraint>)
-
-    All transformed Disjuncts will have a pointer to the block their transformed
+    indicating the order in which the disjuncts were relaxed. All
+    transformed Disjuncts will have a pointer to the block their transformed
     constraints are on, and all transformed Disjunctions will have a
     pointer to the corresponding 'Or' or 'ExactlyOne' constraint.
 
@@ -247,18 +261,6 @@ def _transform_disjunct(self, obj, bigM, transBlock):
 
         relaxationBlock = self._get_disjunct_transformation_block(obj, transBlock)
 
-        # we will keep a map of constraints (hashable, ha!) to a tuple to
-        # indicate what their M value is and where it came from, of the form:
-        # ((lower_value, lower_source, lower_key), (upper_value, upper_source,
-        # upper_key)), where the first tuple is the information for the lower M,
-        # the second tuple is the info for the upper M, source is the Suffix or
-        # argument dictionary and None if the value was calculated, and key is
-        # the key in the Suffix or argument dictionary, and None if it was
-        # calculated. (Note that it is possible the lower or upper is
-        # user-specified and the other is not, hence the need to store
-        # information for both.)
-        relaxationBlock.bigm_src = {}
-
         # This is crazy, but if the disjunction has been previously
         # relaxed, the disjunct *could* be deactivated.  This is a big
         # deal for Hull, as it uses the component_objects /
@@ -278,8 +280,8 @@ def _transform_constraint(
     ):
         # add constraint to the transformation block, we'll transform it there.
         transBlock = disjunct._transformation_block()
-        bigm_src = transBlock.bigm_src
-        constraintMap = transBlock._constraintMap
+        bigm_src = transBlock.private_data().bigm_src
+        constraint_map = transBlock.private_data('pyomo.gdp')
 
         disjunctionRelaxationBlock = transBlock.parent_block()
 
@@ -346,7 +348,7 @@ def _transform_constraint(
             bigm_src[c] = (lower, upper)
 
             self._add_constraint_expressions(
-                c, i, M, disjunct.binary_indicator_var, newConstraint, constraintMap
+                c, i, M, disjunct.binary_indicator_var, newConstraint, constraint_map
             )
 
             # deactivate because we relaxed
@@ -409,7 +411,7 @@ def _update_M_from_suffixes(self, constraint, suffix_list, lower, upper):
     def get_m_value_src(self, constraint):
         transBlock = _get_constraint_transBlock(constraint)
         ((lower_val, lower_source, lower_key), (upper_val, upper_source, upper_key)) = (
-            transBlock.bigm_src[constraint]
+            transBlock.private_data().bigm_src[constraint]
         )
 
         if (
@@ -464,7 +466,7 @@ def get_M_value_src(self, constraint):
         transBlock = _get_constraint_transBlock(constraint)
         # This is a KeyError if it fails, but it is also my fault if it
         # fails... (That is, it's a bug in the mapping.)
-        return transBlock.bigm_src[constraint]
+        return transBlock.private_data().bigm_src[constraint]
 
     def get_M_value(self, constraint):
         """Returns the M values used to transform constraint. Return is a tuple:
@@ -479,7 +481,7 @@ def get_M_value(self, constraint):
         transBlock = _get_constraint_transBlock(constraint)
         # This is a KeyError if it fails, but it is also my fault if it
         # fails... (That is, it's a bug in the mapping.)
-        lower, upper = transBlock.bigm_src[constraint]
+        lower, upper = transBlock.private_data().bigm_src[constraint]
         return (lower[0], upper[0])
 
     def get_all_M_values_by_constraint(self, model):
@@ -499,9 +501,8 @@ def get_all_M_values_by_constraint(self, model):
             # First check if it was transformed at all.
             if transBlock is not None:
                 # If it was transformed with BigM, we get the M values.
-                if hasattr(transBlock, 'bigm_src'):
-                    for cons in transBlock.bigm_src:
-                        m_values[cons] = self.get_M_value(cons)
+                for cons in transBlock.private_data().bigm_src:
+                    m_values[cons] = self.get_M_value(cons)
         return m_values
 
     def get_largest_M_value(self, model):

pyomo/gdp/plugins/bigm_mixin.py

@@ -232,7 +232,7 @@ def _estimate_M(self, expr, constraint):
         return tuple(M)
 
     def _add_constraint_expressions(
-        self, c, i, M, indicator_var, newConstraint, constraintMap
+        self, c, i, M, indicator_var, newConstraint, constraint_map
     ):
         # Since we are both combining components from multiple blocks and using
         # local names, we need to make sure that the first index for
@@ -253,8 +253,10 @@ def _add_constraint_expressions(
                 )
             M_expr = M[0] * (1 - indicator_var)
             newConstraint.add((name, i, 'lb'), c.lower <= c.body - M_expr)
-            constraintMap['transformedConstraints'][c] = [newConstraint[name, i, 'lb']]
-            constraintMap['srcConstraints'][newConstraint[name, i, 'lb']] = c
+            constraint_map.transformed_constraints[c].append(
+                newConstraint[name, i, 'lb']
+            )
+            constraint_map.src_constraint[newConstraint[name, i, 'lb']] = c
         if c.upper is not None:
             if M[1] is None:
                 raise GDP_Error(
@@ -263,13 +265,7 @@ def _add_constraint_expressions(
                 )
             M_expr = M[1] * (1 - indicator_var)
             newConstraint.add((name, i, 'ub'), c.body - M_expr <= c.upper)
-            transformed = constraintMap['transformedConstraints'].get(c)
-            if transformed is not None:
-                constraintMap['transformedConstraints'][c].append(
-                    newConstraint[name, i, 'ub']
-                )
-            else:
-                constraintMap['transformedConstraints'][c] = [
-                    newConstraint[name, i, 'ub']
-                ]
-            constraintMap['srcConstraints'][newConstraint[name, i, 'ub']] = c
+            constraint_map.transformed_constraints[c].append(
+                newConstraint[name, i, 'ub']
+            )
+            constraint_map.src_constraint[newConstraint[name, i, 'ub']] = c

pyomo/gdp/plugins/binary_multiplication.py

@@ -121,7 +121,7 @@ def _transform_disjunct(self, obj, transBlock):
     def _transform_constraint(self, obj, disjunct):
         # add constraint to the transformation block, we'll transform it there.
         transBlock = disjunct._transformation_block()
-        constraintMap = transBlock._constraintMap
+        constraint_map = transBlock.private_data('pyomo.gdp')
 
         disjunctionRelaxationBlock = transBlock.parent_block()
 
@@ -137,14 +137,14 @@ def _transform_constraint(self, obj, disjunct):
                 continue
 
             self._add_constraint_expressions(
-                c, i, disjunct.binary_indicator_var, newConstraint, constraintMap
+                c, i, disjunct.binary_indicator_var, newConstraint, constraint_map
             )
 
             # deactivate because we relaxed
             c.deactivate()
 
     def _add_constraint_expressions(
-        self, c, i, indicator_var, newConstraint, constraintMap
+        self, c, i, indicator_var, newConstraint, constraint_map
     ):
         # Since we are both combining components from multiple blocks and using
         # local names, we need to make sure that the first index for
@@ -156,21 +156,21 @@ def _add_constraint_expressions(
         # over the constraint indices, but I don't think it matters a lot.)
         unique = len(newConstraint)
         name = c.local_name + "_%s" % unique
-        transformed = constraintMap['transformedConstraints'][c] = []
+        transformed = constraint_map.transformed_constraints[c]
 
         lb, ub = c.lower, c.upper
         if (c.equality or lb is ub) and lb is not None:
             # equality
             newConstraint.add((name, i, 'eq'), (c.body - lb) * indicator_var == 0)
             transformed.append(newConstraint[name, i, 'eq'])
-            constraintMap['srcConstraints'][newConstraint[name, i, 'eq']] = c
+            constraint_map.src_constraint[newConstraint[name, i, 'eq']] = c
         else:
             # inequality
             if lb is not None:
                 newConstraint.add((name, i, 'lb'), 0 <= (c.body - lb) * indicator_var)
                 transformed.append(newConstraint[name, i, 'lb'])
-                constraintMap['srcConstraints'][newConstraint[name, i, 'lb']] = c
+                constraint_map.src_constraint[newConstraint[name, i, 'lb']] = c
             if ub is not None:
                 newConstraint.add((name, i, 'ub'), (c.body - ub) * indicator_var <= 0)
                 transformed.append(newConstraint[name, i, 'ub'])
-                constraintMap['srcConstraints'][newConstraint[name, i, 'ub']] = c
+                constraint_map.src_constraint[newConstraint[name, i, 'ub']] = c

pyomo/gdp/plugins/gdp_to_mip_transformation.py

@@ -11,7 +11,8 @@
 
 from functools import wraps
 
-from pyomo.common.collections import ComponentMap
+from pyomo.common.autoslots import AutoSlots
+from pyomo.common.collections import ComponentMap, DefaultComponentMap
 from pyomo.common.log import is_debug_set
 from pyomo.common.modeling import unique_component_name
 
@@ -48,6 +49,17 @@
 from weakref import ref as weakref_ref
 
 
+class _GDPTransformationData(AutoSlots.Mixin):
+    __slots__ = ('src_constraint', 'transformed_constraints')
+
+    def __init__(self):
+        self.src_constraint = ComponentMap()
+        self.transformed_constraints = DefaultComponentMap(list)
+
+
+Block.register_private_data_initializer(_GDPTransformationData, scope='pyomo.gdp')
+
+
 class GDP_to_MIP_Transformation(Transformation):
     """
     Base class for transformations from GDP to MIP
@@ -243,14 +255,7 @@ def _get_disjunct_transformation_block(self, disjunct, transBlock):
         relaxationBlock = relaxedDisjuncts[len(relaxedDisjuncts)]
 
         relaxationBlock.transformedConstraints = Constraint(Any)
-
         relaxationBlock.localVarReferences = Block()
-        # add the map that will link back and forth between transformed
-        # constraints and their originals.
-        relaxationBlock._constraintMap = {
-            'srcConstraints': ComponentMap(),
-            'transformedConstraints': ComponentMap(),
-        }
 
         # add mappings to source disjunct (so we'll know we've relaxed)
         disjunct._transformation_block = weakref_ref(relaxationBlock)