Merge branch 'main' into nonlinear-to-pwl

pyomo/contrib/pynumero/interfaces/pyomo_grey_box_nlp.py

@@ -31,6 +31,7 @@
 )
 from pyomo.contrib.pynumero.interfaces.external_grey_box import ExternalGreyBoxBlock
 from pyomo.contrib.pynumero.interfaces.nlp_projections import ProjectedNLP
+from pyomo.core.base.suffix import SuffixFinder
 
 
 # Todo: make some of the numpy arrays not writable from __init__
@@ -226,13 +227,15 @@ def __init__(self, pyomo_model):
         else:
             need_scaling = True
 
-        self._primals_scaling = np.ones(self.n_primals())
-        scaling_suffix = pyomo_model.component('scaling_factor')
-        if scaling_suffix and scaling_suffix.ctype is pyo.Suffix:
-            need_scaling = True
-            for i, v in enumerate(self._pyomo_model_var_datas):
-                if v in scaling_suffix:
-                    self._primals_scaling[i] = scaling_suffix[v]
+        scaling_finder = SuffixFinder(
+            'scaling_factor', default=1.0, context=self._pyomo_model
+        )
+        self._primals_scaling = np.fromiter(
+            (scaling_finder.find(v) for v in self._pyomo_model_var_datas),
+            count=self.n_primals(),
+            dtype=float,
+        )
+        need_scaling = bool(scaling_finder.all_suffixes)
 
         self._constraints_scaling = BlockVector(len(nlps))
         for i, nlp in enumerate(nlps):

pyomo/contrib/pynumero/interfaces/pyomo_nlp.py

@@ -26,6 +26,7 @@
 from ..sparse.block_matrix import BlockMatrix
 from pyomo.contrib.pynumero.interfaces.ampl_nlp import AslNLP
 from pyomo.contrib.pynumero.interfaces.nlp import NLP
+from pyomo.core.base.suffix import SuffixFinder
 from .external_grey_box import ExternalGreyBoxBlock
 
 
@@ -297,35 +298,41 @@ def get_inequality_constraint_indices(self, constraints):
 
     # overloaded from NLP
     def get_obj_scaling(self):
-        obj = self.get_pyomo_objective()
-        scaling_suffix = self._pyomo_model.component('scaling_factor')
-        if scaling_suffix and scaling_suffix.ctype is pyo.Suffix:
-            if obj in scaling_suffix:
-                return scaling_suffix[obj]
-            return 1.0
-        return None
+        scaling_finder = SuffixFinder(
+            'scaling_factor', default=1.0, context=self._pyomo_model
+        )
+        val = scaling_finder.find(self.get_pyomo_objective())
+        if not scaling_finder.all_suffixes:
+            return None
+        return val
 
     # overloaded from NLP
     def get_primals_scaling(self):
-        scaling_suffix = self._pyomo_model.component('scaling_factor')
-        if scaling_suffix and scaling_suffix.ctype is pyo.Suffix:
-            primals_scaling = np.ones(self.n_primals())
-            for i, v in enumerate(self.get_pyomo_variables()):
-                if v in scaling_suffix:
-                    primals_scaling[i] = scaling_suffix[v]
-            return primals_scaling
-        return None
+        scaling_finder = SuffixFinder(
+            'scaling_factor', default=1.0, context=self._pyomo_model
+        )
+        primals_scaling = np.fromiter(
+            (scaling_finder.find(v) for v in self.get_pyomo_variables()),
+            count=self.n_primals(),
+            dtype=float,
+        )
+        if not scaling_finder.all_suffixes:
+            return None
+        return primals_scaling
 
     # overloaded from NLP
     def get_constraints_scaling(self):
