removed CSR_work from expr

Author: dance858

include/affine.h

@@ -2,6 +2,7 @@
 #define AFFINE_H
 
 #include "expr.h"
+#include "subexpr.h"
 #include "utils/CSR_Matrix.h"
 
 /* Helper function to initialize a linear operator expr (can be used with derived

include/elementwise_univariate.h

@@ -2,6 +2,7 @@
 #define ELEMENTWISE_H
 
 #include "expr.h"
+#include "subexpr.h"
 
 /* Helper function to initialize an elementwise expr (can be used with derived types)
  */

include/expr.h

@@ -8,11 +8,8 @@
 
 #define JAC_IDXS_NOT_SET -1
 
-/* Forward declarations */
-struct expr;
-struct int_double_pair;
-
 /* Function pointer types */
+struct expr;
 typedef void (*forward_fn)(struct expr *node, const double *u);
 typedef void (*jacobian_init_fn)(struct expr *node);
 typedef void (*wsum_hess_init_fn)(struct expr *node);
@@ -50,54 +47,13 @@ typedef struct expr
     // ------------------------------------------------------------------------
     //                      other things
     // ------------------------------------------------------------------------
-    CSR_Matrix *CSR_work;
     is_affine_fn is_affine;
     local_jacobian_fn local_jacobian;   /* used by elementwise univariate atoms*/
     local_wsum_hess_fn local_wsum_hess; /* used by elementwise univariate atoms*/
     free_type_data_fn free_type_data;   /* Cleanup for type-specific fields */
 
 } expr;
 
-/* Type-specific expression structures that "inherit" from expr */
-
-/* Linear operator: y = A * x */
-typedef struct linear_op_expr
-{
-    expr base; /* MUST be first member for casting to work */
-    CSC_Matrix *A_csc;
-    CSR_Matrix *A_csr;
-} linear_op_expr;
-
-/* Power: y = x^p */
-typedef struct power_expr
-{
-    expr base; /* MUST be first member for casting to work */
-    int p;
-} power_expr;
-
-/* Quadratic form: y = x'*Q*x */
-typedef struct quad_form_expr
-{
-    expr base; /* MUST be first member for casting to work */
-    CSR_Matrix *Q;
-} quad_form_expr;
-
-/* Sum reduction along an axis */
-typedef struct sum_expr
-{
-    expr base; /* MUST be first member for casting to work */
-    int axis;
-    struct int_double_pair *int_double_pairs; /* for sorting jacobian entries */
-} sum_expr;
-
-/* Horizontal stack (concatenate) */
-typedef struct hstack_expr
-{
-    expr base; /* MUST be first member for casting to work */
-    expr **args;
-    int n_args;
-} hstack_expr;
-
 expr *new_expr(int d1, int d2, int n_vars);
 void free_expr(expr *node);
 

include/other.h

@@ -2,6 +2,7 @@
 #define OTHER_H
 
 #include "expr.h"
+#include "subexpr.h"
 #include "utils/CSR_Matrix.h"
 
 /* Helper function to initialize a quad_form expr (can be used with derived types) */

include/subexpr.h

@@ -0,0 +1,51 @@
+#ifndef SUBEXPR_H
+#define SUBEXPR_H
+
+#include "expr.h"
+#include "utils/CSC_Matrix.h"
+#include "utils/CSR_Matrix.h"
+
+/* Forward declaration */
+struct int_double_pair;
+
+/* Type-specific expression structures that "inherit" from expr */
+
+/* Linear operator: y = A * x */
+typedef struct linear_op_expr
+{
+    expr base;
+    CSC_Matrix *A_csc;
+    CSR_Matrix *A_csr;
+} linear_op_expr;
+
+/* Power: y = x^p */
+typedef struct power_expr
+{
+    expr base;
+    int p;
+} power_expr;
+
+/* Quadratic form: y = x'*Q*x */
+typedef struct quad_form_expr
+{
+    expr base;
+    CSR_Matrix *Q;
+} quad_form_expr;
+
+/* Sum reduction along an axis */
+typedef struct sum_expr
+{
+    expr base;
+    int axis;
+    struct int_double_pair *int_double_pairs; /* for sorting jacobian entries */
+} sum_expr;
+
+/* Horizontal stack (concatenate) */
+typedef struct hstack_expr
+{
+    expr base;
+    expr **args;
+    int n_args;
+} hstack_expr;
+
+#endif /* SUBEXPR_H */

