Merge branch 'main' into adds-more-affine-atoms

# Conflicts:
#	include/atoms/affine.h
#	src/atoms/affine/diag_mat.c
#	src/atoms/affine/upper_tri.c
#	tests/all_tests.c

Author: Transurgeon

CMakeLists.txt

@@ -4,8 +4,8 @@ set(CMAKE_C_STANDARD 99)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 set(DIFF_ENGINE_VERSION_MAJOR 0)
-set(DIFF_ENGINE_VERSION_MINOR 1)
-set(DIFF_ENGINE_VERSION_PATCH 5)
+set(DIFF_ENGINE_VERSION_MINOR 3)
+set(DIFF_ENGINE_VERSION_PATCH 0)
 set(DIFF_ENGINE_VERSION "${DIFF_ENGINE_VERSION_MAJOR}.${DIFF_ENGINE_VERSION_MINOR}.${DIFF_ENGINE_VERSION_PATCH}")
 add_compile_definitions(DIFF_ENGINE_VERSION="${DIFF_ENGINE_VERSION}")
 
@@ -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)
 

include/affine.h include/atoms/affine.hinclude/affine.h renamed to include/atoms/affine.h

@@ -22,17 +22,16 @@
 #include "subexpr.h"
 #include "utils/CSR_Matrix.h"
 
-expr *new_linear(expr *u, const CSR_Matrix *A, const double *b);
-
 expr *new_add(expr *left, expr *right);
 expr *new_neg(expr *child);
 
 expr *new_sum(expr *child, int axis);
 expr *new_hstack(expr **args, int n_args, int n_vars);
+expr *new_vstack(expr **args, int n_args, int n_vars);
 expr *new_promote(expr *child, int d1, int d2);
 expr *new_trace(expr *child);
 
-expr *new_constant(int d1, int d2, int n_vars, const double *values);
+expr *new_parameter(int d1, int d2, int param_id, int n_vars, const double *values);
 expr *new_variable(int d1, int d2, int var_id, int n_vars);
 
 expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs);
@@ -43,4 +42,27 @@ expr *new_diag_mat(expr *child);
 expr *new_upper_tri(expr *child);
 expr *new_transpose(expr *child);
 
+/* Left matrix multiplication: A @ f(x) where A is a constant or parameter
+ * sparse matrix. param_node is NULL for fixed constants. */
+expr *new_left_matmul(expr *param_node, expr *u, const CSR_Matrix *A);
+
+/* Left matrix multiplication: A @ f(x) where A is a constant or parameter
+ * dense matrix (row-major, m x n). Uses CBLAS for efficient computation. */
+expr *new_left_matmul_dense(expr *param_node, expr *u, int m, int n,
+                            const double *data);
+
+/* Right matrix multiplication: f(x) @ A where A is a constant or parameter
+ * matrix. */
+expr *new_right_matmul(expr *param_node, expr *u, const CSR_Matrix *A);
+
+expr *new_right_matmul_dense(expr *param_node, expr *u, int m, int n,
+                             const double *data);
+
+/* Scalar multiplication: a * f(x) where a comes from param_node */
+expr *new_scalar_mult(expr *param_node, expr *child);
+
+/* Vector elementwise multiplication: a . f(x) where a comes from
+ * param_node */
+expr *new_vector_mult(expr *param_node, expr *child);
+
 #endif /* AFFINE_H */

include/atoms/bivariate_full_dom.h

@@ -0,0 +1,11 @@
+#ifndef BIVARIATE_FULL_DOM_H
+#define BIVARIATE_FULL_DOM_H
+
+#include "expr.h"
+
+expr *new_elementwise_mult(expr *left, expr *right);
+
+/* Matrix multiplication: Z = X @ Y */
+expr *new_matmul(expr *x, expr *y);
+
+#endif /* BIVARIATE_FULL_DOM_H */

include/atoms/bivariate_restricted_dom.h

@@ -0,0 +1,11 @@
+#ifndef BIVARIATE_RESTRICTED_DOM_H
+#define BIVARIATE_RESTRICTED_DOM_H
+
+#include "expr.h"
+
+expr *new_quad_over_lin(expr *left, expr *right);
+expr *new_rel_entr_vector_args(expr *left, expr *right);
+expr *new_rel_entr_first_arg_scalar(expr *left, expr *right);
+expr *new_rel_entr_second_arg_scalar(expr *left, expr *right);
+
+#endif /* BIVARIATE_RESTRICTED_DOM_H */

include/atoms/elementwise_full_dom.h

@@ -0,0 +1,34 @@
+#ifndef ELEMENTWISE_FULL_DOM_H
+#define ELEMENTWISE_FULL_DOM_H
+
+#include "expr.h"
+
+/* Helper function to initialize an elementwise expr
+ * (can be used with derived types) */
+void init_elementwise(expr *node, expr *child);
+
+expr *new_exp(expr *child);
+expr *new_sin(expr *child);
+expr *new_cos(expr *child);
+expr *new_sinh(expr *child);
+expr *new_tanh(expr *child);
+expr *new_asinh(expr *child);
+expr *new_logistic(expr *child);
+expr *new_power(expr *child, double p);
+expr *new_xexp(expr *child);
+expr *new_normal_cdf(expr *child);
+
+/* the jacobian and wsum_hess for elementwise full domain
+   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);
+void wsum_hess_init_elementwise(expr *node);
+void eval_wsum_hess_elementwise(expr *node, const double *w);
+expr *new_elementwise(expr *child);
+
+/* no elementwise atoms are affine according to our
+   convention, so we can have a common implementation */
+bool is_affine_elementwise(const expr *node);
+
+#endif /* ELEMENTWISE_FULL_DOM_H */

