adds constraint forward as well

Author: Transurgeon

include/problem.h

@@ -18,12 +18,13 @@ typedef struct problem
     CSR_Matrix *stacked_jac;
 } problem;
 
-/* Takes ownership of objective and constraints - caller should not free them */
+/* Retains objective and constraints (shared ownership with caller) */
 problem *new_problem(expr *objective, expr **constraints, int n_constraints);
 void problem_allocate(problem *prob, const double *u);
 void free_problem(problem *prob);
 
 double problem_forward(problem *prob, const double *u);
+double *problem_constraint_forward(problem *prob, const double *u);
 double *problem_gradient(problem *prob, const double *u);
 CSR_Matrix *problem_jacobian(problem *prob, const double *u);
 

python/bindings.c

@@ -388,6 +388,56 @@ static PyObject *py_problem_forward(PyObject *self, PyObject *args)
     return Py_BuildValue("(dO)", obj_val, constraint_vals);
 }
 
+static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    double *constraint_vals =
+        problem_constraint_forward(prob, (const double *) PyArray_DATA(u_array));
+
+    PyObject *out = NULL;
+    if (prob->total_constraint_size > 0)
+    {
+        npy_intp size = prob->total_constraint_size;
+        out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+        if (!out)
+        {
+            Py_DECREF(u_array);
+            return NULL;
+        }
+        memcpy(PyArray_DATA((PyArrayObject *) out), constraint_vals,
+               size * sizeof(double));
+    }
+    else
+    {
+        npy_intp size = 0;
+        out = PyArray_SimpleNew(1, &size, NPY_DOUBLE);
+    }
+
+    Py_DECREF(u_array);
+    return out;
+}
+
 static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
 {
     PyObject *prob_capsule;
@@ -500,6 +550,7 @@ static PyMethodDef DNLPMethods[] = {
     {"make_problem", py_make_problem, METH_VARARGS, "Create problem from objective and constraints"},
     {"problem_allocate", py_problem_allocate, METH_VARARGS, "Allocate problem resources"},
     {"problem_forward", py_problem_forward, METH_VARARGS, "Evaluate objective and constraints"},
+    {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS, "Evaluate constraints only"},
     {"problem_gradient", py_problem_gradient, METH_VARARGS, "Compute objective gradient"},
     {"problem_jacobian", py_problem_jacobian, METH_VARARGS, "Compute constraint jacobian"},
     {NULL, NULL, 0, NULL}};

python/convert.py

@@ -132,6 +132,10 @@ def forward(self, u: np.ndarray) -> tuple[float, np.ndarray]:
         """Evaluate objective and constraints. Returns (obj_value, constraint_values)."""
         return diffengine.problem_forward(self._capsule, u)
 
+    def constraint_forward(self, u: np.ndarray) -> np.ndarray:
+        """Evaluate constraints only. Returns constraint_values array."""
+        return diffengine.problem_constraint_forward(self._capsule, u)
+
     def gradient(self, u: np.ndarray) -> np.ndarray:
         """Compute gradient of objective. Returns gradient array."""
         return diffengine.problem_gradient(self._capsule, u)

python/tests/test_problem_convert.py

@@ -67,6 +67,29 @@ def test_problem_jacobian_lowlevel():
     assert np.allclose(jac.toarray(), expected_jac)
 
 
+def test_problem_constraint_forward_lowlevel():
+    """Test problem_constraint_forward for constraint values only (low-level)."""
+    n_vars = 2
+    x = diffengine.make_variable(n_vars, 1, 0, n_vars)
+
+    log_obj = diffengine.make_log(x)
+    objective = diffengine.make_sum(log_obj, -1)
+
+    log_c = diffengine.make_log(x)
+    exp_c = diffengine.make_exp(x)
+    constraints = [log_c, exp_c]
+
+    prob = diffengine.make_problem(objective, constraints)
+    u = np.array([2.0, 4.0])
+    diffengine.problem_allocate(prob, u)
+
+    constraint_vals = diffengine.problem_constraint_forward(prob, u)
+
+    # Expected: [log(2), log(4), exp(2), exp(4)]
+    expected = np.concatenate([np.log(u), np.exp(u)])
+    assert np.allclose(constraint_vals, expected)
+
+
 def test_problem_no_constraints_lowlevel():
     """Test Problem with no constraints (low-level)."""
     n_vars = 3
@@ -320,11 +343,37 @@ def test_problem_repeated_evaluations():
     assert np.allclose(grad2, 1.0 / u2)
 
 
+def test_problem_constraint_forward():
+    """Test Problem.constraint_forward for constraint values only."""
+    x = cp.Variable(2)
+    obj = cp.sum(cp.log(x))
+    constraints = [
+        cp.log(x),
+        cp.exp(x),
+    ]
+
+    cvxpy_prob = cp.Problem(cp.Minimize(obj), constraints)
+    prob = Problem(cvxpy_prob)
+
+    u = np.array([2.0, 4.0])
+    prob.allocate(u)
+
+    # Test constraint_forward
+    constraint_vals = prob.constraint_forward(u)
+    expected = np.concatenate([np.log(u), np.exp(u)])
+    assert np.allclose(constraint_vals, expected)
+
+    # Verify it gives same result as forward's constraint values
+    _, forward_constraint_vals = prob.forward(u)
+    assert np.allclose(constraint_vals, forward_constraint_vals)
+
+
 if __name__ == "__main__":
     # Low-level tests
     test_problem_forward_lowlevel()
     test_problem_gradient_lowlevel()
     test_problem_jacobian_lowlevel()
+    test_problem_constraint_forward_lowlevel()
     test_problem_no_constraints_lowlevel()
     # Problem class tests
     test_problem_single_constraint()
@@ -334,4 +383,5 @@ def test_problem_repeated_evaluations():
     test_problem_no_constraints_convert()
     test_problem_larger_scale()
     test_problem_repeated_evaluations()
+    test_problem_constraint_forward()
     print("All problem tests passed!")

src/problem.c

@@ -63,17 +63,19 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
     problem *prob = (problem *) calloc(1, sizeof(problem));
     if (!prob) return NULL;
 
-    /* Take ownership of objective (no retain - caller transfers ownership) */
+    /* Retain objective (shared ownership with caller) */
     prob->objective = objective;
+    expr_retain(objective);
 
-    /* Copy constraints array (take ownership, no retain) */
+    /* Copy and retain constraints array */
     prob->n_constraints = n_constraints;
     if (n_constraints > 0)
     {
         prob->constraints = (expr **) malloc(n_constraints * sizeof(expr *));
         for (int i = 0; i < n_constraints; i++)
         {
             prob->constraints[i] = constraints[i];
+            expr_retain(constraints[i]);
         }
     }
     else
@@ -139,14 +141,11 @@ void free_problem(problem *prob)
     free(prob->gradient_values);
     free_csr_matrix(prob->stacked_jac);
 
-    /* Free expression trees with shared visited set to handle node sharing */
-    VisitedSet visited;
-    visited_init(&visited);
-
-    free_expr_tree_visited(prob->objective, &visited);
+    /* Release expression references (decrements refcount) */
+    free_expr(prob->objective);
     for (int i = 0; i < prob->n_constraints; i++)
     {
-        free_expr_tree_visited(prob->constraints[i], &visited);
+        free_expr(prob->constraints[i]);
     }
     free(prob->constraints);
 
@@ -173,6 +172,21 @@ double problem_forward(problem *prob, const double *u)
     return obj_val;
 }
 
