[WIP] efficient assembly of problem hessian

docs/problem_struct_design.md

@@ -37,7 +37,7 @@ typedef struct problem
     /* Pre-allocated storage */
     double *constraint_values;
     double *gradient_values;   /* Dense gradient array */
-    CSR_Matrix *stacked_jac;
+    CSR_Matrix *jacobian;
 } problem;
 
 problem *new_problem(expr *objective, expr **constraints, int n_constraints);
@@ -85,16 +85,16 @@ void problem_allocate(problem *prob, const double *u)
     }
 
     /* Allocate stacked jacobian with total_constraint_size rows */
-    prob->stacked_jac = alloc_csr(prob->total_constraint_size, prob->n_vars, total_nnz);
+    prob->jacobian = alloc_csr(prob->total_constraint_size, prob->n_vars, total_nnz);
 
     /* Note: The actual nnz may be smaller after evaluation due to
-     * cancellations. Update stacked_jac->nnz after problem_jacobian(). */
+     * cancellations. Update jacobian->nnz after problem_jacobian(). */
 }
 ```
 
 ### `free_problem`
 - Call `free_expr` on objective and all constraints (decrements refcount)
-- Free allocated arrays and stacked_jac
+- Free allocated arrays and jacobian
 
 ### `problem_forward`
 ```c
@@ -127,7 +127,7 @@ double problem_forward(problem *prob, const double *u)
 - Stack CSR matrices vertically:
   - Total rows = `total_constraint_size`
   - Copy row pointers with offset, copy col indices and values
-- Lazy allocate/reallocate `stacked_jac` based on total nnz
+- Lazy allocate/reallocate `jacobian` based on total nnz
 
 ---
 
@@ -236,6 +236,6 @@ class Problem:
 - **Two-phase init**: `new_problem` creates struct, `problem_allocate` allocates arrays
   - Constraint values array: size = `total_constraint_size`
   - Jacobian: initialize all constraint jacobians first, count total nnz, allocate CSR
-  - The allocated nnz may be a slight overestimate; update `stacked_jac->nnz` after evaluation
+  - The allocated nnz may be a slight overestimate; update `jacobian->nnz` after evaluation
 - **Hessian**: Deferred - not allocated in this design (to be added later)
 - **API**: Returns internal pointers (caller should NOT free)

include/problem.h

@@ -17,7 +17,10 @@ typedef struct problem
     double *gradient_values;
 
     /* Allocated by problem_init_derivatives */
-    CSR_Matrix *stacked_jac;
+    CSR_Matrix *jacobian;
+    CSR_Matrix *lagrange_hessian;
+    int *hess_idx_map; /* Maps all wsum_hess nnz to lagrange_hessian (obj +
+                          constraints) */
 } problem;
 
 /* Retains objective and constraints (shared ownership with caller) */
@@ -29,5 +32,6 @@ double problem_objective_forward(problem *prob, const double *u);
 void problem_constraint_forward(problem *prob, const double *u);
 void problem_gradient(problem *prob);
 void problem_jacobian(problem *prob);
+void problem_hessian(problem *prob, double obj_w, const double *w);
 
 #endif

include/utils/CSR_Matrix.h

@@ -93,7 +93,8 @@ void sum_all_rows_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A, CSR_Matrix
 // void sum_all_rows_csr_fill_values(const CSR_Matrix *A, CSR_Matrix *C,
 //                                  const int *idx_map);
 
-/* Fill accumulator for summing rows using precomputed idx_map for each nnz of A*/
+/* Fill accumulator for summing rows using precomputed idx_map for each nnz of A.
+   Must memset accumulator to zero before calling. */
 void idx_map_accumulator(const CSR_Matrix *A, const int *idx_map,
                          double *accumulator);
 

python/problem/jacobian.h

@@ -33,7 +33,7 @@ static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
 
     problem_jacobian(prob);
 
-    CSR_Matrix *jac = prob->stacked_jac;
+    CSR_Matrix *jac = prob->jacobian;
     npy_intp nnz = jac->nnz;
     npy_intp m_plus_1 = jac->m + 1;
 

src/affine/sum.c

@@ -100,6 +100,7 @@ static void eval_jacobian(expr *node)
 
     /* we have precomputed an idx map between the nonzeros of the child's jacobian
        and this node's jacobian, so we just accumulate accordingly */
+    memset(node->jacobian->x, 0, node->jacobian->nnz * sizeof(double));
     idx_map_accumulator(x->jacobian, ((sum_expr *) node)->idx_map,
                         node->jacobian->x);
 }

src/affine/variable.c

@@ -25,6 +25,19 @@ static void eval_jacobian(expr *node)
     (void) node;
 }
 
