kron_left.c

/*
 * Copyright 2026 Daniel Cederberg and William Zhang
 *
 * This file is part of the DNLP-differentiation-engine project.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "bivariate.h"
#include "subexpr.h"
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/*
 * Kronecker product: Z = kron(C, X) where C is a constant (m x n) matrix
 * and X is a variable expression of shape (p x q).
 *
 * Output Z has shape (m*p, n*q), stored column-major as vec(Z) of length m*p*n*q.
 *
 * Key identity: Z[i*p+k, j*q+l] = C[i,j] * X[k,l]
 * In column-major: vec(Z)[r] where r = (j*q+l)*(m*p) + i*p + k
 *   depends on vec(X)[s] where s = l*p + k, with coefficient C[i,j].
 *
 * Jacobian structure: each row r of J_Z is a scaled copy of row s of J_X.
 * Only rows where C[i,j] != 0 are non-trivial.
 */

/* ------------------------------------------------------------------ */
/*                          Forward pass                               */
/* ------------------------------------------------------------------ */
static void forward(expr *node, const double *u)
{
    kron_left_expr *kn = (kron_left_expr *) node;
    expr *child = node->left;
    CSR_Matrix *C = kn->C;
    int m = kn->m, p = kn->p, q = kn->q;
    int mp = m * p;

    child->forward(child, u);

    /* Zero output first */
    memset(node->value, 0, (size_t) node->size * sizeof(double));

    /* For each nonzero C[i,j], fill block: Z[i*p+k, j*q+l] = C[i,j] * X[k,l] */
    for (int i = 0; i < m; i++)
    {
        for (int idx = C->p[i]; idx < C->p[i + 1]; idx++)
        {
            int j = C->i[idx];
            double cij = C->x[idx];

            for (int l = 0; l < q; l++)
            {
                int z_col_start = (j * q + l) * mp + i * p;
                int x_col_start = l * p;
                for (int k = 0; k < p; k++)
                {
                    node->value[z_col_start + k] =
                        cij * child->value[x_col_start + k];
                }
            }
        }
    }
}

/* ------------------------------------------------------------------ */
/*                          Affine check                               */
/* ------------------------------------------------------------------ */
static bool is_affine(const expr *node)
{
    return node->left->is_affine(node->left);
}

/* ------------------------------------------------------------------ */
/*                      Jacobian initialization                        */
/* ------------------------------------------------------------------ */
static void jacobian_init(expr *node)
{
    kron_left_expr *kn = (kron_left_expr *) node;
    expr *child = node->left;
    CSR_Matrix *C = kn->C;
    int m = kn->m, p = kn->p, q = kn->q;
    int mp = m * p;
    int out_size = node->size; /* m * p * n * q */

    /* Initialize child's Jacobian */
    child->jacobian_init(child);
    CSR_Matrix *Jchild = child->jacobian;

    /*
     * Count total nnz: for each output row r corresponding to nonzero C[i,j],
     * the nnz equals the nnz of child's Jacobian row s = l*p + k.
     */
    int total_nnz = 0;
    for (int i = 0; i < m; i++)
    {
        int row_nnz_C = C->p[i + 1] - C->p[i];
        if (row_nnz_C == 0) continue;

        for (int l = 0; l < q; l++)
        {
            for (int k = 0; k < p; k++)
            {
                int s = l * p + k; /* child row */
                int child_row_nnz = Jchild->p[s + 1] - Jchild->p[s];
                total_nnz += row_nnz_C * child_row_nnz;
            }
        }
    }

    /* Allocate Jacobian */
    node->jacobian = new_csr_matrix(out_size, node->n_vars, total_nnz);

    /* Allocate row_map and row_scale arrays for eval_jacobian */
    kn->row_map = (int *) malloc((size_t) out_size * sizeof(int));
    kn->row_scale_idx = (int *) malloc((size_t) out_size * sizeof(int));

    /*
     * Fill Jacobian sparsity pattern.
     * Iterate in column-major order over Z:
     *   r = (j*q + l) * mp + i*p + k
     *   maps to child row s = l*p + k, with scale C[i,j]
     */
    int nnz_idx = 0;
    for (int r = 0; r < out_size; r++)
    {
        node->jacobian->p[r] = nnz_idx;

        /* Decode r: r = col_Z * mp + row_Z */
        int row_Z = r % mp; /* i*p + k */
        int col_Z = r / mp; /* j*q + l */
        int i = row_Z / p;
        int k = row_Z % p;
        int j = col_Z / q;
        int l = col_Z % q;
        int s = l * p + k; /* child Jacobian row */

        kn->row_map[r] = s;

        /* Find C[i,j] in CSR */
        double cij = 0.0;
        int cij_found = 0;
        for (int idx = C->p[i]; idx < C->p[i + 1]; idx++)
        {
            if (C->i[idx] == j)
            {
                cij = C->x[idx];
                kn->row_scale_idx[r] = idx;
                cij_found = 1;
                break;
            }
        }

        if (!cij_found || cij == 0.0)
        {
            kn->row_scale_idx[r] = -1; /* mark as zero row */
            continue;
        }

        /* Copy child's column indices for row s */
        int child_start = Jchild->p[s];
        int child_end = Jchild->p[s + 1];
        int child_row_nnz = child_end - child_start;

        memcpy(node->jacobian->i + nnz_idx, Jchild->i + child_start,
               (size_t) child_row_nnz * sizeof(int));
        nnz_idx += child_row_nnz;
    }
    node->jacobian->p[out_size] = nnz_idx;
    node->jacobian->nnz = nnz_idx;
    assert(nnz_idx == total_nnz);
}

