[WIP] Add upper_tri and diag_mat affine atoms (#54)

* Add upper_tri and diag_mat affine atoms

Both are thin wrappers that compute flat column-major index arrays
and delegate to new_index. diag_mat extracts the diagonal of a
square matrix into a vector. upper_tri extracts strict upper
triangular elements (excluding diagonal).

Also removes a duplicate new_diag_vec declaration in affine.h.

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

* Fix upper_tri to use row-major ordering to match CVXPY

CVXPY's upper_tri returns elements row-by-row (i outer, j inner),
which differs from the engine's typical column-major convention.
Swap the loop nesting to iterate rows-then-columns for compatibility.
Add a 4x4 forward test that distinguishes the two orderings.

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

* Remove SPDX headers and redundant 3x3 upper_tri test

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

* Upgrade upper_tri jacobian and hessian tests to 4x4

The 3x3 tests produced identical indices for both row-major and
column-major orderings, so they didn't actually verify the fix.
4x4 indices [4,8,12,9,13,14] differ from column-major [4,8,9,12,13,14].

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

* Run clang-format on changed files

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

---------

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

Author: Transurgeon

include/atoms/affine.h

@@ -38,6 +38,8 @@ expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs);
 expr *new_reshape(expr *child, int d1, int d2);
 expr *new_broadcast(expr *child, int target_d1, int target_d2);
 expr *new_diag_vec(expr *child);
+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

src/atoms/affine/diag_mat.c

@@ -0,0 +1,40 @@
+/*
+ * 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 "atoms/affine.h"
+#include <assert.h>
+#include <stdlib.h>
+
+/* Extract diagonal from a square matrix into a column vector.
+ * For an (n, n) matrix in column-major order, diagonal element i
+ * is at flat index i * (n + 1). */
+
+expr *new_diag_mat(expr *child)
+{
+    assert(child->d1 == child->d2);
+    int n = child->d1;
+
+    int *indices = (int *) malloc((size_t) n * sizeof(int));
+    for (int i = 0; i < n; i++)
+    {
+        indices[i] = i * (n + 1);
+    }
+
+    expr *node = new_index(child, n, 1, indices, n);
+    free(indices);
+    return node;
+}

src/atoms/affine/upper_tri.c

@@ -0,0 +1,58 @@
+/*
+ * 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 "atoms/affine.h"
+#include <assert.h>
+#include <stdlib.h>
+
+/* Extract strict upper triangular elements (excluding diagonal)
+ * from a square matrix, in ROW-MAJOR order to match CVXPY.
+ *
+ * NOTE: This is an exception to the engine's column-major
+ * convention. CVXPY's upper_tri iterates row-by-row across
+ * columns (i outer, j inner), so we do the same here for
+ * compatibility.
+ *
+ * For an (n, n) column-major matrix, element (i, j) with
+ * i < j is at flat index j * n + i.
+ * Output has n * (n - 1) / 2 elements. */
+
+expr *new_upper_tri(expr *child)
+{
+    assert(child->d1 == child->d2);
+    int n = child->d1;
+    int n_elems = n * (n - 1) / 2;
+
+    int *indices = NULL;
+    if (n_elems > 0)
+    {
+        indices = (int *) malloc((size_t) n_elems * sizeof(int));
+        int k = 0;
+        for (int i = 0; i < n; i++)
+        {
+            for (int j = i + 1; j < n; j++)
+            {
+                indices[k++] = j * n + i;
+            }
+        }
+        assert(k == n_elems);
+    }
+
+    expr *node = new_index(child, n_elems, 1, indices, n_elems);
+    free(indices);
+    return node;
+}

tests/all_tests.c

@@ -7,12 +7,14 @@
 #ifndef PROFILE_ONLY
 #include "forward_pass/affine/test_add.h"
 #include "forward_pass/affine/test_broadcast.h"
+#include "forward_pass/affine/test_diag_mat.h"
 #include "forward_pass/affine/test_hstack.h"
 #include "forward_pass/affine/test_left_matmul_dense.h"
 #include "forward_pass/affine/test_linear_op.h"
 #include "forward_pass/affine/test_neg.h"
 #include "forward_pass/affine/test_promote.h"
 #include "forward_pass/affine/test_sum.h"
