Affine shortcut

TODO.md

@@ -1,8 +1,4 @@
-3. more tests for chain rule elementwise univariate hessian
 4. in the refactor, add consts
-6. AX where X is a matrix. Can that happen? How is that canonicalized? Maybe it can't happen.
-7. Must be able to compute jacobian and hessian of A @ phi(x), so linear operator needs other code! This requires new infrastructure, I think.
-8. Shortcut hessians of affine stuff?
 10. For performance reasons, is it useful to have a dense matmul with A and B as dense matrices?
 11. right matmul, add broadcasting logic as in left matmul
 

include/problem.h

@@ -33,6 +33,10 @@ typedef struct problem
     int *hess_idx_map; /* Maps all wsum_hess nnz to lagrange_hessian (obj +
                           constraints) */
 
+    /* for the affine shortcut we keep track of the first time the jacobian and
+     * hessian are called */
+    bool jacobian_called;
+
     /* Statistics for performance measurement */
     stats stats;
 } problem;

src/bivariate/const_scalar_mult.c

@@ -75,14 +75,20 @@ static void eval_wsum_hess(expr *node, const double *w)
     }
 }
 
+static bool is_affine(const expr *node)
+{
+    /* Affine iff the child is affine */
+    return node->left->is_affine(node->left);
+}
+
 expr *new_const_scalar_mult(double a, expr *child)
 {
     const_scalar_mult_expr *mult_node =
         (const_scalar_mult_expr *) calloc(1, sizeof(const_scalar_mult_expr));
     expr *node = &mult_node->base;
 
     init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init,
-              eval_jacobian, NULL, NULL);
+              eval_jacobian, is_affine, NULL);
     node->wsum_hess_init = wsum_hess_init;
     node->eval_wsum_hess = eval_wsum_hess;
     node->left = child;

src/bivariate/const_vector_mult.c

@@ -90,14 +90,20 @@ static void free_type_data(expr *node)
     free(vnode->a);
 }
 
+static bool is_affine(const expr *node)
+{
+    /* Affine iff the child is affine */
+    return node->left->is_affine(node->left);
+}
+
 expr *new_const_vector_mult(const double *a, expr *child)
 {
     const_vector_mult_expr *vnode =
         (const_vector_mult_expr *) calloc(1, sizeof(const_vector_mult_expr));
     expr *node = &vnode->base;
 
     init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init,
-              eval_jacobian, NULL, free_type_data);
+              eval_jacobian, is_affine, free_type_data);
     node->wsum_hess_init = wsum_hess_init;
     node->eval_wsum_hess = eval_wsum_hess;
     node->left = child;

src/bivariate/multiply.c

@@ -184,14 +184,20 @@ static void free_type_data(expr *node)
     free_csr_matrix(((elementwise_mult_expr *) node)->CSR_work2);
 }
 
+static bool is_affine(const expr *node)
+{
+    (void) node;
+    return false;
+}
+
 expr *new_elementwise_mult(expr *left, expr *right)
 {
     elementwise_mult_expr *mul_node =
         (elementwise_mult_expr *) calloc(1, sizeof(elementwise_mult_expr));
     expr *node = &mul_node->base;
 
     init_expr(node, left->d1, left->d2, left->n_vars, forward, jacobian_init,
-              eval_jacobian, NULL, free_type_data);
+              eval_jacobian, is_affine, free_type_data);
     node->wsum_hess_init = wsum_hess_init;
     node->eval_wsum_hess = eval_wsum_hess;
     node->left = left;

src/bivariate/quad_over_lin.c

@@ -292,6 +292,12 @@ static void eval_wsum_hess(expr *node, const double *w)
     }
 }
 
+static bool is_affine(const expr *node)
+{
+    (void) node;
+    return false;
+}
+
 expr *new_quad_over_lin(expr *left, expr *right)
 {
     assert((right->d2 == 1 && right->d2 == 1)); /* right must be scalar*/
@@ -305,5 +311,6 @@ expr *new_quad_over_lin(expr *left, expr *right)
     node->eval_jacobian = eval_jacobian;
     node->wsum_hess_init = wsum_hess_init;
     node->eval_wsum_hess = eval_wsum_hess;
+    node->is_affine = is_affine;
     return node;
 }

src/bivariate/rel_entr.c

@@ -169,6 +169,12 @@ static void eval_wsum_hess_vector_args(expr *node, const double *w)
     }
 }
 
