add parameter support for right matmul

Author: dance858

include/bivariate.h

@@ -34,8 +34,8 @@ expr *new_matmul(expr *x, expr *y);
    Only the forward pass possibly updates the parameter. */
 expr *new_left_matmul(expr *param_node, expr *child, const CSR_Matrix *A);
 
-/* Right matrix multiplication: f(x) @ A where A is a fixed parameter matrix */
-expr *new_right_matmul(expr *u, const CSR_Matrix *A);
+/* Right matrix multiplication: f(x) @ A where A comes from a parameter node. */
+expr *new_right_matmul(expr *param_node, expr *u, const CSR_Matrix *A);
 
 /* Scalar multiplication: a * f(x) where a comes from a parameter node */
 expr *new_scalar_mult(expr *param_node, expr *child);

include/subexpr.h

@@ -129,6 +129,7 @@ typedef struct left_matmul_expr
     int *csc_to_csr_workspace;
     int *AT_iwork;      /* work for computing AT values from A */
     expr *param_source; /* parameter node; A/AT values are refreshed from this */
+    void (*refresh_param_values)(struct left_matmul_expr *lin_node);
 } left_matmul_expr;
 
 /* Right matrix multiplication: y = f(x) * A where f(x) is an expression.

src/bivariate/left_matmul.c

@@ -48,6 +48,7 @@
 */
 
 #include "utils/utils.h"
+#include <assert.h>
 #include <string.h>
 
 /* Refresh A and AT values from param_source.
@@ -60,6 +61,8 @@ static void refresh_param_values(left_matmul_expr *lin_node)
     if (!param || param->has_been_refreshed) return;
     param->has_been_refreshed = true;
 
+    assert(param->param_id != PARAM_FIXED);
+
     /* update values of A */
     memcpy(lin_node->A->x, lin_node->param_source->value,
            lin_node->A->nnz * sizeof(double));
@@ -74,7 +77,7 @@ static void forward(expr *node, const double *u)
     left_matmul_expr *lin_node = (left_matmul_expr *) node;
 
     /* possibly refresh A and AT */
-    refresh_param_values(lin_node);
+    if (lin_node->refresh_param_values) lin_node->refresh_param_values(lin_node);
 
     /* child's forward pass */
     node->left->forward(node->left, u);
@@ -205,6 +208,7 @@ expr *new_left_matmul(expr *param_node, expr *child, const CSR_Matrix *A)
     lin_node->A = new_csr(A);
     lin_node->AT = transpose(lin_node->A, lin_node->AT_iwork);
     lin_node->param_source = param_node;
+    lin_node->refresh_param_values = refresh_param_values;
 
     if (param_node) expr_retain(param_node);
 

src/bivariate/right_matmul.c

@@ -20,26 +20,49 @@
 #include "subexpr.h"
 #include "utils/CSR_Matrix.h"
 #include "utils/linalg_sparse_matmuls.h"
+#include <assert.h>
 #include <stdlib.h>
+#include <string.h>
+
+/* Refresh AT and A values from param_source for right matmul.
+   param_source stores values in CSR order for the original A. */
+static void refresh_param_values(left_matmul_expr *lin_node)
+{
+    parameter_expr *param = (parameter_expr *) lin_node->param_source;
+
+    if (!param || param->has_been_refreshed) return;
+    param->has_been_refreshed = true;
+
+    assert(param->param_id != PARAM_FIXED);
+
+    /* update values of original A (stored in lin_node->AT) */
+    memcpy(lin_node->AT->x, lin_node->param_source->value,
+           lin_node->AT->nnz * sizeof(double));
+
+    /* update values of A^T (stored in lin_node->A) */
+    AT_fill_values(lin_node->AT, lin_node->A, lin_node->AT_iwork);
+}
 
 /* This file implements the atom 'right_matmul' corresponding to the operation y =
-   f(x) @ A, where A is a given matrix and f(x) is an arbitrary expression.
+    f(x) @ A, where A is a given matrix and f(x) is an arbitrary expression.
     We implement this by expressing right matmul in terms of left matmul and
     transpose: f(x) @ A = (A^T @ f(x)^T)^T. */
-expr *new_right_matmul(expr *u, const CSR_Matrix *A)
+expr *new_right_matmul(expr *param_node, expr *u, const CSR_Matrix *A)
 {
     /* We can express right matmul using left matmul and transpose:
        u @ A = (A^T @ u^T)^T. */
     int *work_transpose = (int *) malloc(A->n * sizeof(int));
     CSR_Matrix *AT = transpose(A, work_transpose);
-
-    /* Parameter stores CSR data order (same as AT->x) */
     expr *u_transpose = new_transpose(u);
-    expr *param_node = new_parameter(AT->nnz, 1, PARAM_FIXED, u->n_vars, AT->x);
     expr *left_matmul_node = new_left_matmul(param_node, u_transpose, AT);
     expr *node = new_transpose(left_matmul_node);
 
+    /* functionality for parameter */
+    left_matmul_expr *left_matmul_data = (left_matmul_expr *) left_matmul_node;
+    free(left_matmul_data->AT_iwork);
+    left_matmul_data->AT_iwork = work_transpose;
+    left_matmul_data->refresh_param_values = refresh_param_values;
+
     free_csr_matrix(AT);
-    free(work_transpose);
     return node;
 }

