Merge pull request Pyomo#3348 from jsiirola/suffix-finder-context

Add `context` option to `SuffixFinder`

pyomo/core/base/suffix.py

@@ -19,6 +19,7 @@
 from pyomo.common.modeling import NOTSET
 from pyomo.common.pyomo_typing import overload
 from pyomo.common.timing import ConstructionTimer
+from pyomo.core.base.block import BlockData
 from pyomo.core.base.component import ActiveComponent, ModelComponentFactory
 from pyomo.core.base.disable_methods import disable_methods
 from pyomo.core.base.initializer import Initializer
@@ -409,7 +410,7 @@ class AbstractSuffix(Suffix):
 
 
 class SuffixFinder(object):
-    def __init__(self, name, default=None):
+    def __init__(self, name, default=None, context=None):
         """This provides an efficient utility for finding suffix values on a
         (hierarchical) Pyomo model.
 
@@ -424,11 +425,26 @@ def __init__(self, name, default=None):
             Default value to return from `.find()` if no matching Suffix
             is found.
 
+        context: BlockData
+
+            The root of the Block hierarchy to use when searching for
+            Suffix components.  Suffixes outside this hierarchy will not
+            be interrogated and components that are queried (with
+            :py:meth:`find(component_data)` will return the default
+            value.
+
         """
         self.name = name
         self.default = default
         self.all_suffixes = []
-        self._suffixes_by_block = {None: []}
+        self._context = context
+        self._suffixes_by_block = ComponentMap()
+        self._suffixes_by_block[self._context] = []
+        if context is not None:
+            s = context.component(name)
+            if s is not None and s.ctype is Suffix and s.active:
+                self._suffixes_by_block[context].append(s)
+                self.all_suffixes.append(s)
 
     def find(self, component_data):
         """Find suffix value for a given component data object in model tree
@@ -458,7 +474,17 @@ def find(self, component_data):
 
         """
         # Walk parent tree and search for suffixes
-        suffixes = self._get_suffix_list(component_data.parent_block())
+        if isinstance(component_data, BlockData):
+            _block = component_data
+        else:
+            _block = component_data.parent_block()
+        try:
+            suffixes = self._get_suffix_list(_block)
+        except AttributeError:
+            # Component was outside the context (eventually parent
+            # becomes None and parent.parent_block() raises an
+            # AttributeError): we will return the default value
+            return self.default
         # Pass 1: look for the component_data, working root to leaf
         for s in suffixes:
             if component_data in s:

pyomo/core/plugins/transform/scaling.py

@@ -82,15 +82,10 @@ def _create_using(self, original_model, **kwds):
         self._apply_to(scaled_model, **kwds)
         return scaled_model
 
-    def _get_float_scaling_factor(self, component):
-        if self._suffix_finder is None:
-            self._suffix_finder = SuffixFinder('scaling_factor', 1.0)
-        return self._suffix_finder.find(component)
-
     def _apply_to(self, model, rename=True):
         # create a map of component to scaling factor
         component_scaling_factor_map = ComponentMap()
-        self._suffix_finder = SuffixFinder('scaling_factor', 1.0)
+        self._suffix_finder = SuffixFinder('scaling_factor', 1.0, model)
 
         # if the scaling_method is 'user', get the scaling parameters from the suffixes
         if self._scaling_method == 'user':

pyomo/core/tests/transform/test_scaling.py

@@ -13,7 +13,7 @@
 import pyomo.common.unittest as unittest
 import pyomo.environ as pyo
 from pyomo.opt.base.solvers import UnknownSolver
-from pyomo.core.plugins.transform.scaling import ScaleModel
+from pyomo.core.plugins.transform.scaling import ScaleModel, SuffixFinder
 
 
 class TestScaleModelTransformation(unittest.TestCase):
@@ -600,6 +600,13 @@ def con_rule(m, i):
         self.assertAlmostEqual(pyo.value(model.zcon), -8, 4)
 
     def test_get_float_scaling_factor_top_level(self):
+        # Note: the transformation used to have a private method for
+        # finding suffix values (which this method tested).  The
+        # transformation now leverages the SuffixFinder.  To ensure that
+        # the SuffixFinder behaves in the same way as the original local
+        # method, we preserve these tests, but directly test the
+        # SuffixFinder
+
         m = pyo.ConcreteModel()
         m.scaling_factor = pyo.Suffix(direction=pyo.Suffix.EXPORT)
 
@@ -616,17 +623,23 @@ def test_get_float_scaling_factor_top_level(self):
         m.scaling_factor[m.v1] = 0.1
         m.scaling_factor[m.b1.v2] = 0.2
 
