Merge pull request #66 from SparseDifferentiation/quad-form-chain-rule-jacobian

[Ready for review] Quad form chain rule

Author: Transurgeon

include/problem.h

@@ -36,6 +36,8 @@ typedef struct
     int nnz_affine;
     int nnz_nonlinear; /* jacobian of nonlinear constraints */
     int nnz_hessian;
+    int n_vars;
+    int total_constraint_size;
 } Diff_engine_stats;
 
 typedef struct problem

include/subexpr.h

@@ -48,6 +48,7 @@ typedef struct quad_form_expr
 {
     expr base;
     CSR_Matrix *Q;
+    CSC_Matrix *QJf; /* Q * J_f in CSC (for chain rule hessian) */
 } quad_form_expr;
 
 /* Sum reduction along an axis */

include/utils/CSC_Matrix.h

@@ -29,30 +29,37 @@ CSC_Matrix *new_csc_matrix(int m, int n, int nnz);
 /* Free a CSC matrix */
 void free_csc_matrix(CSC_Matrix *matrix);
 
-CSC_Matrix *csr_to_csc(const CSR_Matrix *A);
-
-/* Allocate sparsity pattern for C = A^T D A for diagonal D */
+/* Fill sparsity of C = A^T D A for diagonal D */
 CSR_Matrix *ATA_alloc(const CSC_Matrix *A);
 
-/* Allocate sparsity pattern for C = B^T D A for diagonal D */
+/* Fill sparsity of C = B^T D A for diagonal D */
 CSR_Matrix *BTA_alloc(const CSC_Matrix *A, const CSC_Matrix *B);
 
-/* Compute values for C = A^T D A. C must have precomputed sparsity pattern  */
+/* Fill sparsity of C = BA, where B is symmetric. */
+CSC_Matrix *symBA_alloc(const CSR_Matrix *B, const CSC_Matrix *A);
+
+/* Compute values for C = A^T D A (null d corresponds to D as identity) */
 void ATDA_fill_values(const CSC_Matrix *A, const double *d, CSR_Matrix *C);
 
-/* Compute values for C = B^T D A. C must have precomputed sparsity pattern  */
+/* Compute values for C = B^T D A (null d corresonds to D as identity) */
 void BTDA_fill_values(const CSC_Matrix *A, const CSC_Matrix *B, const double *d,
                       CSR_Matrix *C);
 
-/* C = z^T A where A is in CSC format and C is assumed to have one row.
- * C must have column indices pre-computed. Fills in values of C only.
- */
-void csc_matvec_fill_values(const CSC_Matrix *A, const double *z, CSR_Matrix *C);
+/* Fill values of C = BA. The matrix B does not have to be symmetric */
+void BA_fill_values(const CSR_Matrix *B, const CSC_Matrix *A, CSC_Matrix *C);
+
+/* Fill values of C = x^T A. The matrix C must have filled sparsity. */
+void yTA_fill_values(const CSC_Matrix *A, const double *x, CSR_Matrix *C);
+
+/* Count nonzero columns of a CSC matrix */
+int count_nonzero_cols_csc(const CSC_Matrix *A);
 
-CSC_Matrix *csr_to_csc_fill_sparsity(const CSR_Matrix *A, int *iwork);
+/* convert from CSR to CSC format */
+CSC_Matrix *csr_to_csc_alloc(const CSR_Matrix *A, int *iwork);
 void csr_to_csc_fill_values(const CSR_Matrix *A, CSC_Matrix *C, int *iwork);
 
-CSR_Matrix *csc_to_csr_fill_sparsity(const CSC_Matrix *A, int *iwork);
+/* convert from CSC to CSR format */
+CSR_Matrix *csc_to_csr_alloc(const CSC_Matrix *A, int *iwork);
 void csc_to_csr_fill_values(const CSC_Matrix *A, CSR_Matrix *C, int *iwork);
 
 #endif /* CSC_MATRIX_H */

src/affine/left_matmul.c

@@ -88,12 +88,12 @@ static void jacobian_init_impl(expr *node)
 
     /* initialize child's jacobian and precompute sparsity of its CSC */
     jacobian_init(x);
-    lnode->Jchild_CSC = csr_to_csc_fill_sparsity(x->jacobian, node->work->iwork);
+    lnode->Jchild_CSC = csr_to_csc_alloc(x->jacobian, node->work->iwork);
 
     /* precompute sparsity of this node's jacobian in CSC and CSR */
     lnode->J_CSC = lnode->A->block_left_mult_sparsity(lnode->A, lnode->Jchild_CSC,
                                                       lnode->n_blocks);
-    node->jacobian = csc_to_csr_fill_sparsity(lnode->J_CSC, lnode->csc_to_csr_work);
+    node->jacobian = csc_to_csr_alloc(lnode->J_CSC, lnode->csc_to_csr_work);
 }
 
 static void eval_jacobian(expr *node)

src/bivariate_restricted_dom/quad_over_lin.c

@@ -164,8 +164,7 @@ static void eval_jacobian(expr *node)
         }
 
         /* chain rule (no derivative wrt y) using CSC format */
-        csc_matvec_fill_values(x->work->jacobian_csc, node->work->dwork,
-                               node->jacobian);
+        yTA_fill_values(x->work->jacobian_csc, node->work->dwork, node->jacobian);
 
         /* insert derivative wrt y at right place (for correctness this assumes
            that y does not appear in the numerator, but this will always be

src/expr.c

@@ -51,7 +51,7 @@ void jacobian_csc_init(expr *node)
     }
     node->work->csc_work = (int *) malloc(node->n_vars * sizeof(int));
     node->work->jacobian_csc =
-        csr_to_csc_fill_sparsity(node->jacobian, node->work->csc_work);
+        csr_to_csc_alloc(node->jacobian, node->work->csc_work);
 }
 
 void free_expr(expr *node)