diff --git a/pyomo/contrib/pynumero/interfaces/pyomo_grey_box_nlp.py b/pyomo/contrib/pynumero/interfaces/pyomo_grey_box_nlp.py index e6ed40e9974..66cf99ea862 100644 --- a/pyomo/contrib/pynumero/interfaces/pyomo_grey_box_nlp.py +++ b/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): diff --git a/pyomo/contrib/pynumero/interfaces/pyomo_nlp.py b/pyomo/contrib/pynumero/interfaces/pyomo_nlp.py index e12d0cf568b..725435619ad 100644 --- a/pyomo/contrib/pynumero/interfaces/pyomo_nlp.py +++ b/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() diff --git a/pyomo/contrib/pynumero/interfaces/tests/test_nlp.py b/pyomo/contrib/pynumero/interfaces/tests/test_nlp.py index 4f735e06de7..a291ef1151a 100644 --- a/pyomo/contrib/pynumero/interfaces/tests/test_nlp.py +++ b/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()