+        _finder = SuffixFinder('scaling_factor', 1.0, m)
+
         # SF should be 0.1 from top level
-        sf = ScaleModel()._get_float_scaling_factor(m.v1)
-        assert sf == float(0.1)
+        self.assertEqual(_finder.find(m.v1), 0.1)
         # SF should be 0.1 from top level, lower level ignored
-        sf = ScaleModel()._get_float_scaling_factor(m.b1.v2)
-        assert sf == float(0.2)
+        self.assertEqual(_finder.find(m.b1.v2), 0.2)
         # No SF, should return 1
-        sf = ScaleModel()._get_float_scaling_factor(m.b1.b2.v3)
-        assert sf == 1.0
+        self.assertEqual(_finder.find(m.b1.b2.v3), 1.0)
 
     def test_get_float_scaling_factor_local_level(self):
+        # Note: the transformation used to have a private method for
+        # finding suffix values (which this method tested).  The
+        # transformation now leverages the SuffixFinder.  To ensure that
+        # the SuffixFinder behaves in the same way as the original local
+        # method, we preserve these tests, but directly test the
+        # SuffixFinder
+
         m = pyo.ConcreteModel()
         m.scaling_factor = pyo.Suffix(direction=pyo.Suffix.EXPORT)
 
@@ -647,15 +660,21 @@ def test_get_float_scaling_factor_local_level(self):
         # Add an intermediate scaling factor - this should take priority
         m.b1.scaling_factor[m.b1.b2.v3] = 0.4
 
+        _finder = SuffixFinder('scaling_factor', 1.0, m)
+
         # Should get SF from local levels
-        sf = ScaleModel()._get_float_scaling_factor(m.v1)
-        assert sf == float(0.1)
-        sf = ScaleModel()._get_float_scaling_factor(m.b1.v2)
-        assert sf == float(0.2)
-        sf = ScaleModel()._get_float_scaling_factor(m.b1.b2.v3)
-        assert sf == float(0.4)
+        self.assertEqual(_finder.find(m.v1), 0.1)
+        self.assertEqual(_finder.find(m.b1.v2), 0.2)
+        self.assertEqual(_finder.find(m.b1.b2.v3), 0.4)
 
     def test_get_float_scaling_factor_intermediate_level(self):
+        # Note: the transformation used to have a private method for
+        # finding suffix values (which this method tested).  The
+        # transformation now leverages the SuffixFinder.  To ensure that
+        # the SuffixFinder behaves in the same way as the original local
+        # method, we preserve these tests, but directly test the
+        # SuffixFinder
+
         m = pyo.ConcreteModel()
         m.scaling_factor = pyo.Suffix(direction=pyo.Suffix.EXPORT)
 
@@ -680,15 +699,14 @@ def test_get_float_scaling_factor_intermediate_level(self):
 
         m.b1.b2.b3.scaling_factor[m.b1.b2.b3.v3] = 0.4
 
+        _finder = SuffixFinder('scaling_factor', 1.0, m)
+
         # v1 should be unscaled as SF set below variable level
-        sf = ScaleModel()._get_float_scaling_factor(m.v1)
-        assert sf == 1.0
+        self.assertEqual(_finder.find(m.v1), 1.0)
         # v2 should get SF from b1 level
-        sf = ScaleModel()._get_float_scaling_factor(m.b1.b2.b3.v2)
-        assert sf == float(0.2)
+        self.assertEqual(_finder.find(m.b1.b2.b3.v2), 0.2)
         # v2 should get SF from highest level, ignoring b3 level
-        sf = ScaleModel()._get_float_scaling_factor(m.b1.b2.b3.v3)
-        assert sf == float(0.3)
+        self.assertEqual(_finder.find(m.b1.b2.b3.v3), 0.3)
 
 
 if __name__ == "__main__":

pyomo/core/tests/unit/test_suffix.py

@@ -1795,47 +1795,77 @@ def test_suffix_finder(self):
         m.b1.b2 = Block()
         m.b1.b2.v3 = Var([0])
 
-        _suffix_finder = SuffixFinder('suffix')
-
         # Add Suffixes
         m.suffix = Suffix(direction=Suffix.EXPORT)
         # No suffix on b1 - make sure we can handle missing suffixes
         m.b1.b2.suffix = Suffix(direction=Suffix.EXPORT)
 
+        _suffix_finder = SuffixFinder('suffix')
+        _suffix_b1_finder = SuffixFinder('suffix', context=m.b1)
+        _suffix_b2_finder = SuffixFinder('suffix', context=m.b1.b2)
+
         # Check for no suffix value
