relative entropy hessian

Author: dance858

src/bivariate/rel_entr.c

@@ -1,4 +1,5 @@
 #include "bivariate.h"
+#include <assert.h>
 #include <math.h>
 #include <stdlib.h>
 
@@ -28,6 +29,8 @@ static void jacobian_init_vectors_args(expr *node)
 
     expr *x = node->left;
     expr *y = node->right;
+    assert(x->var_id != NOT_A_VARIABLE && y->var_id != NOT_A_VARIABLE);
+    assert(x->var_id != y->var_id);
 
     /* if x has lower variable idx than y it should appear first */
     if (x->var_id < y->var_id)
@@ -76,6 +79,97 @@ static void eval_jacobian_vector_args(expr *node)
     }
 }
 
+static void wsum_hess_init_vector_args(expr *node)
+{
+    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 4 * node->d1);
+    expr *x = node->left;
+    expr *y = node->right;
+
+    int i, var1_id, var2_id;
+
+    if (x->var_id < y->var_id)
+    {
+        var1_id = x->var_id;
+        var2_id = y->var_id;
+    }
+    else
+    {
+        var1_id = y->var_id;
+        var2_id = x->var_id;
+    }
+
+    /* var1 rows of Hessian */
+    for (i = 0; i < node->d1; i++)
+    {
+        node->wsum_hess->p[var1_id + i] = 2 * i;
+        node->wsum_hess->i[2 * i] = var1_id + i;
+        node->wsum_hess->i[2 * i + 1] = var2_id + i;
+    }
+
+    int nnz = 2 * node->d1;
+
+    /* rows between var1 and var2 */
+    for (i = var1_id + node->d1; i < var2_id; i++)
+    {
+        node->wsum_hess->p[i] = nnz;
+    }
+
+    /* var2 rows of Hessian */
+    for (i = 0; i < node->d1; i++)
+    {
+        node->wsum_hess->p[var2_id + i] = nnz + 2 * i;
+    }
+    memcpy(node->wsum_hess->i + nnz, node->wsum_hess->i, nnz * sizeof(int));
+
+    /* remaining rows */
+    for (i = var2_id + node->d1; i <= node->n_vars; i++)
+    {
+        node->wsum_hess->p[i] = 4 * node->d1;
+    }
+}
+
+static void eval_wsum_hess_vector_args(expr *node, const double *w)
+{
+
+    int i, x_id, y_id;
+    double *x = node->left->value;
+    double *y = node->right->value;
+    double *hess = node->wsum_hess->x;
+
+    if (node->left->var_id < node->right->var_id)
+    {
+        for (i = 0; i < node->d1; i++)
+        {
+            hess[2 * i] = w[i] / x[i];
+            hess[2 * i + 1] = -w[i] / y[i];
+        }
+
+        hess += 2 * node->d1;
+
+        for (i = 0; i < node->d1; i++)
+        {
+            hess[2 * i] = -w[i] / y[i];
+            hess[2 * i + 1] = w[i] * x[i] / (y[i] * y[i]);
+        }
+    }
+    else
+    {
+        for (i = 0; i < node->d1; i++)
+        {
+            hess[2 * i] = w[i] * x[i] / (y[i] * y[i]);
+            hess[2 * i + 1] = -w[i] / y[i];
+        }
+
+        hess += 2 * node->d1;
+
+        for (i = 0; i < node->d1; i++)
+        {
+            hess[2 * i] = -w[i] / y[i];
+            hess[2 * i + 1] = w[i] / x[i];
+        }
+    }
+}
+
 expr *new_rel_entr_vector_args(expr *left, expr *right)
 {
     expr *node = new_expr(left->d1, 1, left->n_vars);
@@ -86,6 +180,8 @@ expr *new_rel_entr_vector_args(expr *left, expr *right)
     node->forward = forward_vector_args;
     node->jacobian_init = jacobian_init_vectors_args;
     node->eval_jacobian = eval_jacobian_vector_args;
+    node->wsum_hess_init = wsum_hess_init_vector_args;
+    node->eval_wsum_hess = eval_wsum_hess_vector_args;
     // node->is_affine = is_affine_elementwise;
     // node->local_jacobian = local_jacobian;
     return node;

tests/all_tests.c

@@ -21,15 +21,16 @@
 #include "jacobian_tests/test_sum.h"
 #include "utils/test_csc_matrix.h"
 #include "utils/test_csr_matrix.h"
-#include "wsum_hess/test_entr.h"
-#include "wsum_hess/test_exp.h"
-#include "wsum_hess/test_hyperbolic.h"
-#include "wsum_hess/test_log.h"
-#include "wsum_hess/test_logistic.h"
-#include "wsum_hess/test_power.h"
+#include "wsum_hess/elementwise/test_entr.h"
+#include "wsum_hess/elementwise/test_exp.h"
+#include "wsum_hess/elementwise/test_hyperbolic.h"
+#include "wsum_hess/elementwise/test_log.h"
+#include "wsum_hess/elementwise/test_logistic.h"
+#include "wsum_hess/elementwise/test_power.h"
+#include "wsum_hess/elementwise/test_trig.h"
+#include "wsum_hess/elementwise/test_xexp.h"
+#include "wsum_hess/test_rel_entr.h"
 #include "wsum_hess/test_sum.h"
-#include "wsum_hess/test_trig.h"
-#include "wsum_hess/test_xexp.h"
 
 int main(void)
 {
@@ -95,6 +96,8 @@ int main(void)
     mu_run_test(test_wsum_hess_sum_log_linear, tests_run);
     mu_run_test(test_wsum_hess_sum_log_axis0, tests_run);
     mu_run_test(test_wsum_hess_sum_log_axis1, tests_run);
+    mu_run_test(test_wsum_hess_rel_entr_1, tests_run);
+    mu_run_test(test_wsum_hess_rel_entr_2, tests_run);
 
     printf("\n--- Utility Tests ---\n");
     mu_run_test(test_diag_csr_mult, tests_run);