From 23d316e857b35e81d990114a918f215e5e3a556e Mon Sep 17 00:00:00 2001
From: Bernard Knueven <Bernard.Knueven@nrel.gov>
Date: Tue, 18 Jun 2024 21:52:52 -0600
Subject: [PATCH] finding scaling_factors in PyomoNLP

---
 .../contrib/pynumero/interfaces/pyomo_nlp.py  | 43 +++++++++----------
 1 file changed, 21 insertions(+), 22 deletions(-)

diff --git a/pyomo/contrib/pynumero/interfaces/pyomo_nlp.py b/pyomo/contrib/pynumero/interfaces/pyomo_nlp.py
index e12d0cf568b..a793419efc0 100644
--- a/pyomo/contrib/pynumero/interfaces/pyomo_nlp.py
+++ b/pyomo/contrib/pynumero/interfaces/pyomo_nlp.py
@@ -20,6 +20,7 @@
 from pyomo.common.tempfiles import TempfileManager
 from pyomo.opt import WriterFactory
 import pyomo.core.base as pyo
+from pyomo.core.base.suffix import SuffixFinder
 from pyomo.common.collections import ComponentMap
 from pyomo.common.env import CtypesEnviron
 from pyomo.solvers.amplfunc_merge import amplfunc_merge
@@ -298,34 +299,32 @@ 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
+        val = SuffixFinder('scaling_factor').find(obj)
+        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_suffix_finder = SuffixFinder('scaling_factor')
+        primals_scaling = np.ones(self.n_primals())
+        ret = None
+        for i, v in enumerate(self.get_pyomo_variables()):
+            val = scaling_suffix_finder.find(v)
+            if val is not None:
+                primals_scaling[i] = val
+                ret = primals_scaling
+        return ret
 
     # 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_suffix_finder = SuffixFinder('scaling_factor')
+        constraints_scaling = np.ones(self.n_constraints())
+        ret = None
+        for i, c in enumerate(self.get_pyomo_constraints()):
+            val = scaling_suffix_finder.find(c)
+            if val is not None:
+                constraints_scaling[i] = val
+                ret = constraints_scaling
+        return ret
 
     def extract_subvector_grad_objective(self, pyomo_variables):
         """Compute the gradient of the objective and return the entries