src/bivariate/scalar_mult.c

@@ -32,8 +32,7 @@ static void forward(expr *node, const double *u)
     child->forward(child, u);
 
     /* local forward pass: multiply each element by scalar a */
-    scalar_mult_expr *sn = (scalar_mult_expr *) node;
-    double a = sn->param_source->value[0];
+    double a = ((scalar_mult_expr *) node)->param_source->value[0];
     for (int i = 0; i < node->size; i++)
     {
         node->value[i] = a * child->value[i];
@@ -56,8 +55,7 @@ static void jacobian_init(expr *node)
 static void eval_jacobian(expr *node)
 {
     expr *child = node->left;
-    scalar_mult_expr *sn = (scalar_mult_expr *) node;
-    double a = sn->param_source->value[0];
+    double a = ((scalar_mult_expr *) node)->param_source->value[0];
 
     /* evaluate child */
     child->eval_jacobian(child);
@@ -87,8 +85,7 @@ static void eval_wsum_hess(expr *node, const double *w)
     expr *x = node->left;
     x->eval_wsum_hess(x, w);
 
-    scalar_mult_expr *sn = (scalar_mult_expr *) node;
-    double a = sn->param_source->value[0];
+    double a = ((scalar_mult_expr *) node)->param_source->value[0];
     for (int j = 0; j < x->wsum_hess->nnz; j++)
     {
         node->wsum_hess->x[j] = a * x->wsum_hess->x[j];

src/bivariate/vector_mult.c

@@ -27,8 +27,10 @@
 static void forward(expr *node, const double *u)
 {
     expr *child = node->left;
-    vector_mult_expr *vn = (vector_mult_expr *) node;
-    const double *a = vn->param_source->value;
+    // vector_mult_expr *vn = (vector_mult_expr *) node;
+    // const double *a = vn->param_source->value;
+
+    const double *a = ((vector_mult_expr *) node)->param_source->value;
 
     /* child's forward pass */
     child->forward(child, u);
@@ -56,8 +58,10 @@ static void jacobian_init(expr *node)
 static void eval_jacobian(expr *node)
 {
     expr *x = node->left;
-    vector_mult_expr *vn = (vector_mult_expr *) node;
-    const double *a = vn->param_source->value;
+    // vector_mult_expr *vn = (vector_mult_expr *) node;
+    // const double *a = vn->param_source->value;
+
+    const double *a = ((vector_mult_expr *) node)->param_source->value;
 
     /* evaluate x */
     x->eval_jacobian(x);
@@ -90,8 +94,10 @@ static void wsum_hess_init(expr *node)
 static void eval_wsum_hess(expr *node, const double *w)
 {
     expr *x = node->left;
-    vector_mult_expr *vn = (vector_mult_expr *) node;
-    const double *a = vn->param_source->value;
+    // vector_mult_expr *vn = (vector_mult_expr *) node;
+    // const double *a = vn->param_source->value;
+
+    const double *a = ((vector_mult_expr *) node)->param_source->value;
 
     /* scale weights w by a */
     for (int i = 0; i < node->size; i++)

tests/all_tests.c

@@ -284,6 +284,7 @@ int main(void)
     mu_run_test(test_param_scalar_mult_problem, tests_run);
     mu_run_test(test_param_vector_mult_problem, tests_run);
     mu_run_test(test_param_left_matmul_problem, tests_run);
+    mu_run_test(test_param_right_matmul_problem, tests_run);
 #endif /* PROFILE_ONLY */
 
 #ifdef PROFILE_ONLY

tests/jacobian_tests/test_right_matmul.h

@@ -27,7 +27,7 @@ const char *test_jacobian_right_matmul_log()
     memcpy(A->x, A_x, 4 * sizeof(double));
 
     expr *log_x = new_log(x);
-    expr *log_x_A = new_right_matmul(log_x, A);
+    expr *log_x_A = new_right_matmul(NULL, log_x, A);
 
     log_x_A->forward(log_x_A, x_vals);
     log_x_A->jacobian_init(log_x_A);
@@ -76,7 +76,7 @@ const char *test_jacobian_right_matmul_log_vector()
     memcpy(A->x, A_x, 4 * sizeof(double));
 
     expr *log_x = new_log(x);
-    expr *log_x_A = new_right_matmul(log_x, A);
+    expr *log_x_A = new_right_matmul(NULL, log_x, A);
 
     log_x_A->forward(log_x_A, x_vals);
     log_x_A->jacobian_init(log_x_A);

tests/problem/test_param_prob.h

@@ -254,4 +254,98 @@ const char *test_param_left_matmul_problem(void)
     return 0;
 }
 
+/*
+ * Test 4: right_param_matmul in constraint
+ *
+ * Problem: minimize sum(x), subject to x @ A, x size 1x2, A is 2x2
+ *   A is a 2x2 matrix parameter (param_id=0, size=4, CSR data order)
+ *   A = [[1,2],[3,4]] → CSR data order theta = [1,2,3,4]
+ *
+ * At x=[1,2]:
+ *   constraint_values = [1*1+2*3, 1*2+2*4] = [7, 10]
+ *   jacobian = [[1,3],[2,4]] = A^T
+ *
+ * After update A = [[5,6],[7,8]] → theta = [5,6,7,8]:
+ *   constraint_values = [1*5+2*7, 1*6+2*8] = [19, 22]
+ *   jacobian = [[5,7],[6,8]] = A^T
+ */
+const char *test_param_right_matmul_problem(void)
+{
+    int n_vars = 2;
+
+    /* Objective: sum(x) */
+    expr *x_obj = new_variable(1, 2, 0, n_vars);
+    expr *objective = new_sum(x_obj, -1);
+
+    /* Constraint: x @ A */
+    expr *x_con = new_variable(1, 2, 0, n_vars);
+    expr *A_param = new_parameter(2, 2, 0, n_vars, NULL);
+
+    /* Dense 2x2 CSR with placeholder zeros (values refreshed from A_param) */
+    CSR_Matrix *A = new_csr_matrix(2, 2, 4);
+    int Ap[3] = {0, 2, 4};
+    int Ai[4] = {0, 1, 0, 1};
+    double Ax[4] = {0.0, 0.0, 0.0, 0.0};
+    memcpy(A->p, Ap, 3 * sizeof(int));
+    memcpy(A->i, Ai, 4 * sizeof(int));
+    memcpy(A->x, Ax, 4 * sizeof(double));
+
+    expr *constraint = new_right_matmul(A_param, x_con, A);
+    free_csr_matrix(A);
+
+    expr *constraints[1] = {constraint};
+
+    /* Create problem */
+    problem *prob = new_problem(objective, constraints, 1, true);
+
+    expr *param_nodes[1] = {A_param};
+    problem_register_params(prob, param_nodes, 1);
+    problem_init_derivatives(prob);
+
+    /* Set A = [[1,2],[3,4]], CSR data order: [1,2,3,4] */
+    double theta[4] = {1.0, 2.0, 3.0, 4.0};
+    problem_update_params(prob, theta);
+
+    double u[2] = {1.0, 2.0};
+    problem_constraint_forward(prob, u);
+    problem_jacobian(prob);
+
+    double expected_cv[2] = {7.0, 10.0};
+    mu_assert("constraint values wrong (A1)",
+              cmp_double_array(prob->constraint_values, expected_cv, 2));
+
+    CSR_Matrix *jac = prob->jacobian;
+    mu_assert("jac rows wrong", jac->m == 2);
+    mu_assert("jac cols wrong", jac->n == 2);
+
+    /* Dense jacobian = [[1,3],[2,4]] = A^T, CSR: row 0 → cols 0,1 vals 1,3;
+     *                                         row 1 → cols 0,1 vals 2,4 */
+    int expected_p[3] = {0, 2, 4};
+    mu_assert("jac->p wrong (A1)", cmp_int_array(jac->p, expected_p, 3));
+
+    int expected_i[4] = {0, 1, 0, 1};
+    mu_assert("jac->i wrong (A1)", cmp_int_array(jac->i, expected_i, 4));
+
+    double expected_x[4] = {1.0, 3.0, 2.0, 4.0};
+    mu_assert("jac->x wrong (A1)", cmp_double_array(jac->x, expected_x, 4));
+
+    /* Update A = [[5,6],[7,8]], CSR data order: [5,6,7,8] */
+    double theta2[4] = {5.0, 6.0, 7.0, 8.0};
+    problem_update_params(prob, theta2);
+
+    problem_constraint_forward(prob, u);
+    problem_jacobian(prob);
+
+    double expected_cv2[2] = {19.0, 22.0};
+    mu_assert("constraint values wrong (A2)",
+              cmp_double_array(prob->constraint_values, expected_cv2, 2));
+
+    double expected_x2[4] = {5.0, 7.0, 6.0, 8.0};
+    mu_assert("jac->x wrong (A2)", cmp_double_array(jac->x, expected_x2, 4));
+
+    free_problem(prob);
+
+    return 0;
+}
+
 #endif /* TEST_PARAM_PROB_H */

tests/wsum_hess/test_right_matmul.h

@@ -33,7 +33,7 @@ const char *test_wsum_hess_right_matmul()
     memcpy(A->x, A_x, 4 * sizeof(double));
 
     expr *log_x = new_log(x);
-    expr *log_x_A = new_right_matmul(log_x, A);
+    expr *log_x_A = new_right_matmul(NULL, log_x, A);
 
     log_x_A->forward(log_x_A, x_vals);
     log_x_A->jacobian_init(log_x_A);
@@ -83,7 +83,7 @@ const char *test_wsum_hess_right_matmul_vector()
     memcpy(A->x, A_x, 4 * sizeof(double));
 
     expr *log_x = new_log(x);
-    expr *log_x_A = new_right_matmul(log_x, A);
+    expr *log_x_A = new_right_matmul(NULL, log_x, A);
 
     log_x_A->forward(log_x_A, x_vals);
     log_x_A->jacobian_init(log_x_A);