jacobian of sum

Author: dance858

include/affine.h

@@ -0,0 +1,15 @@
+#ifndef AFFINE_H
+#define AFFINE_H
+
+#include "expr.h"
+#include "utils/CSR_Matrix.h"
+
+expr *new_linear(expr *u, const CSR_Matrix *A);
+
+expr *new_add(expr *left, expr *right);
+expr *new_sum(expr *child, int axis);
+
+expr *new_constant(int d1, int d2, const double *values);
+expr *new_variable(int d1, int d2, int var_id, int n_vars);
+
+#endif /* AFFINE_H */

include/affine/add.h

@@ -9,6 +9,7 @@
 
 /* Forward declarations */
 struct expr;
+struct int_double_pair;
 
 /* Function pointer types */
 typedef void (*forward_fn)(struct expr *node, const double *u);
@@ -23,15 +24,17 @@ typedef struct expr
     // ------------------------------------------------------------------------
     //                         general quantities
     // ------------------------------------------------------------------------
-    int m;
+    int d1, d2, size;
     int n_vars;
     int var_id;
     int refcount;
     struct expr *left;
     struct expr *right;
     double *dwork;
     int *iwork;
-    int p; /* power of power expression */
+    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 */
 
     // ------------------------------------------------------------------------
     //                     forward pass related quantities
@@ -52,7 +55,7 @@ typedef struct expr
 
 } expr;
 
-expr *new_expr(int m, int n_vars);
+expr *new_expr(int d1, int d2, int n_vars);
 void free_expr(expr *node);
 
 /* Reference counting helpers */

include/affine/constant.h

@@ -2,6 +2,9 @@
 #define CSR_MATRIX_H
 #include <stdbool.h>
 
+/* forward declaration */
+struct int_double_pair;
+
 /* CSR (Compressed Sparse Row) Matrix Format
  *
  * For an m x n matrix with nnz nonzeros:
@@ -59,11 +62,24 @@ void sum_csr_matrices(const CSR_Matrix *A, const CSR_Matrix *B, CSR_Matrix *C);
 void sum_scaled_csr_matrices(const CSR_Matrix *A, const CSR_Matrix *B, CSR_Matrix *C,
                              const double *d1, const double *d2);
 
+/* Sum all rows of A into a single row matrix C */
+void sum_all_rows_csr(const CSR_Matrix *A, CSR_Matrix *C,
+                      struct int_double_pair *pairs);
+
+/* Sum blocks of rows of A into a matrix C */
+void sum_block_of_rows_csr(const CSR_Matrix *A, CSR_Matrix *C,
+                           struct int_double_pair *pairs, int row_block_size);
+
+/* Sum evenly spaced rows of A into a matrix C */
+void sum_evenly_spaced_rows_csr(const CSR_Matrix *A, CSR_Matrix *C,
+                                struct int_double_pair *pairs, int row_spacing);
+
 /* Count number of columns with nonzero entries */
 int count_nonzero_cols(const CSR_Matrix *A, bool *col_nz);
 
 /* inserts 'idx' into array 'arr' in sorted order, and moves the other elements */
 void insert_idx(int idx, int *arr, int len);
 
 CSR_Matrix *transpose(const CSR_Matrix *A, int *iwork);
+
 #endif /* CSR_MATRIX_H */

include/affine/linear_op.h

@@ -0,0 +1,17 @@
+
+#ifndef INT_DOUBLE_PAIR_H
+#define INT_DOUBLE_PAIR_H
+
+typedef struct int_double_pair
+{
+    int col;
+    double val;
+} int_double_pair;
+
+int_double_pair *new_int_double_pair_array(int size);
+void set_int_double_pair_array(int_double_pair *pair, int *ints, double *doubles,
+                               int size);
+void free_int_double_pair_array(int_double_pair *array);
+void sort_int_double_pair_array(int_double_pair *array, int size);
+
+#endif /* INT_DOUBLE_PAIR_H */

include/affine/variable.h

@@ -1,4 +1,4 @@
-#include "affine/add.h"
+#include "affine.h"
 
 static void add_forward(expr *node, const double *u)
 {
@@ -7,7 +7,7 @@ static void add_forward(expr *node, const double *u)
     node->right->forward(node->right, u);
 
     /* add left and right values */
-    for (int i = 0; i < node->m; i++)
+    for (int i = 0; i < node->size; i++)
     {
         node->value[i] = node->left->value[i] + node->right->value[i];
     }
@@ -21,7 +21,7 @@ static void jacobian_init(expr *node)
 
     /* we never have to store more than the sum of children's nnz */
     int nnz_max = node->left->jacobian->nnz + node->right->jacobian->nnz;
-    node->jacobian = new_csr_matrix(node->m, node->n_vars, nnz_max);
+    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz_max);
 }
 
 static void eval_jacobian(expr *node)
@@ -42,9 +42,10 @@ static bool is_affine(expr *node)
 expr *new_add(expr *left, expr *right)
 {
     if (!left || !right) return NULL;
-    if (left->m != right->m) return NULL;
+    if (left->d1 != right->d1) return NULL;
+    if (left->d2 != right->d2) return NULL;
 
-    expr *node = new_expr(left->m, left->n_vars);
+    expr *node = new_expr(left->d1, left->d2, left->n_vars);
     if (!node) return NULL;
 
     node->left = left;

include/expr.h

@@ -1,4 +1,4 @@
-#include "affine/constant.h"
+#include "affine.h"
 #include <string.h>
 
 static void forward(expr *node, const double *u)
@@ -14,13 +14,13 @@ static bool is_affine(expr *node)
     return true; /* constant is affine */
 }
 
-expr *new_constant(int m, const double *values)
+expr *new_constant(int d1, int d2, const double *values)
 {
-    expr *node = new_expr(m, 0);
+    expr *node = new_expr(d1, d2, node->n_vars);
     if (!node) return NULL;
 
     /* Copy constant values */
-    memcpy(node->value, values, m * sizeof(double));
+    memcpy(node->value, values, d1 * d2 * sizeof(double));
 
     node->forward = forward;
     node->is_affine = is_affine;