small edits, most substantial thing removed unnecessary loop

Author: dance858

include/affine.h

@@ -18,6 +18,6 @@ expr *new_trace(expr *child);
 expr *new_constant(int d1, int d2, int n_vars, const double *values);
 expr *new_variable(int d1, int d2, int var_id, int n_vars);
 
-expr *new_index(expr *child, const int *indices, int n_selected);
+expr *new_index(expr *child, const int *indices, int n_idxs);
 
 #endif /* AFFINE_H */

include/subexpr.h

@@ -109,14 +109,13 @@ typedef struct const_vector_mult_expr
     double *a; /* length equals node->size */
 } const_vector_mult_expr;
 
-/* Index/slicing: y = child[indices] where indices is a list of flattened positions */
+/* Index/slicing: y = child[indices] where indices is a list of flat positions */
 typedef struct index_expr
 {
     expr base;
     int *indices;        /* Flattened indices to select (owned, copied) */
-    int n_selected;      /* Number of selected elements */
-    bool has_duplicates; /* True if indices contain duplicates (affects Hessian path) */
-    double *parent_w;    /* Scatter buffer for wsum_hess (size = child->size) */
+    int n_idxs;          /* Number of selected elements */
+    bool has_duplicates; /* True if indices have duplicates (affects Hessian path) */
 } index_expr;
 
 #endif /* SUBEXPR_H */

python/atoms/index.h

@@ -6,7 +6,8 @@
 #include "affine.h"
 #include "common.h"
 
-/* Index/slicing: y = child[indices] where indices is a list of flattened positions */
+/* Index/slicing: y = child[indices] where indices is a list of flattened positions
+ */
 static PyObject *py_make_index(PyObject *self, PyObject *args)
 {
     PyObject *child_capsule;
@@ -17,26 +18,26 @@ static PyObject *py_make_index(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    expr *child = (expr *)PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
     if (!child)
     {
         PyErr_SetString(PyExc_ValueError, "invalid child capsule");
         return NULL;
     }
 
     /* Convert indices array to int32 */
-    PyArrayObject *indices_array =
-        (PyArrayObject *)PyArray_FROM_OTF(indices_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
+    PyArrayObject *indices_array = (PyArrayObject *) PyArray_FROM_OTF(
+        indices_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY);
 
     if (!indices_array)
     {
         return NULL;
     }
 
-    int n_selected = (int)PyArray_SIZE(indices_array);
-    int *indices_data = (int *)PyArray_DATA(indices_array);
+    int n_idxs = (int) PyArray_SIZE(indices_array);
+    int *indices_data = (int *) PyArray_DATA(indices_array);
 
-    expr *node = new_index(child, indices_data, n_selected);
+    expr *node = new_index(child, indices_data, n_idxs);
 
     Py_DECREF(indices_array);
 

src/affine/index.c

@@ -1,19 +1,17 @@
-// SPDX-License-Identifier: Apache-2.0
-
 #include "affine.h"
 #include "subexpr.h"
 #include <stdlib.h>
 #include <string.h>
 
-/* Index/slicing: y = child[indices] where indices is a list of flattened positions */
+/* Index/slicing: y = child[indices] where indices is a list of flat positions */
 
 /* Check if indices array contains duplicates using a bitmap.
  * Returns true if duplicates exist, false otherwise. */
-static bool check_for_duplicates(const int *indices, int n_selected, int max_idx)
+static bool check_for_duplicates(const int *indices, int n_idxs, int max_idx)
 {
-    bool *seen = (bool *)calloc(max_idx, sizeof(bool));
+    bool *seen = (bool *) calloc(max_idx, sizeof(bool));
     bool has_dup = false;
-    for (int i = 0; i < n_selected && !has_dup; i++)
+    for (int i = 0; i < n_idxs && !has_dup; i++)
     {
         if (seen[indices[i]])
         {
@@ -27,115 +25,109 @@ static bool check_for_duplicates(const int *indices, int n_selected, int max_idx
 
 static void forward(expr *node, const double *u)
 {
-    expr *child = node->left;
-    index_expr *idx = (index_expr *)node;
+    expr *x = node->left;
+    index_expr *idx = (index_expr *) node;
 
     /* child's forward pass */
-    child->forward(child, u);
+    x->forward(x, u);
 
     /* gather selected elements */
-    for (int i = 0; i < idx->n_selected; i++)
+    for (int i = 0; i < idx->n_idxs; i++)
     {
-        node->value[i] = child->value[idx->indices[i]];
+        node->value[i] = x->value[idx->indices[i]];
     }
 }
 
 static void jacobian_init(expr *node)
 {
-    expr *child = node->left;
-    index_expr *idx = (index_expr *)node;
-
-    child->jacobian_init(child);
-    CSR_Matrix *J_child = child->jacobian;
-
-    /* count nnz */
-    int nnz = 0;
-    for (int i = 0; i < idx->n_selected; i++)
-    {
-        int row = idx->indices[i];
-        nnz += J_child->p[row + 1] - J_child->p[row];
-    }
+    expr *x = node->left;
+    index_expr *idx = (index_expr *) node;
+    x->jacobian_init(x);
 
-    node->jacobian = new_csr_matrix(idx->n_selected, node->n_vars, nnz);
-    CSR_Matrix *J = node->jacobian;
+    CSR_Matrix *Jx = x->jacobian;
+    CSR_Matrix *J = new_csr_matrix(node->size, node->n_vars, Jx->nnz);
 
-    /* fill p and i arrays in one pass */
+    /* set sparsity pattern */
     J->p[0] = 0;
-    for (int i = 0; i < idx->n_selected; i++)
+    for (int i = 0; i < idx->n_idxs; i++)
     {
         int row = idx->indices[i];
-        int src = J_child->p[row];
-        int len = J_child->p[row + 1] - src;
-        memcpy(J->i + J->p[i], J_child->i + src, len * sizeof(int));
+        int len = Jx->p[row + 1] - Jx->p[row];
+        memcpy(J->i + J->p[i], Jx->i + Jx->p[row], len * sizeof(int));
         J->p[i + 1] = J->p[i] + len;
     }
+
+    node->jacobian = J;
 }
 
 static void eval_jacobian(expr *node)
 {
-    expr *child = node->left;
-    index_expr *idx = (index_expr *)node;
-
-    child->eval_jacobian(child);
+    expr *x = node->left;
+    index_expr *idx = (index_expr *) node;
+    x->eval_jacobian(x);
 
     CSR_Matrix *J = node->jacobian;
-    CSR_Matrix *J_child = child->jacobian;
+    CSR_Matrix *Jx = x->jacobian;
 
-    for (int i = 0; i < idx->n_selected; i++)
+    for (int i = 0; i < idx->n_idxs; i++)
     {
         int row = idx->indices[i];
         int len = J->p[i + 1] - J->p[i];
-        memcpy(J->x + J->p[i], J_child->x + J_child->p[row], len * sizeof(double));
+        memcpy(J->x + J->p[i], Jx->x + Jx->p[row], len * sizeof(double));
     }
 }
 
 static void wsum_hess_init(expr *node)
 {
-    expr *child = node->left;
-    index_expr *idx = (index_expr *)node;
+    expr *x = node->left;
 
     /* initialize child's wsum_hess */
-    child->wsum_hess_init(child);
-
-    /* allocate scatter buffer (zeroed) */
-    idx->parent_w = (double *)calloc(child->size, sizeof(double));
-
-    /* wsum_hess inherits from child (affine has no local Hessian) */
-    /* We need to allocate our own to avoid aliasing issues */
-    CSR_Matrix *child_hess = child->wsum_hess;
-    node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz);
-    memcpy(node->wsum_hess->p, child_hess->p, (child_hess->m + 1) * sizeof(int));
-    memcpy(node->wsum_hess->i, child_hess->i, child_hess->nnz * sizeof(int));
+    x->wsum_hess_init(x);
+
+    /* for setting weight vector to evaluate hessian of child */
+    node->dwork = (double *) calloc(x->size, sizeof(double));
+
+    /* in the implementation of eval_wsum_hess we evaluate the
+       child's hessian with a weight vector that has w[i] = 0
+       if i is not included in idx->indices. This can lead to
+       many numerical zeros in child->wsum_hess that are actually
+       structural zeros, but we do not try to exploit that sparsity
+       right now. */
+    CSR_Matrix *H_child = x->wsum_hess;
+    node->wsum_hess = new_csr_matrix(H_child->m, H_child->n, H_child->nnz);
+    memcpy(node->wsum_hess->p, H_child->p, (H_child->m + 1) * sizeof(int));
+    memcpy(node->wsum_hess->i, H_child->i, H_child->nnz * sizeof(int));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
 {
     expr *child = node->left;
-    index_expr *idx = (index_expr *)node;
+    index_expr *idx = (index_expr *) node;
 
     if (idx->has_duplicates)
     {
-        /* slow path: must zero and accumulate for repeated indices */
-        memset(idx->parent_w, 0, child->size * sizeof(double));
-        for (int i = 0; i < idx->n_selected; i++)
+        /* zero and accumulate for repeated indices */
+        memset(node->dwork, 0, child->size * sizeof(double));
+        for (int i = 0; i < idx->n_idxs; i++)
         {
-            idx->parent_w[idx->indices[i]] += w[i];
+            node->dwork[idx->indices[i]] += w[i];
         }
     }
     else
     {
-        /* fast path: direct write (no memset needed, no accumulation) */
-        for (int i = 0; i < idx->n_selected; i++)
+        /* direct write (no memset needed, no accumulation) */
+        for (int i = 0; i < idx->n_idxs; i++)
         {
-            idx->parent_w[idx->indices[i]] = w[i];
+            node->dwork[idx->indices[i]] = w[i];
         }
     }
 
     /* delegate to child */
-    child->eval_wsum_hess(child, idx->parent_w);
+    child->eval_wsum_hess(child, node->dwork);
 
     /* copy values from child */
-    memcpy(node->wsum_hess->x, child->wsum_hess->x, child->wsum_hess->nnz * sizeof(double));
+    memcpy(node->wsum_hess->x, child->wsum_hess->x,
+           child->wsum_hess->nnz * sizeof(double));
 }
 
 static bool is_affine(const expr *node)
@@ -145,44 +137,37 @@ static bool is_affine(const expr *node)
 
 static void free_type_data(expr *node)
 {
-    index_expr *idx = (index_expr *)node;
+    index_expr *idx = (index_expr *) node;
     if (idx->indices)
     {
         free(idx->indices);
         idx->indices = NULL;
     }
-    if (idx->parent_w)
-    {
-        free(idx->parent_w);
-        idx->parent_w = NULL;
-    }
 }
 
-expr *new_index(expr *child, const int *indices, int n_selected)
+expr *new_index(expr *child, const int *indices, int n_idxs)
 {
     /* allocate type-specific struct */
-    index_expr *idx = (index_expr *)calloc(1, sizeof(index_expr));
+    index_expr *idx = (index_expr *) calloc(1, sizeof(index_expr));
     expr *node = &idx->base;
 
-    /* output shape is (n_selected, 1) - flattened */
-    init_expr(node, n_selected, 1, child->n_vars, forward, jacobian_init,
-              eval_jacobian, is_affine, free_type_data);
+    /* output shape is (n_idxs, 1) - flattened */
+    init_expr(node, n_idxs, 1, child->n_vars, forward, jacobian_init, eval_jacobian,
+              is_affine, free_type_data);
 
-    node->wsum_hess_init = wsum_hess_init;
-    node->eval_wsum_hess = eval_wsum_hess;
     node->left = child;
     expr_retain(child);
 
+    node->wsum_hess_init = wsum_hess_init;
+    node->eval_wsum_hess = eval_wsum_hess;
+
     /* copy indices */
-    idx->indices = (int *)malloc(n_selected * sizeof(int));
-    memcpy(idx->indices, indices, n_selected * sizeof(int));
-    idx->n_selected = n_selected;
+    idx->indices = (int *) malloc(n_idxs * sizeof(int));
+    memcpy(idx->indices, indices, n_idxs * sizeof(int));
+    idx->n_idxs = n_idxs;
 
     /* detect duplicates for Hessian optimization */
-    idx->has_duplicates = check_for_duplicates(indices, n_selected, child->size);
-
-    /* parent_w allocated lazily in wsum_hess_init */
-    idx->parent_w = NULL;
+    idx->has_duplicates = check_for_duplicates(indices, n_idxs, child->size);
 
     return node;
 }

src/other/quad_form.c

@@ -27,7 +27,6 @@ static void forward(expr *node, const double *u)
 
 static void jacobian_init(expr *node)
 {
-    CSR_Matrix *Q = ((quad_form_expr *) node)->Q;
     assert(node->left->var_id != NOT_A_VARIABLE);
     assert(node->left->d2 == 1);
     expr *x = node->left;

tests/jacobian_tests/test_index.h

@@ -83,8 +83,10 @@ const char *test_index_jacobian_of_log(void)
     double expected_x[2] = {1.0, 0.25};
     int expected_i[2] = {0, 2};
 
-    mu_assert("index of log jac vals", cmp_double_array(idx->jacobian->x, expected_x, 2));
-    mu_assert("index of log jac cols", cmp_int_array(idx->jacobian->i, expected_i, 2));
+    mu_assert("index of log jac vals",
+              cmp_double_array(idx->jacobian->x, expected_x, 2));
+    mu_assert("index of log jac cols",
+              cmp_int_array(idx->jacobian->i, expected_i, 2));
 
     free_expr(idx);
     return 0;
@@ -106,8 +108,12 @@ const char *test_index_jacobian_repeated(void)
     int expected_p[3] = {0, 1, 2};
     int expected_i[2] = {0, 0}; /* Both reference col 0 */
 
-    mu_assert("index repeated jac vals", cmp_double_array(idx->jacobian->x, expected_x, 2));
-    mu_assert("index repeated jac i", cmp_int_array(idx->jacobian->i, expected_i, 2));
+    mu_assert("index repeated jac vals",
+              cmp_double_array(idx->jacobian->x, expected_x, 2));
+    mu_assert("index repeated row ptr",
+              cmp_int_array(idx->jacobian->p, expected_p, 4));
+    mu_assert("index repeated jac i",
+              cmp_int_array(idx->jacobian->i, expected_i, 2));
 
     free_expr(idx);
     return 0;