trace

Author: dance858

TODO.md

@@ -2,9 +2,5 @@
 10. For performance reasons, is it useful to have a dense matmul with A and B as dense matrices?
 11. right matmul, add broadcasting logic as in left matmul. Is this necessary?
 
-Going through all atoms to see that sparsity pattern is computed in initialization of jacobian:
-2. trace - not ok
-
 Going through all atoms to see that sparsity pattern is computed in initialization of hessian:
-2. hstack - not ok
 3. trace - not ok

include/subexpr.h

@@ -40,11 +40,11 @@ typedef struct sum_expr
     int *idx_map; /* maps child nnz to summed-row positions */
 } sum_expr;
 
-/* Trace-like reduction: sums entries spaced by child->d1 */
+/* trace */
 typedef struct trace_expr
 {
     expr base;
-    struct int_double_pair *int_double_pairs; /* for sorting jacobian entries */
+    int *idx_map; /* maps child nnz to summed-row positions */
 } trace_expr;
 
 /* Product of all entries */

include/utils/CSR_Matrix.h

@@ -106,6 +106,8 @@ void sum_all_rows_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A, CSR_Matrix
    Must memset accumulator to zero before calling. */
 void idx_map_accumulator(const CSR_Matrix *A, const int *idx_map,
                          double *accumulator);
+void idx_map_accumulator_with_spacing(const CSR_Matrix *A, const int *idx_map,
+                                      double *accumulator, int spacing);
 
 /* Sum blocks of rows of A into a matrix C */
 void sum_block_of_rows_csr(const CSR_Matrix *A, CSR_Matrix *C,
@@ -141,6 +143,11 @@ void sum_evenly_spaced_rows_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A,
 void sum_spaced_rows_into_row_csr(const CSR_Matrix *A, CSR_Matrix *C,
                                   struct int_double_pair *pairs, int offset,
                                   int spacing);
+/* Fills the sparsity and index map for summing spaced rows into a row matrix */
+void sum_spaced_rows_into_row_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A,
+                                                            CSR_Matrix *C,
+                                                            int spacing, int *iwork,
+                                                            int *idx_map);
 
 /* Count number of columns with nonzero entries */
 int count_nonzero_cols(const CSR_Matrix *A, bool *col_nz);

python/atoms/trace.h

@@ -0,0 +1,30 @@
+#ifndef ATOM_TRACE_H
+
+#define ATOM_TRACE_H
+
+#include "common.h"
+
+static PyObject *py_make_trace(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    if (!PyArg_ParseTuple(args, "O", &child_capsule))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+    expr *node = new_trace(child);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create trace node");
+        return NULL;
+    }
+    expr_retain(node); /* Capsule owns a reference */
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif // ATOM_TRACE_H

python/bindings.c

@@ -40,6 +40,7 @@
 #include "atoms/sum.h"
 #include "atoms/tan.h"
 #include "atoms/tanh.h"
+#include "atoms/trace.h"
 #include "atoms/variable.h"
 #include "atoms/xexp.h"
 
