refactor hstack to precompute sparsity pattern + new init_expr to reuse functioanlity

Author: dance858

include/affine.h

@@ -16,9 +16,6 @@ expr *new_add(expr *left, expr *right);
 /* Helper function to initialize a sum expr (can be used with derived types) */
 void init_sum(expr *node, expr *child, int d1);
 
-/* Helper function to initialize an hstack expr (can be used with derived types) */
-void init_hstack(expr *node, int d1, int d2, int n_vars);
-
 expr *new_sum(expr *child, int axis);
 expr *new_hstack(expr **args, int n_args, int n_vars);
 

include/elementwise_univariate.h

@@ -2,7 +2,6 @@
 #define ELEMENTWISE_H
 
 #include "expr.h"
-#include "subexpr.h"
 
 /* Helper function to initialize an elementwise expr (can be used with derived types)
  */
@@ -27,11 +26,11 @@ expr *new_xexp(expr *child);
 void jacobian_init_elementwise(expr *node);
 void eval_jacobian_elementwise(expr *node);
 void wsum_hess_init_elementwise(expr *node);
-void eval_wsum_hess_elementwise(expr *node, double *w);
+void eval_wsum_hess_elementwise(expr *node, const double *w);
 expr *new_elementwise(expr *child);
 
 /* no elementwise atoms are affine according to our convention,
    so we can have a common implementation */
-bool is_affine_elementwise(expr *node);
+bool is_affine_elementwise(const expr *node);
 
 #endif /* ELEMENTWISE_UNIVARIATE_H */

include/expr.h

@@ -14,10 +14,10 @@ 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);
 typedef void (*eval_jacobian_fn)(struct expr *node);
-typedef void (*wsum_hess_fn)(struct expr *node, double *w);
+typedef void (*wsum_hess_fn)(struct expr *node, const double *w);
 typedef void (*local_jacobian_fn)(struct expr *node, double *out);
-typedef void (*local_wsum_hess_fn)(struct expr *node, double *out, double *w);
-typedef bool (*is_affine_fn)(struct expr *node);
+typedef void (*local_wsum_hess_fn)(struct expr *node, double *out, const double *w);
+typedef bool (*is_affine_fn)(const struct expr *node);
 typedef void (*free_type_data_fn)(struct expr *node);
 
 /* Base expression node structure - contains only common fields */
@@ -54,6 +54,8 @@ typedef struct expr
 
 } expr;
 
+void init_expr(expr *node, int d1, int d2, int n_vars);
+
 expr *new_expr(int d1, int d2, int n_vars);
 void free_expr(expr *node);
 

src/affine/add.c

@@ -45,7 +45,7 @@ static void wsum_hess_init(expr *node)
     node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, nnz_max);
 }
 
-static void eval_wsum_hess(expr *node, double *w)
+static void eval_wsum_hess(expr *node, const double *w)
 {
     /* evaluate children's wsum_hess */
     node->left->eval_wsum_hess(node->left, w);
@@ -55,20 +55,14 @@ static void eval_wsum_hess(expr *node, double *w)
     sum_csr_matrices(node->left->wsum_hess, node->right->wsum_hess, node->wsum_hess);
 }
 
-static bool is_affine(expr *node)
+static bool is_affine(const expr *node)
 {
     return node->left->is_affine(node->left) && node->right->is_affine(node->right);
 }
 
 expr *new_add(expr *left, expr *right)
 {
-    if (!left || !right) return NULL;
-    if (left->d1 != right->d1) return NULL;
-    if (left->d2 != right->d2) return NULL;
-
     expr *node = new_expr(left->d1, left->d2, left->n_vars);
-    if (!node) return NULL;
-
     node->left = left;
     node->right = right;
     expr_retain(left);

src/affine/constant.c

@@ -8,20 +8,16 @@ static void forward(expr *node, const double *u)
     (void) u;
 }
 
-static bool is_affine(expr *node)
+static bool is_affine(const expr *node)
 {
     (void) node;
-    return true; /* constant is affine */
+    return true;
 }
 
 expr *new_constant(int d1, int d2, int n_vars, const double *values)
 {
     expr *node = new_expr(d1, d2, n_vars);
-    if (!node) return NULL;
-
-    /* Copy constant values */
     memcpy(node->value, values, d1 * d2 * sizeof(double));
-
     node->forward = forward;
     node->is_affine = is_affine;
 

src/affine/hstack.c

@@ -36,25 +36,18 @@ static void jacobian_init(expr *node)
     }
 
     node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz);
