more mental preparation

Author: dance858

src/bivariate_full_dom/multiply.c

@@ -128,23 +128,29 @@ static void wsum_hess_init_impl(expr *node)
     else
     {
         /* chain rule: the Hessian is in this case given by
-           wsum_hess = term1 + term1^T + term2 + term3 where
+           wsum_hess = C + C^T + term2 + term3 where
 
-           * term1 = J_{g2}^T diag(w) J_{g1}
+           * C = J_{g2}^T diag(w) J_{g1}
            * term2 = sum_k w_k g2_k H_{g1_k}
            * term3 = sum_k w_k g1_k H_{g2_k}
 
-            The two last terms are nonzero only if g1 and g2 are nonlinear.
+           The two last terms are nonzero only if g1 and g2 are nonlinear.
+           Here, we view multiply as the composition h(x) = f(g1(x), g2(x)) where f
+           is the elementwise multiplication operator, and g1 and g2 are the left and
+           right child nodes.
         */
 
+        /* prepare sparsity pattern of csc conversion */
+        jacobian_csc_init(x);
+        jacobian_csc_init(y);
+
         /* both are linear operators */
-        CSC_Matrix *A = ((linear_op_expr *) x)->A_csc;
-        CSC_Matrix *B = ((linear_op_expr *) y)->A_csc;
+        CSC_Matrix *Jg1 = x->work->jacobian_csc;
+        CSC_Matrix *Jg2 = y->work->jacobian_csc;
 
         /* Allocate workspace for Hessian computation */
         elementwise_mult_expr *mul_node = (elementwise_mult_expr *) node;
-        CSR_Matrix *C; /* C = B^T diag(w) A */
-        C = BTA_alloc(A, B);
+        CSR_Matrix *C = BTA_alloc(Jg1, Jg2); /* C = Jg2^T diag(w) Jg1 */
         node->work->iwork = (int *) malloc(C->m * sizeof(int));
 
         CSR_Matrix *CT = AT_alloc(C, node->work->iwork);
@@ -174,14 +180,42 @@ static void eval_wsum_hess(expr *node, const double *w)
     }
     else
     {
+        // ----------------------------------------------------------------------
+        //            convert Jacobians of children to CSC format
+        //      (we only need to do this once if the child is affine)
+        //      TODO: what if we have parameters? Should we set jacobian_csc_filled
+        //      to false whenever parameters change value?
+        // ----------------------------------------------------------------------
+        if (!x->work->jacobian_csc_filled)
+        {
+            csr_to_csc_fill_values(x->jacobian, x->work->jacobian_csc,
+                                   x->work->csc_work);
+
+            if (x->is_affine(x))
+            {
+                x->work->jacobian_csc_filled = true;
+            }
+        }
+
+        if (!y->work->jacobian_csc_filled)
+        {
+            csr_to_csc_fill_values(y->jacobian, y->work->jacobian_csc,
+                                   y->work->csc_work);
+
+            if (y->is_affine(y))
+            {
+                y->work->jacobian_csc_filled = true;
+            }
+        }
+
         /* both are linear operators */
-        CSC_Matrix *A = ((linear_op_expr *) x)->A_csc;
-        CSC_Matrix *B = ((linear_op_expr *) y)->A_csc;
+        CSC_Matrix *Jg1 = x->work->jacobian_csc;
+        CSC_Matrix *Jg2 = y->work->jacobian_csc;
         CSR_Matrix *C = ((elementwise_mult_expr *) node)->CSR_work1;
         CSR_Matrix *CT = ((elementwise_mult_expr *) node)->CSR_work2;
 
         /* Compute C = B^T diag(w) A */
-        BTDA_fill_values(A, B, w, C);
+        BTDA_fill_values(Jg1, Jg2, w, C);
 
         /* Compute CT = C^T = A^T diag(w) B */
         AT_fill_values(C, CT, node->work->iwork);

tests/wsum_hess/affine/test_const_scalar_mult.h

@@ -27,6 +27,7 @@ const char *test_wsum_hess_const_scalar_mult_log_vector(void)
     y->forward(y, u_vals);
 
     /* Initialize and evaluate weighted Hessian with w = [1.0, 0.5, 0.25] */
+    jacobian_init(y);
     wsum_hess_init(y);
     double w[3] = {1.0, 0.5, 0.25};
     y->eval_wsum_hess(y, w);
@@ -72,6 +73,7 @@ const char *test_wsum_hess_const_scalar_mult_log_matrix(void)
     y->forward(y, u_vals);
 
     /* Initialize and evaluate weighted Hessian with w = [1.0, 1.0, 1.0, 1.0] */
+    jacobian_init(y);
     wsum_hess_init(y);
     double w[4] = {1.0, 1.0, 1.0, 1.0};
     y->eval_wsum_hess(y, w);

tests/wsum_hess/affine/test_const_vector_mult.h

@@ -27,6 +27,7 @@ const char *test_wsum_hess_const_vector_mult_log_vector(void)
     y->forward(y, u_vals);
 
     /* Initialize and evaluate weighted Hessian with w = [1.0, 0.5, 0.25] */
+    jacobian_init(y);
     wsum_hess_init(y);
     double w[3] = {1.0, 0.5, 0.25};
     y->eval_wsum_hess(y, w);
@@ -70,6 +71,7 @@ const char *test_wsum_hess_const_vector_mult_log_matrix(void)
     y->forward(y, u_vals);
 
     /* Initialize and evaluate weighted Hessian with w = [1.0, 1.0, 1.0, 1.0] */
+    jacobian_init(y);
     wsum_hess_init(y);
     double w[4] = {1.0, 1.0, 1.0, 1.0};
     y->eval_wsum_hess(y, w);

tests/wsum_hess/affine/test_hstack.h

@@ -135,6 +135,7 @@ const char *test_wsum_hess_hstack_matrix(void)
     expr *hstack_node = new_hstack(args, 4, 18);
 
     hstack_node->forward(hstack_node, u_vals);
+    jacobian_init(hstack_node);
     wsum_hess_init(hstack_node);
     hstack_node->eval_wsum_hess(hstack_node, w);
 

tests/wsum_hess/affine/test_transpose.h

@@ -18,6 +18,7 @@ const char *test_wsum_hess_transpose(void)
 
     double u[8] = {1, 3, 2, 4, 5, 7, 6, 8};
     XYT->forward(XYT, u);
+    jacobian_init(XYT);
     wsum_hess_init(XYT);
     double w[4] = {1, 2, 3, 4};
     XYT->eval_wsum_hess(XYT, w);

tests/wsum_hess/bivariate_full_dom/test_matmul.h

@@ -44,6 +44,7 @@ const char *test_wsum_hess_matmul(void)
 
     /* Forward pass and Hessian initialization */
     Z->forward(Z, u_vals);
+    jacobian_init(Z);
     wsum_hess_init(Z);
     Z->eval_wsum_hess(Z, w);
 
@@ -144,6 +145,7 @@ const char *test_wsum_hess_matmul_yx(void)
 
     /* Forward pass and Hessian initialization */
     Z->forward(Z, u_vals);
+    jacobian_init(Z);
     wsum_hess_init(Z);
     Z->eval_wsum_hess(Z, w);
 

tests/wsum_hess/bivariate_full_dom/test_multiply.h

@@ -22,6 +22,7 @@ const char *test_wsum_hess_multiply_1(void)
     expr *node = new_elementwise_mult(x, y);
 
     node->forward(node, u_vals);
+    jacobian_init(node);
     wsum_hess_init(node);
     node->eval_wsum_hess(node, w);
 
@@ -79,6 +80,7 @@ const char *test_wsum_hess_multiply_sparse_random(void)
     mult_node->forward(mult_node, u_vals);
 
     /* Initialize and evaluate Hessian */
+    jacobian_init(mult_node);
     wsum_hess_init(mult_node);
     double w[5] = {0.50646339, 0.44756224, 0.67295241, 0.16424956, 0.03031469};
     mult_node->eval_wsum_hess(mult_node, w);
@@ -160,6 +162,7 @@ const char *test_wsum_hess_multiply_linear_ops(void)
     mult_node->forward(mult_node, u_vals);
 
     /* Initialize Hessian structure */
+    jacobian_init(mult_node);
     wsum_hess_init(mult_node);
 
     /* Evaluate Hessian with weights */
@@ -207,8 +210,9 @@ const char *test_wsum_hess_multiply_2(void)
     expr *y = new_variable(3, 1, 3, 12);
     expr *node = new_elementwise_mult(x, y);
 
-    node->forward(node, u_vals);
+    jacobian_init(node);
     wsum_hess_init(node);
+    node->forward(node, u_vals);
     node->eval_wsum_hess(node, w);
 
     int expected_p[13] = {0, 0, 0, 0, 1, 2, 3, 3, 3, 4, 5, 6, 6};

tests/wsum_hess/bivariate_restricted_dom/test_quad_over_lin.h

@@ -17,6 +17,7 @@ const char *test_wsum_hess_quad_over_lin_xy(void)
     expr *node = new_quad_over_lin(x, y);
 
     node->forward(node, u_vals);
+    jacobian_init(node);
     wsum_hess_init(node);
     node->eval_wsum_hess(node, &w);
 
@@ -46,6 +47,7 @@ const char *test_wsum_hess_quad_over_lin_yx(void)
     expr *node = new_quad_over_lin(x, y);
 
     node->forward(node, u_vals);
+    jacobian_init(node);
     wsum_hess_init(node);
     node->eval_wsum_hess(node, &w);
 

tests/wsum_hess/bivariate_restricted_dom/test_rel_entr.h

@@ -21,6 +21,7 @@ const char *test_wsum_hess_rel_entr_1(void)
     expr *node = new_rel_entr_vector_args(x, y);
 
     node->forward(node, u_vals);
+    jacobian_init(node);
     wsum_hess_init(node);
     node->eval_wsum_hess(node, w);
 
@@ -52,6 +53,7 @@ const char *test_wsum_hess_rel_entr_2(void)
     expr *node = new_rel_entr_vector_args(x, y);
 
     node->forward(node, u_vals);
+    jacobian_init(node);
     wsum_hess_init(node);
     node->eval_wsum_hess(node, w);
 
@@ -83,6 +85,7 @@ const char *test_wsum_hess_rel_entr_matrix(void)
     expr *node = new_rel_entr_vector_args(x, y);
 
     node->forward(node, u_vals);
+    jacobian_init(node);
     wsum_hess_init(node);
     node->eval_wsum_hess(node, w);
 

tests/wsum_hess/bivariate_restricted_dom/test_rel_entr_scalar_vector.h

@@ -16,6 +16,7 @@ const char *test_wsum_hess_rel_entr_scalar_vector(void)
     expr *node = new_rel_entr_first_arg_scalar(x, y);
 
     node->forward(node, u_vals);
+    jacobian_init(node);
     wsum_hess_init(node);
     node->eval_wsum_hess(node, w);