replace right matmul with left matmul

Author: dance858

include/utils/CSR_Matrix.h

@@ -46,9 +46,8 @@ CSR_Matrix *kron_identity_csr(const CSR_Matrix *A, int p);
 void csr_matvec(const CSR_Matrix *A, const double *x, double *y, int col_offset);
 void csr_matvec_wo_offset(const CSR_Matrix *A, const double *x, double *y);
 
-/* C = z^T A is assumed to have one row. C must have column indices pre-computed
-and transposed matrix AT must be provided. Fills in values of C only.
- */
+/* Computes values of the row matrix C = z^T A (column indices must have been
+   pre-computed) and transposed matrix AT must be provided) */
 void csr_matvec_fill_values(const CSR_Matrix *AT, const double *z, CSR_Matrix *C);
 
 /* Insert value into CSR matrix A with just one row at col_idx. Assumes that A

src/bivariate/right_matmul.c

@@ -15,137 +15,29 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include "affine.h"
 #include "bivariate.h"
 #include "subexpr.h"
+#include "utils/CSR_Matrix.h"
 #include "utils/linalg_sparse_matmuls.h"
 #include <stdlib.h>
 
 /* 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.
-   Here, f(x) can be a vector-valued expression and a matrix-valued
-   expression. The dimensions are f(x) - p x n, A - n x q, y - p x q.
-   Note that here A does not have global column indices but it is a local matrix.
-   This is an important distinction compared to linear_op_expr.
-
-   * To compute the forward pass: vec(y) = B @ vec(f(x)),
-     where B = A^T kron I_p is a Kronecker product of size (p*q) x (p*n).
-
-   * To compute the Jacobian: J_y = B @ J_f(x), where J_f(x) is the
-     Jacobian of f(x) of size (p*n) x n_vars.
-
-    * To compute the contribution to the Lagrange Hessian: we form
-    w_child = B^T @ w and then evaluate the hessian of f(x).
-*/
-
-static void forward(expr *node, const double *u)
-{
-    expr *x = node->left;
-
-    /* child's forward pass */
-    node->left->forward(node->left, u);
-
-    /* y = x * A, vec(y) = B @ vec(x) */
-    csr_matvec_wo_offset(((right_matmul_expr *) node)->B, x->value, node->value);
-}
-
-static bool is_affine(const expr *node)
-{
-    return node->left->is_affine(node->left);
-}
-
-static void free_type_data(expr *node)
-{
-    right_matmul_expr *right_node = (right_matmul_expr *) node;
-    free_csr_matrix(right_node->B);
-    free_csr_matrix(right_node->BT);
-    if (right_node->CSC_work)
-    {
-        free_csc_matrix(right_node->CSC_work);
-    }
-    right_node->B = NULL;
-    right_node->BT = NULL;
-    right_node->CSC_work = NULL;
-}
-
-static void jacobian_init(expr *node)
-{
-    expr *x = node->left;
-    right_matmul_expr *right_node = (right_matmul_expr *) node;
-
-    /* initialize child's jacobian and precompute sparsity of its transpose */
-    x->jacobian_init(x);
-    right_node->CSC_work = csr_to_csc_fill_sparsity(x->jacobian, node->iwork);
-
-    /* precompute sparsity of this node's jacobian */
-    node->jacobian = csr_csc_matmul_alloc(right_node->B, right_node->CSC_work);
-}
-
-static void eval_jacobian(expr *node)
-{
-    expr *x = node->left;
-    right_matmul_expr *right_node = (right_matmul_expr *) node;
-
-    /* evaluate child's jacobian and convert to CSC*/
-    x->eval_jacobian(x);
-    csr_to_csc_fill_values(x->jacobian, right_node->CSC_work, node->iwork);
-
-    /* compute this node's jacobian */
-    csr_csc_matmul_fill_values(right_node->B, right_node->CSC_work, node->jacobian);
-}
-
-static void wsum_hess_init(expr *node)
-{
-    /* initialize child's hessian */
-    expr *x = node->left;
-    x->wsum_hess_init(x);
-
-    /* allocate this node's hessian with the same sparsity as child's */
-    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, x->wsum_hess->nnz);
-    memcpy(node->wsum_hess->p, x->wsum_hess->p, (node->n_vars + 1) * sizeof(int));
-    memcpy(node->wsum_hess->i, x->wsum_hess->i, x->wsum_hess->nnz * sizeof(int));
-
-    /* Allocate workspace for B^T @ w */
-    node->dwork = (double *) malloc(x->d1 * x->d2 * sizeof(double));
-}
-
-static void eval_wsum_hess(expr *node, const double *w)
-{
-    /* Compute B^T @ w, where B = A^T ⊗ I_p */
-    right_matmul_expr *right_node = (right_matmul_expr *) node;
-
-    /* B^T @ w computes the weights for the child expression */
-    csr_matvec_wo_offset(right_node->BT, w, node->dwork);
-
-    /* Propagate to child */
-    node->left->eval_wsum_hess(node->left, node->dwork);
-    memcpy(node->wsum_hess->x, node->left->wsum_hess->x,
-           node->wsum_hess->nnz * sizeof(double));
-}
-
+    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)
 {
-    /* Allocate the type-specific struct */
-    right_matmul_expr *right_matmul_node =
-        (right_matmul_expr *) calloc(1, sizeof(right_matmul_expr));
-    expr *node = &right_matmul_node->base;
-
-    /* Output dimensions: u is p x n, A is n x q, output is p x q */
-    int p = u->d1;
-    int q = A->n;
+    /* 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);
 
-    init_expr(node, p, q, u->n_vars, forward, jacobian_init, eval_jacobian,
-              is_affine, wsum_hess_init, eval_wsum_hess, free_type_data);
-    node->left = u;
-    expr_retain(u);
+    expr *u_transpose = new_transpose(u);
+    expr *left_matmul = new_left_matmul(u_transpose, AT);
+    expr *node = new_transpose(left_matmul);
 
-    /* create B = A^T kron I_p and its transpose */
-    node->iwork = (int *) malloc(node->n_vars * sizeof(int));
-    CSR_Matrix *AT = transpose(A, node->iwork);
-    right_matmul_node->B = kron_identity_csr(AT, p);
-    right_matmul_node->BT = kron_identity_csr(A, p);
     free_csr_matrix(AT);
-
-    right_matmul_node->CSC_work = NULL;
-
+    free(work_transpose);
     return node;
 }