include/utils/CSC_Matrix.h

@@ -39,4 +39,9 @@ CSR_Matrix *ATA_alloc(const CSC_Matrix *A);
  */
 void ATDA_values(const CSC_Matrix *A, 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);
+
 #endif /* CSC_MATRIX_H */

src/affine/hstack.c

@@ -108,7 +108,6 @@ void init_hstack(expr *node, int d1, int d2, int n_vars)
     node->value = (double *) calloc(node->size, sizeof(double));
     node->jacobian = NULL;
     node->wsum_hess = NULL;
-    node->CSR_work = NULL;
     node->jacobian_init = jacobian_init;
     node->wsum_hess_init = NULL;
     node->eval_jacobian = eval_jacobian;

src/affine/linear_op.c

@@ -39,7 +39,6 @@ void init_linear_op(expr *node, expr *child, int d1, int d2)
     node->value = (double *) calloc(node->size, sizeof(double));
     node->jacobian = NULL;
     node->wsum_hess = NULL;
-    node->CSR_work = NULL;
     node->jacobian_init = NULL;
     node->wsum_hess_init = NULL;
     node->eval_jacobian = NULL;

src/affine/sum.c

@@ -161,7 +161,6 @@ void init_sum(expr *node, expr *child, int d1)
     node->value = (double *) calloc(node->size, sizeof(double));
     node->jacobian = NULL;
     node->wsum_hess = NULL;
-    node->CSR_work = NULL;
     node->jacobian_init = jacobian_init;
     node->wsum_hess_init = wsum_hess_init;
     node->eval_jacobian = eval_jacobian;

src/bivariate/quad_over_lin.c

@@ -1,4 +1,6 @@
 #include "bivariate.h"
+#include "subexpr.h"
+#include "utils/CSC_Matrix.h"
 #include <assert.h>
 #include <math.h>
 #include <stdlib.h>
@@ -58,14 +60,15 @@ static void jacobian_init(expr *node)
             }
         }
     }
-    else /* left node is not a variable */
+    else /* left node is not a variable (guaranteed to be a linear operator) */
     {
+        linear_op_expr *lin_x = (linear_op_expr *) x;
         node->dwork = (double *) malloc(x->d1 * sizeof(double));
 
         /* compute required allocation and allocate jacobian */
         bool *col_nz = (bool *) calloc(
             node->n_vars, sizeof(bool)); /* TODO: could use iwork here instead*/
-        int nonzero_cols = count_nonzero_cols(x->jacobian, col_nz);
+        int nonzero_cols = count_nonzero_cols(lin_x->base.jacobian, col_nz);
         node->jacobian = new_csr_matrix(1, node->n_vars, nonzero_cols + 1);
 
         /* precompute column indices */
@@ -88,11 +91,8 @@ static void jacobian_init(expr *node)
         node->jacobian->p[0] = 0;
         node->jacobian->p[1] = node->jacobian->nnz;
 
-        /* store A^T of child's A to simplify chain rule computation */
-        node->iwork = (int *) malloc(x->jacobian->n * sizeof(int));
-        node->CSR_work = transpose(x->jacobian, node->iwork);
-
         /* find position where y should be inserted */
+        node->iwork = (int *) malloc(sizeof(int));
         for (int j = 0; j < node->jacobian->nnz; j++)
         {
             if (node->jacobian->i[j] == y->var_id)
@@ -132,14 +132,16 @@ static void eval_jacobian(expr *node)
     }
     else /* x is not a variable */
     {
+        CSC_Matrix *A_csc = ((linear_op_expr *) x)->A_csc;
+
         /* local jacobian */
         for (int j = 0; j < x->d1; j++)
         {
             node->dwork[j] = (2.0 * x->value[j]) / y->value[0];
         }
 
-        /* chain rule (no derivative wrt y) */
-        csr_matvec_fill_values(node->CSR_work, node->dwork, node->jacobian);
+        /* chain rule (no derivative wrt y) using CSC format */
+        csc_matvec_fill_values(A_csc, node->dwork, node->jacobian);
 
         /* insert derivative wrt y at right place (for correctness this assumes
            that y does not appear in the denominator, but this will always be