more infrastructure prepratation for hess_vec

Author: dance858

include/expr.h

@@ -1,6 +1,7 @@
 #ifndef EXPR_H
 #define EXPR_H
 
+#include "utils/CSC_Matrix.h"
 #include "utils/CSR_Matrix.h"
 #include <stdbool.h>
 #include <stddef.h>
@@ -18,6 +19,8 @@ typedef void (*eval_jacobian_fn)(struct expr *node);
 typedef void (*eval_local_jacobian_fn)(struct expr *node, double *out);
 typedef bool (*is_affine_fn)(struct expr *node);
 
+/* TODO: implement proper polymorphism */
+
 /* Expression node structure */
 typedef struct expr
 {
@@ -36,7 +39,8 @@ typedef struct expr
     int *iwork;
     struct int_double_pair *int_double_pairs; /* for sorting jacobian entries */
     int p;                                    /* power of power expression */
-    int axis; /* axis for sum or similar operations */
+    int axis;      /* axis for sum or similar operations */
+    CSR_Matrix *Q; /* Q for quad_form */
 
     // ------------------------------------------------------------------------
     //                     forward pass related quantities
@@ -48,13 +52,16 @@ typedef struct expr
     //                      jacobian related quantities
     // ------------------------------------------------------------------------
     CSR_Matrix *jacobian;
-    CSR_Matrix *Q;
     CSR_Matrix *CSR_work;
     jacobian_init_fn jacobian_init;
     eval_jacobian_fn eval_jacobian;
     eval_local_jacobian_fn eval_local_jacobian;
     is_affine_fn is_affine;
 
+    // for every linear operator we store A in CSR and CSC
+    CSC_Matrix *A_csc;
+    CSR_Matrix *A_csr;
+
 } expr;
 
 expr *new_expr(int d1, int d2, int n_vars);

include/utils/CSC_Matrix.h

@@ -29,8 +29,9 @@ CSC_Matrix *new_csc_matrix(int m, int n, int nnz);
 /* Free a CSC matrix */
 void free_csc_matrix(CSC_Matrix *matrix);
 
-/* Allocate sparsity pattern for C = A^T D A or C = A^T A
- */
+CSC_Matrix *csr_to_csc(const CSR_Matrix *A);
+
+/* Allocate sparsity pattern for C = A^T D A for diagonal D */
 CSR_Matrix *ATA_alloc(const CSC_Matrix *A);
 
 /* Compute values for C = A^T D A

src/affine/linear_op.c

@@ -27,8 +27,13 @@ expr *new_linear(expr *u, const CSR_Matrix *A)
     node->is_affine = is_affine;
 
     /* allocate jacobian and copy A into it */
+    // TODO: this should eventually be removed
     node->jacobian = new_csr_matrix(A->m, A->n, A->nnz);
     copy_csr_matrix(A, node->jacobian);
 
+    node->A_csr = new_csr_matrix(A->m, A->n, A->nnz);
+    copy_csr_matrix(A, node->A_csr);
+    node->A_csc = csr_to_csc(A);
+
     return node;
 }

src/elementwise_univariate/common.c

@@ -36,7 +36,7 @@ void eval_jacobian_elementwise(expr *node)
     else
     {
         node->eval_local_jacobian(node, node->dwork);
-        diag_csr_mult(node->dwork, child->jacobian, node->jacobian);
+        diag_csr_mult(node->dwork, child->A_csr, node->jacobian);
     }
 }
 

src/expr.c

@@ -59,6 +59,8 @@ void free_expr(expr *node)
     free(node->value);
     free_csr_matrix(node->jacobian);
     free_csr_matrix(node->CSR_work);
+    free_csr_matrix(node->A_csr);
+    free_csc_matrix(node->A_csc);
     free(node->dwork);
     free(node->iwork);
     free_int_double_pair_array(node->int_double_pairs);

src/utils/CSC_Matrix.c

@@ -84,9 +84,6 @@ CSR_Matrix *ATA_alloc(const CSC_Matrix *A)
 
     /* Allocate C and symmetrize it */
     CSR_Matrix *C = new_csr_matrix(n, n, nnz);
-
-    /* TODO: do we need to symmetrize here? If we are a bit careful with symmetry
-       throughout the implementation we can skip this step. */
     symmetrize_csr(Cp, Ci->data, n, C);
 
     /* free workspace */
@@ -152,3 +149,49 @@ void ATDA_values(const CSC_Matrix *A, const double *d, CSR_Matrix *C)
         }
     }
 }
