started refactoring polymorphism

Author: dance858

README.md

@@ -3,4 +3,5 @@
 3. more tests for chain rule elementwise univariate hessian
 4. in the refactor, add consts
 5. multiply with one constant vector/scalar argument
-6. AX where X is a matrix. Can that happen? How is that canonicalized? Maybe it can't happen.
+6. AX where X is a matrix. Can that happen? How is that canonicalized? Maybe it can't happen.
+7. Must be able to compute jacobian and hessian of A @ phi(x), so linear operator needs other code?

include/affine.h

@@ -4,9 +4,20 @@
 #include "expr.h"
 #include "utils/CSR_Matrix.h"
 
+/* Helper function to initialize a linear operator expr (can be used with derived
+ * types) */
+void init_linear_op(expr *node, expr *child, int d1, int d2);
+
 expr *new_linear(expr *u, const CSR_Matrix *A);
 
 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

@@ -3,6 +3,10 @@
 
 #include "expr.h"
 
+/* Helper function to initialize an elementwise expr (can be used with derived types)
+ */
+void init_elementwise(expr *node, expr *child);
+
 expr *new_exp(expr *child);
 expr *new_log(expr *child);
 expr *new_entr(expr *child);

include/expr.h

@@ -21,55 +21,82 @@ typedef void (*wsum_hess_fn)(struct expr *node, 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 (*free_type_data_fn)(struct expr *node);
 
-/* TODO: implement proper polymorphism */
-
-/* Expression node structure */
+/* Base expression node structure - contains only common fields */
 typedef struct expr
 {
     // ------------------------------------------------------------------------
     //                         general quantities
     // ------------------------------------------------------------------------
-    int d1, d2, size;
-    int n_vars;
-    int var_id;
-    int refcount;
+    int d1, d2, size, n_vars, refcount, var_id;
     struct expr *left;
     struct expr *right;
-    struct expr **args; /* hstack can have multiple arguments */
-    int n_args;
     double *dwork;
     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 */
-    CSR_Matrix *Q; /* Q for quad_form */
 
     // ------------------------------------------------------------------------
-    //                     forward pass related quantities
+    //                     oracle related quantities
     // ------------------------------------------------------------------------
     double *value;
-    forward_fn forward;
-
-    // ------------------------------------------------------------------------
-    //                      jacobian related quantities
-    // ------------------------------------------------------------------------
     CSR_Matrix *jacobian;
     CSR_Matrix *wsum_hess;
-    CSR_Matrix *CSR_work;
+    forward_fn forward;
     jacobian_init_fn jacobian_init;
     wsum_hess_init_fn wsum_hess_init;
     eval_jacobian_fn eval_jacobian;
     wsum_hess_fn eval_wsum_hess;
-    local_jacobian_fn local_jacobian;
-    local_wsum_hess_fn local_wsum_hess;
+
+    // ------------------------------------------------------------------------
+    //                      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 */
 
-    // for every linear operator we store A in CSR and CSC
+} 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;
 
-} 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

@@ -4,6 +4,9 @@
 #include "expr.h"
 #include "utils/CSR_Matrix.h"
 
+/* Helper function to initialize a quad_form expr (can be used with derived types) */
+void init_quad_form(expr *node, expr *child);
+
 expr *new_quad_form(expr *child, CSR_Matrix *Q);
 
 #endif /* OTHER_H */

src/affine/hstack.c

@@ -1,49 +1,55 @@
 #include "affine.h"
 #include <assert.h>
