more jacobians

Author: dance858

include/bivariate.h

@@ -0,0 +1,13 @@
+#ifndef BIVARIATE_H
+#define BIVARIATE_H
+
+#include "expr.h"
+
+expr *new_elementwise_mult(expr *left, expr *right);
+expr *new_rel_entr_vector_args(expr *left, expr *right);
+expr *new_quad_over_lin(expr *left, expr *right);
+
+expr *new_rel_entr_first_arg_scalar(expr *left, expr *right);
+expr *new_rel_entr_second_arg_scalar(expr *left, expr *right);
+
+#endif /* BIVARIATE_H */

include/elementwise_univariate.h

@@ -9,6 +9,12 @@ expr *new_entr(expr *child);
 expr *new_sin(expr *child);
 expr *new_cos(expr *child);
 expr *new_tan(expr *child);
+expr *new_sinh(expr *child);
+expr *new_tanh(expr *child);
+expr *new_asinh(expr *child);
+expr *new_atanh(expr *child);
+expr *new_logistic(expr *child);
+expr *new_power(expr *child, int p);
 
 /* the jacobian for elementwise atoms are always initialized in the
    same way and implement the chain rule in the same way */
@@ -19,4 +25,4 @@ void eval_jacobian_elementwise(expr *node);
    so we can have a common implementation */
 bool is_affine_elementwise(expr *node);
 
-#endif /* ELEMENTWISE_H */
+#endif /* ELEMENTWISE_UNIVARIATE_H */

include/expr.h

@@ -26,9 +26,12 @@ typedef struct expr
     int m;
     int n_vars;
     int var_id;
+    int refcount;
     struct expr *left;
     struct expr *right;
     double *dwork;
+    int *iwork;
+    int p; /* power of power expression */
 
     // ------------------------------------------------------------------------
     //                     forward pass related quantities
@@ -40,6 +43,8 @@ typedef struct expr
     //                      jacobian related quantities
     // ------------------------------------------------------------------------
     CSR_Matrix *jacobian;
+    CSR_Matrix *Q;
+    CSR_Matrix *CSR_work;
     jacobian_init_fn jacobian_init;
     eval_jacobian_fn eval_jacobian;
     eval_local_jacobian_fn eval_local_jacobian;
@@ -50,4 +55,7 @@ typedef struct expr
 expr *new_expr(int m, int n_vars);
 void free_expr(expr *node);
 
+/* Reference counting helpers */
+void expr_retain(expr *node);
+
 #endif /* EXPR_H */

include/other.h

@@ -0,0 +1,9 @@
+#ifndef OTHER_H
+#define OTHER_H
+
+#include "expr.h"
+#include "utils/CSR_Matrix.h"
+
+expr *new_quad_form(expr *child, CSR_Matrix *Q);
+
+#endif /* OTHER_H */

include/utils/CSR_Matrix.h

@@ -1,5 +1,6 @@
 #ifndef CSR_MATRIX_H
 #define CSR_MATRIX_H
+#include <stdbool.h>
 
 /* CSR (Compressed Sparse Row) Matrix Format
  *
@@ -30,9 +31,20 @@ void free_csr_matrix(CSR_Matrix *matrix);
 /* Copy CSR matrix A to C */
 void copy_csr_matrix(const CSR_Matrix *A, CSR_Matrix *C);
 
-/* matvec y = Ax, where A indices minus col_offset gives x indices */
+/* matvec y = Ax, where A indices minus col_offset gives x indices. Returns y as
+ * dense. */
 void csr_matvec(const CSR_Matrix *A, const double *x, double *y, int col_offset);
 
+/* C = z^T A is assumed to have one row. C must have column indices pre-computed
+and transposed matrix AT must be provided. Fills in values of C only.
+ */
+void csr_matvec_fill_values(const CSR_Matrix *AT, const double *z, CSR_Matrix *C);
+
+/* Insert value into CSR matrix A with just one row at col_idx. Assumes that A
+has enough space and that A does not have an element at col_idx. It does update
+nnz. */
+void csr_insert_value(CSR_Matrix *A, int col_idx, double value);
+
 /* Compute C = diag(d) * A where d is an array and A, C are CSR matrices
  * d must have length m
  * C must be pre-allocated with same dimensions as A */