+
+CSC_Matrix *csr_to_csc(const CSR_Matrix *A)
+{
+    CSC_Matrix *C = new_csc_matrix(A->m, A->n, A->nnz);
+
+    int i, j, start;
+    int *count = malloc(A->n * sizeof(int));
+
+    memset(count, 0, A->n * sizeof(int));
+
+    // -------------------------------------------------------------------
+    //              compute nnz in each column of A
+    // -------------------------------------------------------------------
+    for (i = 0; i < A->m; ++i)
+    {
+        for (j = A->p[i]; j < A->p[i + 1]; ++j)
+        {
+            count[A->i[j]]++;
+        }
+    }
+
+    // ------------------------------------------------------------------
+    //                      compute column pointers
+    // ------------------------------------------------------------------
+    C->p[0] = 0;
+    for (i = 0; i < A->n; ++i)
+    {
+        C->p[i + 1] = C->p[i] + count[i];
+        count[i] = C->p[i];
+    }
+
+    // ------------------------------------------------------------------
+    //                         fill matrix
+    // ------------------------------------------------------------------
+    for (i = 0; i < A->m; ++i)
+    {
+        for (j = A->p[i]; j < A->p[i + 1]; ++j)
+        {
+            C->x[count[A->i[j]]] = A->x[j];
+            C->i[count[A->i[j]]] = i;
+            count[A->i[j]]++;
+        }
+    }
+
+    return C;
+}

src/utils/CSR_Matrix.c

@@ -5,6 +5,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+
 CSR_Matrix *new_csr_matrix(int m, int n, int nnz)
 {
     CSR_Matrix *matrix = (CSR_Matrix *) malloc(sizeof(CSR_Matrix));
@@ -423,7 +424,7 @@ CSR_Matrix *transpose(const CSR_Matrix *A, int *iwork)
     }
 
     // ------------------------------------------------------------------