-}
-
-static void eval_jacobian(expr *node)
-{
-    hstack_expr *hnode = (hstack_expr *) node;
 
-    /* evaluate children's jacobians */
+    /* precompute sparsity pattern of this jacobian's node */
     int row_offset = 0;
     CSR_Matrix *A = node->jacobian;
     A->nnz = 0;
 
     for (int i = 0; i < hnode->n_args; i++)
     {
         expr *child = hnode->args[i];
-        child->eval_jacobian(child);
         CSR_Matrix *B = child->jacobian;
 
-        /* copy columns and values */
-        memcpy(A->x + A->nnz, B->x, B->nnz * sizeof(double));
+        /* copy columns */
         memcpy(A->i + A->nnz, B->i, B->nnz * sizeof(int));
 
         /* set row pointers */
@@ -69,9 +62,27 @@ static void eval_jacobian(expr *node)
     A->p[node->size] = A->nnz;
 }
 
-static bool is_affine(expr *node)
+static void eval_jacobian(expr *node)
 {
     hstack_expr *hnode = (hstack_expr *) node;
+    CSR_Matrix *A = node->jacobian;
+    A->nnz = 0;
+
+    for (int i = 0; i < hnode->n_args; i++)
+    {
+        expr *child = hnode->args[i];
+        child->eval_jacobian(child);
+
+        /* copy values */
+        memcpy(A->x + A->nnz, child->jacobian->x,
+               child->jacobian->nnz * sizeof(double));
+        A->nnz += child->jacobian->nnz;
+    }
+}
+
+static bool is_affine(const expr *node)
+{
+    const hstack_expr *hnode = (const hstack_expr *) node;
 
     for (int i = 0; i < hnode->n_args; i++)
     {
@@ -92,33 +103,6 @@ static void free_type_data(expr *node)
     }
 }
 
-/* Helper function to initialize an hstack expr */
-void init_hstack(expr *node, int d1, int d2, int n_vars)
-{
-    node->d1 = d1;
-    node->d2 = d2;
-    node->size = d1 * d2;
-    node->n_vars = n_vars;
-    node->var_id = -1;
-    node->refcount = 1;
-    node->left = NULL;
-    node->right = NULL;
-    node->dwork = NULL;
-    node->iwork = NULL;
-    node->value = (double *) calloc(node->size, sizeof(double));
-    node->jacobian = NULL;
-    node->wsum_hess = NULL;
-    node->jacobian_init = jacobian_init;
-    node->wsum_hess_init = NULL;
-    node->eval_jacobian = eval_jacobian;
-    node->eval_wsum_hess = NULL;
-    node->local_jacobian = NULL;
-    node->local_wsum_hess = NULL;
-    node->forward = forward;
-    node->is_affine = is_affine;
-    node->free_type_data = free_type_data;
-}
-
 expr *new_hstack(expr **args, int n_args, int n_vars)
 {
     /* compute second dimension */
@@ -129,20 +113,18 @@ expr *new_hstack(expr **args, int n_args, int n_vars)
     }
 
     /* Allocate the type-specific struct */
-    hstack_expr *hnode = (hstack_expr *) malloc(sizeof(hstack_expr));
-    if (!hnode) return NULL;
-
+    hstack_expr *hnode = (hstack_expr *) calloc(1, sizeof(hstack_expr));
     expr *node = &hnode->base;
 
-    /* Initialize base hstack fields */
-    init_hstack(node, args[0]->d1, d2, n_vars);
+    /* Initialize basic fields */
+    init_expr(node, args[0]->d1, d2, n_vars);
 
-    /* Check if allocation succeeded */
-    if (!node->value)
-    {
-        free(hnode);
-        return NULL;
-    }
+    /* Set function pointers */
+    node->forward = forward;
+    node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
+    node->is_affine = is_affine;
+    node->free_type_data = free_type_data;
 
     /* Set type-specific fields */
     hnode->args = args;

src/affine/linear_op.c

@@ -11,7 +11,7 @@ static void forward(expr *node, const double *u)
     csr_matvec(node->jacobian, x->value, node->value, x->var_id);
 }
 
-static bool is_affine(expr *node)
+static bool is_affine(const expr *node)
 {
     return node->left->is_affine(node->left);
 }

src/affine/sum.c

@@ -109,7 +109,7 @@ static void wsum_hess_init(expr *node)
     node->dwork = malloc(x->size * sizeof(double));
 }
 
-static void eval_wsum_hess(expr *node, double *w)
+static void eval_wsum_hess(expr *node, const double *w)
 {
     expr *x = node->left;
     sum_expr *snode = (sum_expr *) node;
@@ -134,7 +134,7 @@ static void eval_wsum_hess(expr *node, double *w)
     copy_csr_matrix(x->wsum_hess, node->wsum_hess);
 }
 
-static bool is_affine(expr *node)
+static bool is_affine(const expr *node)
 {
     return node->left->is_affine(node->left);
 }

src/affine/variable.c

@@ -20,7 +20,7 @@ static void jacobian_init(expr *node)
     node->jacobian->p[size] = size;
 }
 
-static bool is_affine(expr *node)
+static bool is_affine(const expr *node)
 {
     (void) node;
     return true;

src/elementwise_univariate/common.c

@@ -1,4 +1,6 @@
 #include "elementwise_univariate.h"
+#include "expr.h"
+#include "subexpr.h"
 #include <stdlib.h>
 
 void jacobian_init_elementwise(expr *node)
@@ -72,7 +74,7 @@ void wsum_hess_init_elementwise(expr *node)
     }
 }
 
-void eval_wsum_hess_elementwise(expr *node, double *w)
+void eval_wsum_hess_elementwise(expr *node, const double *w)
 {
     expr *child = node->left;
 
@@ -89,7 +91,7 @@ void eval_wsum_hess_elementwise(expr *node, double *w)
     }
 }
 
-bool is_affine_elementwise(expr *node)
+bool is_affine_elementwise(const expr *node)
 {
     (void) node;
     return false;