hessian for elementwise univariate atoms

Author: dance858

README.md

@@ -0,0 +1,2 @@
+1. power should be double
+2. can we reuse calculations, like in hessian of logistic

include/elementwise_univariate.h

@@ -15,11 +15,15 @@ expr *new_asinh(expr *child);
 expr *new_atanh(expr *child);
 expr *new_logistic(expr *child);
 expr *new_power(expr *child, int p);
+expr *new_xexp(expr *child);
 
-/* the jacobian for elementwise atoms are always initialized in the
-   same way and implement the chain rule in the same way */
+/* the jacobian and wsum_hess for elementwise univariate atoms are always
+   initialized in the same way and implement the chain rule in the same way */
 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);
+expr *new_elementwise(expr *child);
 
 /* no elementwise atoms are affine according to our convention,
    so we can have a common implementation */

include/expr.h

@@ -15,8 +15,11 @@ struct int_double_pair;
 /* Function pointer types */
 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 (*eval_local_jacobian_fn)(struct expr *node, double *out);
+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);
 
 /* TODO: implement proper polymorphism */
@@ -52,10 +55,14 @@ typedef struct expr
     //                      jacobian related quantities
     // ------------------------------------------------------------------------
     CSR_Matrix *jacobian;
+    CSR_Matrix *wsum_hess;
     CSR_Matrix *CSR_work;
     jacobian_init_fn jacobian_init;
+    wsum_hess_init_fn wsum_hess_init;
     eval_jacobian_fn eval_jacobian;
-    eval_local_jacobian_fn eval_local_jacobian;
+    wsum_hess_fn eval_wsum_hess;
+    local_jacobian_fn local_jacobian;
+    local_wsum_hess_fn local_wsum_hess;
     is_affine_fn is_affine;
 
     // for every linear operator we store A in CSR and CSC

src/bivariate/rel_entr.c

@@ -87,7 +87,7 @@ expr *new_rel_entr_vector_args(expr *left, expr *right)
     node->jacobian_init = jacobian_init_vectors_args;
     node->eval_jacobian = eval_jacobian_vector_args;
     // node->is_affine = is_affine_elementwise;
-    // node->eval_local_jacobian = eval_local_jacobian;
+    // node->local_jacobian = local_jacobian;
     return node;
 }
 

src/elementwise_univariate/common.c

@@ -5,7 +5,7 @@ void jacobian_init_elementwise(expr *node)
 {
     expr *child = node->left;
 
-    // if the variable is a child
+    /* if the variable is a child */
     if (child->var_id != -1)
     {
         node->jacobian = new_csr_matrix(node->size, node->n_vars, node->size);
@@ -16,7 +16,7 @@ void jacobian_init_elementwise(expr *node)
         }
         node->jacobian->p[node->size] = node->size;
     }
-    // otherwise it should be a linear operator
+    /* otherwise it should be a linear operator */
     else
     {
         node->jacobian = new_csr_matrix(child->jacobian->m, child->jacobian->n,
@@ -31,17 +31,73 @@ void eval_jacobian_elementwise(expr *node)
 
     if (child->var_id != -1)
     {
-        node->eval_local_jacobian(node, node->jacobian->x);
+        node->local_jacobian(node, node->jacobian->x);
     }
     else
     {
-        node->eval_local_jacobian(node, node->dwork);
+        node->local_jacobian(node, node->dwork);
         diag_csr_mult(node->dwork, child->A_csr, node->jacobian);
     }
 }
 
+void wsum_hess_init_elementwise(expr *node)
+{
+    expr *child = node->left;
+    int id = child->var_id;
+    int i;
+
+    /* if the variable is a child*/
+    if (id != -1)
+    {
+        node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, node->size);
+
+        for (i = 0; i < node->size; i++)
+        {
+            node->wsum_hess->p[id + i] = i;
+            node->wsum_hess->i[i] = id + i;
+        }
+
+        for (i = id + node->size; i <= node->n_vars; i++)
+        {
+            node->wsum_hess->p[i] = node->size;
+        }
+    }
+    /* otherwise it will be a linear operator */
+    else
+    {
+        node->wsum_hess = ATA_alloc(child->A_csc);
+    }
+}
+
+void eval_wsum_hess_elementwise(expr *node, double *w)
+{
+    expr *child = node->left;
+
+    if (child->var_id != -1)
+    {
+        node->local_wsum_hess(node, node->wsum_hess->x, w);
+    }
+    else
+    {
+        node->local_wsum_hess(node, node->dwork, w);
+        ATDA_values(child->A_csc, node->dwork, node->wsum_hess);
+    }
+}
+
 bool is_affine_elementwise(expr *node)
 {
     (void) node;
     return false;
 }