+static void wsum_hess_init(expr *node)
+{
+    /* Variables have zero Hessian */
+    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 0);
+}
+
+static void wsum_hess_eval(expr *node, const double *w)
+{
+    /* Variables have zero Hessian */
+    (void) node;
+    (void) w;
+}
+
 static bool is_affine(const expr *node)
 {
     (void) node;
@@ -39,6 +52,8 @@ expr *new_variable(int d1, int d2, int var_id, int n_vars)
     node->is_affine = is_affine;
     node->jacobian_init = jacobian_init;
     node->eval_jacobian = eval_jacobian;
+    node->wsum_hess_init = wsum_hess_init;
+    node->eval_wsum_hess = wsum_hess_eval;
 
     return node;
 }

src/problem.c

@@ -2,6 +2,20 @@
 #include <stdlib.h>
 #include <string.h>
 
+// TODO: move this into common file
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+
+/* forward declaration */
+static void problem_lagrange_hess_fill_sparsity(problem *prob, int *iwork);
+
+/* Helper function to compare integers for qsort */
+static int compare_int_asc(const void *a, const void *b)
+{
+    int ia = *((const int *) a);
+    int ib = *((const int *) b);
+    return (ia > ib) - (ia < ib);
+}
+
 problem *new_problem(expr *objective, expr **constraints, int n_constraints)
 {
     problem *prob = (problem *) calloc(1, sizeof(problem));
@@ -49,35 +63,129 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints)
     prob->gradient_values = (double *) calloc(prob->n_vars, sizeof(double));
 
     /* Derivative structures allocated by problem_init_derivatives */
-    prob->stacked_jac = NULL;
+    prob->jacobian = NULL;
 
     return prob;
 }
 
 void problem_init_derivatives(problem *prob)
 {
-    /* 1. Initialize objective jacobian */
+    // -------------------------------------------------------------------------------
+    //                           Jacobian structure
+    // -------------------------------------------------------------------------------
     prob->objective->jacobian_init(prob->objective);
-
-    /* 2. Initialize constraint jacobians and count total nnz */
-    int total_nnz = 0;
+    int nnz = 0;
     for (int i = 0; i < prob->n_constraints; i++)
     {
         expr *c = prob->constraints[i];
         c->jacobian_init(c);
-        total_nnz += c->jacobian->nnz;
+        nnz += c->jacobian->nnz;
     }
 
-    /* 3. Allocate stacked jacobian */
-    if (prob->total_constraint_size > 0)
+    prob->jacobian = new_csr_matrix(prob->total_constraint_size, prob->n_vars, nnz);
+
+    // -------------------------------------------------------------------------------
+    //                        Lagrange Hessian structure
+    // -------------------------------------------------------------------------------
+    prob->objective->wsum_hess_init(prob->objective);
+    nnz = prob->objective->wsum_hess->nnz;
+
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        prob->constraints[i]->wsum_hess_init(prob->constraints[i]);
+        nnz += prob->constraints[i]->wsum_hess->nnz;
+    }
+
+    prob->lagrange_hessian = new_csr_matrix(prob->n_vars, prob->n_vars, nnz);
+    prob->hess_idx_map = (int *) malloc(nnz * sizeof(int));
+    int *iwork = (int *) malloc(MAX(nnz, prob->n_vars) * sizeof(int));
+    problem_lagrange_hess_fill_sparsity(prob, iwork);
+    free(iwork);
+}
+
+static void problem_lagrange_hess_fill_sparsity(problem *prob, int *iwork)
+{
+    expr **constrs = prob->constraints;
+    int *cols = iwork;
+    int *col_to_pos = iwork; /* reused after qsort */
+    int nnz = 0;
+    CSR_Matrix *H_obj = prob->objective->wsum_hess;
+    CSR_Matrix *H_c;
+    CSR_Matrix *H = prob->lagrange_hessian;
+    H->p[0] = 0;
+
+    // ----------------------------------------------------------------------
+    //                      Fill sparsity pattern
+    // ----------------------------------------------------------------------
+    for (int row = 0; row < H->m; row++)
     {
-        prob->stacked_jac =
-            new_csr_matrix(prob->total_constraint_size, prob->n_vars, total_nnz);
+        /* gather columns from objective hessian */
+        int count = H_obj->p[row + 1] - H_obj->p[row];
+        memcpy(cols, H_obj->i + H_obj->p[row], count * sizeof(int));
+
+        /* gather columns from constraint hessians */
+        for (int c_idx = 0; c_idx < prob->n_constraints; c_idx++)
+        {
+            H_c = constrs[c_idx]->wsum_hess;
+            int c_len = H_c->p[row + 1] - H_c->p[row];
+            memcpy(cols + count, H_c->i + H_c->p[row], c_len * sizeof(int));
+            count += c_len;
+        }
+
+        /* find unique columns */
+        qsort(cols, count, sizeof(int), compare_int_asc);
+        int prev_col = -1;
+        for (int j = 0; j < count; j++)
+        {
+            if (cols[j] != prev_col)
+            {
+                H->i[nnz] = cols[j];
+                nnz++;
+                prev_col = cols[j];
+            }
+        }
+
+        H->p[row + 1] = nnz;
     }
 
-    /* TODO: 4. Initialize objective wsum_hess */
+    H->nnz = nnz;
+
+    // ----------------------------------------------------------------------
+    //                           Build idx map
+    // ----------------------------------------------------------------------
+    int idx_offset = 0;
+
+    /* map objective hessian entries */
+    for (int row = 0; row < H->m; row++)
+    {
+        for (int idx = H->p[row]; idx < H->p[row + 1]; idx++)
+        {
+            col_to_pos[H->i[idx]] = idx;
+        }
+
+        for (int j = H_obj->p[row]; j < H_obj->p[row + 1]; j++)
+        {
+            prob->hess_idx_map[idx_offset++] = col_to_pos[H_obj->i[j]];
+        }
+    }
 
-    /* TODO: 5. Initialize constraint wsum_hess */
+    /* map constraint hessian entries */
+    for (int c_idx = 0; c_idx < prob->n_constraints; c_idx++)
+    {
+        H_c = constrs[c_idx]->wsum_hess;
+        for (int row = 0; row < H->m; row++)
+        {
+            for (int idx = H->p[row]; idx < H->p[row + 1]; idx++)
+            {
+                col_to_pos[H->i[idx]] = idx;
+            }
+
+            for (int j = H_c->p[row]; j < H_c->p[row + 1]; j++)
+            {
+                prob->hess_idx_map[idx_offset++] = col_to_pos[H_c->i[j]];
+            }
+        }
+    }
 }
 
 void free_problem(problem *prob)
