refactor of affine atoms done

Author: dance858

include/affine.h

@@ -9,9 +9,6 @@ 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);
-
 expr *new_sum(expr *child, int axis);
 expr *new_hstack(expr **args, int n_args, int n_vars);
 

src/affine/linear_op.c

@@ -4,6 +4,7 @@
 static void forward(expr *node, const double *u)
 {
     expr *x = node->left;
+
     /* child's forward pass */
     node->left->forward(node->left, u);
 
@@ -19,42 +20,30 @@ static bool is_affine(const expr *node)
 static void free_type_data(expr *node)
 {
     linear_op_expr *lin_node = (linear_op_expr *) node;
-    free_csr_matrix(lin_node->A_csr);
+    /* memory pointing to by A_csr will already be freed when the
+       jacobian is freed, so free_csr_matrix(lin_node->A_csr) should
+       be commented out */
+    // free_csr_matrix(lin_node->A_csr);
     free_csc_matrix(lin_node->A_csc);
+    lin_node->A_csr = NULL;
+    lin_node->A_csc = NULL;
 }
 
 expr *new_linear(expr *u, const CSR_Matrix *A)
 {
     /* Allocate the type-specific struct */
     linear_op_expr *lin_node = (linear_op_expr *) calloc(1, sizeof(linear_op_expr));
-    if (!lin_node) return NULL;
-
     expr *node = &lin_node->base;
-
-    /* Initialize base fields */
     init_expr(node, A->m, 1, u->n_vars, forward, NULL, NULL, is_affine,
               free_type_data);
-
-    /* Set left child */
     node->left = u;
     expr_retain(u);
 
-    /* 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);
-
     /* 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);
+    node->jacobian = lin_node->A_csr;
 
     return node;
 }

src/affine/sum.c

@@ -145,35 +145,6 @@ static void free_type_data(expr *node)
     free_int_double_pair_array(snode->int_double_pairs);
 }
 
-/* Helper function to initialize a sum expr */
-void init_sum(expr *node, expr *child, int d1)
-{
-    node->d1 = d1;
-    node->d2 = 1;
-    node->size = d1 * 1;
-    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->jacobian_init = jacobian_init;
-    node->wsum_hess_init = wsum_hess_init;
-    node->eval_jacobian = eval_jacobian;
-    node->eval_wsum_hess = eval_wsum_hess;
-    node->local_jacobian = NULL;
-    node->local_wsum_hess = NULL;
-    node->is_affine = is_affine;
-    node->forward = forward;
-    node->free_type_data = free_type_data;
-
-    expr_retain(child);
-}
-
 expr *new_sum(expr *child, int axis)
 {
     int d1 = 0;
@@ -195,20 +166,16 @@ expr *new_sum(expr *child, int axis)
     }
 
     /* Allocate the type-specific struct */
-    sum_expr *snode = (sum_expr *) malloc(sizeof(sum_expr));
-    if (!snode) return NULL;
-
+    sum_expr *snode = (sum_expr *) calloc(1, sizeof(sum_expr));
     expr *node = &snode->base;
+    init_expr(node, d1, 1, child->n_vars, forward, jacobian_init, eval_jacobian,
+              is_affine, free_type_data);
+    node->left = child;
+    expr_retain(child);
 
-    /* Initialize base sum fields */
-    init_sum(node, child, d1);
-
-    /* Check if allocation succeeded */
-    if (!node->value)
-    {
-        free(snode);
-        return NULL;
-    }
+    /* hessian function pointers */
+    node->wsum_hess_init = wsum_hess_init;
+    node->eval_wsum_hess = eval_wsum_hess;
 
     /* Set type-specific fields */
     snode->axis = axis;

src/affine/variable.c

@@ -29,13 +29,10 @@ static bool is_affine(const expr *node)
 expr *new_variable(int d1, int d2, int var_id, int n_vars)
 {
     expr *node = new_expr(d1, d2, n_vars);
-    if (!node) return NULL;
-
     node->forward = forward;
     node->var_id = var_id;
     node->is_affine = is_affine;
     node->jacobian_init = jacobian_init;
-    // node->jacobian = NULL;
 
     return node;
 }

src/expr.c

@@ -38,13 +38,20 @@ void free_expr(expr *node)
     /* recursively free children */
     free_expr(node->left);
     free_expr(node->right);
+    node->left = NULL;
+    node->right = NULL;
 
     /* free value array and jacobian */
     free(node->value);
     free_csr_matrix(node->jacobian);
     free_csr_matrix(node->wsum_hess);
     free(node->dwork);
     free(node->iwork);
+    node->value = NULL;
+    node->jacobian = NULL;
+    node->wsum_hess = NULL;
+    node->dwork = NULL;
+    node->iwork = NULL;
 
     /* free type-specific data */
     if (node->free_type_data)