+
+expr *new_elementwise(expr *child)
+{
+    expr *node = new_expr(child->d1, child->d2, child->n_vars);
+    node->left = child;
+    expr_retain(child);
+    node->is_affine = is_affine_elementwise;
+    node->jacobian_init = jacobian_init_elementwise;
+    node->eval_jacobian = eval_jacobian_elementwise;
+    node->wsum_hess_init = wsum_hess_init_elementwise;
+    node->eval_wsum_hess = eval_wsum_hess_elementwise;
+}

src/elementwise_univariate/entr.c

@@ -15,7 +15,7 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void eval_local_jacobian(expr *node, double *vals)
+static void local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
     for (int j = 0; j < node->size; j++)
@@ -24,15 +24,21 @@ static void eval_local_jacobian(expr *node, double *vals)
     }
 }
 
+static void local_wsum_hess(expr *node, double *out, double *w)
+{
+    double *x = node->left->value;
+
+    for (int j = 0; j < node->size; j++)
+    {
+        out[j] = -w[j] / x[j];
+    }
+}
+
 expr *new_entr(expr *child)
 {
-    expr *node = new_expr(child->d1, child->d2, child->n_vars);
-    node->left = child;
-    expr_retain(child);
+    expr *node = new_elementwise(child);
     node->forward = forward;
-    node->jacobian_init = jacobian_init_elementwise;
-    node->eval_jacobian = eval_jacobian_elementwise;
-    node->is_affine = is_affine_elementwise;
-    node->eval_local_jacobian = eval_local_jacobian;
+    node->local_jacobian = local_jacobian;
+    node->local_wsum_hess = local_wsum_hess;
     return node;
 }

src/elementwise_univariate/exp.c

@@ -14,25 +14,25 @@ static void forward(expr *node, const double *u)
     }
 }
 
-static void eval_local_jacobian(expr *node, double *vals)
+static void local_jacobian(expr *node, double *vals)
 {
     memcpy(vals, node->value, node->size * sizeof(double));
 }
 
-expr *new_exp(expr *child)
+static void local_wsum_hess(expr *node, double *out, double *w)
 {
-    if (!child) return NULL;
-
-    expr *node = new_expr(child->d1, child->d2, child->n_vars);
-    if (!node) return NULL;
+    double *x = node->left->value;
+    for (int j = 0; j < node->size; j++)
+    {
+        out[j] = w[j] * exp(x[j]);
+    }
+}
 
-    node->left = child;
-    expr_retain(child);
+expr *new_exp(expr *child)
+{
+    expr *node = new_elementwise(child);
     node->forward = forward;
-    node->is_affine = is_affine_elementwise;
-    node->jacobian_init = jacobian_init_elementwise;
-    node->eval_jacobian = eval_jacobian_elementwise;
-    node->eval_local_jacobian = eval_local_jacobian;
-
+    node->local_jacobian = local_jacobian;
+    node->local_wsum_hess = local_wsum_hess;
     return node;
 }

src/elementwise_univariate/hyperbolic.c

@@ -1,4 +1,5 @@
 #include "elementwise_univariate.h"
+#include <assert.h>
 #include <math.h>
 
 /* ----------------------- sinh ----------------------- */
@@ -11,7 +12,7 @@ static void sinh_forward(expr *node, const double *u)
     }
 }
 
-static void sinh_eval_local_jacobian(expr *node, double *vals)
+static void sinh_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
     for (int j = 0; j < node->size; j++)
@@ -20,16 +21,21 @@ static void sinh_eval_local_jacobian(expr *node, double *vals)
     }
 }
 
+static void sinh_local_wsum_hess(expr *node, double *out, double *w)
+{
+    double *x = node->left->value;
+    for (int j = 0; j < node->size; j++)
+    {
+        out[j] = w[j] * sinh(x[j]);
+    }
+}
+
 expr *new_sinh(expr *child)
 {
-    expr *node = new_expr(child->d1, child->d2, child->n_vars);
-    node->left = child;
-    expr_retain(child);
+    expr *node = new_elementwise(child);
     node->forward = sinh_forward;
-    node->jacobian_init = jacobian_init_elementwise;
-    node->eval_jacobian = eval_jacobian_elementwise;
-    node->eval_local_jacobian = sinh_eval_local_jacobian;
-    node->is_affine = is_affine_elementwise;
+    node->local_jacobian = sinh_local_jacobian;
+    node->local_wsum_hess = sinh_local_wsum_hess;
     return node;
 }
 