+static bool is_affine(const expr *node)
+{
+    (void) node;
+    return false;
+}
+
 expr *new_rel_entr_vector_args(expr *left, expr *right)
 {
     expr *node = new_expr(left->d1, left->d2, left->n_vars);
@@ -181,5 +187,6 @@ expr *new_rel_entr_vector_args(expr *left, expr *right)
     node->eval_jacobian = eval_jacobian_vector_args;
     node->wsum_hess_init = wsum_hess_init_vector_args;
     node->eval_wsum_hess = eval_wsum_hess_vector_args;
+    node->is_affine = is_affine;
     return node;
 }

src/bivariate/rel_entr_scalar_vector.c

@@ -191,6 +191,12 @@ static void eval_wsum_hess_scalar_vector(expr *node, const double *w)
     }
 }
 
+static bool is_affine(const expr *node)
+{
+    (void) node;
+    return false;
+}
+
 expr *new_rel_entr_first_arg_scalar(expr *left, expr *right)
 {
     assert(left->d1 == 1 && left->d2 == 1);
@@ -204,5 +210,6 @@ expr *new_rel_entr_first_arg_scalar(expr *left, expr *right)
     node->eval_jacobian = eval_jacobian_scalar_vector;
     node->wsum_hess_init = wsum_hess_init_scalar_vector;
     node->eval_wsum_hess = eval_wsum_hess_scalar_vector;
+    node->is_affine = is_affine;
     return node;
 }

src/bivariate/rel_entr_vector_scalar.c

@@ -192,6 +192,14 @@ static void eval_wsum_hess_vector_scalar(expr *node, const double *w)
     }
 }
 
+static bool is_affine(const expr *node)
+{
+    (void) node;
+    return false;
+}
+
+/* TODO: modify relative entropy atoms to use init_expr initialization? Also make
+ * init_expr take in the hessian functions?*/
 expr *new_rel_entr_second_arg_scalar(expr *left, expr *right)
 {
     assert(right->d1 == 1 && right->d2 == 1);
@@ -205,5 +213,6 @@ expr *new_rel_entr_second_arg_scalar(expr *left, expr *right)
     node->eval_jacobian = eval_jacobian_vector_scalar;
     node->wsum_hess_init = wsum_hess_init_vector_scalar;
     node->eval_wsum_hess = eval_wsum_hess_vector_scalar;
+    node->is_affine = is_affine;
     return node;
 }

src/other/quad_form.c

@@ -177,14 +177,20 @@ static void free_type_data(expr *node)
     qnode->Q = NULL;
 }
 
+static bool is_affine(const expr *node)
+{
+    (void) node;
+    return false;
+}
+
 expr *new_quad_form(expr *left, CSR_Matrix *Q)
 {
     assert(left->d1 == 1 || left->d2 == 1); /* left must be a vector */
     quad_form_expr *qnode = (quad_form_expr *) calloc(1, sizeof(quad_form_expr));
     expr *node = &qnode->base;
 
-    init_expr(node, 1, 1, left->n_vars, forward, jacobian_init, eval_jacobian, NULL,
-              free_type_data);
+    init_expr(node, 1, 1, left->n_vars, forward, jacobian_init, eval_jacobian,
+              is_affine, free_type_data);
     node->left = left;
     expr_retain(left);
     node->wsum_hess_init = wsum_hess_init;

src/problem.c

@@ -17,6 +17,7 @@ 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;
@@ -99,6 +100,7 @@ void problem_init_derivatives(problem *prob)
     }
 
     prob->lagrange_hessian = new_csr_matrix(prob->n_vars, prob->n_vars, nnz);
+    memset(prob->lagrange_hessian->x, 0, nnz * sizeof(double)); /* affine shortcut */
     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);
@@ -274,13 +276,22 @@ 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))
+        {
+            /* skip evaluation for affine constraints after first call */
+            nnz_offset += c->jacobian->nnz;
+            continue;
+        }
+
         c->eval_jacobian(c);
         memcpy(J->x + nnz_offset, c->jacobian->x, c->jacobian->nnz * sizeof(double));
         nnz_offset += c->jacobian->nnz;
@@ -289,6 +300,7 @@ void problem_jacobian(problem *prob)
     /* update actual nnz (may be less than allocated) */
     J->nnz = nnz_offset;
 
+    prob->jacobian_called = true;
     clock_gettime(CLOCK_MONOTONIC, &timer.end);
     prob->stats.time_eval_jacobian += GET_ELAPSED_SECONDS(timer);
 }
@@ -308,6 +320,12 @@ void problem_hessian(problem *prob, double obj_w, const double *w)
     expr **constrs = prob->constraints;
     for (int i = 0; i < prob->n_constraints; i++)
     {
+        if (constrs[i]->is_affine(constrs[i]))
+        {
+            /* skip evaluation for affine constraints */
+            offset += constrs[i]->size;
+            continue;
+        }
         constrs[i]->eval_wsum_hess(constrs[i], w + offset);
         offset += constrs[i]->size;
     }