+#include <stdlib.h>
 #include <string.h>
 
 static void forward(expr *node, const double *u)
 {
+    hstack_expr *hnode = (hstack_expr *) node;
 
     /* children's forward passes */
-    for (int i = 0; i < node->n_args; i++)
+    for (int i = 0; i < hnode->n_args; i++)
     {
-        node->args[i]->forward(node->args[i], u);
+        hnode->args[i]->forward(hnode->args[i], u);
     }
 
     /* concatenate values horizontally */
     int offset = 0;
-    for (int i = 0; i < node->n_args; i++)
+    for (int i = 0; i < hnode->n_args; i++)
     {
-        expr *child = node->args[i];
+        expr *child = hnode->args[i];
         memcpy(node->value + offset, child->value, child->size * sizeof(double));
         offset += child->size;
     }
 }
 
 static void jacobian_init(expr *node)
 {
+    hstack_expr *hnode = (hstack_expr *) node;
+
     /* initialize children's jacobians */
     int nnz = 0;
-    for (int i = 0; i < node->n_args; i++)
+    for (int i = 0; i < hnode->n_args; i++)
     {
-        node->args[i]->jacobian_init(node->args[i]);
-        nnz += node->args[i]->jacobian->nnz;
+        hnode->args[i]->jacobian_init(hnode->args[i]);
+        nnz += hnode->args[i]->jacobian->nnz;
     }
 
     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 */
     int row_offset = 0;
     CSR_Matrix *A = node->jacobian;
     A->nnz = 0;
 
-    for (int i = 0; i < node->n_args; i++)
+    for (int i = 0; i < hnode->n_args; i++)
     {
-        expr *child = node->args[i];
+        expr *child = hnode->args[i];
         child->eval_jacobian(child);
         CSR_Matrix *B = child->jacobian;
 
@@ -65,16 +71,55 @@ static void eval_jacobian(expr *node)
 
 static bool is_affine(expr *node)
 {
-    for (int i = 0; i < node->n_args; i++)
+    hstack_expr *hnode = (hstack_expr *) node;
+
+    for (int i = 0; i < hnode->n_args; i++)
     {
-        if (!node->args[i]->is_affine(node->args[i]))
+        if (!hnode->args[i]->is_affine(hnode->args[i]))
         {
             return false;
         }
     }
     return true;
 }
 
+static void free_type_data(expr *node)
+{
+    hstack_expr *hnode = (hstack_expr *) node;
+    for (int i = 0; i < hnode->n_args; i++)
+    {
+        free_expr(hnode->args[i]);
+    }
+}
+
+/* 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->CSR_work = 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 */
@@ -84,20 +129,30 @@ expr *new_hstack(expr **args, int n_args, int n_vars)
         d2 += args[i]->d2;
     }
 
-    expr *node = new_expr(args[0]->d1, d2, n_vars);
-    if (!node) return NULL;
-    node->args = args;
-    node->n_args = n_args;
+    /* Allocate the type-specific struct */
+    hstack_expr *hnode = (hstack_expr *) malloc(sizeof(hstack_expr));
+    if (!hnode) return NULL;
+
+    expr *node = &hnode->base;
+
+    /* Initialize base hstack fields */
+    init_hstack(node, args[0]->d1, d2, n_vars);
+
+    /* Check if allocation succeeded */
+    if (!node->value)
+    {
+        free(hnode);
+        return NULL;
+    }
+
+    /* Set type-specific fields */
+    hnode->args = args;
+    hnode->n_args = n_args;
 
     for (int i = 0; i < n_args; i++)
     {
         expr_retain(args[i]);
     }
 
-    node->forward = forward;
-    node->is_affine = is_affine;
-    node->jacobian_init = jacobian_init;
-    node->eval_jacobian = eval_jacobian;
-
     return node;
 }

src/affine/linear_op.c

@@ -16,24 +16,70 @@ static bool is_affine(expr *node)
     return node->left->is_affine(node->left);
 }
 
-expr *new_linear(expr *u, const CSR_Matrix *A)
+static void free_type_data(expr *node)
 {
-    expr *node = new_expr(A->m, 1, u->n_vars);
-    if (!node) return NULL;
+    linear_op_expr *lin_node = (linear_op_expr *) node;
+    free_csr_matrix(lin_node->A_csr);
+    free_csc_matrix(lin_node->A_csc);
+}
 
-    node->left = u;
-    expr_retain(u);
+/* Helper function to initialize a linear operator expr */
+void init_linear_op(expr *node, expr *child, int d1, int d2)
+{
+    node->d1 = d1;
+    node->d2 = d2;
+    node->size = d1 * d2;
+    node->n_vars = child->n_vars;
+    node->var_id = -1;
+    node->refcount = 1;
+    node->left = child;
+    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->CSR_work = NULL;
+    node->jacobian_init = NULL;
+    node->wsum_hess_init = NULL;
+    node->eval_jacobian = NULL;
+    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_retain(child);
+}
+
+expr *new_linear(expr *u, const CSR_Matrix *A)
+{
+    /* Allocate the type-specific struct */
+    linear_op_expr *lin_node = (linear_op_expr *) malloc(sizeof(linear_op_expr));
+    if (!lin_node) return NULL;
+
+    expr *node = &lin_node->base;
+
+    /* Initialize base linear operator fields */
+    init_linear_op(node, u, A->m, 1);
+
+    /* Check if allocation succeeded */
+    if (!node->value)
+    {
+        free(lin_node);
+        return NULL;
+    }
 
     /* 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);
+    /* Initialize type-specific fields */
+    lin_node->A_csr = new_csr_matrix(A->m, A->n, A->nnz);
+    copy_csr_matrix(A, lin_node->A_csr);
+    lin_node->A_csc = csr_to_csc(A);
 
     return node;
 }