include/atoms/elementwise_restricted_dom.h

@@ -0,0 +1,18 @@
+#ifndef ELEMENTWISE_RESTRICTED_DOM_H
+#define ELEMENTWISE_RESTRICTED_DOM_H
+
+#include "expr.h"
+
+/* Shared init functions for restricted domain atoms
+ * (variable-child only, no linear operator support) */
+void jacobian_init_restricted(expr *node);
+void wsum_hess_init_restricted(expr *node);
+bool is_affine_restricted(const expr *node);
+expr *new_restricted(expr *child);
+
+expr *new_log(expr *child);
+expr *new_entr(expr *child);
+expr *new_atanh(expr *child);
+expr *new_tan(expr *child);
+
+#endif /* ELEMENTWISE_RESTRICTED_DOM_H */

include/other.h include/atoms/non_elementwise_full_dom.hinclude/other.h renamed to include/atoms/non_elementwise_full_dom.h

@@ -15,8 +15,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#ifndef OTHER_H
-#define OTHER_H
+#ifndef NON_ELEMENTWISE_FULL_DOM_H
+#define NON_ELEMENTWISE_FULL_DOM_H
 
 #include "expr.h"
 #include "subexpr.h"
@@ -33,4 +33,4 @@ expr *new_prod_axis_zero(expr *child);
 /* product of entries along axis=1 (rowwise products) */
 expr *new_prod_axis_one(expr *child);
 
-#endif /* OTHER_H */
+#endif /* NON_ELEMENTWISE_FULL_DOM_H */

include/bivariate.h

@@ -39,7 +39,24 @@ typedef void (*local_wsum_hess_fn)(struct expr *node, double *out, const double
 typedef bool (*is_affine_fn)(const struct expr *node);
 typedef void (*free_type_data_fn)(struct expr *node);
 
-/* Base expression node structure - contains only common fields */
+/* Workspace for derivative computation */
+typedef struct
+{
+    double *dwork;
+    int *iwork;
+    CSC_Matrix *jacobian_csc;
+    int *csc_work; /* for CSR-CSC conversion */
+
+    /* jacobian_csc_filled is only used for affine functions to avoid redundant
+       conversions. Could become relevant for non-affine functions if we start
+       supporting common subexpressions on the Python side. */
+    bool jacobian_csc_filled;
+    double *local_jac_diag; /* cached f'(g(x)) diagonal */
+    CSR_Matrix *hess_term1; /* Jg^T D Jg workspace */
+    CSR_Matrix *hess_term2; /* child wsum_hess workspace */
+} Expr_Work;
+
+/* Base expression node structure */
 typedef struct expr
 {
     // ------------------------------------------------------------------------
@@ -48,8 +65,6 @@ typedef struct expr
     int d1, d2, size, n_vars, refcount, var_id;
     struct expr *left;
     struct expr *right;
-    double *dwork;
-    int *iwork;
 
     // ------------------------------------------------------------------------
     //                     oracle related quantities
@@ -58,8 +73,8 @@ typedef struct expr
     CSR_Matrix *jacobian;
     CSR_Matrix *wsum_hess;
     forward_fn forward;
-    jacobian_init_fn jacobian_init;
-    wsum_hess_init_fn wsum_hess_init;
+    jacobian_init_fn jacobian_init_impl;
+    wsum_hess_init_fn wsum_hess_init_impl;
     eval_jacobian_fn eval_jacobian;
     wsum_hess_fn eval_wsum_hess;
 
@@ -70,6 +85,13 @@ typedef struct expr
     local_jacobian_fn local_jacobian;   /* used by elementwise univariate atoms*/
     local_wsum_hess_fn local_wsum_hess; /* used by elementwise univariate atoms*/
     free_type_data_fn free_type_data;   /* Cleanup for type-specific fields */
+    Expr_Work *work;                    /* derivative workspace */
+    /* Set to true on all nodes by problem_update_params() via
+       expr_set_needs_refresh(). Atoms that cache parameter data
+       (e.g. left_matmul_dense) check this flag before their forward
+       pass: if true, they refresh their cached matrices from
+       param_source->value and clear the flag to false. */
+    bool needs_parameter_refresh;
 
     // name of node just for debugging - should be removed later
     char name[32];
@@ -83,6 +105,18 @@ void init_expr(expr *node, int d1, int d2, int n_vars, forward_fn forward,
 
 void free_expr(expr *node);
 
+/* Guarded init: skips if already initialized (safe for DAGs
+ * where a node may be visited through multiple parents). */
+void jacobian_init(expr *node);
+void wsum_hess_init(expr *node);
+
+/* Initialize CSC form of the Jacobian from the CSR Jacobian.
+ * Must be called after jacobian_init. */
+void jacobian_csc_init(expr *node);
+
+/* Recursively set needs_parameter_refresh on node and all children */
+void expr_set_needs_refresh(expr *node);
+
 /* Reference counting helpers */
 void expr_retain(expr *node);