-        scaling_suffix = self._pyomo_model.component('scaling_factor')
-        if scaling_suffix and scaling_suffix.ctype is pyo.Suffix:
-            constraints_scaling = np.ones(self.n_constraints())
-            for i, c in enumerate(self.get_pyomo_constraints()):
-                if c in scaling_suffix:
-                    constraints_scaling[i] = scaling_suffix[c]
-            return constraints_scaling
-        return None
+        scaling_finder = SuffixFinder(
+            'scaling_factor', default=1.0, context=self._pyomo_model
+        )
+        constraints_scaling = np.fromiter(
+            (scaling_finder.find(v) for v in self.get_pyomo_constraints()),
+            count=self.n_constraints(),
+            dtype=float,
+        )
+        if not scaling_finder.all_suffixes:
+            return None
+        return constraints_scaling
 
     def extract_subvector_grad_objective(self, pyomo_variables):
         """Compute the gradient of the objective and return the entries
@@ -606,13 +613,15 @@ def __init__(self, pyomo_model):
         else:
             need_scaling = True
 
-        self._primals_scaling = np.ones(self.n_primals())
-        scaling_suffix = self._pyomo_nlp._pyomo_model.component('scaling_factor')
-        if scaling_suffix and scaling_suffix.ctype is pyo.Suffix:
-            need_scaling = True
-            for i, v in enumerate(self.get_pyomo_variables()):
-                if v in scaling_suffix:
-                    self._primals_scaling[i] = scaling_suffix[v]
+        scaling_finder = SuffixFinder(
+            'scaling_factor', default=1.0, context=self._pyomo_model
+        )
+        self._primals_scaling = np.fromiter(
+            (scaling_finder.find(v) for v in self.get_pyomo_variables()),
+            count=self.n_primals(),
+            dtype=float,
+        )
+        need_scaling = bool(scaling_finder.all_suffixes)
 
         self._constraints_scaling = []
         pyomo_nlp_scaling = self._pyomo_nlp.get_constraints_scaling()

pyomo/contrib/pynumero/interfaces/tests/test_nlp.py

@@ -699,6 +699,42 @@ def test_indices_methods(self):
         dense_hess = hess.todense()
         self.assertTrue(np.array_equal(dense_hess, expected_hess))
 
+    def test_subblock_scaling(self):
+        m = pyo.ConcreteModel()
+        m.b = b = pyo.Block()
+        b.x = pyo.Var(bounds=(5e-17, 5e-16), initialize=1e-16)
+        b.scaling_factor = pyo.Suffix(direction=pyo.Suffix.EXPORT)
+        b.scaling_factor[b.x] = 1e16
+
+        b.c = pyo.Constraint(rule=b.x == 1e-16)
+        b.scaling_factor[b.c] = 1e16
+
+        b.o = pyo.Objective(expr=b.x)
+        b.scaling_factor[b.o] = 1e16
+
+        nlp = PyomoNLP(m)
+
+        assert nlp.get_obj_scaling() == 1e16
+        assert nlp.get_primals_scaling()[0] == 1e16
+        assert nlp.get_constraints_scaling()[0] == 1e16
+
+    def test_subblock_no_scaling(self):
+        m = pyo.ConcreteModel()
+        m.b = pyo.Block()
+        m.b.x = pyo.Var([1, 2], initialize={1: 100, 2: 20})
+
+        # Components so we don't have an empty NLP
+        m.b.eq = pyo.Constraint(expr=m.b.x[1] * m.b.x[2] == 2000)
+        m.b.obj = pyo.Objective(expr=m.b.x[1] ** 2 + m.b.x[2] ** 2)
+
+        m.scaling_factor = pyo.Suffix(direction=pyo.Suffix.EXPORT)
+        m.scaling_factor[m.b.x[1]] = 1e-2
+        m.scaling_factor[m.b.x[2]] = 1e-1
+
+        nlp = PyomoNLP(m.b)
+        scaling = nlp.get_primals_scaling()
+        assert scaling is None
+
     def test_no_objective(self):
         m = pyo.ConcreteModel()
         m.x = pyo.Var()

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__":