+#include "forward_pass/affine/test_upper_tri.h"
 #include "forward_pass/affine/test_variable_parameter.h"
 #include "forward_pass/affine/test_vstack.h"
 #include "forward_pass/bivariate_full_dom/test_matmul.h"
@@ -23,6 +25,7 @@
 #include "forward_pass/other/test_prod_axis_one.h"
 #include "forward_pass/other/test_prod_axis_zero.h"
 #include "jacobian_tests/affine/test_broadcast.h"
+#include "jacobian_tests/affine/test_diag_mat.h"
 #include "jacobian_tests/affine/test_hstack.h"
 #include "jacobian_tests/affine/test_index.h"
 #include "jacobian_tests/affine/test_left_matmul.h"
@@ -33,6 +36,7 @@
 #include "jacobian_tests/affine/test_sum.h"
 #include "jacobian_tests/affine/test_trace.h"
 #include "jacobian_tests/affine/test_transpose.h"
+#include "jacobian_tests/affine/test_upper_tri.h"
 #include "jacobian_tests/affine/test_vector_mult.h"
 #include "jacobian_tests/affine/test_vstack.h"
 #include "jacobian_tests/bivariate_full_dom/test_elementwise_mult.h"
@@ -61,6 +65,7 @@
 #include "utils/test_linalg_utils_matmul_chain_rule.h"
 #include "utils/test_matrix.h"
 #include "wsum_hess/affine/test_broadcast.h"
+#include "wsum_hess/affine/test_diag_mat.h"
 #include "wsum_hess/affine/test_hstack.h"
 #include "wsum_hess/affine/test_index.h"
 #include "wsum_hess/affine/test_left_matmul.h"
@@ -69,6 +74,7 @@
 #include "wsum_hess/affine/test_sum.h"
 #include "wsum_hess/affine/test_trace.h"
 #include "wsum_hess/affine/test_transpose.h"
+#include "wsum_hess/affine/test_upper_tri.h"
 #include "wsum_hess/affine/test_vector_mult.h"
 #include "wsum_hess/affine/test_vstack.h"
 #include "wsum_hess/bivariate_full_dom/test_matmul.h"
@@ -128,6 +134,8 @@ int main(void)
     mu_run_test(test_forward_prod_axis_one, tests_run);
     mu_run_test(test_matmul, tests_run);
     mu_run_test(test_left_matmul_dense, tests_run);
+    mu_run_test(test_diag_mat_forward, tests_run);
+    mu_run_test(test_upper_tri_forward_4x4, tests_run);
 
     printf("\n--- Jacobian Tests ---\n");
     mu_run_test(test_neg_jacobian, tests_run);
@@ -208,6 +216,10 @@ int main(void)
     mu_run_test(test_jacobian_right_matmul_log_vector, tests_run);
     mu_run_test(test_jacobian_matmul, tests_run);
     mu_run_test(test_jacobian_transpose, tests_run);
+    mu_run_test(test_diag_mat_jacobian_variable, tests_run);
+    mu_run_test(test_diag_mat_jacobian_of_log, tests_run);
+    mu_run_test(test_upper_tri_jacobian_variable, tests_run);
+    mu_run_test(test_upper_tri_jacobian_of_log, tests_run);
 
     printf("\n--- Weighted Sum of Hessian Tests ---\n");
     mu_run_test(test_wsum_hess_log, tests_run);
@@ -275,6 +287,8 @@ int main(void)
     mu_run_test(test_wsum_hess_trace_log_variable, tests_run);
     mu_run_test(test_wsum_hess_trace_composite, tests_run);
     mu_run_test(test_wsum_hess_transpose, tests_run);
+    mu_run_test(test_wsum_hess_diag_mat_log, tests_run);
+    mu_run_test(test_wsum_hess_upper_tri_log, tests_run);
     mu_run_test(test_wsum_hess_exp_sum, tests_run);
     mu_run_test(test_wsum_hess_exp_sum_mult, tests_run);
     mu_run_test(test_wsum_hess_exp_sum_matmul, tests_run);

