From bcade49abd70bf576aebb1d355880444b64527d1 Mon Sep 17 00:00:00 2001 From: dance858 Date: Tue, 13 Oct 2026 09:01:10 +0200 Subject: [PATCH] problem: build the Lagrange Hessian with sum_n_matrices_* Replace problem_lagrange_hess_fill_sparsity, hess_idx_map and the accumulator loop with sum_n_matrices_alloc / fill_zero / fill_add. lagrange_hessian stays a CSR_matrix *, now the CSR view of hess_sum. Each Hessian is added right after it is evaluated, so constraints that share a subexpression no longer sum overwritten values. Affine constraints, whose Hessian is zero, skip the add. problem.c no longer uses accumulator, compose_csr_idx_map_for_spd or csr_kernels.h. Co-Authored-By: Claude Opus 5.5 --- include/problem.h | 4 +- src/problem.c | 171 +++++------------------------------ tests/all_tests.c | 1 + tests/problem/test_problem.h | 39 ++++++++ 4 files changed, 65 insertions(+), 150 deletions(-) diff --git a/include/problem.h b/include/problem.h index 27cf224f..b9f28777 100644 --- a/include/problem.h +++ b/include/problem.h @@ -60,8 +60,8 @@ typedef struct problem /* allocated by problem_init_derivatives */ CSR_matrix *jacobian; - CSR_matrix *lagrange_hessian; - int *hess_idx_map; /* maps all wsum_hess nnz to lagrange_hessian */ + CSR_matrix *lagrange_hessian; /* CSR view of hess_sum; do not free */ + matrix *hess_sum; /* sum of all wsum_hess */ COO_matrix *jacobian_coo; COO_matrix *lagrange_hessian_coo; /* lower triangular part stored in COO */ diff --git a/src/problem.c b/src/problem.c index b65d19eb..fe5b859b 100644 --- a/src/problem.c +++ b/src/problem.c @@ -16,9 +16,8 @@ * limitations under the License. */ #include "problem.h" -#include "int_utils.h" -#include "sparsediff_linalg/csr_kernels.h" #include "sparsediff_linalg/matrix.h" +#include "sparsediff_linalg/matrix_sum.h" #include "sparsediff_linalg/tracked_alloc.h" #include "subexpr.h" #include @@ -26,9 +25,6 @@ #include #include -/* forward declaration */ -static void problem_lagrange_hess_fill_sparsity(problem *prob, int *iwork); - problem *new_problem(expr *objective, expr **constraints, int n_constraints, bool verbose) { @@ -83,104 +79,6 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints, return prob; } -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->to_csr(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++) - { - /* 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->to_csr(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 */ - sort_int_array(cols, count); - 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; - } - - H->nnz = nnz; - - // ---------------------------------------------------------------------- - // Build idx map - // ---------------------------------------------------------------------- - int idx_offset = 0; - - /* map objective hessian entries */ - int obj_start = idx_offset; - 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]]; - } - } - if (prob->objective->wsum_hess->kind == MATRIX_STACKED_PD) - { - compose_csr_idx_map_for_spd(prob->objective->wsum_hess, H_obj, - prob->hess_idx_map + obj_start); - } - - /* map constraint hessian entries */ - for (int c_idx = 0; c_idx < prob->n_constraints; c_idx++) - { - H_c = constrs[c_idx]->wsum_hess->to_csr(constrs[c_idx]->wsum_hess); - int c_start = idx_offset; - 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]]; - } - } - if (constrs[c_idx]->wsum_hess->kind == MATRIX_STACKED_PD) - { - compose_csr_idx_map_for_spd(constrs[c_idx]->wsum_hess, H_c, - prob->hess_idx_map + c_start); - } - } -} - void problem_init_jacobian(problem *prob) { if (prob->jacobian != NULL) return; @@ -254,28 +152,22 @@ void problem_init_hessian(problem *prob) // ------------------------------------------------------------------------------- // Lagrange Hessian structure // ------------------------------------------------------------------------------- + int n_ops = 1 + prob->n_constraints; + matrix **ops = (matrix **) sp_malloc(n_ops * sizeof(matrix *)); wsum_hess_init(prob->objective); - int nnz = prob->objective->wsum_hess->nnz; - + ops[0] = prob->objective->wsum_hess; for (int i = 0; i < prob->n_constraints; i++) { wsum_hess_init(prob->constraints[i]); - nnz += prob->constraints[i]->wsum_hess->nnz; + ops[i + 1] = prob->constraints[i]->wsum_hess; } - - int hess_nnz_ub = MIN(nnz, sat_mul_int(prob->n_vars, prob->n_vars)); - prob->lagrange_hessian = new_CSR_matrix(prob->n_vars, prob->n_vars, hess_nnz_ub); - - prob->hess_idx_map = (int *) sp_malloc(nnz * sizeof(int)); - int *iwork = (int *) sp_malloc(MAX(nnz, prob->n_vars) * sizeof(int)); - problem_lagrange_hess_fill_sparsity(prob, iwork); - CSR_trim(prob->lagrange_hessian); + prob->hess_sum = sum_n_matrices_alloc(ops, n_ops); + sp_free(ops); + prob->lagrange_hessian = prob->hess_sum->to_csr(prob->hess_sum); prob->stats.nnz_hessian = prob->lagrange_hessian->nnz; - sp_free(iwork); /* affine shortcut */ - memset(prob->lagrange_hessian->x, 0, - prob->lagrange_hessian->nnz * sizeof(double)); + sum_n_matrices_fill_zero(prob->hess_sum); clock_gettime(CLOCK_MONOTONIC, &timer.end); prob->stats.time_init_derivatives += GET_ELAPSED_SECONDS(timer); @@ -393,10 +285,9 @@ void free_problem(problem *prob) sp_free(prob->constraint_values); sp_free(prob->gradient_values); free_CSR_matrix(prob->jacobian); - free_CSR_matrix(prob->lagrange_hessian); + free_matrix(prob->hess_sum); free_COO_matrix(prob->jacobian_coo); free_COO_matrix(prob->lagrange_hessian_coo); - sp_free(prob->hess_idx_map); sp_free(prob->constraint_jac_seen); /* Release expression references (decrements refcount) */ @@ -559,44 +450,28 @@ void problem_hessian(problem *prob, double obj_w, const double *w) clock_gettime(CLOCK_MONOTONIC, &timer.start); // ------------------------------------------------------------------------ - // evaluate hessian of objective and constraints + // evaluate hessians of objective and constraints and sum them // ------------------------------------------------------------------------ + /* add each hessian right after evaluating it: the objective and the + constraints may share a subexpression, whose later evaluation would + overwrite earlier values */ + matrix *H = prob->hess_sum; + sum_n_matrices_fill_zero(H); expr *obj = prob->objective; eval_wsum_hess(obj, &obj_w); + sum_n_matrices_fill_add(obj->wsum_hess, 0, H); int offset = 0; - expr **constrs = prob->constraints; for (int i = 0; i < prob->n_constraints; i++) { - if (constrs[i]->is_affine(constrs[i])) + expr *c = prob->constraints[i]; + /* an affine constraint's hessian is zero: skip it */ + if (!c->is_affine(c)) { - /* skip evaluation for affine constraints */ - offset += constrs[i]->size; - continue; + eval_wsum_hess(c, w + offset); + sum_n_matrices_fill_add(c->wsum_hess, i + 1, H); } - 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 */ - accumulator(obj->wsum_hess->x, obj->wsum_hess->nnz, idx_map, H->x); - offset = obj->wsum_hess->nnz; - - /* accumulate constraint functions */ - for (int i = 0; i < prob->n_constraints; i++) - { - matrix *c_hess = constrs[i]->wsum_hess; - accumulator(c_hess->x, c_hess->nnz, idx_map + offset, H->x); - offset += c_hess->nnz; + offset += c->size; } clock_gettime(CLOCK_MONOTONIC, &timer.end); diff --git a/tests/all_tests.c b/tests/all_tests.c index f7ec62c2..768d8b56 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -407,6 +407,7 @@ int main(void) mu_run_test(test_problem_jacobian_spd_sum_constraint, tests_run); mu_run_test(test_problem_constraint_forward, tests_run); mu_run_test(test_problem_hessian, tests_run); + mu_run_test(test_problem_hessian_shared_subexpr, tests_run); mu_run_test(test_problem_hessian_sum_exp_left_matmul_dense_transpose, tests_run); printf("\n--- Parameter Tests ---\n"); diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h index 55d4ce5c..3b2c59c5 100644 --- a/tests/problem/test_problem.h +++ b/tests/problem/test_problem.h @@ -347,6 +347,45 @@ const char *test_problem_hessian(void) return 0; } +/* Constraints c1 = log(x) and c2 = exp(log(x)) share the node log(x), so + * evaluating c2's hessian re-evaluates log(x) with c2's weights. The + * Lagrange hessian must still use c1's values for c1: + * L = sum(x) + w[0:3]^T log(x) + w[3:6]^T exp(log(x)), + * H = diag(-w[i] / x[i]^2), since exp(log(x)) = x has a zero hessian. + * Summing after all evaluations would give diag(-w[3+i] / x[i]) instead. + */ +const char *test_problem_hessian_shared_subexpr(void) +{ + int n_vars = 3; + expr *x = new_variable(3, 1, 0, n_vars); + expr *objective = new_sum(x, -1); + expr *log_x = new_log(x); + expr *exp_log_x = new_exp(log_x); + expr *constraints[2] = {log_x, exp_log_x}; + + problem *prob = new_problem(objective, constraints, 2, false); + problem_init_derivatives(prob); + + double u[3] = {1.0, 2.0, 3.0}; + problem_objective_forward(prob, u); + problem_constraint_forward(prob, u); + problem_jacobian(prob); + + double w[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + problem_hessian(prob, 1.0, w); + + CSR_matrix *H = prob->lagrange_hessian; + int expected_p[4] = {0, 1, 2, 3}; + int expected_i[3] = {0, 1, 2}; + double expected_x[3] = {-1.0, -2.0 / 4.0, -3.0 / 9.0}; + mu_assert("H->p wrong", cmp_int_array(H->p, expected_p, 4)); + mu_assert("H->i wrong", cmp_int_array(H->i, expected_i, 3)); + mu_assert("H->x wrong", cmp_double_array(H->x, expected_x, 3)); + + free_problem(prob); + return 0; +} + /* Mirror the Python pipeline: build the user-reported sum(exp(A @ X.T)) as a problem objective (no constraints), call problem_hessian, then read prob->lagrange_hessian. Uses the exact A and X values the user