@@ -87,7 +195,9 @@ void free_problem(problem *prob)
     /* Free allocated arrays */
     free(prob->constraint_values);
     free(prob->gradient_values);
-    free_csr_matrix(prob->stacked_jac);
+    free_csr_matrix(prob->jacobian);
+    free_csr_matrix(prob->lagrange_hessian);
+    free(prob->hess_idx_map);
 
     /* Release expression references (decrements refcount) */
     free_expr(prob->objective);
@@ -141,14 +251,15 @@ void problem_gradient(problem *prob)
 
 void problem_jacobian(problem *prob)
 {
-    CSR_Matrix *stacked = prob->stacked_jac;
+    CSR_Matrix *stacked = prob->jacobian;
 
     /* Initialize row pointers */
     stacked->p[0] = 0;
 
     int row_offset = 0;
     int nnz_offset = 0;
 
+    /* TODO: here we eventually want to evaluate affine jacobians only once */
     for (int i = 0; i < prob->n_constraints; i++)
     {
         expr *c = prob->constraints[i];
@@ -176,3 +287,41 @@ void problem_jacobian(problem *prob)
     /* Update actual nnz (may be less than allocated) */
     stacked->nnz = nnz_offset;
 }
+
+void problem_hessian(problem *prob, double obj_w, const double *w)
+{
+    // ------------------------------------------------------------------------
+    //             evaluate hessian of objective and constraints
+    // ------------------------------------------------------------------------
+    expr *obj = prob->objective;
+    expr **constrs = prob->constraints;
+
+    obj->eval_wsum_hess(obj, &obj_w);
+
+    int offset = 0;
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        constrs[i]->eval_wsum_hess(constrs[i], w + offset);
+        offset += constrs[i]->size;
+    }
+
+    // ------------------------------------------------------------------------
+    //           assemble Lagrange hessian using index map
+    // ------------------------------------------------------------------------
+    CSR_Matrix *H = prob->lagrange_hessian;
+    int *idx_map = prob->hess_idx_map;
+
+    /* zero out hessian before adding contribution from obj and constraints */
+    memset(H->x, 0, H->nnz * sizeof(double));
+
+    /* accumulate objective function */
+    idx_map_accumulator(obj->wsum_hess, idx_map, H->x);
+    offset = obj->wsum_hess->nnz;
+
+    /* accumulate constraint functions */
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        idx_map_accumulator(constrs[i]->wsum_hess, idx_map + offset, H->x);
+        offset += constrs[i]->wsum_hess->nnz;
+    }
+}

src/utils/CSR_Matrix.c

@@ -675,7 +675,7 @@ void sum_evenly_spaced_rows_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A,
 void idx_map_accumulator(const CSR_Matrix *A, const int *idx_map,
                          double *accumulator)
 {
-    memset(accumulator, 0, A->nnz * sizeof(double));
+    // memset(accumulator, 0, A->nnz * sizeof(double));
 
     for (int row = 0; row < A->m; row++)
     {

tests/all_tests.c

@@ -172,6 +172,7 @@ int main(void)
     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);
+    mu_run_test(test_problem_hessian, tests_run);
 
     printf("\n=== All %d tests passed ===\n", tests_run);