/* ------------------------------------------------------------------ */
/*                      Jacobian evaluation                            */
/* ------------------------------------------------------------------ */
static void eval_jacobian(expr *node)
{
    kron_left_expr *kn = (kron_left_expr *) node;
    expr *child = node->left;
    CSR_Matrix *C = kn->C;
    CSR_Matrix *Jchild = child->jacobian;
    CSR_Matrix *J = node->jacobian;
    int out_size = node->size;

    /* Evaluate child's Jacobian */
    child->eval_jacobian(child);

    /* Fill values: each active row r copies child row s scaled by C[i,j] */
    for (int r = 0; r < out_size; r++)
    {
        int j_start = J->p[r];
        int j_end = J->p[r + 1];
        int row_nnz = j_end - j_start;
        if (row_nnz == 0) continue;

        int s = kn->row_map[r];
        int c_idx = kn->row_scale_idx[r];
        double cij = C->x[c_idx];

        int child_start = Jchild->p[s];

        for (int t = 0; t < row_nnz; t++)
        {
            J->x[j_start + t] = cij * Jchild->x[child_start + t];
        }
    }
}

/* ------------------------------------------------------------------ */
/*                  Weighted-sum Hessian initialization                 */
/* ------------------------------------------------------------------ */
static void wsum_hess_init(expr *node)
{
    expr *child = node->left;

    /* Initialize child's Hessian */
    child->wsum_hess_init(child);

    /* kron_left is linear in X, so Hessian has same sparsity as child */
    node->wsum_hess =
        new_csr_matrix(node->n_vars, node->n_vars, child->wsum_hess->nnz);
    memcpy(node->wsum_hess->p, child->wsum_hess->p,
           (size_t) (node->n_vars + 1) * sizeof(int));
    memcpy(node->wsum_hess->i, child->wsum_hess->i,
           (size_t) child->wsum_hess->nnz * sizeof(int));

    /* Allocate workspace for reverse-mode weight accumulation */
    node->dwork = (double *) calloc((size_t) child->size, sizeof(double));
}

/* ------------------------------------------------------------------ */
/*                  Weighted-sum Hessian evaluation                     */
/* ------------------------------------------------------------------ */
static void eval_wsum_hess(expr *node, const double *w)
{
    kron_left_expr *kn = (kron_left_expr *) node;
    expr *child = node->left;
    CSR_Matrix *C = kn->C;
    int m = kn->m, p = kn->p, q = kn->q;
    int mp = m * p;
    int child_size = child->size;

    /*
     * Reverse mode: w_child[s] = sum over active r mapping to s of (C[i,j] * w[r])
     * This is the adjoint of the forward pass.
     */
    memset(node->dwork, 0, (size_t) child_size * sizeof(double));

    for (int i = 0; i < m; i++)
    {
        for (int idx = C->p[i]; idx < C->p[i + 1]; idx++)
        {
            int j = C->i[idx];
            double cij = C->x[idx];

            for (int l = 0; l < q; l++)
            {
                for (int k = 0; k < p; k++)
                {
                    int r = (j * q + l) * mp + i * p + k;
                    int s = l * p + k;
                    node->dwork[s] += cij * w[r];
                }
            }
        }
    }

    /* Delegate to child */
    child->eval_wsum_hess(child, node->dwork);
    memcpy(node->wsum_hess->x, child->wsum_hess->x,
           (size_t) node->wsum_hess->nnz * sizeof(double));
}

/* ------------------------------------------------------------------ */
/*                          Cleanup                                    */
/* ------------------------------------------------------------------ */
static void free_type_data(expr *node)
{
    kron_left_expr *kn = (kron_left_expr *) node;
    free_csr_matrix(kn->C);
    free(kn->row_map);
    free(kn->row_scale_idx);
    kn->C = NULL;
    kn->row_map = NULL;
    kn->row_scale_idx = NULL;
}

/* ------------------------------------------------------------------ */
/*                        Constructor                                  */
/* ------------------------------------------------------------------ */
expr *new_kron_left(expr *child, const CSR_Matrix *C, int p, int q)
{
    int m = C->m;
    int n = C->n;

    /* Verify child dimensions */
    if (child->size != p * q)
    {
        fprintf(stderr,
                "Error in new_kron_left: child size %d != p*q = %d*%d = %d\n",
                child->size, p, q, p * q);
        exit(1);
    }

    /* Output: kron(C, X) has shape (m*p, n*q) */
    int d1 = m * p;
    int d2 = n * q;

    kron_left_expr *kn = (kron_left_expr *) calloc(1, sizeof(kron_left_expr));
    expr *node = &kn->base;
    init_expr(node, d1, d2, child->n_vars, forward, jacobian_init, eval_jacobian,
              is_affine, wsum_hess_init, eval_wsum_hess, free_type_data);

    node->left = child;
    expr_retain(child);

    kn->m = m;
    kn->n = n;
    kn->p = p;
    kn->q = q;
    kn->C = new_csr(C);
    kn->row_map = NULL;
    kn->row_scale_idx = NULL;

    return node;
}