Merge origin/main into parameter-support-v2

Resolve conflicts from quad-form chain rule PRs (#65, #66):
- Keep parameter_expr struct from this branch
- Update new_left_matmul calls in new tests to use 3-arg signature (NULL param_node)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Transurgeon

include/problem.h

@@ -36,6 +36,8 @@ typedef struct
     int nnz_affine;
     int nnz_nonlinear; /* jacobian of nonlinear constraints */
     int nnz_hessian;
+    int n_vars;
+    int total_constraint_size;
 } Diff_engine_stats;
 
 typedef struct problem

include/subexpr.h

@@ -40,12 +40,11 @@ typedef struct parameter_expr
     bool has_been_refreshed;
 } parameter_expr;
 
-/* Linear operator: y = A * x + b */
+/* Linear operator: y = A * x + b
+ * The matrix A is stored as node->jacobian (CSR). */
 typedef struct linear_op_expr
 {
     expr base;
-    CSC_Matrix *A_csc;
-    CSR_Matrix *A_csr;
     double *b; /* constant offset vector (NULL if no offset) */
 } linear_op_expr;
 
@@ -61,6 +60,7 @@ typedef struct quad_form_expr
 {
     expr base;
     CSR_Matrix *Q;
+    CSC_Matrix *QJf; /* Q * J_f in CSC (for chain rule hessian) */
 } quad_form_expr;
 
 /* Sum reduction along an axis */
@@ -110,8 +110,12 @@ typedef struct hstack_expr
 typedef struct elementwise_mult_expr
 {
     expr base;
-    CSR_Matrix *CSR_work1;
-    CSR_Matrix *CSR_work2;
+    CSR_Matrix *CSR_work1; /* C  = Jg2^T diag(w) Jg1 */
+    CSR_Matrix *CSR_work2; /* CT = C^T */
+    int *idx_map_C;        /* C[j]  -> wsum_hess pos */
+    int *idx_map_CT;       /* CT[j] -> wsum_hess pos */
+    int *idx_map_Hx;       /* x->wsum_hess[j] -> pos */
+    int *idx_map_Hy;       /* y->wsum_hess[j] -> pos */
 } elementwise_mult_expr;
 
 /* Left matrix multiplication: y = A * f(x) where f(x) is an expression. Note that

include/utils/CSC_Matrix.h

@@ -29,30 +29,37 @@ CSC_Matrix *new_csc_matrix(int m, int n, int nnz);
 /* Free a CSC matrix */
 void free_csc_matrix(CSC_Matrix *matrix);
 
-CSC_Matrix *csr_to_csc(const CSR_Matrix *A);
-
-/* Allocate sparsity pattern for C = A^T D A for diagonal D */
+/* Fill sparsity of C = A^T D A for diagonal D */
 CSR_Matrix *ATA_alloc(const CSC_Matrix *A);
 
-/* Allocate sparsity pattern for C = B^T D A for diagonal D */
+/* Fill sparsity of C = B^T D A for diagonal D */
 CSR_Matrix *BTA_alloc(const CSC_Matrix *A, const CSC_Matrix *B);
 
-/* Compute values for C = A^T D A. C must have precomputed sparsity pattern  */
+/* Fill sparsity of C = BA, where B is symmetric. */
+CSC_Matrix *symBA_alloc(const CSR_Matrix *B, const CSC_Matrix *A);
+
+/* Compute values for C = A^T D A (null d corresponds to D as identity) */
 void ATDA_fill_values(const CSC_Matrix *A, const double *d, CSR_Matrix *C);
 
-/* Compute values for C = B^T D A. C must have precomputed sparsity pattern  */
+/* Compute values for C = B^T D A (null d corresonds to D as identity) */
 void BTDA_fill_values(const CSC_Matrix *A, const CSC_Matrix *B, const double *d,
                       CSR_Matrix *C);
 
-/* C = z^T A where A is in CSC format and C is assumed to have one row.
- * C must have column indices pre-computed. Fills in values of C only.
- */
-void csc_matvec_fill_values(const CSC_Matrix *A, const double *z, CSR_Matrix *C);
+/* Fill values of C = BA. The matrix B does not have to be symmetric */
+void BA_fill_values(const CSR_Matrix *B, const CSC_Matrix *A, CSC_Matrix *C);
+
+/* Fill values of C = x^T A. The matrix C must have filled sparsity. */
+void yTA_fill_values(const CSC_Matrix *A, const double *x, CSR_Matrix *C);
+
+/* Count nonzero columns of a CSC matrix */
+int count_nonzero_cols_csc(const CSC_Matrix *A);
 
-CSC_Matrix *csr_to_csc_fill_sparsity(const CSR_Matrix *A, int *iwork);
+/* convert from CSR to CSC format */
+CSC_Matrix *csr_to_csc_alloc(const CSR_Matrix *A, int *iwork);
 void csr_to_csc_fill_values(const CSR_Matrix *A, CSC_Matrix *C, int *iwork);
 
-CSR_Matrix *csc_to_csr_fill_sparsity(const CSC_Matrix *A, int *iwork);
+/* convert from CSC to CSR format */
+CSR_Matrix *csc_to_csr_alloc(const CSC_Matrix *A, int *iwork);
 void csc_to_csr_fill_values(const CSC_Matrix *A, CSR_Matrix *C, int *iwork);
 
 #endif /* CSC_MATRIX_H */

