Better abstraction for keeping track of which constraints have been evaluated

include/problem.h

@@ -65,9 +65,10 @@ typedef struct problem
     COO_matrix *jacobian_coo;
     COO_matrix *lagrange_hessian_coo; /* lower triangular part stored in COO */
 
-    /* for the affine shortcut we keep track of the first time the jacobian and
-     * hessian are called */
-    bool jacobian_called;
+    /* per-constraint jacobian values_version last copied into the aggregated
+     * jacobian; the copy is skipped when the version is unchanged (affine
+     * constraints after their first eval) */
+    uint64_t *constraint_jac_seen;
 
     /* statistics for performance measurement */
     Diff_engine_stats stats;

src/problem.c

@@ -43,7 +43,6 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints,
     prob->objective = objective;
     expr_retain(objective);
     prob->n_vars = objective->n_vars;
-    prob->jacobian_called = false;
 
     /* constraints array */
     prob->total_constraint_size = 0;
@@ -192,11 +191,18 @@ void problem_init_jacobian(problem *prob)
     //                           Jacobian structure
     // -------------------------------------------------------------------------------
     jacobian_init(prob->objective);
+    if (prob->n_constraints > 0)
+    {
+        prob->constraint_jac_seen =
+            (uint64_t *) sp_malloc(prob->n_constraints * sizeof(uint64_t));
+    }
     int nnz = 0;
     for (int i = 0; i < prob->n_constraints; i++)
     {
         expr *c = prob->constraints[i];
         jacobian_init(c);
+        /* deliberately stale so the first problem_jacobian call copies */
+        prob->constraint_jac_seen[i] = c->jacobian->values_version - 1;
         CSR_matrix *Jc = c->jacobian->to_csr(c->jacobian);
         nnz += Jc->nnz;
 
@@ -384,6 +390,7 @@ void free_problem(problem *prob)
     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) */
     free_expr(prob->objective);
@@ -450,9 +457,6 @@ void problem_update_params(problem *prob, const double *theta)
     {
         expr_set_needs_refresh(prob->constraints[i]);
     }
-
-    /* Force re-evaluation of affine Jacobians on next call */
-    prob->jacobian_called = false;
 }
 
 double problem_objective_forward(problem *prob, const double *u)
@@ -513,28 +517,27 @@ void problem_jacobian(problem *prob)
 {
     Timer timer;
     clock_gettime(CLOCK_MONOTONIC, &timer.start);
-    bool first_call = !prob->jacobian_called;
-
     CSR_matrix *J = prob->jacobian;
     int nnz_offset = 0;
 
     for (int i = 0; i < prob->n_constraints; i++)
     {
         expr *c = prob->constraints[i];
-        if (!first_call && c->is_affine(c))
+        eval_jacobian(c);
+
+        /* copy only when the constraint's jacobian values actually changed
+           (an affine constraint's eval is a no-op after its first call and
+           leaves the version untouched until the next parameter update) */
+        if (prob->constraint_jac_seen[i] != c->jacobian->values_version)
         {
-            /* skip evaluation for affine constraints after first call */
-            nnz_offset += c->jacobian->nnz;
-            continue;
+            memcpy(J->x + nnz_offset, c->jacobian->x,
+                   c->jacobian->nnz * sizeof(double));
+            prob->constraint_jac_seen[i] = c->jacobian->values_version;
         }
-
-        eval_jacobian(c);
-        memcpy(J->x + nnz_offset, c->jacobian->x, c->jacobian->nnz * sizeof(double));
         nnz_offset += c->jacobian->nnz;
     }
 
     assert(nnz_offset == J->nnz);
-    prob->jacobian_called = true;
     clock_gettime(CLOCK_MONOTONIC, &timer.end);
     prob->stats.time_eval_jacobian += GET_ELAPSED_SECONDS(timer);
 }

tests/all_tests.c

@@ -556,6 +556,7 @@ int main(void)
     mu_run_test(test_param_right_matmul_rectangular, tests_run);
     mu_run_test(test_param_shared_left_matmul_problem, tests_run);
     mu_run_test(test_param_fixed_skip_in_update, tests_run);
+    mu_run_test(test_problem_jacobian_memcpy_skip, tests_run);
     mu_run_test(test_param_scalar_mult_problem_with_constant, tests_run);
     mu_run_test(test_param_convolve_problem, tests_run);
 

tests/problem/test_param_prob.h

@@ -618,4 +618,59 @@ const char *test_param_convolve_problem(void)
     return 0;
 }
 
+/* problem_jacobian copies a constraint's block into the aggregated jacobian
+   only when the constraint's values_version moved: an affine constraint's
+   slice is not rewritten after its first copy (proven by poisoning it), and
+   a parameter update re-arms the copy. */
+const char *test_problem_jacobian_memcpy_skip(void)
+{
+    int n = 2;
+    expr *x = new_variable(2, 1, 0, n);
+
+    /* objective: sum(log(x)) */
+    expr *objective = new_sum(new_log(x), -1);
+
+    /* constraint 1 (affine, parameterized): a * x */
+    double theta[1] = {3.0};
+    expr *a_param = new_parameter(1, 1, 0, n, theta);
+    expr *scaled = new_scalar_mult(a_param, x);
+
+    /* constraint 2 (non-affine): exp(x) */
+    expr *exp_c = new_exp(x);
+
+    expr *constraints[2] = {scaled, exp_c};
+    problem *prob = new_problem(objective, constraints, 2, false);
+    expr *param_nodes[1] = {a_param};
+    problem_register_params(prob, param_nodes, 1);
+    problem_init_derivatives(prob);
+
+    double u[2] = {1.0, 2.0};
+    problem_constraint_forward(prob, u);
+    problem_jacobian(prob);
+
+    double expected1[4] = {3.0, 3.0, exp(1.0), exp(2.0)};
+    mu_assert("first jacobian wrong",
+              cmp_double_array(prob->jacobian->x, expected1, 4));
+
+    /* poison the affine slice; a re-eval must not rewrite it */
+    prob->jacobian->x[0] = 42.0;
+    prob->jacobian->x[1] = 42.0;
+    problem_jacobian(prob);
+    mu_assert("affine slice must not be recopied", prob->jacobian->x[0] == 42.0);
+    mu_assert("affine slice must not be recopied", prob->jacobian->x[1] == 42.0);
+    mu_assert("non-affine slice must be recopied", prob->jacobian->x[2] == exp(1.0));
+
+    /* a parameter update re-arms the copy */
+    theta[0] = 5.0;
+    problem_update_params(prob, theta);
+    problem_constraint_forward(prob, u);
+    problem_jacobian(prob);
+    double expected2[4] = {5.0, 5.0, exp(1.0), exp(2.0)};
+    mu_assert("jacobian after update wrong",
+              cmp_double_array(prob->jacobian->x, expected2, 4));
+
+    free_problem(prob);
+    return 0;
+}
+
 #endif /* TEST_PARAM_PROB_H */