tests/forward_pass/affine/test_diag_mat.h

@@ -0,0 +1,29 @@
+#include <stdio.h>
+
+#include "atoms/affine.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_diag_mat_forward(void)
+{
+    /* 3x3 matrix variable (column-major): [1,2,3,4,5,6,7,8,9]
+     * Matrix:  1 4 7
+     *          2 5 8
+     *          3 6 9
+     * Diagonal: (0,0)=1, (1,1)=5, (2,2)=9 */
+    double u[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0};
+    expr *var = new_variable(3, 3, 0, 9);
+    expr *dm = new_diag_mat(var);
+
+    mu_assert("diag_mat d1", dm->d1 == 3);
+    mu_assert("diag_mat d2", dm->d2 == 1);
+
+    dm->forward(dm, u);
+
+    double expected[3] = {1.0, 5.0, 9.0};
+    mu_assert("diag_mat forward", cmp_double_array(dm->value, expected, 3));
+
+    free_expr(dm);
+    return 0;
+}

tests/forward_pass/affine/test_upper_tri.h

@@ -0,0 +1,41 @@
+#include <stdio.h>
+
+#include "atoms/affine.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_upper_tri_forward_4x4(void)
+{
+    /* 4x4 matrix variable (column-major): [1..16]
+     * Matrix:  1  5   9  13
+     *          2  6  10  14
+     *          3  7  11  15
+     *          4  8  12  16
+     * Upper tri in row-major order (matching CVXPY):
+     *   Row 0: (0,1)=5,  (0,2)=9,  (0,3)=13
+     *   Row 1: (1,2)=10, (1,3)=14
+     *   Row 2: (2,3)=15
+     * Flat indices: 4, 8, 12, 9, 13, 14
+     *
+     * NOTE: column-major order would give [4, 8, 9, 12, 13, 14]
+     * instead. This test verifies the row-major ordering. */
+    double u[16];
+    for (int k = 0; k < 16; k++)
+    {
+        u[k] = (double) (k + 1);
+    }
+    expr *var = new_variable(4, 4, 0, 16);
+    expr *ut = new_upper_tri(var);
+
+    mu_assert("upper_tri 4x4 d1", ut->d1 == 6);
+    mu_assert("upper_tri 4x4 d2", ut->d2 == 1);
+
+    ut->forward(ut, u);
+
+    double expected[6] = {5.0, 9.0, 13.0, 10.0, 14.0, 15.0};
+    mu_assert("upper_tri forward 4x4", cmp_double_array(ut->value, expected, 6));
+
+    free_expr(ut);
+    return 0;
+}

tests/jacobian_tests/affine/test_diag_mat.h

@@ -0,0 +1,61 @@
+#include <stdio.h>
+
+#include "atoms/affine.h"
+#include "atoms/elementwise_restricted_dom.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_diag_mat_jacobian_variable(void)
+{
+    /* diag_mat of a 2x2 variable (4 vars total)
+     * Diagonal indices in column-major: [0, 3]
+     * Jacobian is 2x4 CSR: row 0 has col 0, row 1 has col 3 */
+    double u[4] = {1.0, 2.0, 3.0, 4.0};
+    expr *var = new_variable(2, 2, 0, 4);
+    expr *dm = new_diag_mat(var);
+
+    dm->forward(dm, u);
+    jacobian_init(dm);
+    dm->eval_jacobian(dm);
+
+    double expected_x[2] = {1.0, 1.0};
+    int expected_p[3] = {0, 1, 2};
+    int expected_i[2] = {0, 3};
+
+    mu_assert("diag_mat jac vals", cmp_double_array(dm->jacobian->x, expected_x, 2));
+    mu_assert("diag_mat jac p", cmp_int_array(dm->jacobian->p, expected_p, 3));
+    mu_assert("diag_mat jac i", cmp_int_array(dm->jacobian->i, expected_i, 2));
+
+    free_expr(dm);
+    return 0;
+}
+
+const char *test_diag_mat_jacobian_of_log(void)
+{
+    /* diag_mat(log(X)) where X is 2x2 variable
+     * X = [[1, 3], [2, 4]] (column-major: [1, 2, 3, 4])
+     * Diagonal: x[0]=1, x[3]=4
+     * d/dx log at diagonal positions:
+     * Row 0: 1/1 = 1.0 at col 0
+     * Row 1: 1/4 = 0.25 at col 3 */
+    double u[4] = {1.0, 2.0, 3.0, 4.0};
+    expr *var = new_variable(2, 2, 0, 4);
+    expr *log_node = new_log(var);
+    expr *dm = new_diag_mat(log_node);
+
+    dm->forward(dm, u);
+    jacobian_init(dm);
+    dm->eval_jacobian(dm);
+
+    double expected_x[2] = {1.0, 0.25};
+    int expected_i[2] = {0, 3};
+
+    mu_assert("diag_mat log jac vals",
+              cmp_double_array(dm->jacobian->x, expected_x, 2));
+    mu_assert("diag_mat log jac cols",
+              cmp_int_array(dm->jacobian->i, expected_i, 2));
+
+    free_expr(dm);
+    return 0;
+}

