elementwise jacobian

Author: dance858

CMakeLists.txt

@@ -2,19 +2,17 @@ cmake_minimum_required(VERSION 3.10)
 project(DNLP_Diff_Engine C)
 
 set(CMAKE_C_STANDARD 99)
+set(CMAKE_BUILD_TYPE Debug)
+
+# Debug-friendly flags
+add_compile_options(-g -O0)
 
 # Include directories
 include_directories(${PROJECT_SOURCE_DIR}/include)
+include_directories(${PROJECT_SOURCE_DIR}/tests)
 
-# Source files
-set(SOURCES
-    src/expr.c
-    src/elementwise/exp.c
-    src/elementwise/log.c
-    src/affine/variable.c
-    src/affine/constant.c
-    src/affine/add.c
-)
+# Source files - automatically gather all .c files from src/
+file(GLOB_RECURSE SOURCES "src/*.c")
 
 # Create library
 add_library(dnlp_diff ${SOURCES})
@@ -23,15 +21,10 @@ target_link_libraries(dnlp_diff m)
 # Enable testing
 enable_testing()
 
-# Single test executable that links all test source files
-add_executable(test_all
-    tests/forward_pass/all_tests.c
-    tests/forward_pass/test_helpers.c
-    tests/forward_pass/affine/test_variable_constant.c
-    tests/forward_pass/affine/test_add.c
-    tests/forward_pass/elementwise/test_exp.c
-    tests/forward_pass/elementwise/test_log.c
-    tests/forward_pass/composite/test_composite.c
+# Single test executable combining all tests
+add_executable(all_tests
+    tests/all_tests.c
+    tests/test_helpers.c
 )
-target_link_libraries(test_all dnlp_diff)
-add_test(NAME AllTests COMMAND test_all)
+target_link_libraries(all_tests dnlp_diff)
+add_test(NAME AllTests COMMAND all_tests)

README.md

@@ -1,91 +0,0 @@
-# DNLP Diff Engine
-
-A minimalistic C-based expression tree system for nonlinear programming with automatic differentiation capabilities.
-
-## Structure
-
-```
-.
-├── include/              # Header files
-│   ├── expr.h           # Base expression node definition
-│   ├── affine/          # Affine operations
-│   │   ├── variable.h
-│   │   ├── constant.h
-│   │   └── add.h
-│   └── elementwise/     # Elementwise operations
-│       ├── exp.h
-│       └── log.h
-├── src/
-│   ├── expr.c           # Base node implementation
-│   ├── affine/          # Affine operations
-│   │   ├── variable.c
-│   │   ├── constant.c
-│   │   └── add.c
-│   └── elementwise/     # Elementwise operations
-│       ├── exp.c
-│       └── log.c
-├── tests/
-│   └── forward_pass/    # Forward pass tests
-│       └── test_forward_pass.c
-└── CMakeLists.txt       # Build configuration
-```
-
-## Features
-
-- **Expression nodes**: Each node has dimension `m`, preallocated value memory, and up to two children
-- **Leaf nodes**: Variable (reads from input `u`) and Constant (fixed values)
-- **Affine operations**: Addition
-- **Elementwise operations**: Exponential (`exp`) and Logarithm (`log`)
-- **Forward pass**: Compute output values through the expression tree
-
-## Building
-
-```bash
-mkdir build && cd build
-cmake ..
-make
-```
-
-## Running Tests
-
-```bash
-# Run the test executable
-./test_forward_pass
-
-# Or use CTest
-ctest --verbose
-```
-
-## Example Usage
-
-```c
-#include "expr.h"
-#include "affine/variable.h"
-#include "affine/constant.h"
-#include "affine/add.h"
-#include "elementwise/exp.h"
-#include "elementwise/log.h"
-
-// Build expression tree: log(exp(x) + c)
-double u[2] = {1.0, 2.0};
-double c[2] = {1.0, 1.0};
-
-expr* var = new_variable(2);
-expr* exp_node = new_exp(var);
-expr* const_node = new_constant(2, c);
-expr* sum = new_add(exp_node, const_node);
-expr* log_node = new_log(sum);
-
-// Compute forward pass
-log_node->forward(log_node, u);
-
-// Result is in log_node->value
-```
-
-## Design
-
-- Each node allocates its own output memory
-- Forward pass recursively evaluates children before computing node output
-- Clean separation: affine operations in `src/affine/`, elementwise in `src/elementwise/`
-- Memory management: nodes must be freed manually (children are not auto-freed)
-# DNLP-diff-engine

include/affine/linear_op.h

@@ -0,0 +1,10 @@
+#ifndef LINEAR_OP_H
+#define LINEAR_OP_H
+
+#include "expr.h"
+#include "utils/CSR_Matrix.h"
+
+/* linear operator f(u) = Au */
+expr *new_linear(expr *u, const CSR_Matrix *A);
+
+#endif /* LINEAR_OP_H */

include/affine/variable.h

@@ -3,6 +3,6 @@
 
 #include "expr.h"
 
-expr *new_variable(int m);
+expr *new_variable(int m, int var_id, int n_vars);
 
 #endif /* VARIABLE_H */

include/elementwise.h

@@ -0,0 +1,22 @@
+#ifndef ELEMENTWISE_H
+#define ELEMENTWISE_H
+
+#include "expr.h"
+
+expr *new_exp(expr *child);
+expr *new_log(expr *child);
+expr *new_entr(expr *child);
+expr *new_sin(expr *child);
+expr *new_cos(expr *child);
+expr *new_tan(expr *child);
+
+/* the jacobian for elementwise 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);
+
+/* no elementwise atoms are affine according to our convention,
+   so we can have a common implementation */
+bool is_affine_elementwise(expr *node);
+
+#endif /* ELEMENTWISE_H */

include/elementwise_univariate.h

@@ -1,26 +1,53 @@
 #ifndef EXPR_H
 #define EXPR_H
 
+#include "utils/CSR_Matrix.h"
+#include <stdbool.h>
 #include <stddef.h>
 
-/* Forward declaration */
+#define JAC_IDXS_NOT_SET -1
+
+/* Forward declarations */
 struct expr;
 
-/* Function pointer type for forward pass computation */
+/* Function pointer types */
 typedef void (*forward_fn)(struct expr *node, const double *u);
+typedef void (*jacobian_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 bool (*is_affine_fn)(struct expr *node);
 
 /* Expression node structure */
 typedef struct expr
 {
-    int m;              /* Output dimension */
-    double *value;      /* Preallocated output value array */
-    struct expr *left;  /* Left child (can be NULL) */
-    struct expr *right; /* Right child (can be NULL) */
-    forward_fn forward; /* Forward pass function */
+    // ------------------------------------------------------------------------
+    //                         general quantities
+    // ------------------------------------------------------------------------
+    int m;
+    int n_vars;
+    int var_id;
+    struct expr *left;
+    struct expr *right;
+    double *dwork;
+
+    // ------------------------------------------------------------------------
+    //                     forward pass related quantities
+    // ------------------------------------------------------------------------
+    double *value;
+    forward_fn forward;
+
+    // ------------------------------------------------------------------------
+    //                      jacobian related quantities
+    // ------------------------------------------------------------------------
+    CSR_Matrix *jacobian;
+    jacobian_init_fn jacobian_init;
+    eval_jacobian_fn eval_jacobian;
+    eval_local_jacobian_fn eval_local_jacobian;
+    is_affine_fn is_affine;
+
 } expr;
 
-/* Memory management functions */
-expr *new_expr(int m);
+expr *new_expr(int m, int n_vars);
 void free_expr(expr *node);
 
 #endif /* EXPR_H */

include/expr.h

@@ -0,0 +1,16 @@
+#ifndef MATRIX_H
+#define MATRIX_H
+
+/* Coordinate (COO) format sparse matrix */
+typedef struct COO_Matrix
+{
+    int nnz;
+    int *rows;
+    int *cols;
+    double *vals;
+} COO_Matrix;
+
+COO_Matrix *new_coo_matrix(int nnz);
+void free_coo_matrix(COO_Matrix *matrix);
+
+#endif /* MATRIX_H */

include/utils/COO_Matrix.h

@@ -0,0 +1,46 @@
+#ifndef CSR_MATRIX_H
+#define CSR_MATRIX_H
+
+/* CSR (Compressed Sparse Row) Matrix Format
+ *
+ * For an m x n matrix with nnz nonzeros:
+ * - p: array of size (m + 1) indicating start of each row
+ * - i: array of size nnz containing column indices
+ * - x: array of size nnz containing values
+ * - m: number of rows
+ * - n: number of columns
+ * - nnz: number of nonzero entries
+ */
+typedef struct CSR_Matrix
+{
+    int *p;
+    int *i;
+    double *x;
+    int m;
+    int n;
+    int nnz;
+} CSR_Matrix;
+
+/* Allocate a new CSR matrix with given dimensions and nnz */
+CSR_Matrix *new_csr_matrix(int m, int n, int nnz);
+
+/* Free a CSR matrix */
+void free_csr_matrix(CSR_Matrix *matrix);
+
+/* Copy CSR matrix A to C */
+void copy_csr_matrix(const CSR_Matrix *A, CSR_Matrix *C);
+
+/* matvec y = Ax, where A indices minus col_offset gives x indices */
+void csr_matvec(const CSR_Matrix *A, const double *x, double *y, int col_offset);
+
+/* Compute C = diag(d) * A where d is an array and A, C are CSR matrices
+ * d must have length m
+ * C must be pre-allocated with same dimensions as A */
+void diag_csr_mult(const double *d, const CSR_Matrix *A, CSR_Matrix *C);
+
+/* Compute C = A + B where A, B, C are CSR matrices
+ * A and B must have same dimensions
+ * C must be pre-allocated with sufficient nnz capacity */
+void sum_csr_matrices(const CSR_Matrix *A, const CSR_Matrix *B, CSR_Matrix *C);
+
+#endif /* CSR_MATRIX_H */

include/utils/CSR_Matrix.h

@@ -13,17 +13,46 @@ static void add_forward(expr *node, const double *u)
     }
 }
 
+static void jacobian_init(expr *node)
+{
+    /* initialize children's jacobians */
+    node->left->jacobian_init(node->left);
+    node->right->jacobian_init(node->right);
+
+    /* we never have to store more than the sum of children's nnz */
+    int nnz_max = node->left->jacobian->nnz + node->right->jacobian->nnz;
+    node->jacobian = new_csr_matrix(node->m, node->n_vars, nnz_max);
+}
+
+static void eval_jacobian(expr *node)
+{
+    /* evaluate children's jacobians */
+    node->left->eval_jacobian(node->left);
+    node->right->eval_jacobian(node->right);
+
+    /* sum children's jacobians */
+    sum_csr_matrices(node->left->jacobian, node->right->jacobian, node->jacobian);
+}
+
+static bool is_affine(expr *node)
+{
+    return node->left->is_affine(node->left) && node->right->is_affine(node->right);
+}
+
 expr *new_add(expr *left, expr *right)
 {
     if (!left || !right) return NULL;
-    if (left->m != right->m) return NULL; /* Dimension mismatch */
+    if (left->m != right->m) return NULL;
 
-    expr *node = new_expr(left->m);
+    expr *node = new_expr(left->m, left->n_vars);
     if (!node) return NULL;
 
     node->left = left;
     node->right = right;
     node->forward = add_forward;
+    node->is_affine = is_affine;
+    node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
 
     return node;
 }