+double *problem_constraint_forward(problem *prob, const double *u)
+{
+    /* Evaluate constraints only and copy values */
+    int offset = 0;
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr *c = prob->constraints[i];
+        c->forward(c, u);
+        memcpy(prob->constraint_values + offset, c->value, c->size * sizeof(double));
+        offset += c->size;
+    }
+
+    return prob->constraint_values;
+}
+
 double *problem_gradient(problem *prob, const double *u)
 {
     /* Jacobian on objective */

tests/all_tests.c

@@ -125,6 +125,7 @@ int main(void)
     mu_run_test(test_problem_gradient, tests_run);
     mu_run_test(test_problem_jacobian, tests_run);
     mu_run_test(test_problem_jacobian_multi, tests_run);
+    mu_run_test(test_problem_constraint_forward, tests_run);
 
     printf("\n=== All %d tests passed ===\n", tests_run);
 

tests/problem/test_problem.h

@@ -37,8 +37,12 @@ const char *test_problem_new_free(void)
     mu_assert("n_constraints wrong", prob->n_constraints == 1);
     mu_assert("total_constraint_size wrong", prob->total_constraint_size == 3);
 
-    /* Free problem (owns and frees all expressions) */
+    /* Free problem and expressions (shared ownership) */
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_x);
+    free_expr(x);
+    free_expr(x_constraint);
 
     return 0;
 }