-        assert _suffix_finder.find(m.b1.b2.v3[0]) == None
+        self.assertEqual(_suffix_finder.find(m.b1.b2.v3[0]), None)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.b2.v3[0]), None)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.b2.v3[0]), None)
 
         # Check finding default values
         # Add a default at the top level
         m.suffix[None] = 1
-        assert _suffix_finder.find(m.b1.b2.v3[0]) == 1
+        self.assertEqual(_suffix_finder.find(m.b1.b2.v3[0]), 1)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.b2.v3[0]), None)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.b2.v3[0]), None)
 
         # Add a default suffix at a lower level
         m.b1.b2.suffix[None] = 2
-        assert _suffix_finder.find(m.b1.b2.v3[0]) == 2
+        self.assertEqual(_suffix_finder.find(m.b1.b2.v3[0]), 2)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.b2.v3[0]), 2)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.b2.v3[0]), 2)
 
         # Check for container at lowest level
         m.b1.b2.suffix[m.b1.b2.v3] = 3
-        assert _suffix_finder.find(m.b1.b2.v3[0]) == 3
+        self.assertEqual(_suffix_finder.find(m.b1.b2.v3[0]), 3)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.b2.v3[0]), 3)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.b2.v3[0]), 3)
 
         # Check for container at top level
         m.suffix[m.b1.b2.v3] = 4
-        assert _suffix_finder.find(m.b1.b2.v3[0]) == 4
+        self.assertEqual(_suffix_finder.find(m.b1.b2.v3[0]), 4)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.b2.v3[0]), 3)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.b2.v3[0]), 3)
 
         # Check for specific values at lowest level
         m.b1.b2.suffix[m.b1.b2.v3[0]] = 5
-        assert _suffix_finder.find(m.b1.b2.v3[0]) == 5
+        self.assertEqual(_suffix_finder.find(m.b1.b2.v3[0]), 5)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.b2.v3[0]), 5)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.b2.v3[0]), 5)
 
         # Check for specific values at top level
         m.suffix[m.b1.b2.v3[0]] = 6
-        assert _suffix_finder.find(m.b1.b2.v3[0]) == 6
+        self.assertEqual(_suffix_finder.find(m.b1.b2.v3[0]), 6)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.b2.v3[0]), 5)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.b2.v3[0]), 5)
 
         # Make sure we don't find default suffixes at lower levels
-        assert _suffix_finder.find(m.b1.v2) == 1
+        self.assertEqual(_suffix_finder.find(m.b1.v2), 1)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.v2), None)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.v2), None)
 
         # Make sure we don't find specific suffixes at lower levels
         m.b1.b2.suffix[m.v1] = 5
-        assert _suffix_finder.find(m.v1) == 1
+        self.assertEqual(_suffix_finder.find(m.v1), 1)
+        self.assertEqual(_suffix_b1_finder.find(m.v1), None)
+        self.assertEqual(_suffix_b2_finder.find(m.v1), None)
+
+        # Make sure we can look up Blocks and that they will match
+        # suffixes that they hold
+        self.assertEqual(_suffix_finder.find(m.b1.b2), 2)
+        self.assertEqual(_suffix_b1_finder.find(m.b1.b2), 2)
+        self.assertEqual(_suffix_b2_finder.find(m.b1.b2), 2)
+
+        self.assertEqual(_suffix_finder.find(m.b1), 1)
+        self.assertEqual(_suffix_b1_finder.find(m.b1), None)
+        self.assertEqual(_suffix_b2_finder.find(m.b1), None)
 
 
 if __name__ == "__main__":

pyomo/repn/plugins/nl_writer.py

@@ -510,8 +510,8 @@ def compile(self, column_order, row_order, obj_order, model_id):
 class CachingNumericSuffixFinder(SuffixFinder):
     scale = True
 
-    def __init__(self, name, default=None):
-        super().__init__(name, default)
+    def __init__(self, name, default=None, context=None):
+        super().__init__(name, default, context)
         self.suffix_cache = {}
 
     def __call__(self, obj):
@@ -646,7 +646,7 @@ def write(self, model):
         # Data structures to support variable/constraint scaling
         #
         if self.config.scale_model and 'scaling_factor' in suffix_data:
-            scaling_factor = CachingNumericSuffixFinder('scaling_factor', 1)
+            scaling_factor = CachingNumericSuffixFinder('scaling_factor', 1, model)
             scaling_cache = scaling_factor.suffix_cache
             del suffix_data['scaling_factor']
         else: