Combine Objective and Constraint functions

src/condor/contrib.py

@@ -912,6 +912,7 @@ def __call__(self, *args, **kwargs):
         # don't have access to instance yet.
         # print(cls, "__call__")
         wrapper_object = super().__call__(*args, **kwargs)
+        wrapper_object.condor_model.wrapper = wrapper_object
         return wrapper_object.condor_model
 
 

src/condor/implementations/iterative.py

@@ -507,69 +507,144 @@ class ScipyCG(ScipyMinimizeBase):
     method_string = "CG"
 
 
+class ObjectiveConstraintData:
+    def __init__(self, x, p, obj_expr):
+        self.x = x
+        self.p = p
+        self.obj_expr = obj_expr
+        self.constr_exprs = []
+        self.constr_dicts = []
+        self._output_cache = (None, None)
+        self._output_jac_cache = (None, None)
+
+    @property
+    def constraints(self):
+        yield from self.constr_dicts
+
+    def add_constraint(self, expr, type):
+        self.constr_exprs.append(expr)
+        self.constr_dicts.append({"type": type})
+
+    def construct(self):
+        obj_constr_expr = concat([self.obj_expr, *self.constr_exprs])
+        self.obj_constr_func = expression_to_operator(
+            [self.x, self.p],
+            obj_constr_expr,
+            "obj_constr",
+        )
+        self.obj_constr_jac = expression_to_operator(
+            [self.x, self.p],
+            jacobian(obj_constr_expr, self.x),
+            "obj_constr_jac",
+        )
+        for i, d in enumerate(self.constr_dicts):
+            d["fun"] = self.constraint_func(i)
+            d["jac"] = self.constraint_jac(i)
+
+    def _get_output(self, x, p):
+        cx, cout = self._output_cache
+        if all(x == cx):
+            return cout
+        else:
+            out = self.obj_constr_func(x, p)
+            out_arr = out.toarray().squeeze()
+            self._output_cache = (x.copy(), out_arr.copy())
+            return out_arr
+
+    def _get_output_jac(self, x, p):
+        cx, cout = self._output_jac_cache
+        if all(x == cx):
+            return cout
+        else:
+            out = self.obj_constr_jac(x, p)
+            out_arr = out.toarray().squeeze()
+            self._output_jac_cache = (x.copy(), out_arr.copy())
+            return out_arr
+
+    def objective_func(self):
+        def obj(x, p):
+            return self._get_output(x, p)[0]
+
+        return obj
+
+    def constraint_func(self, i):
+        def constr(x, p):
+            return self._get_output(x, p)[i + 1]
+
+        return constr
+
+    def objective_jac(self):
+        def obj(x, p):
+            return self._get_output_jac(x, p)[0].T
+
+        return obj
+
+    def constraint_jac(self, i):
+        def constr(*x):
+            return self._get_output_jac(*x)[i + 1].T
+
+        return constr
+
+
 class ScipySLSQP(ScipyMinimizeBase):
     method_string = "SLSQP"
 
     def construct(self, model, *args, **kwargs):
         super().construct(model, *args, **kwargs)
-        self.equality_con_exprs = []
-        self.inequality_con_exprs = []
-        self.con = []
+
+        obj_constr_data = ObjectiveConstraintData(self.x, self.p, self.f)
         for g, lbg, ubg in zip(self.g_split, self.lbg, self.ubg):
             if lbg == ubg:
-                self.equality_con_exprs.append(g - lbg)
+                obj_constr_data.add_constraint(g - lbg, "eq")
             else:
                 if lbg > -np.inf:
-                    self.inequality_con_exprs.append(g - lbg)
+                    obj_constr_data.add_constraint(g - lbg, "ineq")
                 if ubg < np.inf:
-                    self.inequality_con_exprs.append(ubg - g)
-
-        if self.equality_con_exprs:
-            self.equality_con_expr = concat(self.equality_con_exprs)
-            self.eq_g_func = expression_to_operator(
-                [self.x, self.p],
-                self.equality_con_expr,
-                f"{model.__name__}_equality_constraint",
-            )
-            self.eq_g_jac_func = expression_to_operator(
-                [self.x, self.p],
-                jacobian(self.equality_con_expr, self.x),
-                f"{model.__name__}_equality_constraint_jac",
-            )
-            self.con.append(
-                dict(
-                    type="eq",
-                    fun=lambda x, p: self.eq_g_func(x, p).toarray().squeeze().T,
-                    jac=self.eq_g_jac_func,
-                )
-            )
+                    obj_constr_data.add_constraint(ubg - g, "ineq")
 
-        if self.inequality_con_exprs:
-            self.inequality_con_expr = concat(self.inequality_con_exprs)
-            self.ineq_g_func = expression_to_operator(
-                [self.x, self.p],
-                self.inequality_con_expr,
-                f"{model.__name__}_inequality_constraint",
-            )
-            self.ineq_g_jac_func = expression_to_operator(
-                [self.x, self.p],
-                jacobian(self.inequality_con_expr, self.x),
-                f"{model.__name__}_inequality_constraint_jac",
-            )
-            self.con.append(
-                dict(
-                    type="ineq",
-                    fun=lambda x, p: self.ineq_g_func(x, p).toarray().squeeze().T,
-                    jac=self.ineq_g_jac_func,
-                )
-            )
+        obj_constr_data.construct()
+        self.con = list(obj_constr_data.constraints)
+        self.obj = obj_constr_data.objective_func()
+        self.obj_jac = obj_constr_data.objective_jac()
 
     def prepare_constraints(self, extra_args):
         scipy_constraints = self.con
         for con in scipy_constraints:
             con["args"] = extra_args
         return scipy_constraints
 
+    def run_optimizer(self, model_instance):
+        extra_args = (self.eval_p,) if self.has_p else ([],)
+
+        scipy_constraints = self.prepare_constraints(extra_args)
+
+        if self.init_callback is not None:
+            self.init_callback(
+                model_instance.parameter,
+                {**self.options, "lbx": self.lbx, "ubx": self.ubx},
+            )
+
+        min_out = minimize(
+            self.obj,
+            self.x0,
+            jac=self.obj_jac,
+            method=self.method_string,
+            args=extra_args,
+            constraints=scipy_constraints,
+            bounds=np.vstack([self.lbx, self.ubx]).T,
+            # tol = 1E-9,
+            # options=dict(disp=True),
+            options=self.options,
+            callback=SciPyIterCallbackWrapper.create_or_none(
+                self.model, model_instance.parameter, self.iter_callback
+            ),
+        )
+
+        model_instance.bind_field(self.model.variable.wrap(min_out.x))
+        model_instance.objective = min_out.fun
+        self.x0 = min_out.x
+        self.stats = model_instance._stats = min_out
+
 
 class ScipyTrustConstr(ScipyMinimizeBase):
     method_string = "trust-constr"