@@ -70,6 +71,8 @@ static PyMethodDef DNLPMethods[] = {
     {"make_exp", py_make_exp, METH_VARARGS, "Create exp node"},
     {"make_index", py_make_index, METH_VARARGS, "Create index node"},
     {"make_add", py_make_add, METH_VARARGS, "Create add node"},
+    {"make_trace", py_make_trace, METH_VARARGS,
+     "Create trace node from an expr capsule (make_trace(child))"},
     {"make_hstack", py_make_hstack, METH_VARARGS,
      "Create hstack node from list of expr capsules and n_vars (make_hstack([e1, "
      "e2, ...], n_vars))"},

src/affine/sum.c

@@ -73,6 +73,11 @@ static void jacobian_init(expr *node)
     node->iwork = malloc(MAX(node->jacobian->n, x->jacobian->nnz) * sizeof(int));
     snode->idx_map = malloc(x->jacobian->nnz * sizeof(int));
 
+    /* the idx_map array maps each nonzero entry j in x->jacobian
+       to the corresponding index in the output row matrix C. Specifically, for
+       each nonzero entry j in A, idx_map[j] gives the position in C->x where
+       the value from x->jacobian->x[j] should be accumulated. */
+
     if (axis == -1)
     {
         sum_all_rows_csr_fill_sparsity_and_idx_map(x->jacobian, node->jacobian,

src/affine/trace.c

@@ -1,8 +1,11 @@
 #include "affine.h"
 #include "utils/int_double_pair.h"
+#include "utils/utils.h"
 #include <assert.h>
+#include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <utils/iVec.h>
 
 static void forward(expr *node, const double *u)
 {
@@ -25,21 +28,39 @@ static void forward(expr *node, const double *u)
 static void jacobian_init(expr *node)
 {
     expr *x = node->left;
+    assert(x->d1 == x->d2);
 
     /* initialize child's jacobian */
     x->jacobian_init(x);
 
-    /* count total nnz in the rows of the jacobian that should be summed */
+    // ---------------------------------------------------------------
+    //    count total nnz and allocate matrix with sufficient space
+    // ---------------------------------------------------------------
     const CSR_Matrix *A = x->jacobian;
     int total_nnz = 0;
     int row_spacing = x->d1 + 1;
+
     for (int row = 0; row < A->m; row += row_spacing)
     {
-        total_nnz += (A->p[row + 1] - A->p[row]);
+        total_nnz += A->p[row + 1] - A->p[row];
     }
 
     node->jacobian = new_csr_matrix(1, node->n_vars, total_nnz);
-    ((trace_expr *) node)->int_double_pairs = new_int_double_pair_array(total_nnz);
+
+    // ---------------------------------------------------------------
+    // fill sparsity pattern and idx_map
+    // ---------------------------------------------------------------
+    trace_expr *tnode = (trace_expr *) node;
+    node->iwork = malloc(MAX(node->jacobian->n, total_nnz) * sizeof(int));
+
+    /* the idx_map array maps each nonzero entry j in the original matrix A (from the
+       selected, evenly spaced rows) to the corresponding index in the output row
+       matrix C. Specifically, for each nonzero entry j in A (from the selected
+       rows), idx_map[j] gives the position in C->x where the value from A->x[j]
+       should be accumulated. */
+    tnode->idx_map = malloc(x->jacobian->nnz * sizeof(int));
+    sum_spaced_rows_into_row_csr_fill_sparsity_and_idx_map(
+        A, node->jacobian, row_spacing, node->iwork, tnode->idx_map);
 }
 
 static void eval_jacobian(expr *node)
@@ -51,8 +72,9 @@ static void eval_jacobian(expr *node)
     x->eval_jacobian(x);
 
     /* local jacobian */
-    sum_spaced_rows_into_row_csr(x->jacobian, node->jacobian,
-                                 tnode->int_double_pairs, 0, x->d1 + 1);
+    memset(node->jacobian->x, 0, node->jacobian->nnz * sizeof(double));
+    idx_map_accumulator_with_spacing(x->jacobian, tnode->idx_map, node->jacobian->x,
+                                     x->d1 + 1);
 }
 
 /* Placeholders for Hessian-related functions */
@@ -66,16 +88,12 @@ static void wsum_hess_init(expr *node)
     node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, x->wsum_hess->nnz);
     node->dwork = (double *) calloc(x->size, sizeof(double));
 
-    /* TODO: here we could copy over sparsity pattern once we have checked
-    that all atoms fill their sparsity pattern in the init functions. Perhaps
-    we should only take sparsity pattern of rows that are summed? Not the rows
-    which will get zero weight in the hessian. That would be very cool.
-    But must eval_wsum_hess then also ignore contributions with zero weight? that
-    would be bad. */
-
-    /* lets conclude that the hessian can be made more sophisticated */
-
-    /* but perhaps let's keep it as simple as possible! */
+    /* We copy over the sparsity pattern from the child. This also includes the
+       contribution to wsum_hess of entries of the child that will always have
+       zero weight in eval_wsum_hess. We do this for simplicity. But the Hessian
+       can for sure be made more sophisticated. */
+    memcpy(node->wsum_hess->p, x->wsum_hess->p, (x->n_vars + 1) * sizeof(int));
+    memcpy(node->wsum_hess->i, x->wsum_hess->i, x->wsum_hess->nnz * sizeof(int));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -90,14 +108,7 @@ static void eval_wsum_hess(expr *node, const double *w)
 
     x->eval_wsum_hess(x, node->dwork);
 
-    /* TODO: here we only need to copy over values once we have filled the sparsity
-     * pattern in wsum_hess_init*/
-    node->wsum_hess->nnz = x->wsum_hess->nnz;
-    memcpy(node->wsum_hess->p, x->wsum_hess->p, sizeof(int) * (node->n_vars + 1));
-    memcpy(node->wsum_hess->i, x->wsum_hess->i, sizeof(int) * x->wsum_hess->nnz);
     memcpy(node->wsum_hess->x, x->wsum_hess->x, sizeof(double) * x->wsum_hess->nnz);
-
-    /* This might contain many many zeros actually! Hmm...*/
 }
 
 static bool is_affine(const expr *node)
@@ -107,27 +118,21 @@ static bool is_affine(const expr *node)
 
 static void free_type_data(expr *node)
 {
-    trace_expr *tnode = (trace_expr *) node;
-    free_int_double_pair_array(tnode->int_double_pairs);
+    if (node)
+    {
+        trace_expr *tnode = (trace_expr *) node;
+        free(tnode->idx_map);
+    }
 }
 
 expr *new_trace(expr *child)
 {
-    /* Output is a single scalar */
-    int d1 = 1;
-
     trace_expr *tnode = (trace_expr *) calloc(1, sizeof(trace_expr));
     expr *node = &tnode->base;
-    init_expr(node, d1, 1, child->n_vars, forward, jacobian_init, eval_jacobian,
+    init_expr(node, 1, 1, child->n_vars, forward, jacobian_init, eval_jacobian,
               is_affine, wsum_hess_init, eval_wsum_hess, free_type_data);
     node->left = child;
     expr_retain(child);
 
-    /* Initialize type-specific fields */
-    tnode->int_double_pairs = NULL;
-
-    // just for debugging, should be removed
-    strcpy(node->name, "trace");
-
     return node;
 }

src/utils/CSR_Matrix.c

@@ -754,6 +754,19 @@ void idx_map_accumulator(const CSR_Matrix *A, const int *idx_map,
     }
 }
 
+void idx_map_accumulator_with_spacing(const CSR_Matrix *A, const int *idx_map,
+                                      double *accumulator, int spacing)
+{
+    /* don't forget to initialze accumulator to 0 before calling this */
+    for (int row = 0; row < A->m; row += spacing)
+    {
+        for (int j = A->p[row]; j < A->p[row + 1]; j++)
+        {
+            accumulator[idx_map[j]] += A->x[j];
+        }
+    }
+}
+
 void sum_spaced_rows_into_row_csr(const CSR_Matrix *A, CSR_Matrix *C,
                                   struct int_double_pair *pairs, int offset,
                                   int spacing)
@@ -1058,7 +1071,8 @@ void sum_all_rows_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A, CSR_Matrix
     C->nnz = unique_nnz;
 
     // -------------------------------------------------------------------
-    //         Map child values to summed-row positions
+    //  Map child values to summed-row positions. col_to_pos maps
+    //  column indices to positions in C's row.
     // -------------------------------------------------------------------
     int *col_to_pos = iwork;
     for (int idx = 0; idx < unique_nnz; idx++)
@@ -1090,3 +1104,61 @@ void sum_all_rows_csr_fill_values(const CSR_Matrix *A, CSR_Matrix *C,
     }
 }
 */
