Merge pull request #3295 from bknueven/pyomonlp-scaling

PyomoNLP scaling factors on sub-blocks

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()