Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions include/problem.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 */

Expand Down
171 changes: 23 additions & 148 deletions src/problem.c
Original file line number Diff line number Diff line change
Expand Up @@ -16,19 +16,15 @@
* 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 <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* 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)
{
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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) */
Expand Down Expand Up @@ -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);
Expand Down
1 change: 1 addition & 0 deletions tests/all_tests.c
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
39 changes: 39 additions & 0 deletions tests/problem/test_problem.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading