jacobian checker of expression

Author: dance858

CMakeLists.txt

@@ -103,6 +103,7 @@ if(NOT SKBUILD)
     add_executable(all_tests
         tests/all_tests.c
         tests/test_helpers.c
+        tests/numerical_diff.c
     )
     target_link_libraries(all_tests dnlp_diff)
 

src/utils/dense_matrix.c

@@ -156,15 +156,16 @@ static void dense_block_left_mult_values(const Matrix *A, const CSC_Matrix *J,
             }
             else
             {
-                /* scatter sparse J col into dense vector and then compute A @ j_dense */
+                /* scatter sparse J col into dense vector and then compute A @
+                 * j_dense */
                 memset(j_dense, 0, n * sizeof(double));
                 for (s = start; s < end; s++)
                 {
                     j_dense[J->i[s] - block_start] = J->x[s];
                 }
 
-                cblas_dgemv(CblasRowMajor, CblasNoTrans, m, n, 1.0, dm->x, n, j_dense, 1,
-                            0.0, C->x + i, 1);
+                cblas_dgemv(CblasRowMajor, CblasNoTrans, m, n, 1.0, dm->x, n,
+                            j_dense, 1, 0.0, C->x + i, 1);
             }
         }
     }

tests/all_tests.c

@@ -7,12 +7,12 @@
 #include "forward_pass/affine/test_add.h"
 #include "forward_pass/affine/test_broadcast.h"
 #include "forward_pass/affine/test_hstack.h"
-#include "forward_pass/affine/test_vstack.h"
 #include "forward_pass/affine/test_linear_op.h"
 #include "forward_pass/affine/test_neg.h"
 #include "forward_pass/affine/test_promote.h"
 #include "forward_pass/affine/test_sum.h"
 #include "forward_pass/affine/test_variable_constant.h"
+#include "forward_pass/affine/test_vstack.h"
 #include "forward_pass/composite/test_composite.h"
 #include "forward_pass/elementwise/test_exp.h"
 #include "forward_pass/elementwise/test_log.h"
@@ -27,7 +27,6 @@
 #include "jacobian_tests/test_const_vector_mult.h"
 #include "jacobian_tests/test_elementwise_mult.h"
 #include "jacobian_tests/test_hstack.h"
-#include "jacobian_tests/test_vstack.h"
 #include "jacobian_tests/test_index.h"
 #include "jacobian_tests/test_left_matmul.h"
 #include "jacobian_tests/test_log.h"
@@ -46,6 +45,8 @@
 #include "jacobian_tests/test_sum.h"
 #include "jacobian_tests/test_trace.h"
 #include "jacobian_tests/test_transpose.h"
+#include "jacobian_tests/test_vstack.h"
+#include "numerical_diff/test_numerical_diff.h"
 #include "problem/test_problem.h"
 #include "utils/test_cblas.h"
 #include "utils/test_coo_matrix.h"
@@ -66,7 +67,6 @@
 #include "wsum_hess/test_const_scalar_mult.h"
 #include "wsum_hess/test_const_vector_mult.h"
 #include "wsum_hess/test_hstack.h"
-#include "wsum_hess/test_vstack.h"
 #include "wsum_hess/test_index.h"
 #include "wsum_hess/test_left_matmul.h"
 #include "wsum_hess/test_matmul.h"
@@ -83,6 +83,7 @@
 #include "wsum_hess/test_sum.h"
 #include "wsum_hess/test_trace.h"
 #include "wsum_hess/test_transpose.h"
+#include "wsum_hess/test_vstack.h"
 #endif /* PROFILE_ONLY */
 
 #ifdef PROFILE_ONLY
@@ -297,6 +298,9 @@ int main(void)
     mu_run_test(test_dense_matrix_trans, tests_run);
     mu_run_test(test_sparse_vs_dense_mult_vec_blocks, tests_run);
 
+    printf("\n--- Numerical Diff Tests ---\n");
+    mu_run_test(test_check_jacobian_composite_log, tests_run);
+
     printf("\n--- Problem Struct Tests ---\n");
     mu_run_test(test_problem_new_free, tests_run);
     mu_run_test(test_problem_objective_forward, tests_run);

tests/numerical_diff.c

@@ -0,0 +1,99 @@
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "numerical_diff.h"
+
+#define ABS_TOL 1e-6
+#define REL_TOL 1e-6
+
+static int is_close(double a, double b)
+{
+    return fabs(a - b) <= fmax(ABS_TOL, REL_TOL * fmax(fabs(a), fabs(b)));
+}
+
+static void csr_to_dense(const CSR_Matrix *A, double *dense)
+{
+    for (int row = 0; row < A->m; row++)
+    {
+        for (int idx = A->p[row]; idx < A->p[row + 1]; idx++)
+        {
+            dense[row * A->n + A->i[idx]] = A->x[idx];
+        }
+    }
+}
+
+double *numerical_jacobian(expr *node, const double *u, double h)
+{
+    int m = node->size;
+    int n = node->n_vars;
+    double inv_2h = 1.0 / (2.0 * h);
+
+    double *J = calloc((size_t) m * n, sizeof(double));
+    double *u_work = malloc(n * sizeof(double));
+    double *f_plus = malloc(m * sizeof(double));
+    double *f_minus = malloc(m * sizeof(double));
+
+    memcpy(u_work, u, n * sizeof(double));
+
+    for (int j = 0; j < n; j++)
+    {
+        u_work[j] = u[j] + h;
+        node->forward(node, u_work);
+        memcpy(f_plus, node->value, m * sizeof(double));
+
+        u_work[j] = u[j] - h;
+        node->forward(node, u_work);
+        memcpy(f_minus, node->value, m * sizeof(double));
+
+        u_work[j] = u[j];
+
+        for (int k = 0; k < m; k++)
+        {
+            J[k * n + j] = (f_plus[k] - f_minus[k]) * inv_2h;
+        }
+    }
+
+    free(f_minus);
+    free(f_plus);
+    free(u_work);
+    return J;
+}
+
+int check_jacobian(expr *node, const double *u, double h)
+{
+    int m = node->size;
+    int n = node->n_vars;
+
+    node->jacobian_init(node);
+    node->forward(node, u);
+    node->eval_jacobian(node);
+
+    double *J_num = numerical_jacobian(node, u, h);
+
+    /* restore expression state after perturbations */
+    node->forward(node, u);
+
+    double *J_analytical = calloc((size_t) m * n, sizeof(double));
+    csr_to_dense(node->jacobian, J_analytical);
+
+    int result = 1;
+    for (int i = 0; i < m * n; i++)
+    {
+        if (!is_close(J_analytical[i], J_num[i]))
+        {
+            int row = i / n;
+            int col = i % n;
+            printf("  check_jacobian FAILED at (%d, %d):"
+                   " analytical=%e, numerical=%e\n",
+                   row, col, J_analytical[i], J_num[i]);
+            result = 0;
+            break;
+        }
+    }
+
+    free(J_analytical);
+    free(J_num);
+    return result;
+}

tests/numerical_diff.h

@@ -0,0 +1,18 @@
+#ifndef NUMERICAL_DIFF_H
+#define NUMERICAL_DIFF_H
+
+#include "expr.h"
+
+#define NUMERICAL_DIFF_DEFAULT_H 1e-7
+
+/* Compute dense numerical Jacobian via central differences.
+ * Returns malloc'd row-major array (node->size x node->n_vars).
+ * Caller must free(). */
+double *numerical_jacobian(expr *node, const double *u, double h);
+
+/* Evaluate analytical Jacobian, compute numerical Jacobian,
+ * and compare. Returns 1 on match, 0 on mismatch.
+ * Prints diagnostic on first failing entry. */
+int check_jacobian(expr *node, const double *u, double h);
+
+#endif /* NUMERICAL_DIFF_H */

tests/numerical_diff/test_numerical_diff.h

@@ -0,0 +1,30 @@
+#include <string.h>
+
+#include "affine.h"
+#include "elementwise_univariate.h"
+#include "minunit.h"
+#include "numerical_diff.h"
+
+const char *test_check_jacobian_composite_log(void)
+{
+    double u_vals[6] = {0, 0, 1, 2, 3, 0};
+
+    CSR_Matrix *A = new_csr_matrix(2, 6, 6);
+    double Ax[6] = {3, 2, 1, 2, 1, 1};
+    int Ai[6] = {2, 3, 4, 2, 3, 4};
+    int Ap[3] = {0, 3, 6};
+    memcpy(A->x, Ax, 6 * sizeof(double));
+    memcpy(A->i, Ai, 6 * sizeof(int));
+    memcpy(A->p, Ap, 3 * sizeof(int));
+
+    expr *u = new_variable(3, 1, 2, 6);
+    expr *Au = new_linear(u, A, NULL);
+    expr *log_node = new_log(Au);
+
+    mu_assert("check_jacobian failed",
+              check_jacobian(log_node, u_vals, NUMERICAL_DIFF_DEFAULT_H));
+
+    free_expr(log_node);
+    free_csr_matrix(A);
+    return 0;
+}