include/utils/CSR_sum.h

@@ -86,4 +86,15 @@ void sum_spaced_rows_into_row_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A,
                                                             int spacing, int *iwork,
                                                             int *idx_map);
 
+/* 4-way sorted merge of CSR matrices A, B, C, D (same dimensions).
+ * Allocates and returns the output CSR with the union sparsity pattern.
+ * Allocates and fills idx_maps[0..3] (one per input, size input->nnz
+ * each) mapping each input entry to its position in the output.
+ * Caller owns the returned CSR and all 4 idx_map arrays. */
+CSR_Matrix *sum_4_csr_fill_sparsity_and_idx_maps(const CSR_Matrix *A,
+                                                 const CSR_Matrix *B,
+                                                 const CSR_Matrix *C,
+                                                 const CSR_Matrix *D,
+                                                 int *idx_maps[4]);
+
 #endif /* CSR_SUM_H */

src/affine/left_matmul.c

@@ -111,12 +111,12 @@ static void jacobian_init_impl(expr *node)
 
     /* initialize child's jacobian and precompute sparsity of its CSC */
     jacobian_init(x);
-    lnode->Jchild_CSC = csr_to_csc_fill_sparsity(x->jacobian, node->work->iwork);
+    lnode->Jchild_CSC = csr_to_csc_alloc(x->jacobian, node->work->iwork);
 
     /* precompute sparsity of this node's jacobian in CSC and CSR */
     lnode->J_CSC = lnode->A->block_left_mult_sparsity(lnode->A, lnode->Jchild_CSC,
                                                       lnode->n_blocks);
-    node->jacobian = csc_to_csr_fill_sparsity(lnode->J_CSC, lnode->csc_to_csr_work);
+    node->jacobian = csc_to_csr_alloc(lnode->J_CSC, lnode->csc_to_csr_work);
 }
 
 static void eval_jacobian(expr *node)

src/affine/linear_op.c

@@ -16,6 +16,7 @@
  * limitations under the License.
  */
 #include "affine.h"
+#include "utils/CSR_Matrix.h"
 #include <assert.h>
 #include <stdlib.h>
 #include <string.h>
@@ -28,8 +29,8 @@ static void forward(expr *node, const double *u)
     /* child's forward pass */
     node->left->forward(node->left, u);
 
-    /* y = A * x */
-    csr_matvec(lin_node->A_csr, x->value, node->value, x->var_id);
+    /* y = A * x (A is stored as node->jacobian) */
+    csr_matvec(node->jacobian, x->value, node->value, x->var_id);
 
     /* y += b (if offset exists) */
     if (lin_node->b != NULL)
@@ -49,29 +50,17 @@ static bool is_affine(const expr *node)
 static void free_type_data(expr *node)
 {
     linear_op_expr *lin_node = (linear_op_expr *) node;
-    /* memory pointing to by A_csr will be freed when the jacobian is freed,
-       so if the jacobian is not null we must not free A_csr. */
-
-    if (!node->jacobian)
-    {
-        free_csr_matrix(lin_node->A_csr);
-    }
-
-    free_csc_matrix(lin_node->A_csc);
-
     if (lin_node->b != NULL)
     {
         free(lin_node->b);
         lin_node->b = NULL;
     }
-
-    lin_node->A_csr = NULL;
-    lin_node->A_csc = NULL;
 }
 
 static void jacobian_init_impl(expr *node)
 {
-    node->jacobian = ((linear_op_expr *) node)->A_csr;
+    /* jacobian is set at construction time — nothing to do */
+    (void) node;
 }
 
 static void eval_jacobian(expr *node)
@@ -80,21 +69,33 @@ static void eval_jacobian(expr *node)
     (void) node;
 }
 
+static void wsum_hess_init_impl(expr *node)
+{
+    /* Linear operator Hessian is always zero */
+    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 0);
+}
+
+static void eval_wsum_hess(expr *node, const double *w)
+{
+    /* Linear operator Hessian is always zero - nothing to evaluate */
+    (void) node;
+    (void) w;
+}
+
 expr *new_linear(expr *u, const CSR_Matrix *A, const double *b)
 {
     assert(u->d2 == 1);
     /* Allocate the type-specific struct */
     linear_op_expr *lin_node = (linear_op_expr *) calloc(1, sizeof(linear_op_expr));
     expr *node = &lin_node->base;
     init_expr(node, A->m, 1, u->n_vars, forward, jacobian_init_impl, eval_jacobian,
-              is_affine, NULL, NULL, free_type_data);
+              is_affine, wsum_hess_init_impl, eval_wsum_hess, free_type_data);
     node->left = u;
     expr_retain(u);
 
-    /* Initialize type-specific fields */
-    lin_node->A_csr = new_csr_matrix(A->m, A->n, A->nnz);
-    copy_csr_matrix(A, lin_node->A_csr);
-    lin_node->A_csc = csr_to_csc(A);
+    /* Store A directly as the jacobian (linear op jacobian is constant) */
+    node->jacobian = new_csr_matrix(A->m, A->n, A->nnz);
+    copy_csr_matrix(A, node->jacobian);
 
     /* Initialize offset (copy b if provided, otherwise NULL) */
     if (b != NULL)