@@ -43,4 +55,15 @@ void diag_csr_mult(const double *d, const CSR_Matrix *A, CSR_Matrix *C);
  * C must be pre-allocated with sufficient nnz capacity */
 void sum_csr_matrices(const CSR_Matrix *A, const CSR_Matrix *B, CSR_Matrix *C);
 
+/* Compute C = diag(d1) * A + diag(d2) * B where A, B, C are CSR matrices */
+void sum_scaled_csr_matrices(const CSR_Matrix *A, const CSR_Matrix *B, CSR_Matrix *C,
+                             const double *d1, const double *d2);
+
+/* 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 */

src/affine/add.c

@@ -49,6 +49,8 @@ expr *new_add(expr *left, expr *right)
 
     node->left = left;
     node->right = right;
+    expr_retain(left);
+    expr_retain(right);
     node->forward = add_forward;
     node->is_affine = is_affine;
     node->jacobian_init = jacobian_init;

src/affine/linear_op.c

@@ -22,6 +22,7 @@ expr *new_linear(expr *u, const CSR_Matrix *A)
     if (!node) return NULL;
 
     node->left = u;
+    expr_retain(u);
     node->forward = forward;
     node->is_affine = is_affine;
 

src/affine/variable.c

@@ -7,6 +7,18 @@ static void forward(expr *node, const double *u)
     memcpy(node->value, u + node->var_id, node->m * sizeof(double));
 }
 
+static void jacobian_init(expr *node)
+{
+    node->jacobian = new_csr_matrix(node->m, node->n_vars, node->m);
+    for (int j = 0; j < node->m; j++)
+    {
+        node->jacobian->p[j] = j;
+        node->jacobian->i[j] = j + node->var_id;
+        node->jacobian->x[j] = 1.0;
+    }
+    node->jacobian->p[node->m] = node->m;
+}
+
 static bool is_affine(expr *node)
 {
     (void) node;
@@ -21,6 +33,8 @@ expr *new_variable(int m, int var_id, int n_vars)
     node->forward = forward;
     node->var_id = var_id;
     node->is_affine = is_affine;
+    node->jacobian_init = jacobian_init;
+    // node->jacobian = NULL;
 
     return node;
 }

src/bivariate/multiply.c

@@ -0,0 +1,66 @@
+#include "bivariate.h"
+#include <math.h>
+#include <stdlib.h>
+
+// ------------------------------------------------------------------------------
+// Implementation of elementwise multiplication when both arguments are vectors.
+// If one argument is a scalar variable, the broadcasting should be represented
+// as a linear operator child node?
+// ------------------------------------------------------------------------------
+static void forward(expr *node, const double *u)
+{
+    expr *x = node->left;
+    expr *y = node->right;
+
+    /* children's forward passes */
+    x->forward(x, u);
+    y->forward(y, u);
+
+    /* local forward pass */
+    for (int i = 0; i < node->m; i++)
+    {
+        node->value[i] = x->value[i] * y->value[i];
+    }
+}
+
+static void jacobian_init(expr *node)
+{
+    /* if a child is a variable we initialize its jacobian for a
+       short chain rule implementation */
+    if (node->left->var_id != -1)
+    {
+        node->left->jacobian_init(node->left);
+    }
+
+    if (node->right->var_id != -1)
+    {
+        node->right->jacobian_init(node->right);
+    }
+
+    node->dwork = (double *) malloc(2 * node->m * sizeof(double));
+    int nnz_max = node->left->jacobian->nnz + node->right->jacobian->nnz;
+    node->jacobian = new_csr_matrix(node->m, node->n_vars, nnz_max);
+}
+
+static void eval_jacobian(expr *node)
+{
+    expr *x = node->left;
+    expr *y = node->right;
+
+    /* chain rule */
+    sum_scaled_csr_matrices(x->jacobian, y->jacobian, node->jacobian, y->value,
+                            x->value);
+}
+
+expr *new_elementwise_mult(expr *left, expr *right)
+{
+    expr *node = new_expr(left->m, left->n_vars);
+    node->left = left;
+    node->right = right;
+    expr_retain(left);
+    expr_retain(right);
+    node->forward = forward;
+    node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
+    return node;
+}

src/bivariate/quad_over_lin.c