tests/jacobian_tests/affine/test_upper_tri.h

@@ -0,0 +1,69 @@
+#include <stdio.h>
+
+#include "atoms/affine.h"
+#include "atoms/elementwise_restricted_dom.h"
+#include "expr.h"
+#include "minunit.h"
+#include "test_helpers.h"
+
+const char *test_upper_tri_jacobian_variable(void)
+{
+    /* upper_tri of a 4x4 variable (16 vars total)
+     * Row-major upper tri indices: [4, 8, 12, 9, 13, 14]
+     * Jacobian is 6x16 CSR: row k has a single 1.0 at col indices[k] */
+    double u[16];
+    for (int k = 0; k < 16; k++)
+    {
+        u[k] = (double) (k + 1);
+    }
+    expr *var = new_variable(4, 4, 0, 16);
+    expr *ut = new_upper_tri(var);
+
+    ut->forward(ut, u);
+    jacobian_init(ut);
+    ut->eval_jacobian(ut);
+
+    double expected_x[6] = {1.0, 1.0, 1.0, 1.0, 1.0, 1.0};
+    int expected_p[7] = {0, 1, 2, 3, 4, 5, 6};
+    int expected_i[6] = {4, 8, 12, 9, 13, 14};
+
+    mu_assert("upper_tri jac vals",
+              cmp_double_array(ut->jacobian->x, expected_x, 6));
+    mu_assert("upper_tri jac p", cmp_int_array(ut->jacobian->p, expected_p, 7));
+    mu_assert("upper_tri jac i", cmp_int_array(ut->jacobian->i, expected_i, 6));
+
+    free_expr(ut);
+    return 0;
+}
+
+const char *test_upper_tri_jacobian_of_log(void)
+{
+    /* upper_tri(log(X)) where X is 4x4 variable
+     * Row-major upper tri indices: [4, 8, 12, 9, 13, 14]
+     * Values at those positions: u[4]=5, u[8]=9, u[12]=13,
+     *                            u[9]=10, u[13]=14, u[14]=15
+     * d/dx log at those positions: 1/5, 1/9, 1/13, 1/10, 1/14, 1/15 */
+    double u[16];
+    for (int k = 0; k < 16; k++)
+    {
+        u[k] = (double) (k + 1);
+    }
+    expr *var = new_variable(4, 4, 0, 16);
+    expr *log_node = new_log(var);
+    expr *ut = new_upper_tri(log_node);
+
+    ut->forward(ut, u);
+    jacobian_init(ut);
+    ut->eval_jacobian(ut);
+
+    double expected_x[6] = {0.2, 1.0 / 9.0, 1.0 / 13.0, 0.1, 1.0 / 14.0, 1.0 / 15.0};
+    int expected_i[6] = {4, 8, 12, 9, 13, 14};
+
+    mu_assert("upper_tri log jac vals",
+              cmp_double_array(ut->jacobian->x, expected_x, 6));
+    mu_assert("upper_tri log jac cols",
+              cmp_int_array(ut->jacobian->i, expected_i, 6));
+
+    free_expr(ut);
+    return 0;
+}