From 2c69d25b8d9e79c2c1afe7d11adeb2d50f601517 Mon Sep 17 00:00:00 2001 From: William Zijie Zhang Date: Tue, 13 Jan 2026 20:47:48 -0500 Subject: [PATCH 1/2] adds matmul shortcut and jacobian_only API for problem --- include/problem.h | 1 + python/bindings.c | 3 +++ src/bivariate/left_matmul.c | 45 ++++++++++++++++++++++++++----- src/bivariate/right_matmul.c | 46 +++++++++++++++++++++++++++----- src/dnlp_diff_engine/__init__.py | 13 ++++++--- src/problem.c | 21 +++++++++++++++ 6 files changed, 113 insertions(+), 16 deletions(-) diff --git a/include/problem.h b/include/problem.h index 8f685b5..bf6a71c 100644 --- a/include/problem.h +++ b/include/problem.h @@ -26,6 +26,7 @@ typedef struct problem /* Retains objective and constraints (shared ownership with caller) */ problem *new_problem(expr *objective, expr **constraints, int n_constraints); void problem_init_derivatives(problem *prob); +void problem_init_jacobian_only(problem *prob); /* Skip Hessian setup */ void free_problem(problem *prob); double problem_objective_forward(problem *prob, const double *u); diff --git a/python/bindings.c b/python/bindings.c index 1042701..4f60890 100644 --- a/python/bindings.c +++ b/python/bindings.c @@ -38,6 +38,7 @@ #include "problem/gradient.h" #include "problem/hessian.h" #include "problem/init_derivatives.h" +#include "problem/init_jacobian_only.h" #include "problem/jacobian.h" #include "problem/make_problem.h" #include "problem/objective_forward.h" @@ -94,6 +95,8 @@ static PyMethodDef DNLPMethods[] = { "Create problem from objective and constraints"}, {"problem_init_derivatives", py_problem_init_derivatives, METH_VARARGS, "Initialize derivative structures"}, + {"problem_init_jacobian_only", py_problem_init_jacobian_only, METH_VARARGS, + "Initialize Jacobian structure only (skip Hessian)"}, {"problem_objective_forward", py_problem_objective_forward, METH_VARARGS, "Evaluate objective only"}, {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS, diff --git a/src/bivariate/left_matmul.c b/src/bivariate/left_matmul.c index 517f48c..151423b 100644 --- a/src/bivariate/left_matmul.c +++ b/src/bivariate/left_matmul.c @@ -64,11 +64,37 @@ static void jacobian_init(expr *node) expr *x = node->left; left_matmul_expr *lin_node = (left_matmul_expr *) node; - /* initialize child's jacobian and precompute sparsity of its transpose */ + /* initialize child's jacobian */ x->jacobian_init(x); - lin_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork); - /* precompute sparsity of this node's jacobian */ + /* Fast path: child is a variable with identity Jacobian. + * For y = A @ x, J_y = A @ J_x where J_x is identity shifted by var_id. + * So J_y is just A with column indices shifted by var_id. + * This avoids O(nnz(A) * n) CSR-CSC multiplication. */ + if (x->var_id != NOT_A_VARIABLE) + { + CSR_Matrix *A = lin_node->A; + node->jacobian = new_csr_matrix(A->m, node->n_vars, A->nnz); + + /* Copy row pointers directly */ + memcpy(node->jacobian->p, A->p, (A->m + 1) * sizeof(int)); + + /* Copy column indices with offset by var_id */ + for (int k = 0; k < A->nnz; k++) + { + node->jacobian->i[k] = A->i[k] + x->var_id; + } + + /* Copy values directly (A is constant, variable Jacobian entries are 1.0) */ + memcpy(node->jacobian->x, A->x, A->nnz * sizeof(double)); + + /* Signal fast path by leaving CSC_work as NULL */ + lin_node->CSC_work = NULL; + return; + } + + /* Slow path: general case - compute A @ J_x via matrix multiplication */ + lin_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork); node->jacobian = csr_csc_matmul_alloc(lin_node->A, lin_node->CSC_work); } @@ -77,11 +103,18 @@ static void eval_jacobian(expr *node) expr *x = node->left; left_matmul_expr *lin_node = (left_matmul_expr *) node; - /* evaluate child's jacobian and convert to CSC */ + /* evaluate child's jacobian */ x->eval_jacobian(x); - csr_to_csc_fill_values(x->jacobian, lin_node->CSC_work, node->iwork); - /* compute this node's jacobian */ + /* Fast path: child is a variable. + * Values already set in jacobian_init (A is constant, variable Jacobian is 1.0) */ + if (lin_node->CSC_work == NULL) + { + return; + } + + /* Slow path: convert to CSC and multiply */ + csr_to_csc_fill_values(x->jacobian, lin_node->CSC_work, node->iwork); csr_csc_matmul_fill_values(lin_node->A, lin_node->CSC_work, node->jacobian); } diff --git a/src/bivariate/right_matmul.c b/src/bivariate/right_matmul.c index 0c349fb..b6aa450 100644 --- a/src/bivariate/right_matmul.c +++ b/src/bivariate/right_matmul.c @@ -54,11 +54,38 @@ static void jacobian_init(expr *node) expr *x = node->left; right_matmul_expr *right_node = (right_matmul_expr *) node; - /* initialize child's jacobian and precompute sparsity of its transpose */ + /* initialize child's jacobian */ x->jacobian_init(x); - right_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork); - /* precompute sparsity of this node's jacobian */ + /* Fast path: child is a variable with identity Jacobian. + * For y = x @ A, we compute vec(y) = B @ vec(x) where B = A^T kron I_p. + * J_y = B @ J_x where J_x is identity shifted by var_id. + * So J_y is just B with column indices shifted by var_id. + * This avoids O(nnz(B) * n) CSR-CSC multiplication. */ + if (x->var_id != NOT_A_VARIABLE) + { + CSR_Matrix *B = right_node->B; + node->jacobian = new_csr_matrix(B->m, node->n_vars, B->nnz); + + /* Copy row pointers directly */ + memcpy(node->jacobian->p, B->p, (B->m + 1) * sizeof(int)); + + /* Copy column indices with offset by var_id */ + for (int k = 0; k < B->nnz; k++) + { + node->jacobian->i[k] = B->i[k] + x->var_id; + } + + /* Copy values directly (B is constant, variable Jacobian entries are 1.0) */ + memcpy(node->jacobian->x, B->x, B->nnz * sizeof(double)); + + /* Signal fast path by leaving CSC_work as NULL */ + right_node->CSC_work = NULL; + return; + } + + /* Slow path: general case - compute B @ J_x via matrix multiplication */ + right_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork); node->jacobian = csr_csc_matmul_alloc(right_node->B, right_node->CSC_work); } @@ -67,11 +94,18 @@ static void eval_jacobian(expr *node) expr *x = node->left; right_matmul_expr *right_node = (right_matmul_expr *) node; - /* evaluate child's jacobian and convert to CSC*/ + /* evaluate child's jacobian */ x->eval_jacobian(x); - csr_to_csc_fill_values(x->jacobian, right_node->CSC_work, node->iwork); - /* compute this node's jacobian */ + /* Fast path: child is a variable. + * Values already set in jacobian_init (B is constant, variable Jacobian is 1.0) */ + if (right_node->CSC_work == NULL) + { + return; + } + + /* Slow path: convert to CSC and multiply */ + csr_to_csc_fill_values(x->jacobian, right_node->CSC_work, node->iwork); csr_csc_matmul_fill_values(right_node->B, right_node->CSC_work, node->jacobian); } diff --git a/src/dnlp_diff_engine/__init__.py b/src/dnlp_diff_engine/__init__.py index a8cda90..92236d8 100644 --- a/src/dnlp_diff_engine/__init__.py +++ b/src/dnlp_diff_engine/__init__.py @@ -68,8 +68,8 @@ def _convert_multiply(expr, children): # multiply has args: [left, right] left_arg, right_arg = expr.args - # Check if left is a constant - if isinstance(left_arg, cp.Constant): + # Check if left is a constant (use is_constant() to handle Promote and other wrappers) + if left_arg.is_constant(): value = np.asarray(left_arg.value, dtype=np.float64) # Scalar constant @@ -83,7 +83,7 @@ def _convert_multiply(expr, children): return _diffengine.make_const_vector_mult(children[1], vector) # Check if right is a constant - elif isinstance(right_arg, cp.Constant): + elif right_arg.is_constant(): value = np.asarray(right_arg.value, dtype=np.float64) # Scalar constant @@ -323,10 +323,15 @@ def __init__(self, cvxpy_problem: cp.Problem): self._allocated = False def init_derivatives(self): - """Initialize derivative structures. Must be called before forward/gradient/jacobian.""" + """Initialize derivative structures. Must be called before forward/gradient/jacobian/hessian.""" _diffengine.problem_init_derivatives(self._capsule) self._allocated = True + def init_jacobian_only(self): + """Initialize Jacobian structure only (skip Hessian). Use when only Jacobian is needed.""" + _diffengine.problem_init_jacobian_only(self._capsule) + self._allocated = True + def objective_forward(self, u: np.ndarray) -> float: """Evaluate objective. Returns obj_value float.""" return _diffengine.problem_objective_forward(self._capsule, u) diff --git a/src/problem.c b/src/problem.c index 90763a2..7c5f3bb 100644 --- a/src/problem.c +++ b/src/problem.c @@ -104,6 +104,27 @@ void problem_init_derivatives(problem *prob) free(iwork); } +void problem_init_jacobian_only(problem *prob) +{ + // ------------------------------------------------------------------------------- + // Jacobian structure only (skip Hessian) + // ------------------------------------------------------------------------------- + prob->objective->jacobian_init(prob->objective); + int nnz = 0; + for (int i = 0; i < prob->n_constraints; i++) + { + expr *c = prob->constraints[i]; + c->jacobian_init(c); + nnz += c->jacobian->nnz; + } + + prob->jacobian = new_csr_matrix(prob->total_constraint_size, prob->n_vars, nnz); + + /* Leave Hessian structures as NULL */ + prob->lagrange_hessian = NULL; + prob->hess_idx_map = NULL; +} + static void problem_lagrange_hess_fill_sparsity(problem *prob, int *iwork) { expr **constrs = prob->constraints; From f9389e0dd55b7d78ea078c571f5a7d204ca377b3 Mon Sep 17 00:00:00 2001 From: William Zijie Zhang Date: Tue, 13 Jan 2026 20:48:07 -0500 Subject: [PATCH 2/2] adds new python problem file for binding --- python/problem/init_jacobian_only.h | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 python/problem/init_jacobian_only.h diff --git a/python/problem/init_jacobian_only.h b/python/problem/init_jacobian_only.h new file mode 100644 index 0000000..646ad71 --- /dev/null +++ b/python/problem/init_jacobian_only.h @@ -0,0 +1,27 @@ +#ifndef PROBLEM_INIT_JACOBIAN_ONLY_H +#define PROBLEM_INIT_JACOBIAN_ONLY_H + +#include "common.h" + +static PyObject *py_problem_init_jacobian_only(PyObject *self, PyObject *args) +{ + PyObject *prob_capsule; + if (!PyArg_ParseTuple(args, "O", &prob_capsule)) + { + return NULL; + } + + problem *prob = + (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME); + if (!prob) + { + PyErr_SetString(PyExc_ValueError, "invalid problem capsule"); + return NULL; + } + + problem_init_jacobian_only(prob); + + Py_RETURN_NONE; +} + +#endif /* PROBLEM_INIT_JACOBIAN_ONLY_H */