wsum of add and sum

Author: dance858

README.md

@@ -1,3 +1,6 @@
 1. power should be double
 2. can we reuse calculations, like in hessian of logistic
-3. more tests for chain rule elementwise univariate hessian
+3. more tests for chain rule elementwise univariate hessian
+4. in the refactor, add consts
+5. multiply with one constant vector/scalar argument
+6. AX where X is a matrix. Can that happen? How is that canonicalized? Maybe it can't happen.

include/utils/mini_numpy.h

@@ -0,0 +1,22 @@
+#ifndef MINI_NUMPY_H
+#define MINI_NUMPY_H
+
+/* Repeat each element of array 'a' 'repeats' times
+ * Example: a = [1, 2], len = 2, repeats = 3
+ *          result = [1, 1, 1, 2, 2, 2]
+ */
+void repeat(double *result, const double *a, int len, int repeats);
+
+/* Tile array 'a' 'tiles' times
+ * Example: a = [1, 2], len = 2, tiles = 3
+ *          result = [1, 2, 1, 2, 1, 2]
+ */
+void tile(double *result, const double *a, int len, int tiles);
+
+/* Fill array with 'size' copies of 'value'
+ * Example: size = 5, value = 3.0
+ *          result = [3.0, 3.0, 3.0, 3.0, 3.0]
+ */
+void scaled_ones(double *result, int size, double value);
+
+#endif /* MINI_NUMPY_H */

src/affine/add.c

@@ -34,6 +34,27 @@ static void eval_jacobian(expr *node)
     sum_csr_matrices(node->left->jacobian, node->right->jacobian, node->jacobian);
 }
 
+static void wsum_hess_init(expr *node)
+{
+    /* initialize children's wsum_hess */
+    node->left->wsum_hess_init(node->left);
+    node->right->wsum_hess_init(node->right);
+
+    /* we never have to store more than the sum of children's nnz */
+    int nnz_max = node->left->wsum_hess->nnz + node->right->wsum_hess->nnz;
+    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, nnz_max);
+}
+
+static void eval_wsum_hess(expr *node, double *w)
+{
+    /* evaluate children's wsum_hess */
+    node->left->eval_wsum_hess(node->left, w);
+    node->right->eval_wsum_hess(node->right, w);
+
+    /* sum children's wsum_hess */
+    sum_csr_matrices(node->left->wsum_hess, node->right->wsum_hess, node->wsum_hess);
+}
+
 static bool is_affine(expr *node)
 {
     return node->left->is_affine(node->left) && node->right->is_affine(node->right);
@@ -56,6 +77,8 @@ expr *new_add(expr *left, expr *right)
     node->is_affine = is_affine;
     node->jacobian_init = jacobian_init;
     node->eval_jacobian = eval_jacobian;
+    node->wsum_hess_init = wsum_hess_init;
+    node->eval_wsum_hess = eval_wsum_hess;
 
     return node;
 }

src/affine/sum.c

@@ -1,5 +1,6 @@
 #include "affine.h"
 #include "utils/int_double_pair.h"
+#include "utils/mini_numpy.h"
 #include <assert.h>
 #include <stdlib.h>
 #include <string.h>
@@ -91,6 +92,40 @@ static void eval_jacobian(expr *node)
     }
 }
 