@@ -0,0 +1,162 @@
+#include "bivariate.h"
+#include <assert.h>
+#include <math.h>
+#include <stdlib.h>
+
+// ------------------------------------------------------------------------------
+// Implementation of quad-over-lin. The second argument will always be a variable
+// that only appears in this node and as the left-hand side of another equality
+// constraint.
+// ------------------------------------------------------------------------------
+static void forward(expr *node, const double *u)
+{
+    expr *x = node->left;
+    expr *y = node->right;
+
+    /* children's forward passes */
+    x->forward(x, u);
+    y->forward(y, u);
+
+    /* local forward pass */
+    node->value[0] = 0.0;
+
+    for (int i = 0; i < x->m; i++)
+    {
+        node->value[0] += x->value[i] * x->value[i];
+    }
+
+    node->value[0] /= y->value[0];
+}
+
+static void jacobian_init(expr *node)
+{
+    expr *x = node->left;
+    expr *y = node->right;
+
+    /* if left node is a variable */
+    if (x->var_id != -1)
+    {
+        node->jacobian = new_csr_matrix(1, node->n_vars, x->m + 1);
+        node->jacobian->p[0] = 0;
+        node->jacobian->p[1] = x->m + 1;
+
+        /* if x has lower idx than y*/
+        if (x->var_id < y->var_id)
+        {
+            for (int j = 0; j < x->m; j++)
+            {
+                node->jacobian->i[j] = x->var_id + j;
+            }
+            node->jacobian->i[x->m] = y->var_id;
+        }
+        else /* y has lower idx than x */
+        {
+            node->jacobian->i[0] = y->var_id;
+            for (int j = 0; j < x->m; j++)
+            {
+                node->jacobian->i[j + 1] = x->var_id + j;
+            }
+        }
+    }
+    else /* left node is not a variable */
+    {
+        node->dwork = (double *) malloc(x->m * 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);
+        node->jacobian = new_csr_matrix(1, node->n_vars, nonzero_cols + 1);
+
+        /* precompute column indices */
+        node->jacobian->nnz = 0;
+        for (int j = 0; j < node->n_vars; j++)
+        {
+            if (col_nz[j])
+            {
+                node->jacobian->i[node->jacobian->nnz] = j;
+                node->jacobian->nnz++;
+            }
+        }
+        assert(nonzero_cols == node->jacobian->nnz);
+
+        free(col_nz);
+
+        /* insert y variable index at correct position */
+        insert_idx(y->var_id, node->jacobian->i, node->jacobian->nnz);
+        node->jacobian->nnz += 1;
+        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 */
+        for (int j = 0; j < node->jacobian->nnz; j++)
+        {
+            if (node->jacobian->i[j] == y->var_id)
+            {
+                node->iwork[0] = j;
+                break;
+            }
+        }
+    }
+}
+
+static void eval_jacobian(expr *node)
+{
+    expr *x = node->left;
+    expr *y = node->right;
+
+    /* if x is a variable */
+    if (x->var_id != -1)
+    {
+        /* if x has lower idx than y*/
+        if (x->var_id < y->var_id)
+        {
+            for (int j = 0; j < x->m; j++)
+            {
+                node->jacobian->x[j] = (2.0 * x->value[j]) / y->value[0];
+            }
+            node->jacobian->x[x->m] = -node->value[0] / y->value[0];
+        }
+        else /* y has lower idx than x */
+        {
+            node->jacobian->x[0] = -node->value[0] / y->value[0];
+            for (int j = 0; j < x->m; j++)
+            {
+                node->jacobian->x[j + 1] = (2.0 * x->value[j]) / y->value[0];
+            }
+        }
+    }
+    else /* x is not a variable */
+    {
+        /* local jacobian */
+        for (int j = 0; j < x->m; 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);
+
+        /* insert derivative wrt y at right place (for correctness this assumes
+           that y does not appear in the denominator, but this will always be
+           the case since y is a new variable for the numerator) */
+        node->jacobian->x[node->iwork[0]] = -node->value[0] / y->value[0];
+    }
+}
+
+expr *new_quad_over_lin(expr *left, expr *right)
+{
+    expr *node = new_expr(left->m, left->n_vars);
+    node->left = left;
+    node->right = right;
+    expr_retain(left);
+    expr_retain(right);
+    node->forward = forward;
+    node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
+    return node;
+}