Added methods to directly add costs and constraints.

src/hippopt/base/optimization_problem.py

@@ -38,11 +38,52 @@ def generate_optimization_objects(
             input_structure=input_structure
         )
 
+    def add_cost(
+        self,
+        expression: cs.MX | Generator[cs.MX, None, None],
+        scaling: float | cs.MX = 1.0,
+    ):
+        if isinstance(expression, types.GeneratorType):
+            for expr in expression:
+                self.add_cost(expr, scaling)
+        else:
+            assert isinstance(expression, cs.MX)
+            if expression.is_op(cs.OP_LE) or expression.is_op(cs.OP_LT):
+                raise ValueError(
+                    "The conversion from an inequality to a cost is not yet supported"
+                )
+            if expression.is_op(cs.OP_EQ):
+                error_expr = expression.dep(0) - expression.dep(1)
+                self._solver.add_cost(scaling * cs.sumsqr(error_expr))
+            else:
+                self._solver.add_cost(scaling * expression)  # noqa
+
+    def add_constraint(
+        self,
+        expression: cs.MX | Generator[cs.MX, None, None],
+        expected_value: float | cs.MX = 0.0,
+    ):
+        if isinstance(expression, types.GeneratorType):
+            for expr in expression:
+                self.add_constraint(expr, expected_value)
+        else:
+            assert isinstance(expression, cs.MX)
+            if (
+                expression.is_op(cs.OP_LE)
+                or expression.is_op(cs.OP_LT)
+                or expression.is_op(cs.OP_EQ)
+            ):
+                self._solver.add_constraint(expression)
+            else:
+                if not expression.is_scalar():
+                    raise ValueError("The input expression is not supported.")
+                self._solver.add_constraint(expression == expected_value)  # noqa
+
     def add_expression(
         self,
         mode: ExpressionType,
         expression: cs.MX | Generator[cs.MX, None, None],
-        expected_value: float = 0.0,
+        **kwargs,
     ):
         if isinstance(expression, types.GeneratorType):
             for expr in expression:
@@ -51,29 +92,10 @@ def add_expression(
             assert isinstance(expression, cs.MX)
             match mode:
                 case ExpressionType.subject_to:
-                    if (
-                        expression.is_op(cs.OP_LE)
-                        or expression.is_op(cs.OP_LT)
-                        or expression.is_op(cs.OP_EQ)
-                    ):
-                        self._solver.add_constraint(expression)
-                    else:
-                        if not expression.is_scalar():
-                            raise ValueError("The input expression is not supported.")
-                        self._solver.add_constraint(
-                            expression == expected_value  # noqa
-                        )
+                    self.add_constraint(expression, **kwargs)
 
                 case ExpressionType.minimize:
-                    if expression.is_op(cs.OP_LE) or expression.is_op(cs.OP_LT):
-                        raise ValueError(
-                            "The conversion from an inequality to a cost is not yet supported"
-                        )
-                    if expression.is_op(cs.OP_EQ):
-                        error_expr = expression.dep(0) - expression.dep(1)
-                        self._solver.add_cost(cs.sumsqr(error_expr))
-                    else:
-                        self._solver.add_cost(expression)
+                    self.add_cost(expression, **kwargs)
                 case _:
                     pass
 

test/test_optimization_problem.py

@@ -144,23 +144,17 @@ def test_opti_solver_with_parameters_and_lists():
         c.append(20.0 * np.random.rand(3) - 10.0)
         initial_guess[j].parameter = c[j]
 
-    problem.add_expression(
-        mode=ExpressionType.minimize,
-        expression=(
-            a[j][k] * cs.power(var[j].composite.variable[k], 2)
-            + b[j][k] * var[j].composite.variable[k]
-            for j in range(len(initial_guess))
-            for k in range(0, 3)
-        ),
+    problem.add_cost(
+        a[j][k] * cs.power(var[j].composite.variable[k], 2)
+        + b[j][k] * var[j].composite.variable[k]
+        for j in range(len(initial_guess))
+        for k in range(0, 3)
     )
 
-    problem.add_expression(
-        mode=ExpressionType.subject_to,
-        expression=(  # noqa
-            var[j].composite.variable[k] >= c[j][k]
-            for j in range(len(initial_guess))
-            for k in range(3)
-        ),
+    problem.add_constraint(
+        var[j].composite.variable[k] >= c[j][k]  # noqa
+        for j in range(len(initial_guess))
+        for k in range(3)
     )
 
     problem.solver().set_initial_guess(initial_guess=initial_guess)
@@ -221,7 +215,9 @@ def test_switch_costs():
         ExpressionType.subject_to, new_variables.x + new_variables.y == a - 1
     )  # noqa
     new_problem.add_expression(
-        ExpressionType.subject_to, new_variables.x * new_variables.x
+        ExpressionType.subject_to,
+        new_variables.x * new_variables.x + 1,
+        expected_value=1,
     )
     output, cost_value = new_problem.solver().solve()
     expected_cost = a * (a - 1) ** 2
@@ -250,7 +246,9 @@ def test_switch_constraints():
     new_problem.add_expression(
         ExpressionType.subject_to, new_variables.x + new_variables.y == a - 1
     )  # noqa
-    new_problem.add_expression(ExpressionType.minimize, new_variables.x == 5)
+    new_problem.add_expression(
+        ExpressionType.minimize, new_variables.x == 5, scaling=1.0
+    )
     output, cost_value = new_problem.solver().solve()
     assert cost_value == pytest.approx(expected=initial_cost_value, rel=0.1)
     assert output.x == pytest.approx(initial_output.x)