+
+/*
+ * Sums evenly spaced rows from A into a single row in C and fills an index map.
+ * A: input CSR matrix
+ * C: output CSR matrix (must have m=1)
+ * spacing: row spacing
+ * iwork: workspace of size at least max(A->n, A->nnz)
+ * idx_map: output index map, size at least A->nnz
+ */
+void sum_spaced_rows_into_row_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A,
+                                                            CSR_Matrix *C,
+                                                            int spacing, int *iwork,
+                                                            int *idx_map)
+{
+    assert(C->m == 1);
+    C->n = A->n;
+
+    /* gather all column indices from the spaced rows */
+    int count = 0;
+    for (int row = 0; row < A->m; row += spacing)
+    {
+        int len = A->p[row + 1] - A->p[row];
+        memcpy(iwork + count, A->i + A->p[row], len * sizeof(int));
+        count += len;
+    }
+
+    /* fill sparsity pattern */
+    sort_int_array(iwork, count);
+    int unique_nnz = 0;
+    int prev_col = -1;
+    for (int i = 0; i < count; i++)
+    {
+        int col = iwork[i];
+        if (col != prev_col)
+        {
+            C->i[unique_nnz++] = col;
+            prev_col = col;
+        }
+    }
+    C->nnz = unique_nnz;
+    C->p[0] = 0;
+    C->p[1] = C->nnz;
+
+    /* fill idx_map */
+    int *col_to_pos = iwork;
+    for (int idx = 0; idx < unique_nnz; idx++)
+    {
+        col_to_pos[C->i[idx]] = idx;
+    }
+
+    for (int row = 0; row < A->m; row += spacing)
+    {
+        for (int j = A->p[row]; j < A->p[row + 1]; j++)
+        {
+            idx_map[j] = col_to_pos[A->i[j]];
+        }
+    }
+}

tests/all_tests.c

@@ -222,12 +222,9 @@ int main(void)
     mu_run_test(test_wsum_hess_broadcast_row, tests_run);
     mu_run_test(test_wsum_hess_broadcast_col, tests_run);
     mu_run_test(test_wsum_hess_broadcast_scalar_to_matrix, tests_run);
-    // This test leads to seg fault
-    // mu_run_test(test_wsum_hess_trace_variable, tests_run);
-
-    // This test fails - not sure how sophisticated we should make
-    // wsum_hess for trace
-    // mu_run_test(test_wsum_hess_trace_composite, tests_run);
+    mu_run_test(test_wsum_hess_trace_variable, tests_run);
+    mu_run_test(test_wsum_hess_trace_log_variable, tests_run);
+    mu_run_test(test_wsum_hess_trace_composite, tests_run);
 
     printf("\n--- Utility Tests ---\n");
     mu_run_test(test_diag_csr_mult, tests_run);
@@ -257,7 +254,7 @@ int main(void)
     mu_run_test(test_problem_constraint_forward, tests_run);
     mu_run_test(test_problem_hessian, tests_run);
 
-      printf("\n=== All %d tests passed ===\n", tests_run);
+    printf("\n=== All %d tests passed ===\n", tests_run);
 
     return 0;
 }