+static void wsum_hess_init(expr *node)
+{
+    expr *x = node->left;
+    /* initialize child's wsum_hess */
+    x->wsum_hess_init(x);
+
+    /* we never have to store more than the child's nnz */
+    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, x->wsum_hess->nnz);
+    node->dwork = malloc(x->size * sizeof(double));
+}
+
+static void eval_wsum_hess(expr *node, double *w)
+{
+    expr *x = node->left;
+
+    if (node->axis == -1)
+    {
+        scaled_ones(node->dwork, x->size, *w);
+    }
+    else if (node->axis == 0)
+    {
+        repeat(node->dwork, w, x->d2, x->d1);
+    }
+    else if (node->axis == 1)
+    {
+        tile(node->dwork, w, x->d1, x->d2);
+    }
+
+    x->eval_wsum_hess(x, node->dwork);
+
+    /* todo: is this copy necessary or can we just change pointers? */
+    copy_csr_matrix(x->wsum_hess, node->wsum_hess);
+}
+
 static bool is_affine(expr *node)
 {
     return node->left->is_affine(node->left);
@@ -125,6 +160,8 @@ expr *new_sum(expr *child, int axis)
     node->is_affine = is_affine;
     node->jacobian_init = jacobian_init;
     node->eval_jacobian = eval_jacobian;
+    node->wsum_hess_init = wsum_hess_init;
+    node->eval_wsum_hess = eval_wsum_hess;
     node->axis = axis;
 
     return node;

src/utils/mini_numpy.c

@@ -0,0 +1,33 @@
+#include "utils/mini_numpy.h"
+
+void repeat(double *result, const double *a, int len, int repeats)
+{
+    int idx = 0;
+    for (int i = 0; i < len; i++)
+    {
+        for (int j = 0; j < repeats; j++)
+        {
+            result[idx++] = a[i];
+        }
+    }
+}
+
+void tile(double *result, const double *a, int len, int tiles)
+{
+    int idx = 0;
+    for (int i = 0; i < tiles; i++)
+    {
+        for (int j = 0; j < len; j++)
+        {
+            result[idx++] = a[j];
+        }
+    }
+}
+
+void scaled_ones(double *result, int size, double value)
+{
+    for (int i = 0; i < size; i++)
+    {
+        result[i] = value;
+    }
+}

tests/all_tests.c

@@ -27,6 +27,7 @@
 #include "wsum_hess/test_log.h"
 #include "wsum_hess/test_logistic.h"
 #include "wsum_hess/test_power.h"
+#include "wsum_hess/test_sum.h"
 #include "wsum_hess/test_trig.h"
 #include "wsum_hess/test_xexp.h"
 
@@ -90,6 +91,9 @@ int main(void)
     mu_run_test(test_wsum_hess_tanh, tests_run);
     mu_run_test(test_wsum_hess_asinh, tests_run);
     mu_run_test(test_wsum_hess_atanh, tests_run);
+    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);
 
     printf("\n--- Utility Tests ---\n");
     mu_run_test(test_diag_csr_mult, tests_run);

tests/test_helpers.c

@@ -5,11 +5,19 @@
 
 #define EPSILON 1e-7
 