@@ -43,7 +49,7 @@ static void tanh_forward(expr *node, const double *u)
     }
 }
 
-static void tanh_eval_local_jacobian(expr *node, double *vals)
+static void tanh_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
     for (int j = 0; j < node->size; j++)
@@ -53,16 +59,22 @@ static void tanh_eval_local_jacobian(expr *node, double *vals)
     }
 }
 
+static void tanh_local_wsum_hess(expr *node, double *out, double *w)
+{
+    double *x = node->left->value;
+    for (int j = 0; j < node->size; j++)
+    {
+        double c = cosh(x[j]);
+        out[j] = w[j] * (-2.0 * tanh(x[j]) / (c * c));
+    }
+}
+
 expr *new_tanh(expr *child)
 {
-    expr *node = new_expr(child->d1, child->d2, child->n_vars);
-    node->left = child;
-    expr_retain(child);
+    expr *node = new_elementwise(child);
     node->forward = tanh_forward;
-    node->jacobian_init = jacobian_init_elementwise;
-    node->eval_jacobian = eval_jacobian_elementwise;
-    node->eval_local_jacobian = tanh_eval_local_jacobian;
-    node->is_affine = is_affine_elementwise;
+    node->local_jacobian = tanh_local_jacobian;
+    node->local_wsum_hess = tanh_local_wsum_hess;
     return node;
 }
 
@@ -76,7 +88,7 @@ static void asinh_forward(expr *node, const double *u)
     }
 }
 
-static void asinh_eval_local_jacobian(expr *node, double *vals)
+static void asinh_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
     for (int j = 0; j < node->size; j++)
@@ -85,16 +97,22 @@ static void asinh_eval_local_jacobian(expr *node, double *vals)
     }
 }
 
+static void asinh_local_wsum_hess(expr *node, double *out, double *w)
+{
+    double *x = node->left->value;
+    for (int j = 0; j < node->size; j++)
+    {
+        double c = 1.0 + x[j] * x[j];
+        out[j] = w[j] * (-x[j]) / pow(c, 1.5);
+    }
+}
+
 expr *new_asinh(expr *child)
 {
-    expr *node = new_expr(child->d1, child->d2, child->n_vars);
-    node->left = child;
-    expr_retain(child);
+    expr *node = new_elementwise(child);
     node->forward = asinh_forward;
-    node->jacobian_init = jacobian_init_elementwise;
-    node->eval_jacobian = eval_jacobian_elementwise;
-    node->eval_local_jacobian = asinh_eval_local_jacobian;
-    node->is_affine = is_affine_elementwise;
+    node->local_jacobian = asinh_local_jacobian;
+    node->local_wsum_hess = asinh_local_wsum_hess;
     return node;
 }
 
@@ -108,7 +126,7 @@ static void atanh_forward(expr *node, const double *u)
     }
 }
 
-static void atanh_eval_local_jacobian(expr *node, double *vals)
+static void atanh_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
     for (int j = 0; j < node->size; j++)
@@ -117,15 +135,21 @@ static void atanh_eval_local_jacobian(expr *node, double *vals)
     }
 }
 
+static void atanh_local_wsum_hess(expr *node, double *out, double *w)
+{
+    double *x = node->left->value;
+    for (int j = 0; j < node->size; j++)
+    {
+        double c = 1.0 - x[j] * x[j];
+        out[j] = w[j] * (2.0 * x[j]) / (c * c);
+    }
+}
+
 expr *new_atanh(expr *child)
 {
-    expr *node = new_expr(child->d1, child->d2, child->n_vars);
-    node->left = child;
-    expr_retain(child);
+    expr *node = new_elementwise(child);
     node->forward = atanh_forward;
-    node->jacobian_init = jacobian_init_elementwise;
-    node->eval_jacobian = eval_jacobian_elementwise;
-    node->eval_local_jacobian = atanh_eval_local_jacobian;
-    node->is_affine = is_affine_elementwise;
+    node->local_jacobian = atanh_local_jacobian;
+    node->local_wsum_hess = atanh_local_wsum_hess;
     return node;
 }