-    //  fill transposed matrix (this is a bottleneck)
+    //  fill transposed matrix
     // ------------------------------------------------------------------
     for (i = 0; i < A->m; ++i)
     {

tests/all_tests.c

@@ -72,6 +72,8 @@ int main(void)
     mu_run_test(test_csr_sum, tests_run);
     mu_run_test(test_csr_sum2, tests_run);
     mu_run_test(test_transpose, tests_run);
+    mu_run_test(test_csr_to_csc1, tests_run);
+    mu_run_test(test_csr_to_csc2, tests_run);
     mu_run_test(test_csr_vecmat_values_sparse, tests_run);
     mu_run_test(test_sum_all_rows_csr, tests_run);
     mu_run_test(test_sum_block_of_rows_csr, tests_run);

tests/utils/test_csc_matrix.h

@@ -7,6 +7,85 @@
 #include "test_helpers.h"
 #include "utils/CSC_Matrix.h"
 
+const char *test_csr_to_csc1()
+{
+    CSR_Matrix *A = new_csr_matrix(4, 5, 5);
+    double Ax[5] = {1.0, 1.0, 3.0, 2.0, 4.0};
+    int Ai[5] = {0, 4, 1, 0, 1};
+    int Ap[5] = {0, 2, 3, 4, 5};
+    memcpy(A->x, Ax, 5 * sizeof(double));
+    memcpy(A->i, Ai, 5 * sizeof(int));
+    memcpy(A->p, Ap, 5 * sizeof(int));
+
+    CSC_Matrix *C = csr_to_csc(A);
+
+    double Cx_correct[5] = {1.0, 2.0, 3.0, 4.0, 1.0};
+    int Ci_correct[5] = {0, 2, 1, 3, 0};
+    int Cp_correct[6] = {0, 2, 4, 4, 4, 5};
+
+    mu_assert("C vals incorrect", cmp_double_array(C->x, Cx_correct, 5));
+    mu_assert("C rows incorrect", cmp_int_array(C->i, Ci_correct, 5));
+    mu_assert("C cols incorrect", cmp_int_array(C->p, Cp_correct, 6));
+
+    free_csr_matrix(A);
+    free_csc_matrix(C);
+
+    return 0;
+}
+
+const char *test_csr_to_csc2()
+{
+    CSR_Matrix *A = new_csr_matrix(20, 30, 120);
+    double Ax[120] = {9, 6, 5, 9, 7, 3, 8, 2, 6, 1, 3, 9, 2, 8, 9, 1, 4, 9, 2, 1,
+                      3, 4, 2, 8, 6, 2, 9, 7, 3, 8, 3, 7, 9, 2, 2, 2, 5, 5, 3, 5,
+                      1, 6, 7, 2, 7, 3, 3, 7, 3, 5, 4, 7, 7, 3, 6, 3, 6, 1, 8, 8,
+                      3, 2, 2, 3, 4, 5, 5, 5, 8, 3, 5, 3, 7, 5, 1, 4, 9, 6, 6, 7,
+                      4, 6, 8, 2, 7, 3, 5, 3, 3, 4, 7, 3, 6, 4, 2, 1, 1, 5, 5, 8,
+                      1, 9, 5, 2, 3, 8, 5, 8, 4, 5, 5, 6, 9, 6, 4, 4, 1, 8, 9, 8};
+    int Ai[120] = {1,  2,  3,  19, 21, 22, 9,  10, 19, 20, 25, 0,  6,  8,  9,
+                   12, 15, 19, 20, 21, 26, 2,  5,  6,  8,  12, 14, 16, 19, 27,
+                   8,  11, 13, 15, 25, 26, 27, 10, 12, 19, 22, 23, 24, 25, 28,
+                   1,  11, 12, 15, 18, 24, 13, 22, 2,  5,  6,  9,  18, 24, 3,
+                   6,  8,  22, 20, 27, 7,  9,  17, 26, 29, 0,  1,  11, 13, 15,
+                   16, 18, 23, 24, 4,  5,  8,  9,  16, 20, 23, 4,  6,  14, 15,
+                   24, 8,  9,  11, 12, 20, 22, 29, 2,  5,  12, 14, 15, 19, 21,
+                   10, 19, 27, 1,  5,  6,  9,  11, 15, 21, 26, 3,  15, 26, 27};
+    int Ap[21] = {0,  6,  11, 21, 30, 37, 45,  51,  53,  59, 63,
+                  65, 70, 79, 86, 91, 98, 105, 108, 116, 120};
+    memcpy(A->x, Ax, 120 * sizeof(double));
+    memcpy(A->i, Ai, 120 * sizeof(int));
+    memcpy(A->p, Ap, 21 * sizeof(int));
+
+    CSC_Matrix *C = csr_to_csc(A);
+
+    double Cx_correct[120] = {
+        9, 5, 9, 3, 3, 4, 6, 4, 3, 5, 5, 8, 1, 7, 5, 2, 6, 4, 8, 5, 2, 8, 3, 3,
+        3, 5, 5, 8, 6, 3, 2, 6, 3, 8, 9, 6, 5, 8, 6, 6, 2, 5, 8, 7, 3, 7, 4, 9,
+        1, 2, 3, 7, 2, 1, 9, 7, 5, 9, 3, 9, 4, 2, 3, 1, 4, 5, 6, 8, 7, 4, 2, 5,
+        5, 1, 9, 9, 6, 9, 3, 5, 2, 5, 1, 2, 3, 7, 1, 7, 1, 3, 4, 3, 1, 7, 2, 1,
+        6, 6, 3, 7, 4, 8, 6, 7, 3, 2, 2, 3, 2, 8, 4, 9, 8, 5, 4, 8, 8, 7, 3, 5};
+    int Ci_correct[120] = {
+        2,  12, 0,  6,  12, 18, 0,  3,  8,  16, 0,  9,  19, 13, 14, 3,  8,  13,
+        16, 18, 2,  3,  8,  9,  14, 18, 11, 2,  3,  4,  9,  13, 15, 1,  2,  8,
+        11, 13, 15, 18, 1,  5,  17, 4,  6,  12, 15, 18, 2,  3,  5,  6,  15, 16,
+        4,  7,  12, 3,  14, 16, 2,  4,  6,  12, 14, 16, 18, 19, 3,  12, 13, 11,
+        6,  8,  12, 0,  1,  2,  3,  5,  16, 17, 1,  2,  10, 13, 15, 0,  2,  16,
+        18, 0,  5,  7,  9,  15, 5,  12, 13, 5,  6,  8,  12, 14, 1,  4,  5,  2,
+        4,  11, 18, 19, 3,  4,  10, 17, 19, 5,  11, 15};
+    int Cp_correct[31] = {0,  2,  6,  10,  13,  15,  20,  26,  27, 33, 40,
+                          43, 48, 54, 57,  60,  68,  71,  72,  75, 82, 87,
+                          91, 96, 99, 104, 107, 112, 117, 118, 120};
+
+    mu_assert("C vals incorrect", cmp_double_array(C->x, Cx_correct, 120));
+    mu_assert("C rows incorrect", cmp_int_array(C->i, Ci_correct, 120));
+    mu_assert("C cols incorrect", cmp_int_array(C->p, Cp_correct, 31));
+
+    free_csr_matrix(A);
+    free_csc_matrix(C);
+
+    return 0;
+}
+
 /* Test ATA_alloc with a simple 3x3 example
  * A is 4x3 (4 rows, 3 columns):
  * [x  0  x]