+#define ABS_TOL 1e-6
+#define REL_TOL 1e-6
+
+int is_equal_double(double a, double b)
+{
+    return fabs(a - b) <= fmax(ABS_TOL, REL_TOL * fmax(fabs(a), fabs(b)));
+}
+
 int cmp_double_array(const double *actual, const double *expected, int size)
 {
     for (int i = 0; i < size; i++)
     {
-        if (fabs(actual[i] - expected[i]) > EPSILON)
+        if (!is_equal_double(actual[i], expected[i]))
         {
             printf("  FAILED: actual[%d] = %f, expected %f\n", i, actual[i],
                    expected[i]);

tests/wsum_hess/test_sum.h

@@ -0,0 +1,134 @@
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "affine.h"
+#include "elementwise_univariate.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_wsum_hess_sum_log_linear()
+{
+    double Ax[6] = {1, 1, 2, 3, 1, -1};
+    int Ai[6] = {0, 1, 0, 1, 0, 1};
+    int Ap[4] = {0, 2, 4, 6};
+    CSR_Matrix *A = new_csr_matrix(3, 2, 6);
+    memcpy(A->x, Ax, 6 * sizeof(double));
+    memcpy(A->i, Ai, 6 * sizeof(int));
+    memcpy(A->p, Ap, 4 * sizeof(int));
+    double x_vals[2] = {2.0, 1.0};
+    double w = 1.5;
+
+    expr *x = new_variable(2, 1, 0, 2);
+    expr *Ax_node = new_linear(x, A);
+    expr *log_node = new_log(Ax_node);
+    expr *sum_node = new_sum(log_node, -1);
+
+    sum_node->forward(sum_node, x_vals);
+    sum_node->jacobian_init(sum_node);
+    // sum_node->eval_jacobian(sum_node);
+    sum_node->wsum_hess_init(sum_node);
+    sum_node->eval_wsum_hess(sum_node, &w);
+
+    double expected_x[4] = {-1.5 * 526.0 / 441.0, 1.5 * 338.0 / 441.0,
+                            1.5 * 338.0 / 441.0, -1.5 * 571.0 / 441.0};
+    int expected_p[3] = {0, 2, 4};
+    int expected_i[4] = {0, 1, 0, 1};
+
+    mu_assert("vals incorrect",
+              cmp_double_array(sum_node->wsum_hess->x, expected_x, 4));
+    mu_assert("rows incorrect",
+              cmp_int_array(sum_node->wsum_hess->p, expected_p, 3));
+    mu_assert("cols incorrect",
+              cmp_int_array(sum_node->wsum_hess->i, expected_i, 4));
+
+    free_expr(sum_node);
+    free_expr(log_node);
+    free_expr(Ax_node);
+    free_csr_matrix(A);
+    free_expr(x);
+
+    return 0;
+}
+
+const char *test_wsum_hess_sum_log_axis0()
+{
+    /* Test: wsum_hess of sum(log(x), axis=0) where x is 3x2
+     * x = [[1, 4],
+     *      [2, 5],
+     *      [3, 6]]
+     */
+
+    double x[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
+    double w[2] = {-1.0, -2.0};
+    expr *x_node = new_variable(3, 2, 0, 6);
+    expr *log_node = new_log(x_node);
+    expr *sum_node = new_sum(log_node, 0);
+
+    sum_node->forward(sum_node, x);
+    sum_node->jacobian_init(sum_node);
+    sum_node->wsum_hess_init(sum_node);
+    sum_node->eval_wsum_hess(sum_node, w);
+
+    /* Expected diagonal values */
+    double expected_x[6] = {-w[0] / (x[0] * x[0]), -w[0] / (x[1] * x[1]),
+                            -w[0] / (x[2] * x[2]), -w[1] / (x[3] * x[3]),
+                            -w[1] / (x[4] * x[4]), -w[1] / (x[5] * x[5])};
+    int expected_p[7] = {0, 1, 2, 3, 4, 5, 6};
+    int expected_i[6] = {0, 1, 2, 3, 4, 5};
+
+    mu_assert("vals incorrect",
+              cmp_double_array(sum_node->wsum_hess->x, expected_x, 6));
+    mu_assert("rows incorrect",
+              cmp_int_array(sum_node->wsum_hess->p, expected_p, 7));
+    mu_assert("cols incorrect",
+              cmp_int_array(sum_node->wsum_hess->i, expected_i, 6));
+
+    free_expr(sum_node);
+    free_expr(log_node);
+    free_expr(x_node);
+
+    return 0;
+}
+
+const char *test_wsum_hess_sum_log_axis1()
+{
+    /* Test: wsum_hess of sum(log(x), axis=1) where x is 3x2
+     * x = [[1, 4],
+     *      [2, 5],
+     *      [3, 6]]
+     */
+
+    double x[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
+    double w[3] = {-1.0, -2.0, -3.0};
+    expr *x_node = new_variable(3, 2, 0, 6);
+    expr *log_node = new_log(x_node);
+    expr *sum_node = new_sum(log_node, 1);
+
+    sum_node->forward(sum_node, x);
+    sum_node->jacobian_init(sum_node);
+    sum_node->wsum_hess_init(sum_node);
+    sum_node->eval_wsum_hess(sum_node, w);
+
+    /* Expected diagonal values */
+    double expected_x[6] = {-w[0] / (x[0] * x[0]), -w[1] / (x[1] * x[1]),
+                            -w[2] / (x[2] * x[2]), -w[0] / (x[3] * x[3]),
+                            -w[1] / (x[4] * x[4]), -w[2] / (x[5] * x[5])};
+    int expected_p[7] = {0, 1, 2, 3, 4, 5, 6};
+    int expected_i[6] = {0, 1, 2, 3, 4, 5};
+
+    mu_assert("vals incorrect",
+              cmp_double_array(sum_node->wsum_hess->x, expected_x, 6));
+    mu_assert("rows incorrect",
+              cmp_int_array(sum_node->wsum_hess->p, expected_p, 7));
+    mu_assert("cols incorrect",
+              cmp_int_array(sum_node->wsum_hess->i, expected_i, 6));
+
+    free_expr(sum_node);
+    free_expr(log_node);
+    free_expr(x_node);
+
+    return 0;
+}