@@ -73,6 +77,10 @@ const char *test_problem_forward(void)
     mu_assert("constraint[2] wrong", fabs(prob->constraint_values[2] - 3.0) < 1e-10);
 
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_x);
+    free_expr(x);
+    free_expr(x_constraint);
 
     return 0;
 }
@@ -99,6 +107,9 @@ const char *test_problem_gradient(void)
     mu_assert("grad[2] wrong", fabs(grad[2] - 0.25) < 1e-10);
 
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_x);
+    free_expr(x);
 
     return 0;
 }
@@ -146,6 +157,11 @@ const char *test_problem_jacobian(void)
     mu_assert("jac->x[1] wrong", fabs(jac->x[1] - 0.25) < 1e-10);
 
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_obj);
+    free_expr(x_obj);
+    free_expr(log_c1);
+    free_expr(x_c1);
 
     return 0;
 }
@@ -221,6 +237,60 @@ const char *test_problem_jacobian_multi(void)
     mu_assert("jac->x[3] wrong", fabs(jac->x[3] - exp(4.0)) < 1e-10);
 
     free_problem(prob);
+    free_expr(objective);
+    free_expr(log_obj);
+    free_expr(log_c1);
+    free_expr(exp_c2);
+    free_expr(x);
+
+    return 0;
+}
+
+/*
+ * Test problem_constraint_forward: evaluate constraints only
+ *   Constraint 1: log(x) -> [log(2), log(4)]
+ *   Constraint 2: exp(x) -> [exp(2), exp(4)]
+ */
+const char *test_problem_constraint_forward(void)
+{
+    int n_vars = 2;
+
+    /* Shared variable */
+    expr *x = new_variable(2, 1, 0, n_vars);
+
+    /* Objective: sum(log(x)) */
+    expr *log_obj = new_log(x);
+    expr *objective = new_sum(log_obj, -1);
+
+    /* Constraint 1: log(x) */
+    expr *log_c1 = new_log(x);
+
+    /* Constraint 2: exp(x) */
+    expr *exp_c2 = new_exp(x);
+
+    expr *constraints[2] = {log_c1, exp_c2};
+
+    problem *prob = new_problem(objective, constraints, 2);
+
+    double u[2] = {2.0, 4.0};
+    problem_allocate(prob, u);
+
+    double *constraint_vals = problem_constraint_forward(prob, u);
+
+    /* Check constraint values:
+     * [log(2), log(4), exp(2), exp(4)]
+     */
+    mu_assert("constraint[0] wrong", fabs(constraint_vals[0] - log(2.0)) < 1e-10);
+    mu_assert("constraint[1] wrong", fabs(constraint_vals[1] - log(4.0)) < 1e-10);
+    mu_assert("constraint[2] wrong", fabs(constraint_vals[2] - exp(2.0)) < 1e-10);
+    mu_assert("constraint[3] wrong", fabs(constraint_vals[3] - exp(4.0)) < 1e-10);
+
+    free_problem(prob);
+    free_expr(objective);
+    free_expr(log_obj);
+    free_expr(log_c1);
+    free_expr(exp_c2);
+    free_expr(x);
 
     return 0;
 }