From 1febd656d8be2a73287e93929ee58477e34a727f Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Sat, 3 Jan 2026 18:13:03 -0500
Subject: [PATCH 1/7] adds initial code for tree conversion from cvxpy

---
 python/bindings.c           |  80 ++++++-
 python/convert.py           | 401 ++++++++++++++++++++++++++++++++++++
 python/test_inverse_data.py |  19 ++
 3 files changed, 499 insertions(+), 1 deletion(-)
 create mode 100644 python/convert.py
 create mode 100644 python/test_inverse_data.py

diff --git a/python/bindings.c b/python/bindings.c
index b4bcd38..2f4cb67 100644
--- a/python/bindings.c
+++ b/python/bindings.c
@@ -10,9 +10,13 @@
 // Capsule name for expr* pointers
 #define EXPR_CAPSULE_NAME "DNLP_EXPR"
 
+static int numpy_initialized = 0;
+
 static int ensure_numpy(void)
 {
-    import_array();
+    if (numpy_initialized) return 0;
+    import_array1(-1);
+    numpy_initialized = 1;
     return 0;
 }
 
@@ -65,6 +69,77 @@ static PyObject *py_make_log(PyObject *self, PyObject *args)
     return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
 }
 
+static PyObject *py_make_exp(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    if (!PyArg_ParseTuple(args, "O", &child_capsule))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_exp(child);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create exp node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+static PyObject *py_make_add(PyObject *self, PyObject *args)
+{
+    PyObject *left_capsule, *right_capsule;
+    if (!PyArg_ParseTuple(args, "OO", &left_capsule, &right_capsule))
+    {
+        return NULL;
+    }
+    expr *left = (expr *) PyCapsule_GetPointer(left_capsule, EXPR_CAPSULE_NAME);
+    expr *right = (expr *) PyCapsule_GetPointer(right_capsule, EXPR_CAPSULE_NAME);
+    if (!left || !right)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_add(left, right);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create add node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+static PyObject *py_make_sum(PyObject *self, PyObject *args)
+{
+    PyObject *child_capsule;
+    int axis;
+    if (!PyArg_ParseTuple(args, "Oi", &child_capsule, &axis))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_sum(child, axis);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create sum node");
+        return NULL;
+    }
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
 static PyObject *py_forward(PyObject *self, PyObject *args)
 {
     PyObject *node_capsule;
@@ -106,6 +181,9 @@ static PyObject *py_forward(PyObject *self, PyObject *args)
 static PyMethodDef DNLPMethods[] = {
     {"make_variable", py_make_variable, METH_VARARGS, "Create variable node"},
     {"make_log", py_make_log, METH_VARARGS, "Create log node"},
+    {"make_exp", py_make_exp, METH_VARARGS, "Create exp node"},
+    {"make_add", py_make_add, METH_VARARGS, "Create add node"},
+    {"make_sum", py_make_sum, METH_VARARGS, "Create sum node"},
     {"forward", py_forward, METH_VARARGS, "Run forward pass and return values"},
     {NULL, NULL, 0, NULL}};
 
diff --git a/python/convert.py b/python/convert.py
new file mode 100644
index 0000000..385b2fc
--- /dev/null
+++ b/python/convert.py
@@ -0,0 +1,401 @@
+import cvxpy as cp
+import numpy as np
+from cvxpy.reductions.inverse_data import InverseData
+import DNLP_diff_engine as diffengine
+
+
+def _chain_add(children):
+    """Chain multiple children with binary adds: a + b + c -> add(add(a, b), c)."""
+    result = children[0]
+    for child in children[1:]:
+        result = diffengine.make_add(result, child)
+    return result
+
+
+# Mapping from CVXPY atom names to C diff engine functions
+ATOM_CONVERTERS = {
+    # Elementwise unary
+    "log": lambda child: diffengine.make_log(child),
+    "exp": lambda child: diffengine.make_exp(child),
+
+    # N-ary (handles 2+ args)
+    "AddExpression": _chain_add,
+
+    # Reductions
+    "Sum": lambda child: diffengine.make_sum(child, -1),  # axis=-1 = sum all
+}
+
+
+def build_variable_dict(variables: list) -> tuple[dict, int]:
+    """
+    Build dictionary mapping CVXPY variable ids to C variables.
+
+    Args:
+        variables: list of CVXPY Variable objects
+
+    Returns:
+        var_dict: {var.id: c_variable} mapping
+        n_vars: total number of scalar variables
+    """
+    id_map, _, n_vars, var_shapes = InverseData.get_var_offsets(variables)
+
+    var_dict = {}
+    for var in variables:
+        offset, size = id_map[var.id]
+        shape = var_shapes[var.id]
+        if len(shape) == 2:
+            d1, d2 = shape[0], shape[1]
+        elif len(shape) == 1:
+            d1, d2 = shape[0], 1
+        else:  # scalar
+            d1, d2 = 1, 1
+        c_var = diffengine.make_variable(d1, d2, offset, n_vars)
+        var_dict[var.id] = c_var
+
+    return var_dict, n_vars
+
+
+def _convert_expr(expr, var_dict: dict):
+    """Convert CVXPY expression using pre-built variable dictionary."""
+    # Base case: variable lookup
+    if isinstance(expr, cp.Variable):
+        return var_dict[expr.id]
+
+    # Recursive case: atoms
+    atom_name = type(expr).__name__
+    if atom_name in ATOM_CONVERTERS:
+        children = [_convert_expr(arg, var_dict) for arg in expr.args]
+        converter = ATOM_CONVERTERS[atom_name]
+        # N-ary ops (like AddExpression) take list, unary ops take single arg
+        if atom_name == "AddExpression":
+            return converter(children)
+        return converter(*children) if len(children) > 1 else converter(children[0])
+
+    raise NotImplementedError(f"Atom '{atom_name}' not supported")
+
+
+def convert_problem(problem: cp.Problem) -> tuple:
+    """
+    Convert CVXPY Problem to C expressions.
+
+    Args:
+        problem: CVXPY Problem object
+
+    Returns:
+        c_objective: C expression for objective
+        c_constraints: list of C expressions for constraints
+    """
+    var_dict, _ = build_variable_dict(problem.variables())
+
+    # Convert objective
+    c_objective = _convert_expr(problem.objective.expr, var_dict)
+
+    # Convert constraints (expression part only for now)
+    c_constraints = []
+    for constr in problem.constraints:
+        c_expr = _convert_expr(constr.expr, var_dict)
+        c_constraints.append(c_expr)
+
+    return c_objective, c_constraints
+
+
+def cvxpy_expr_to_C(expr):
+    """
+    Convert a standalone CVXPY expression to C.
+    Collects variables from the expression and builds variable dict.
+    """
+    # Collect variables
+    variables = []
+    seen_ids = set()
+
+    def collect_vars(e):
+        if isinstance(e, cp.Variable):
+            if e.id not in seen_ids:
+                variables.append(e)
+                seen_ids.add(e.id)
+        elif hasattr(e, 'args'):
+            for arg in e.args:
+                collect_vars(arg)
+
+    collect_vars(expr)
+
+    var_dict, _ = build_variable_dict(variables)
+    return _convert_expr(expr, var_dict)
+
+
+# === Problem-based tests ===
+
+def test_problem_simple_sum_log():
+    """Test converting cp.sum(cp.log(x))."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    c_obj, c_constrs = convert_problem(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0])
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_problem_simple_sum_log passed")
+
+
+def test_problem_two_variables():
+    """Test problem with two variables: sum(log(x + y))."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
+    c_obj, c_constrs = convert_problem(problem)
+
+    # Variables flattened: [x0, x1, y0, y1]
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(np.array([1+3, 2+4])))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_problem_two_variables passed")
+
+
+def test_variable_reuse():
+    """Test that same variable used twice works correctly."""
+    x = cp.Variable(2)
+    # log(x) + exp(x) uses x twice
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
+    c_obj, c_constrs = convert_problem(problem)
+
+    test_values = np.array([1.0, 2.0])
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values) + np.exp(test_values))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_variable_reuse passed")
+
+
+# === Standalone expression tests ===
+
+def test_log_conversion():
+    """Test converting CVXPY log expression to C diff engine."""
+    x = cp.Variable(3)
+    c_expr = cvxpy_expr_to_C(cp.log(x))
+
+    test_values = np.array([1.0, 2.0, 3.0])
+    result = diffengine.forward(c_expr, test_values)
+    expected = np.log(test_values)
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_log_conversion passed")
+
+
+def test_exp_conversion():
+    """Test converting CVXPY exp expression to C diff engine."""
+    x = cp.Variable(3)
+    c_expr = cvxpy_expr_to_C(cp.exp(x))
+
+    test_values = np.array([0.0, 1.0, 2.0])
+    result = diffengine.forward(c_expr, test_values)
+    expected = np.exp(test_values)
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_exp_conversion passed")
+
+
+def test_add_conversion():
+    """Test converting CVXPY addition expression to C diff engine."""
+    x = cp.Variable(3)
+    y = cp.Variable(3)
+    c_expr = cvxpy_expr_to_C(x + y)
+
+    test_values = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
+    result = diffengine.forward(c_expr, test_values)
+    expected = test_values[:3] + test_values[3:]
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_add_conversion passed")
+
+
+def test_composite_expression():
+    """Test converting a composite CVXPY expression: log(exp(x) + y)."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    c_expr = cvxpy_expr_to_C(cp.log(cp.exp(x) + y))
+
+    test_values = np.array([1.0, 2.0, 0.5, 1.0])
+    result = diffengine.forward(c_expr, test_values)
+    expected = np.log(np.exp(test_values[:2]) + test_values[2:])
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_composite_expression passed")
+
+
+# === Complex tests ===
+
+def test_four_variables():
+    """Test problem with 4 variables: sum(log(a + b) + exp(c + d))."""
+    a = cp.Variable(3)
+    b = cp.Variable(3)
+    c = cp.Variable(3)
+    d = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a + b) + cp.exp(c + d))))
+    c_obj, _ = convert_problem(problem)
+
+    # Variables flattened: [a0,a1,a2, b0,b1,b2, c0,c1,c2, d0,d1,d2]
+    a_vals = np.array([1.0, 2.0, 3.0])
+    b_vals = np.array([0.5, 1.0, 1.5])
+    c_vals = np.array([0.1, 0.2, 0.3])
+    d_vals = np.array([0.1, 0.1, 0.1])
+    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
+
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(a_vals + b_vals) + np.exp(c_vals + d_vals))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_four_variables passed")
+
+
+def test_deep_nesting():
+    """Test deeply nested composition: log(exp(log(exp(x))))."""
+    x = cp.Variable(4)
+    expr = cp.log(cp.exp(cp.log(cp.exp(x))))
+    c_expr = cvxpy_expr_to_C(expr)
+
+    test_values = np.array([0.5, 1.0, 1.5, 2.0])
+    result = diffengine.forward(c_expr, test_values)
+    # log(exp(log(exp(x)))) = log(exp(x)) = x (for positive x)
+    expected = np.log(np.exp(np.log(np.exp(test_values))))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_deep_nesting passed")
+
+
+def test_chained_additions():
+    """Test multiple chained additions: x + y + z + w."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    z = cp.Variable(2)
+    w = cp.Variable(2)
+    c_expr = cvxpy_expr_to_C(x + y + z + w)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([3.0, 4.0])
+    z_vals = np.array([5.0, 6.0])
+    w_vals = np.array([7.0, 8.0])
+    test_values = np.concatenate([x_vals, y_vals, z_vals, w_vals])
+
+    result = diffengine.forward(c_expr, test_values)
+    expected = x_vals + y_vals + z_vals + w_vals
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_chained_additions passed")
+
+
+def test_variable_used_multiple_times():
+    """Test variable used 3+ times: log(x) + exp(x) + x."""
+    x = cp.Variable(3)
+    expr = cp.log(x) + cp.exp(x) + x
+    c_expr = cvxpy_expr_to_C(expr)
+
+    test_values = np.array([1.0, 2.0, 3.0])
+    result = diffengine.forward(c_expr, test_values)
+    expected = np.log(test_values) + np.exp(test_values) + test_values
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_variable_used_multiple_times passed")
+
+
+def test_larger_variable_size():
+    """Test with larger variable (100 elements)."""
+    x = cp.Variable(100)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.linspace(0.1, 2.0, 100)
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(np.exp(test_values)))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_larger_variable_size passed")
+
+
+def test_matrix_variable():
+    """Test with 2D matrix variable (3x4)."""
+    X = cp.Variable((3, 4))
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
+    c_obj, _ = convert_problem(problem)
+
+    # Matrix flattened in column-major (Fortran) order by CVXPY
+    test_values = np.arange(1.0, 13.0)  # 12 elements
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_matrix_variable passed")
+
+
+def test_mixed_sizes():
+    """Test with variables of different sizes."""
+    a = cp.Variable(2)
+    b = cp.Variable(5)
+    c = cp.Variable(3)
+
+    # sum of each variable's log, then sum all
+    expr = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))
+    c_expr = cvxpy_expr_to_C(expr)
+
+    a_vals = np.array([1.0, 2.0])
+    b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+    c_vals = np.array([1.0, 2.0, 3.0])
+    test_values = np.concatenate([a_vals, b_vals, c_vals])
+
+    result = diffengine.forward(c_expr, test_values)
+    expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_mixed_sizes passed")
+
+
+def test_complex_objective():
+    """Test complex objective: sum(log(x + y)) + sum(exp(y + z)) + sum(log(z + x))."""
+    x = cp.Variable(3)
+    y = cp.Variable(3)
+    z = cp.Variable(3)
+
+    obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_problem(problem)
+
+    x_vals = np.array([1.0, 2.0, 3.0])
+    y_vals = np.array([0.5, 1.0, 1.5])
+    z_vals = np.array([0.2, 0.3, 0.4])
+    test_values = np.concatenate([x_vals, y_vals, z_vals])
+
+    result = diffengine.forward(c_obj, test_values)
+    expected = (np.sum(np.log(x_vals + y_vals)) +
+                np.sum(np.exp(y_vals + z_vals)) +
+                np.sum(np.log(z_vals + x_vals)))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_complex_objective passed")
+
+
+def test_nested_sums():
+    """Test sum of sum (should collapse to single sum)."""
+    x = cp.Variable(4)
+    # sum(sum(log(x))) - outer sum is over a scalar, so effectively just sum(log(x))
+    expr = cp.sum(cp.log(x))
+    c_expr = cvxpy_expr_to_C(expr)
+
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+    result = diffengine.forward(c_expr, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_nested_sums passed")
+
+
+if __name__ == "__main__":
+    # Problem-based tests
+    test_problem_simple_sum_log()
+    test_problem_two_variables()
+    test_variable_reuse()
+
+    # Standalone expression tests
+    test_log_conversion()
+    test_exp_conversion()
+    test_add_conversion()
+    test_composite_expression()
+
+    # Complex tests
+    test_four_variables()
+    test_deep_nesting()
+    test_chained_additions()
+    test_variable_used_multiple_times()
+    test_larger_variable_size()
+    test_matrix_variable()
+    test_mixed_sizes()
+    test_complex_objective()
+    test_nested_sums()
+
+    print("\nAll tests passed!")
diff --git a/python/test_inverse_data.py b/python/test_inverse_data.py
new file mode 100644
index 0000000..82150ad
--- /dev/null
+++ b/python/test_inverse_data.py
@@ -0,0 +1,19 @@
+import cvxpy as cp
+from cvxpy.reductions.inverse_data import InverseData
+from cvxpy.utilities.coeff_extractor import CoeffExtractor
+
+x = cp.Variable((2,))
+x.value = [1.0, 2.0]
+y = cp.Variable()
+y.value = 4.0
+z = cp.Variable((2,2))
+expr = cp.sum(z) + cp.sum(x) + y + 5 + 2*y + x[0] * 3
+constr = [x + y <= 5, z == 0]
+
+problem = cp.Problem(cp.Minimize(expr), constr)
+inverse_data = InverseData(problem)
+print(inverse_data.var_offsets)
+#extractor = CoeffExtractor(inverse_data, cp.SCIPY_CANON_BACKEND)
+#print(str(constr[0].expr))
+#A = extractor.affine([expr, constr[0].expr])
+#print(A.toarray().reshape((3,4), order='F'))

From c5eca4aaf40b5b0f7fa46006bc2232f9c746d65e Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Sat, 3 Jan 2026 18:14:23 -0500
Subject: [PATCH 2/7] remove test inverse data

---
 python/test_inverse_data.py | 19 -------------------
 1 file changed, 19 deletions(-)
 delete mode 100644 python/test_inverse_data.py

diff --git a/python/test_inverse_data.py b/python/test_inverse_data.py
deleted file mode 100644
index 82150ad..0000000
--- a/python/test_inverse_data.py
+++ /dev/null
@@ -1,19 +0,0 @@
-import cvxpy as cp
-from cvxpy.reductions.inverse_data import InverseData
-from cvxpy.utilities.coeff_extractor import CoeffExtractor
-
-x = cp.Variable((2,))
-x.value = [1.0, 2.0]
-y = cp.Variable()
-y.value = 4.0
-z = cp.Variable((2,2))
-expr = cp.sum(z) + cp.sum(x) + y + 5 + 2*y + x[0] * 3
-constr = [x + y <= 5, z == 0]
-
-problem = cp.Problem(cp.Minimize(expr), constr)
-inverse_data = InverseData(problem)
-print(inverse_data.var_offsets)
-#extractor = CoeffExtractor(inverse_data, cp.SCIPY_CANON_BACKEND)
-#print(str(constr[0].expr))
-#A = extractor.affine([expr, constr[0].expr])
-#print(A.toarray().reshape((3,4), order='F'))

From ca03bccf00db6d88baae7fca793a4a92d942881d Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Sat, 3 Jan 2026 19:54:14 -0500
Subject: [PATCH 3/7] adds basic jacobian bindings and tests

---
 python/bindings.c |  60 +++++++++
 python/convert.py | 325 +++++++++++++++++++++++++++-------------------
 python/example.py |   9 --
 3 files changed, 250 insertions(+), 144 deletions(-)
 delete mode 100644 python/example.py

diff --git a/python/bindings.c b/python/bindings.c
index 2f4cb67..68e4b76 100644
--- a/python/bindings.c
+++ b/python/bindings.c
@@ -178,6 +178,65 @@ static PyObject *py_forward(PyObject *self, PyObject *args)
     return out;
 }
 
+static PyObject *py_jacobian(PyObject *self, PyObject *args)
+{
+    PyObject *node_capsule;
+    PyObject *u_obj;
+    if (!PyArg_ParseTuple(args, "OO", &node_capsule, &u_obj))
+    {
+        return NULL;
+    }
+
+    expr *node = (expr *) PyCapsule_GetPointer(node_capsule, EXPR_CAPSULE_NAME);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid node capsule");
+        return NULL;
+    }
+
+    PyArrayObject *u_array =
+        (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!u_array)
+    {
+        return NULL;
+    }
+
+    // Run forward pass first (required before jacobian)
+    node->forward(node, (const double *) PyArray_DATA(u_array));
+
+    // Initialize and evaluate jacobian
+    node->jacobian_init(node);
+    node->eval_jacobian(node);
+
+    CSR_Matrix *jac = node->jacobian;
+
+    // Create numpy arrays for CSR components
+    npy_intp nnz = jac->nnz;
+    npy_intp m_plus_1 = jac->m + 1;
+
+    PyObject *data = PyArray_SimpleNew(1, &nnz, NPY_DOUBLE);
+    PyObject *indices = PyArray_SimpleNew(1, &nnz, NPY_INT32);
+    PyObject *indptr = PyArray_SimpleNew(1, &m_plus_1, NPY_INT32);
+
+    if (!data || !indices || !indptr)
+    {
+        Py_XDECREF(data);
+        Py_XDECREF(indices);
+        Py_XDECREF(indptr);
+        Py_DECREF(u_array);
+        return NULL;
+    }
+
+    memcpy(PyArray_DATA((PyArrayObject *) data), jac->x, nnz * sizeof(double));
+    memcpy(PyArray_DATA((PyArrayObject *) indices), jac->i, nnz * sizeof(int));
+    memcpy(PyArray_DATA((PyArrayObject *) indptr), jac->p, m_plus_1 * sizeof(int));
+
+    Py_DECREF(u_array);
+
+    // Return tuple: (data, indices, indptr, shape)
+    return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
+}
+
 static PyMethodDef DNLPMethods[] = {
     {"make_variable", py_make_variable, METH_VARARGS, "Create variable node"},
     {"make_log", py_make_log, METH_VARARGS, "Create log node"},
@@ -185,6 +244,7 @@ static PyMethodDef DNLPMethods[] = {
     {"make_add", py_make_add, METH_VARARGS, "Create add node"},
     {"make_sum", py_make_sum, METH_VARARGS, "Create sum node"},
     {"forward", py_forward, METH_VARARGS, "Run forward pass and return values"},
+    {"jacobian", py_jacobian, METH_VARARGS, "Compute jacobian and return CSR components"},
     {NULL, NULL, 0, NULL}};
 
 static struct PyModuleDef dnlp_module = {PyModuleDef_HEAD_INIT, "DNLP_diff_engine",
diff --git a/python/convert.py b/python/convert.py
index 385b2fc..b8f4bc1 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -1,9 +1,16 @@
 import cvxpy as cp
 import numpy as np
+from scipy import sparse
 from cvxpy.reductions.inverse_data import InverseData
 import DNLP_diff_engine as diffengine
 
 
+def get_jacobian(c_expr, values):
+    """Compute jacobian and return as scipy sparse CSR matrix."""
+    data, indices, indptr, shape = diffengine.jacobian(c_expr, values)
+    return sparse.csr_matrix((data, indices, indptr), shape=shape)
+
+
 def _chain_add(children):
     """Chain multiple children with binary adds: a + b + c -> add(add(a, b), c)."""
     result = children[0]
@@ -41,7 +48,7 @@ def build_variable_dict(variables: list) -> tuple[dict, int]:
 
     var_dict = {}
     for var in variables:
-        offset, size = id_map[var.id]
+        offset, _ = id_map[var.id]
         shape = var_shapes[var.id]
         if len(shape) == 2:
             d1, d2 = shape[0], shape[1]
@@ -52,7 +59,7 @@ def build_variable_dict(variables: list) -> tuple[dict, int]:
         c_var = diffengine.make_variable(d1, d2, offset, n_vars)
         var_dict[var.id] = c_var
 
-    return var_dict, n_vars
+    return var_dict
 
 
 def _convert_expr(expr, var_dict: dict):
@@ -85,7 +92,7 @@ def convert_problem(problem: cp.Problem) -> tuple:
         c_objective: C expression for objective
         c_constraints: list of C expressions for constraints
     """
-    var_dict, _ = build_variable_dict(problem.variables())
+    var_dict = build_variable_dict(problem.variables())
 
     # Convert objective
     c_objective = _convert_expr(problem.objective.expr, var_dict)
@@ -99,66 +106,40 @@ def convert_problem(problem: cp.Problem) -> tuple:
     return c_objective, c_constraints
 
 
-def cvxpy_expr_to_C(expr):
-    """
-    Convert a standalone CVXPY expression to C.
-    Collects variables from the expression and builds variable dict.
-    """
-    # Collect variables
-    variables = []
-    seen_ids = set()
-
-    def collect_vars(e):
-        if isinstance(e, cp.Variable):
-            if e.id not in seen_ids:
-                variables.append(e)
-                seen_ids.add(e.id)
-        elif hasattr(e, 'args'):
-            for arg in e.args:
-                collect_vars(arg)
-
-    collect_vars(expr)
-
-    var_dict, _ = build_variable_dict(variables)
-    return _convert_expr(expr, var_dict)
+# === Tests ===
 
-
-# === Problem-based tests ===
-
-def test_problem_simple_sum_log():
+def test_simple_sum_log():
     """Test converting cp.sum(cp.log(x))."""
     x = cp.Variable(3)
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
-    c_obj, c_constrs = convert_problem(problem)
+    c_obj, _ = convert_problem(problem)
 
     test_values = np.array([1.0, 2.0, 3.0])
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(test_values))
     assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_problem_simple_sum_log passed")
+    print("test_simple_sum_log passed")
 
 
-def test_problem_two_variables():
+def test_two_variables():
     """Test problem with two variables: sum(log(x + y))."""
     x = cp.Variable(2)
     y = cp.Variable(2)
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
-    c_obj, c_constrs = convert_problem(problem)
+    c_obj, _ = convert_problem(problem)
 
-    # Variables flattened: [x0, x1, y0, y1]
     test_values = np.array([1.0, 2.0, 3.0, 4.0])
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(np.array([1+3, 2+4])))
     assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_problem_two_variables passed")
+    print("test_two_variables passed")
 
 
 def test_variable_reuse():
     """Test that same variable used twice works correctly."""
     x = cp.Variable(2)
-    # log(x) + exp(x) uses x twice
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
-    c_obj, c_constrs = convert_problem(problem)
+    c_obj, _ = convert_problem(problem)
 
     test_values = np.array([1.0, 2.0])
     result = diffengine.forward(c_obj, test_values)
@@ -167,60 +148,6 @@ def test_variable_reuse():
     print("test_variable_reuse passed")
 
 
-# === Standalone expression tests ===
-
-def test_log_conversion():
-    """Test converting CVXPY log expression to C diff engine."""
-    x = cp.Variable(3)
-    c_expr = cvxpy_expr_to_C(cp.log(x))
-
-    test_values = np.array([1.0, 2.0, 3.0])
-    result = diffengine.forward(c_expr, test_values)
-    expected = np.log(test_values)
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_log_conversion passed")
-
-
-def test_exp_conversion():
-    """Test converting CVXPY exp expression to C diff engine."""
-    x = cp.Variable(3)
-    c_expr = cvxpy_expr_to_C(cp.exp(x))
-
-    test_values = np.array([0.0, 1.0, 2.0])
-    result = diffengine.forward(c_expr, test_values)
-    expected = np.exp(test_values)
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_exp_conversion passed")
-
-
-def test_add_conversion():
-    """Test converting CVXPY addition expression to C diff engine."""
-    x = cp.Variable(3)
-    y = cp.Variable(3)
-    c_expr = cvxpy_expr_to_C(x + y)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
-    result = diffengine.forward(c_expr, test_values)
-    expected = test_values[:3] + test_values[3:]
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_add_conversion passed")
-
-
-def test_composite_expression():
-    """Test converting a composite CVXPY expression: log(exp(x) + y)."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    c_expr = cvxpy_expr_to_C(cp.log(cp.exp(x) + y))
-
-    test_values = np.array([1.0, 2.0, 0.5, 1.0])
-    result = diffengine.forward(c_expr, test_values)
-    expected = np.log(np.exp(test_values[:2]) + test_values[2:])
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_composite_expression passed")
-
-
-# === Complex tests ===
-
 def test_four_variables():
     """Test problem with 4 variables: sum(log(a + b) + exp(c + d))."""
     a = cp.Variable(3)
@@ -230,7 +157,6 @@ def test_four_variables():
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a + b) + cp.exp(c + d))))
     c_obj, _ = convert_problem(problem)
 
-    # Variables flattened: [a0,a1,a2, b0,b1,b2, c0,c1,c2, d0,d1,d2]
     a_vals = np.array([1.0, 2.0, 3.0])
     b_vals = np.array([0.5, 1.0, 1.5])
     c_vals = np.array([0.1, 0.2, 0.3])
@@ -244,26 +170,26 @@ def test_four_variables():
 
 
 def test_deep_nesting():
-    """Test deeply nested composition: log(exp(log(exp(x))))."""
+    """Test deeply nested composition: sum(log(exp(log(exp(x)))))."""
     x = cp.Variable(4)
-    expr = cp.log(cp.exp(cp.log(cp.exp(x))))
-    c_expr = cvxpy_expr_to_C(expr)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(cp.log(cp.exp(x)))))))
+    c_obj, _ = convert_problem(problem)
 
     test_values = np.array([0.5, 1.0, 1.5, 2.0])
-    result = diffengine.forward(c_expr, test_values)
-    # log(exp(log(exp(x)))) = log(exp(x)) = x (for positive x)
-    expected = np.log(np.exp(np.log(np.exp(test_values))))
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(np.exp(np.log(np.exp(test_values)))))
     assert np.allclose(result, expected), f"Expected {expected}, got {result}"
     print("test_deep_nesting passed")
 
 
 def test_chained_additions():
-    """Test multiple chained additions: x + y + z + w."""
+    """Test multiple chained additions: sum(x + y + z + w)."""
     x = cp.Variable(2)
     y = cp.Variable(2)
     z = cp.Variable(2)
     w = cp.Variable(2)
-    c_expr = cvxpy_expr_to_C(x + y + z + w)
+    problem = cp.Problem(cp.Minimize(cp.sum(x + y + z + w)))
+    c_obj, _ = convert_problem(problem)
 
     x_vals = np.array([1.0, 2.0])
     y_vals = np.array([3.0, 4.0])
@@ -271,21 +197,21 @@ def test_chained_additions():
     w_vals = np.array([7.0, 8.0])
     test_values = np.concatenate([x_vals, y_vals, z_vals, w_vals])
 
-    result = diffengine.forward(c_expr, test_values)
-    expected = x_vals + y_vals + z_vals + w_vals
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(x_vals + y_vals + z_vals + w_vals)
     assert np.allclose(result, expected), f"Expected {expected}, got {result}"
     print("test_chained_additions passed")
 
 
 def test_variable_used_multiple_times():
-    """Test variable used 3+ times: log(x) + exp(x) + x."""
+    """Test variable used 3+ times: sum(log(x) + exp(x) + x)."""
     x = cp.Variable(3)
-    expr = cp.log(x) + cp.exp(x) + x
-    c_expr = cvxpy_expr_to_C(expr)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + x)))
+    c_obj, _ = convert_problem(problem)
 
     test_values = np.array([1.0, 2.0, 3.0])
-    result = diffengine.forward(c_expr, test_values)
-    expected = np.log(test_values) + np.exp(test_values) + test_values
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values) + np.exp(test_values) + test_values)
     assert np.allclose(result, expected), f"Expected {expected}, got {result}"
     print("test_variable_used_multiple_times passed")
 
@@ -309,7 +235,6 @@ def test_matrix_variable():
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
     c_obj, _ = convert_problem(problem)
 
-    # Matrix flattened in column-major (Fortran) order by CVXPY
     test_values = np.arange(1.0, 13.0)  # 12 elements
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(test_values))
@@ -322,17 +247,15 @@ def test_mixed_sizes():
     a = cp.Variable(2)
     b = cp.Variable(5)
     c = cp.Variable(3)
-
-    # sum of each variable's log, then sum all
-    expr = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))
-    c_expr = cvxpy_expr_to_C(expr)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))))
+    c_obj, _ = convert_problem(problem)
 
     a_vals = np.array([1.0, 2.0])
     b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
     c_vals = np.array([1.0, 2.0, 3.0])
     test_values = np.concatenate([a_vals, b_vals, c_vals])
 
-    result = diffengine.forward(c_expr, test_values)
+    result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
     assert np.allclose(result, expected), f"Expected {expected}, got {result}"
     print("test_mixed_sizes passed")
@@ -343,7 +266,6 @@ def test_complex_objective():
     x = cp.Variable(3)
     y = cp.Variable(3)
     z = cp.Variable(3)
-
     obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
     problem = cp.Problem(cp.Minimize(obj))
     c_obj, _ = convert_problem(problem)
@@ -361,33 +283,157 @@ def test_complex_objective():
     print("test_complex_objective passed")
 
 
-def test_nested_sums():
-    """Test sum of sum (should collapse to single sum)."""
+def test_log_exp_identity():
+    """Test log(exp(x)) = x identity."""
+    x = cp.Variable(5)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([-1.0, 0.0, 1.0, 2.0, 3.0])
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(test_values)  # log(exp(x)) = x
+    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
+    print("test_log_exp_identity passed")
+
+
+# === Jacobian Tests ===
+
+def test_jacobian_sum_log():
+    """Test jacobian of sum(log(x)). Gradient is [1/x_1, 1/x_2, ...]."""
     x = cp.Variable(4)
-    # sum(sum(log(x))) - outer sum is over a scalar, so effectively just sum(log(x))
-    expr = cp.sum(cp.log(x))
-    c_expr = cvxpy_expr_to_C(expr)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    c_obj, _ = convert_problem(problem)
 
     test_values = np.array([1.0, 2.0, 3.0, 4.0])
-    result = diffengine.forward(c_expr, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_nested_sums passed")
+    jac = get_jacobian(c_obj, test_values)
 
+    # sum(log(x)) is scalar, so jacobian is 1 x n
+    expected = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
+    print("test_jacobian_sum_log passed")
 
-if __name__ == "__main__":
-    # Problem-based tests
-    test_problem_simple_sum_log()
-    test_problem_two_variables()
-    test_variable_reuse()
 
-    # Standalone expression tests
-    test_log_conversion()
-    test_exp_conversion()
-    test_add_conversion()
-    test_composite_expression()
+def test_jacobian_sum_exp():
+    """Test jacobian of sum(exp(x)). Gradient is [exp(x_1), exp(x_2), ...]."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([0.0, 1.0, 2.0])
+    jac = get_jacobian(c_obj, test_values)
+
+    expected = np.exp(test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
+    print("test_jacobian_sum_exp passed")
+
+
+def test_jacobian_four_variables():
+    """Test jacobian of sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
+    a = cp.Variable(2)
+    b = cp.Variable(2)
+    c = cp.Variable(2)
+    d = cp.Variable(2)
+    obj = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.exp(c)) + cp.sum(cp.exp(d))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_problem(problem)
+
+    # Variables: [a0, a1, b0, b1, c0, c1, d0, d1]
+    a_vals = np.array([1.0, 2.0])
+    b_vals = np.array([0.5, 1.0])
+    c_vals = np.array([0.1, 0.2])
+    d_vals = np.array([0.1, 0.1])
+    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
+    jac = get_jacobian(c_obj, test_values)
+
+    # f = sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))
+    df_da = 1.0 / a_vals
+    df_db = 1.0 / b_vals
+    df_dc = np.exp(c_vals)
+    df_dd = np.exp(d_vals)
+    expected = np.concatenate([df_da, df_db, df_dc, df_dd]).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
+    print("test_jacobian_four_variables passed")
+
+
+def test_jacobian_two_variables():
+    """Test jacobian of sum(log(x) + log(y)) with two variables."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y))))
+    c_obj, _ = convert_problem(problem)
+
+    # Variables: [x0, x1, y0, y1]
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+    jac = get_jacobian(c_obj, test_values)
+
+    # f = sum(log(x) + log(y)) = log(x0) + log(x1) + log(y0) + log(y1)
+    # df/dx = [1/x0, 1/x1], df/dy = [1/y0, 1/y1]
+    expected = np.array([[1/1, 1/2, 1/3, 1/4]])
+    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
+    print("test_jacobian_two_variables passed")
 
-    # Complex tests
+
+def test_jacobian_variable_reuse():
+    """Test jacobian when same variable is used multiple times."""
+    x = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([1.0, 2.0])
+    jac = get_jacobian(c_obj, test_values)
+
+    # f = sum(log(x) + exp(x))
+    # df/dx_i = 1/x_i + exp(x_i)
+    expected = (1.0 / test_values + np.exp(test_values)).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
+    print("test_jacobian_variable_reuse passed")
+
+
+def test_jacobian_large_variable():
+    """Test jacobian of sum(log(x)) with larger variable."""
+    x = cp.Variable(10)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.linspace(1.0, 10.0, 10)
+    jac = get_jacobian(c_obj, test_values)
+
+    # f = sum(log(x)), df/dx_i = 1/x_i
+    expected = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
+    print("test_jacobian_large_variable passed")
+
+
+def test_jacobian_complex_objective():
+    """Test jacobian of sum(log(x) + exp(y) + log(z))."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    z = cp.Variable(2)
+    obj = cp.sum(cp.log(x) + cp.exp(y) + cp.log(z))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_problem(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([0.5, 1.0])
+    z_vals = np.array([0.2, 0.3])
+    test_values = np.concatenate([x_vals, y_vals, z_vals])
+    jac = get_jacobian(c_obj, test_values)
+
+    # f = sum(log(x) + exp(y) + log(z))
+    # df/dx_i = 1/x_i, df/dy_i = exp(y_i), df/dz_i = 1/z_i
+    df_dx = 1.0 / x_vals
+    df_dy = np.exp(y_vals)
+    df_dz = 1.0 / z_vals
+    expected = np.concatenate([df_dx, df_dy, df_dz]).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
+    print("test_jacobian_complex_objective passed")
+
+
+if __name__ == "__main__":
+    # Forward pass tests
+    test_simple_sum_log()
+    test_two_variables()
+    test_variable_reuse()
     test_four_variables()
     test_deep_nesting()
     test_chained_additions()
@@ -396,6 +442,15 @@ def test_nested_sums():
     test_matrix_variable()
     test_mixed_sizes()
     test_complex_objective()
-    test_nested_sums()
+    test_log_exp_identity()
+
+    # Jacobian tests
+    test_jacobian_sum_log()
+    test_jacobian_sum_exp()
+    test_jacobian_four_variables()
+    test_jacobian_two_variables()
+    test_jacobian_variable_reuse()
+    test_jacobian_large_variable()
+    test_jacobian_complex_objective()
 
     print("\nAll tests passed!")
diff --git a/python/example.py b/python/example.py
deleted file mode 100644
index 37b70aa..0000000
--- a/python/example.py
+++ /dev/null
@@ -1,9 +0,0 @@
-import numpy as np
-import DNLP_diff_engine as dnlp
-
-x = dnlp.make_variable(3, 1, 0, 3)
-log_x = dnlp.make_log(x)
-
-u = np.array([1.0, 2.0, 3.0])
-out = dnlp.forward(log_x, u)
-print("log(x) forward:", out)

From 9b4076eee9661a61923084b1d76250063cb13621 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Sat, 3 Jan 2026 20:49:22 -0500
Subject: [PATCH 4/7] adding docs for problem struct

---
 docs/problem_struct_design.md | 205 ++++++++++++++++++++++++++++++++++
 1 file changed, 205 insertions(+)
 create mode 100644 docs/problem_struct_design.md

diff --git a/docs/problem_struct_design.md b/docs/problem_struct_design.md
new file mode 100644
index 0000000..7a3cd09
--- /dev/null
+++ b/docs/problem_struct_design.md
@@ -0,0 +1,205 @@
+# Design: C Problem Struct
+
+## Summary
+
+Create a native C `problem` struct that encapsulates objective + constraints, with methods for:
+- `forward(u)` - evaluate objective and all constraints
+- `gradient(u)` - return objective gradient (jacobian.T)
+- `jacobian(u)` - return single stacked CSR matrix of constraint jacobians
+
+## Files to Create/Modify
+
+1. **`include/problem.h`** - New header defining problem struct
+2. **`src/problem.c`** - Implementation
+3. **`python/bindings.c`** - Python bindings for problem
+4. **`python/convert.py`** - Update to return problem capsule
+5. **`tests/problem/test_problem.h`** - C tests
+
+---
+
+## Step 1: Create `include/problem.h`
+
+```c
+#ifndef PROBLEM_H
+#define PROBLEM_H
+
+#include "expr.h"
+#include "utils/CSR_Matrix.h"
+
+typedef struct problem
+{
+    expr *objective;           /* Objective expression (scalar) */
+    expr **constraints;        /* Array of constraint expressions */
+    int n_constraints;
+    int n_vars;
+    int total_constraint_size; /* Sum of all constraint sizes */
+
+    /* Pre-allocated storage */
+    double *constraint_values;
+    double *gradient_values;   /* Dense gradient array */
+    CSR_Matrix *stacked_jac;
+} problem;
+
+problem *new_problem(expr *objective, expr **constraints, int n_constraints);
+void free_problem(problem *prob);
+double problem_forward(problem *prob, const double *u);
+double *problem_gradient(problem *prob, const double *u);
+CSR_Matrix *problem_jacobian(problem *prob, const double *u);
+
+#endif
+```
+
+---
+
+## Step 2: Create `src/problem.c`
+
+Key functions:
+
+### `new_problem`
+- Retain (increment refcount) on objective and all constraints
+- Compute `total_constraint_size = sum(constraints[i]->size)`
+- Pre-allocate `constraint_values` and `gradient_values` arrays
+
+### `free_problem`
+- Call `free_expr` on objective and all constraints (decrements refcount)
+- Free allocated arrays and stacked_jac
+
+### `problem_forward`
+```c
+double problem_forward(problem *prob, const double *u)
+{
+    prob->objective->forward(prob->objective, u);
+    double obj_val = prob->objective->value[0];
+
+    int offset = 0;
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr *c = prob->constraints[i];
+        c->forward(c, u);
+        memcpy(prob->constraint_values + offset, c->value, c->size * sizeof(double));
+        offset += c->size;
+    }
+    return obj_val;
+}
+```
+
+### `problem_gradient`
+- Run forward pass on objective
+- Call jacobian_init + eval_jacobian
+- Objective jacobian is 1 x n_vars row vector
+- Copy sparse row to dense `gradient_values` array
+- Return pointer to internal array
+
+### `problem_jacobian`
+- Forward + jacobian for each constraint
+- Stack CSR matrices vertically:
+  - Total rows = `total_constraint_size`
+  - Copy row pointers with offset, copy col indices and values
+- Lazy allocate/reallocate `stacked_jac` based on total nnz
+
+---
+
+## Step 3: Update `python/bindings.c`
+
+Add capsule and functions:
+
+```c
+#define PROBLEM_CAPSULE_NAME "DNLP_PROBLEM"
+
+static void problem_capsule_destructor(PyObject *capsule) { ... }
+
+static PyObject *py_make_problem(PyObject *self, PyObject *args)
+{
+    PyObject *obj_capsule, *constraints_list;
+    // Parse objective capsule and list of constraint capsules
+    // Extract expr* pointers, call new_problem
+    // Return PyCapsule
+}
+
+static PyObject *py_problem_forward(PyObject *self, PyObject *args)
+{
+    // Returns: (obj_value, constraint_values_array)
+}
+
+static PyObject *py_problem_gradient(PyObject *self, PyObject *args)
+{
+    // Returns: numpy array of size n_vars
+}
+
+static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
+{
+    // Returns: (data, indices, indptr, (m, n)) for scipy CSR
+}
+```
+
+Add to DNLPMethods:
+```c
+{"make_problem", py_make_problem, METH_VARARGS, "Create problem"},
+{"problem_forward", py_problem_forward, METH_VARARGS, "Evaluate problem"},
+{"problem_gradient", py_problem_gradient, METH_VARARGS, "Compute gradient"},
+{"problem_jacobian", py_problem_jacobian, METH_VARARGS, "Compute constraint jacobian"},
+```
+
+---
+
+## Step 4: Update `python/convert.py`
+
+```python
+def convert_problem(problem: cp.Problem):
+    """Convert CVXPY Problem to C problem struct."""
+    var_dict = build_variable_dict(problem.variables())
+
+    c_objective = _convert_expr(problem.objective.expr, var_dict)
+    c_constraints = [_convert_expr(c.expr, var_dict) for c in problem.constraints]
+
+    return diffengine.make_problem(c_objective, c_constraints)
+
+
+class Problem:
+    """Wrapper for C problem struct with clean Python API."""
+
+    def __init__(self, cvxpy_problem: cp.Problem):
+        self._capsule = convert_problem(cvxpy_problem)
+
+    def forward(self, u: np.ndarray) -> tuple[float, np.ndarray]:
+        return diffengine.problem_forward(self._capsule, u)
+
+    def gradient(self, u: np.ndarray) -> np.ndarray:
+        return diffengine.problem_gradient(self._capsule, u)
+
+    def jacobian(self, u: np.ndarray) -> sparse.csr_matrix:
+        data, indices, indptr, shape = diffengine.problem_jacobian(self._capsule, u)
+        return sparse.csr_matrix((data, indices, indptr), shape=shape)
+```
+
+---
+
+## Step 5: Add Tests
+
+### C tests in `tests/problem/test_problem.h`:
+- `test_problem_forward` - verify objective and constraint values
+- `test_problem_gradient` - verify gradient matches manual calculation
+- `test_problem_jacobian_stacking` - verify stacked matrix structure
+
+### Python tests in `convert.py`:
+- `test_problem_forward` - compare with numpy
+- `test_problem_gradient` - gradient of sum(log(x)) = 1/x
+- `test_problem_jacobian` - verify stacked jacobian shape and values
+
+---
+
+## Implementation Order
+
+1. Create `include/problem.h`
+2. Create `src/problem.c` with new_problem, free_problem, problem_forward
+3. Add problem_gradient and problem_jacobian
+4. Add Python bindings to `bindings.c`
+5. Rebuild: `cmake --build build`
+6. Update `convert.py` with Problem class
+7. Add and run tests
+
+## Key Design Notes
+
+- **Memory**: Uses expr refcounting - new_problem retains, free_problem releases
+- **Efficiency**: Pre-allocates arrays, lazy-allocates stacked jacobian
+- **API**: Returns internal pointers (caller should NOT free)

From 55faf29f6a083645f3111ced28298b8ab00af2d1 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Sat, 3 Jan 2026 21:11:54 -0500
Subject: [PATCH 5/7] cleanup tests to combine jac and forward mode

---
 python/convert.py | 343 +++++++++++++++++++++++-----------------------
 1 file changed, 174 insertions(+), 169 deletions(-)

diff --git a/python/convert.py b/python/convert.py
index b8f4bc1..cc34d6e 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -108,48 +108,82 @@ def convert_problem(problem: cp.Problem) -> tuple:
 
 # === Tests ===
 
-def test_simple_sum_log():
-    """Test converting cp.sum(cp.log(x))."""
-    x = cp.Variable(3)
+def test_sum_log():
+    """Test sum(log(x)) forward and jacobian."""
+    x = cp.Variable(4)
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
     c_obj, _ = convert_problem(problem)
 
-    test_values = np.array([1.0, 2.0, 3.0])
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_simple_sum_log passed")
+    assert np.allclose(result, expected)
 
+    # Jacobian: d/dx sum(log(x)) = [1/x_1, 1/x_2, ...]
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_sum_exp():
+    """Test sum(exp(x)) forward and jacobian."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([0.0, 1.0, 2.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.exp(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian: d/dx sum(exp(x)) = [exp(x_1), exp(x_2), ...]
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = np.exp(test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
 
-def test_two_variables():
-    """Test problem with two variables: sum(log(x + y))."""
+
+def test_two_variables_elementwise_add():
+    """Test sum(log(x + y)) - elementwise after add."""
     x = cp.Variable(2)
     y = cp.Variable(2)
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
     c_obj, _ = convert_problem(problem)
 
     test_values = np.array([1.0, 2.0, 3.0, 4.0])
+
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(np.array([1+3, 2+4])))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_two_variables passed")
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
 
 
 def test_variable_reuse():
-    """Test that same variable used twice works correctly."""
+    """Test sum(log(x) + exp(x)) - same variable used twice."""
     x = cp.Variable(2)
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
     c_obj, _ = convert_problem(problem)
 
     test_values = np.array([1.0, 2.0])
+
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(test_values) + np.exp(test_values))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_variable_reuse passed")
+    assert np.allclose(result, expected)
+
+    # Jacobian: d/dx_i = 1/x_i + exp(x_i)
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values + np.exp(test_values)).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
 
 
-def test_four_variables():
-    """Test problem with 4 variables: sum(log(a + b) + exp(c + d))."""
+def test_four_variables_elementwise_add():
+    """Test sum(log(a + b) + exp(c + d)) - elementwise after add."""
     a = cp.Variable(3)
     b = cp.Variable(3)
     c = cp.Variable(3)
@@ -163,27 +197,32 @@ def test_four_variables():
     d_vals = np.array([0.1, 0.1, 0.1])
     test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
 
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(a_vals + b_vals) + np.exp(c_vals + d_vals))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_four_variables passed")
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
 
 
 def test_deep_nesting():
-    """Test deeply nested composition: sum(log(exp(log(exp(x)))))."""
+    """Test sum(log(exp(log(exp(x))))) - deeply nested elementwise."""
     x = cp.Variable(4)
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(cp.log(cp.exp(x)))))))
     c_obj, _ = convert_problem(problem)
 
     test_values = np.array([0.5, 1.0, 1.5, 2.0])
+
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(np.exp(np.log(np.exp(test_values)))))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_deep_nesting passed")
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for nested elementwise compositions not yet supported
 
 
 def test_chained_additions():
-    """Test multiple chained additions: sum(x + y + z + w)."""
+    """Test sum(x + y + z + w) - chained additions."""
     x = cp.Variable(2)
     y = cp.Variable(2)
     z = cp.Variable(2)
@@ -197,53 +236,70 @@ def test_chained_additions():
     w_vals = np.array([7.0, 8.0])
     test_values = np.concatenate([x_vals, y_vals, z_vals, w_vals])
 
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(x_vals + y_vals + z_vals + w_vals)
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_chained_additions passed")
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for sum(add(...)) patterns not yet supported
 
 
 def test_variable_used_multiple_times():
-    """Test variable used 3+ times: sum(log(x) + exp(x) + x)."""
+    """Test sum(log(x) + exp(x) + x) - variable used 3+ times."""
     x = cp.Variable(3)
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + x)))
     c_obj, _ = convert_problem(problem)
 
     test_values = np.array([1.0, 2.0, 3.0])
+
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(test_values) + np.exp(test_values) + test_values)
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_variable_used_multiple_times passed")
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for expressions with sum(variable) not yet supported
 
 
-def test_larger_variable_size():
-    """Test with larger variable (100 elements)."""
+def test_larger_variable():
+    """Test sum(log(x)) with larger variable (100 elements)."""
     x = cp.Variable(100)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
     c_obj, _ = convert_problem(problem)
 
-    test_values = np.linspace(0.1, 2.0, 100)
+    test_values = np.linspace(1.0, 10.0, 100)
+
+    # Forward
     result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(np.exp(test_values)))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_larger_variable_size passed")
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
 
 
 def test_matrix_variable():
-    """Test with 2D matrix variable (3x4)."""
+    """Test sum(log(X)) with 2D matrix variable (3x4)."""
     X = cp.Variable((3, 4))
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
     c_obj, _ = convert_problem(problem)
 
     test_values = np.arange(1.0, 13.0)  # 12 elements
+
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_matrix_variable passed")
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
 
 
 def test_mixed_sizes():
-    """Test with variables of different sizes."""
+    """Test sum(log(a)) + sum(log(b)) + sum(log(c)) with different sized variables."""
     a = cp.Variable(2)
     b = cp.Variable(5)
     c = cp.Variable(3)
@@ -255,80 +311,19 @@ def test_mixed_sizes():
     c_vals = np.array([1.0, 2.0, 3.0])
     test_values = np.concatenate([a_vals, b_vals, c_vals])
 
+    # Forward
     result = diffengine.forward(c_obj, test_values)
     expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_mixed_sizes passed")
-
-
-def test_complex_objective():
-    """Test complex objective: sum(log(x + y)) + sum(exp(y + z)) + sum(log(z + x))."""
-    x = cp.Variable(3)
-    y = cp.Variable(3)
-    z = cp.Variable(3)
-    obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_problem(problem)
-
-    x_vals = np.array([1.0, 2.0, 3.0])
-    y_vals = np.array([0.5, 1.0, 1.5])
-    z_vals = np.array([0.2, 0.3, 0.4])
-    test_values = np.concatenate([x_vals, y_vals, z_vals])
-
-    result = diffengine.forward(c_obj, test_values)
-    expected = (np.sum(np.log(x_vals + y_vals)) +
-                np.sum(np.exp(y_vals + z_vals)) +
-                np.sum(np.log(z_vals + x_vals)))
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_complex_objective passed")
-
-
-def test_log_exp_identity():
-    """Test log(exp(x)) = x identity."""
-    x = cp.Variable(5)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([-1.0, 0.0, 1.0, 2.0, 3.0])
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(test_values)  # log(exp(x)) = x
-    assert np.allclose(result, expected), f"Expected {expected}, got {result}"
-    print("test_log_exp_identity passed")
-
-
-# === Jacobian Tests ===
+    assert np.allclose(result, expected)
 
-def test_jacobian_sum_log():
-    """Test jacobian of sum(log(x)). Gradient is [1/x_1, 1/x_2, ...]."""
-    x = cp.Variable(4)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-    jac = get_jacobian(c_obj, test_values)
-
-    # sum(log(x)) is scalar, so jacobian is 1 x n
-    expected = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
-    print("test_jacobian_sum_log passed")
-
-
-def test_jacobian_sum_exp():
-    """Test jacobian of sum(exp(x)). Gradient is [exp(x_1), exp(x_2), ...]."""
-    x = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([0.0, 1.0, 2.0])
+    # Jacobian
     jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
 
-    expected = np.exp(test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
-    print("test_jacobian_sum_exp passed")
 
-
-def test_jacobian_four_variables():
-    """Test jacobian of sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
+def test_multiple_variables_log_exp():
+    """Test sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
     a = cp.Variable(2)
     b = cp.Variable(2)
     c = cp.Variable(2)
@@ -337,75 +332,74 @@ def test_jacobian_four_variables():
     problem = cp.Problem(cp.Minimize(obj))
     c_obj, _ = convert_problem(problem)
 
-    # Variables: [a0, a1, b0, b1, c0, c1, d0, d1]
     a_vals = np.array([1.0, 2.0])
     b_vals = np.array([0.5, 1.0])
     c_vals = np.array([0.1, 0.2])
     d_vals = np.array([0.1, 0.1])
     test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
-    jac = get_jacobian(c_obj, test_values)
 
-    # f = sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = (np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) +
+                np.sum(np.exp(c_vals)) + np.sum(np.exp(d_vals)))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
     df_da = 1.0 / a_vals
     df_db = 1.0 / b_vals
     df_dc = np.exp(c_vals)
     df_dd = np.exp(d_vals)
-    expected = np.concatenate([df_da, df_db, df_dc, df_dd]).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
-    print("test_jacobian_four_variables passed")
+    expected_jac = np.concatenate([df_da, df_db, df_dc, df_dd]).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
 
 
-def test_jacobian_two_variables():
-    """Test jacobian of sum(log(x) + log(y)) with two variables."""
+def test_two_variables_separate_sums():
+    """Test sum(log(x) + log(y)) - two variables with separate elementwise ops."""
     x = cp.Variable(2)
     y = cp.Variable(2)
     problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y))))
     c_obj, _ = convert_problem(problem)
 
-    # Variables: [x0, x1, y0, y1]
     test_values = np.array([1.0, 2.0, 3.0, 4.0])
-    jac = get_jacobian(c_obj, test_values)
-
-    # f = sum(log(x) + log(y)) = log(x0) + log(x1) + log(y0) + log(y1)
-    # df/dx = [1/x0, 1/x1], df/dy = [1/y0, 1/y1]
-    expected = np.array([[1/1, 1/2, 1/3, 1/4]])
-    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
-    print("test_jacobian_two_variables passed")
-
 
-def test_jacobian_variable_reuse():
-    """Test jacobian when same variable is used multiple times."""
-    x = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
-    c_obj, _ = convert_problem(problem)
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values[:2]) + np.log(test_values[2:]))
+    assert np.allclose(result, expected)
 
-    test_values = np.array([1.0, 2.0])
+    # Jacobian
     jac = get_jacobian(c_obj, test_values)
-
-    # f = sum(log(x) + exp(x))
-    # df/dx_i = 1/x_i + exp(x_i)
-    expected = (1.0 / test_values + np.exp(test_values)).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
-    print("test_jacobian_variable_reuse passed")
+    expected_jac = np.array([[1/1, 1/2, 1/3, 1/4]])
+    assert np.allclose(jac.toarray(), expected_jac)
 
 
-def test_jacobian_large_variable():
-    """Test jacobian of sum(log(x)) with larger variable."""
-    x = cp.Variable(10)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+def test_complex_objective_elementwise_add():
+    """Test sum(log(x + y)) + sum(exp(y + z)) + sum(log(z + x)) - elementwise after add."""
+    x = cp.Variable(3)
+    y = cp.Variable(3)
+    z = cp.Variable(3)
+    obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
+    problem = cp.Problem(cp.Minimize(obj))
     c_obj, _ = convert_problem(problem)
 
-    test_values = np.linspace(1.0, 10.0, 10)
-    jac = get_jacobian(c_obj, test_values)
+    x_vals = np.array([1.0, 2.0, 3.0])
+    y_vals = np.array([0.5, 1.0, 1.5])
+    z_vals = np.array([0.2, 0.3, 0.4])
+    test_values = np.concatenate([x_vals, y_vals, z_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = (np.sum(np.log(x_vals + y_vals)) +
+                np.sum(np.exp(y_vals + z_vals)) +
+                np.sum(np.log(z_vals + x_vals)))
+    assert np.allclose(result, expected)
 
-    # f = sum(log(x)), df/dx_i = 1/x_i
-    expected = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
-    print("test_jacobian_large_variable passed")
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
 
 
-def test_jacobian_complex_objective():
-    """Test jacobian of sum(log(x) + exp(y) + log(z))."""
+def test_complex_objective_no_add():
+    """Test sum(log(x) + exp(y) + log(z)) - multiple elementwise ops without add composition."""
     x = cp.Variable(2)
     y = cp.Variable(2)
     z = cp.Variable(2)
@@ -417,40 +411,51 @@ def test_jacobian_complex_objective():
     y_vals = np.array([0.5, 1.0])
     z_vals = np.array([0.2, 0.3])
     test_values = np.concatenate([x_vals, y_vals, z_vals])
-    jac = get_jacobian(c_obj, test_values)
 
-    # f = sum(log(x) + exp(y) + log(z))
-    # df/dx_i = 1/x_i, df/dy_i = exp(y_i), df/dz_i = 1/z_i
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(x_vals) + np.exp(y_vals) + np.log(z_vals))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
     df_dx = 1.0 / x_vals
     df_dy = np.exp(y_vals)
     df_dz = 1.0 / z_vals
-    expected = np.concatenate([df_dx, df_dy, df_dz]).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected), f"Expected {expected}, got {jac.toarray()}"
-    print("test_jacobian_complex_objective passed")
+    expected_jac = np.concatenate([df_dx, df_dy, df_dz]).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_log_exp_identity():
+    """Test sum(log(exp(x))) = sum(x) identity - nested elementwise."""
+    x = cp.Variable(5)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([-1.0, 0.0, 1.0, 2.0, 3.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(test_values)  # log(exp(x)) = x
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for nested elementwise compositions not yet supported
 
 
 if __name__ == "__main__":
-    # Forward pass tests
-    test_simple_sum_log()
-    test_two_variables()
+    test_sum_log()
+    test_sum_exp()
+    test_two_variables_elementwise_add()
     test_variable_reuse()
-    test_four_variables()
+    test_four_variables_elementwise_add()
     test_deep_nesting()
     test_chained_additions()
     test_variable_used_multiple_times()
-    test_larger_variable_size()
+    test_larger_variable()
     test_matrix_variable()
     test_mixed_sizes()
-    test_complex_objective()
+    test_multiple_variables_log_exp()
+    test_two_variables_separate_sums()
+    test_complex_objective_elementwise_add()
+    test_complex_objective_no_add()
     test_log_exp_identity()
-
-    # Jacobian tests
-    test_jacobian_sum_log()
-    test_jacobian_sum_exp()
-    test_jacobian_four_variables()
-    test_jacobian_two_variables()
-    test_jacobian_variable_reuse()
-    test_jacobian_large_variable()
-    test_jacobian_complex_objective()
-
-    print("\nAll tests passed!")

From d7a3a027a22d5c35bc802c832ea58a11a48a9c16 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Sun, 4 Jan 2026 19:53:49 -0500
Subject: [PATCH 6/7] update problem struct design and move tests

---
 docs/problem_struct_design.md |  40 +++-
 python/convert.py             | 356 ----------------------------------
 2 files changed, 38 insertions(+), 358 deletions(-)

diff --git a/docs/problem_struct_design.md b/docs/problem_struct_design.md
index 7a3cd09..4c9144e 100644
--- a/docs/problem_struct_design.md
+++ b/docs/problem_struct_design.md
@@ -41,6 +41,7 @@ typedef struct problem
 } problem;
 
 problem *new_problem(expr *objective, expr **constraints, int n_constraints);
+void problem_allocate(problem *prob, const double *u);
 void free_problem(problem *prob);
 double problem_forward(problem *prob, const double *u);
 double *problem_gradient(problem *prob, const double *u);
@@ -58,7 +59,38 @@ Key functions:
 ### `new_problem`
 - Retain (increment refcount) on objective and all constraints
 - Compute `total_constraint_size = sum(constraints[i]->size)`
-- Pre-allocate `constraint_values` and `gradient_values` arrays
+- Does NOT allocate storage arrays (use `problem_allocate` separately)
+
+### `problem_allocate`
+Separate function to allocate memory for constraint values and jacobian:
+
+```c
+void problem_allocate(problem *prob, const double *u)
+{
+    /* 1. Allocate constraint values array */
+    prob->constraint_values = malloc(prob->total_constraint_size * sizeof(double));
+
+    /* 2. Allocate jacobian:
+     *    - First, initialize all constraint jacobians
+     *    - Count total nnz across all constraints
+     *    - Allocate CSR matrix with this nnz (may be slight overestimate)
+     */
+    int total_nnz = 0;
+    for (int i = 0; i < prob->n_constraints; i++)
+    {
+        expr *c = prob->constraints[i];
+        c->forward(c, u);
+        c->jacobian_init(c);
+        total_nnz += c->jacobian->nnz;
+    }
+
+    /* Allocate stacked jacobian with total_constraint_size rows */
+    prob->stacked_jac = alloc_csr(prob->total_constraint_size, prob->n_vars, total_nnz);
+
+    /* Note: The actual nnz may be smaller after evaluation due to
+     * cancellations. Update stacked_jac->nnz after problem_jacobian(). */
+}
+```
 
 ### `free_problem`
 - Call `free_expr` on objective and all constraints (decrements refcount)
@@ -201,5 +233,9 @@ class Problem:
 ## Key Design Notes
 
 - **Memory**: Uses expr refcounting - new_problem retains, free_problem releases
-- **Efficiency**: Pre-allocates arrays, lazy-allocates stacked jacobian
+- **Two-phase init**: `new_problem` creates struct, `problem_allocate` allocates arrays
+  - Constraint values array: size = `total_constraint_size`
+  - Jacobian: initialize all constraint jacobians first, count total nnz, allocate CSR
+  - The allocated nnz may be a slight overestimate; update `stacked_jac->nnz` after evaluation
+- **Hessian**: Deferred - not allocated in this design (to be added later)
 - **API**: Returns internal pointers (caller should NOT free)
diff --git a/python/convert.py b/python/convert.py
index cc34d6e..7a1be62 100644
--- a/python/convert.py
+++ b/python/convert.py
@@ -1,5 +1,4 @@
 import cvxpy as cp
-import numpy as np
 from scipy import sparse
 from cvxpy.reductions.inverse_data import InverseData
 import DNLP_diff_engine as diffengine
@@ -104,358 +103,3 @@ def convert_problem(problem: cp.Problem) -> tuple:
         c_constraints.append(c_expr)
 
     return c_objective, c_constraints
-
-
-# === Tests ===
-
-def test_sum_log():
-    """Test sum(log(x)) forward and jacobian."""
-    x = cp.Variable(4)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx sum(log(x)) = [1/x_1, 1/x_2, ...]
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_sum_exp():
-    """Test sum(exp(x)) forward and jacobian."""
-    x = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([0.0, 1.0, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.exp(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx sum(exp(x)) = [exp(x_1), exp(x_2), ...]
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = np.exp(test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_two_variables_elementwise_add():
-    """Test sum(log(x + y)) - elementwise after add."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(np.array([1+3, 2+4])))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_variable_reuse():
-    """Test sum(log(x) + exp(x)) - same variable used twice."""
-    x = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values) + np.exp(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian: d/dx_i = 1/x_i + exp(x_i)
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values + np.exp(test_values)).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_four_variables_elementwise_add():
-    """Test sum(log(a + b) + exp(c + d)) - elementwise after add."""
-    a = cp.Variable(3)
-    b = cp.Variable(3)
-    c = cp.Variable(3)
-    d = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a + b) + cp.exp(c + d))))
-    c_obj, _ = convert_problem(problem)
-
-    a_vals = np.array([1.0, 2.0, 3.0])
-    b_vals = np.array([0.5, 1.0, 1.5])
-    c_vals = np.array([0.1, 0.2, 0.3])
-    d_vals = np.array([0.1, 0.1, 0.1])
-    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(a_vals + b_vals) + np.exp(c_vals + d_vals))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_deep_nesting():
-    """Test sum(log(exp(log(exp(x))))) - deeply nested elementwise."""
-    x = cp.Variable(4)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(cp.log(cp.exp(x)))))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([0.5, 1.0, 1.5, 2.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(np.exp(np.log(np.exp(test_values)))))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for nested elementwise compositions not yet supported
-
-
-def test_chained_additions():
-    """Test sum(x + y + z + w) - chained additions."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    z = cp.Variable(2)
-    w = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(x + y + z + w)))
-    c_obj, _ = convert_problem(problem)
-
-    x_vals = np.array([1.0, 2.0])
-    y_vals = np.array([3.0, 4.0])
-    z_vals = np.array([5.0, 6.0])
-    w_vals = np.array([7.0, 8.0])
-    test_values = np.concatenate([x_vals, y_vals, z_vals, w_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(x_vals + y_vals + z_vals + w_vals)
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for sum(add(...)) patterns not yet supported
-
-
-def test_variable_used_multiple_times():
-    """Test sum(log(x) + exp(x) + x) - variable used 3+ times."""
-    x = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + x)))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values) + np.exp(test_values) + test_values)
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for expressions with sum(variable) not yet supported
-
-
-def test_larger_variable():
-    """Test sum(log(x)) with larger variable (100 elements)."""
-    x = cp.Variable(100)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.linspace(1.0, 10.0, 100)
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_matrix_variable():
-    """Test sum(log(X)) with 2D matrix variable (3x4)."""
-    X = cp.Variable((3, 4))
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.arange(1.0, 13.0)  # 12 elements
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_mixed_sizes():
-    """Test sum(log(a)) + sum(log(b)) + sum(log(c)) with different sized variables."""
-    a = cp.Variable(2)
-    b = cp.Variable(5)
-    c = cp.Variable(3)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))))
-    c_obj, _ = convert_problem(problem)
-
-    a_vals = np.array([1.0, 2.0])
-    b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
-    c_vals = np.array([1.0, 2.0, 3.0])
-    test_values = np.concatenate([a_vals, b_vals, c_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = (1.0 / test_values).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_multiple_variables_log_exp():
-    """Test sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
-    a = cp.Variable(2)
-    b = cp.Variable(2)
-    c = cp.Variable(2)
-    d = cp.Variable(2)
-    obj = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.exp(c)) + cp.sum(cp.exp(d))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_problem(problem)
-
-    a_vals = np.array([1.0, 2.0])
-    b_vals = np.array([0.5, 1.0])
-    c_vals = np.array([0.1, 0.2])
-    d_vals = np.array([0.1, 0.1])
-    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = (np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) +
-                np.sum(np.exp(c_vals)) + np.sum(np.exp(d_vals)))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    df_da = 1.0 / a_vals
-    df_db = 1.0 / b_vals
-    df_dc = np.exp(c_vals)
-    df_dd = np.exp(d_vals)
-    expected_jac = np.concatenate([df_da, df_db, df_dc, df_dd]).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_two_variables_separate_sums():
-    """Test sum(log(x) + log(y)) - two variables with separate elementwise ops."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([1.0, 2.0, 3.0, 4.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(test_values[:2]) + np.log(test_values[2:]))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    expected_jac = np.array([[1/1, 1/2, 1/3, 1/4]])
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_complex_objective_elementwise_add():
-    """Test sum(log(x + y)) + sum(exp(y + z)) + sum(log(z + x)) - elementwise after add."""
-    x = cp.Variable(3)
-    y = cp.Variable(3)
-    z = cp.Variable(3)
-    obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_problem(problem)
-
-    x_vals = np.array([1.0, 2.0, 3.0])
-    y_vals = np.array([0.5, 1.0, 1.5])
-    z_vals = np.array([0.2, 0.3, 0.4])
-    test_values = np.concatenate([x_vals, y_vals, z_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = (np.sum(np.log(x_vals + y_vals)) +
-                np.sum(np.exp(y_vals + z_vals)) +
-                np.sum(np.log(z_vals + x_vals)))
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
-
-
-def test_complex_objective_no_add():
-    """Test sum(log(x) + exp(y) + log(z)) - multiple elementwise ops without add composition."""
-    x = cp.Variable(2)
-    y = cp.Variable(2)
-    z = cp.Variable(2)
-    obj = cp.sum(cp.log(x) + cp.exp(y) + cp.log(z))
-    problem = cp.Problem(cp.Minimize(obj))
-    c_obj, _ = convert_problem(problem)
-
-    x_vals = np.array([1.0, 2.0])
-    y_vals = np.array([0.5, 1.0])
-    z_vals = np.array([0.2, 0.3])
-    test_values = np.concatenate([x_vals, y_vals, z_vals])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(np.log(x_vals) + np.exp(y_vals) + np.log(z_vals))
-    assert np.allclose(result, expected)
-
-    # Jacobian
-    jac = get_jacobian(c_obj, test_values)
-    df_dx = 1.0 / x_vals
-    df_dy = np.exp(y_vals)
-    df_dz = 1.0 / z_vals
-    expected_jac = np.concatenate([df_dx, df_dy, df_dz]).reshape(1, -1)
-    assert np.allclose(jac.toarray(), expected_jac)
-
-
-def test_log_exp_identity():
-    """Test sum(log(exp(x))) = sum(x) identity - nested elementwise."""
-    x = cp.Variable(5)
-    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
-    c_obj, _ = convert_problem(problem)
-
-    test_values = np.array([-1.0, 0.0, 1.0, 2.0, 3.0])
-
-    # Forward
-    result = diffengine.forward(c_obj, test_values)
-    expected = np.sum(test_values)  # log(exp(x)) = x
-    assert np.allclose(result, expected)
-
-    # TODO: Jacobian for nested elementwise compositions not yet supported
-
-
-if __name__ == "__main__":
-    test_sum_log()
-    test_sum_exp()
-    test_two_variables_elementwise_add()
-    test_variable_reuse()
-    test_four_variables_elementwise_add()
-    test_deep_nesting()
-    test_chained_additions()
-    test_variable_used_multiple_times()
-    test_larger_variable()
-    test_matrix_variable()
-    test_mixed_sizes()
-    test_multiple_variables_log_exp()
-    test_two_variables_separate_sums()
-    test_complex_objective_elementwise_add()
-    test_complex_objective_no_add()
-    test_log_exp_identity()

From 624f5f54c089c1ab318cc2ea2c611ebfb4b5a8c4 Mon Sep 17 00:00:00 2001
From: William Zijie Zhang <william@gridmatic.com>
Date: Sun, 4 Jan 2026 19:54:08 -0500
Subject: [PATCH 7/7] adds new folder for tests

---
 python/tests/convert_tests.py | 358 ++++++++++++++++++++++++++++++++++
 1 file changed, 358 insertions(+)
 create mode 100644 python/tests/convert_tests.py

diff --git a/python/tests/convert_tests.py b/python/tests/convert_tests.py
new file mode 100644
index 0000000..0bed90a
--- /dev/null
+++ b/python/tests/convert_tests.py
@@ -0,0 +1,358 @@
+import cvxpy as cp
+import numpy as np
+import DNLP_diff_engine as diffengine
+from convert import convert_problem, get_jacobian
+
+
+def test_sum_log():
+    """Test sum(log(x)) forward and jacobian."""
+    x = cp.Variable(4)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian: d/dx sum(log(x)) = [1/x_1, 1/x_2, ...]
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_sum_exp():
+    """Test sum(exp(x)) forward and jacobian."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([0.0, 1.0, 2.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.exp(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian: d/dx sum(exp(x)) = [exp(x_1), exp(x_2), ...]
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = np.exp(test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_two_variables_elementwise_add():
+    """Test sum(log(x + y)) - elementwise after add."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(np.array([1+3, 2+4])))
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
+
+
+def test_variable_reuse():
+    """Test sum(log(x) + exp(x)) - same variable used twice."""
+    x = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([1.0, 2.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values) + np.exp(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian: d/dx_i = 1/x_i + exp(x_i)
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values + np.exp(test_values)).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_four_variables_elementwise_add():
+    """Test sum(log(a + b) + exp(c + d)) - elementwise after add."""
+    a = cp.Variable(3)
+    b = cp.Variable(3)
+    c = cp.Variable(3)
+    d = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a + b) + cp.exp(c + d))))
+    c_obj, _ = convert_problem(problem)
+
+    a_vals = np.array([1.0, 2.0, 3.0])
+    b_vals = np.array([0.5, 1.0, 1.5])
+    c_vals = np.array([0.1, 0.2, 0.3])
+    d_vals = np.array([0.1, 0.1, 0.1])
+    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(a_vals + b_vals) + np.exp(c_vals + d_vals))
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
+
+
+def test_deep_nesting():
+    """Test sum(log(exp(log(exp(x))))) - deeply nested elementwise."""
+    x = cp.Variable(4)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(cp.log(cp.exp(x)))))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([0.5, 1.0, 1.5, 2.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(np.exp(np.log(np.exp(test_values)))))
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for nested elementwise compositions not yet supported
+
+
+def test_chained_additions():
+    """Test sum(x + y + z + w) - chained additions."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    z = cp.Variable(2)
+    w = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(x + y + z + w)))
+    c_obj, _ = convert_problem(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([3.0, 4.0])
+    z_vals = np.array([5.0, 6.0])
+    w_vals = np.array([7.0, 8.0])
+    test_values = np.concatenate([x_vals, y_vals, z_vals, w_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(x_vals + y_vals + z_vals + w_vals)
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for sum(add(...)) patterns not yet supported
+
+
+def test_variable_used_multiple_times():
+    """Test sum(log(x) + exp(x) + x) - variable used 3+ times."""
+    x = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + x)))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values) + np.exp(test_values) + test_values)
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for expressions with sum(variable) not yet supported
+
+
+def test_larger_variable():
+    """Test sum(log(x)) with larger variable (100 elements)."""
+    x = cp.Variable(100)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.linspace(1.0, 10.0, 100)
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_matrix_variable():
+    """Test sum(log(X)) with 2D matrix variable (3x4)."""
+    X = cp.Variable((3, 4))
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(X))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.arange(1.0, 13.0)  # 12 elements
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_mixed_sizes():
+    """Test sum(log(a)) + sum(log(b)) + sum(log(c)) with different sized variables."""
+    a = cp.Variable(2)
+    b = cp.Variable(5)
+    c = cp.Variable(3)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.log(c))))
+    c_obj, _ = convert_problem(problem)
+
+    a_vals = np.array([1.0, 2.0])
+    b_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+    c_vals = np.array([1.0, 2.0, 3.0])
+    test_values = np.concatenate([a_vals, b_vals, c_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = (1.0 / test_values).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_multiple_variables_log_exp():
+    """Test sum(log(a)) + sum(log(b)) + sum(exp(c)) + sum(exp(d))."""
+    a = cp.Variable(2)
+    b = cp.Variable(2)
+    c = cp.Variable(2)
+    d = cp.Variable(2)
+    obj = cp.sum(cp.log(a)) + cp.sum(cp.log(b)) + cp.sum(cp.exp(c)) + cp.sum(cp.exp(d))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_problem(problem)
+
+    a_vals = np.array([1.0, 2.0])
+    b_vals = np.array([0.5, 1.0])
+    c_vals = np.array([0.1, 0.2])
+    d_vals = np.array([0.1, 0.1])
+    test_values = np.concatenate([a_vals, b_vals, c_vals, d_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = (np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) +
+                np.sum(np.exp(c_vals)) + np.sum(np.exp(d_vals)))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    df_da = 1.0 / a_vals
+    df_db = 1.0 / b_vals
+    df_dc = np.exp(c_vals)
+    df_dd = np.exp(d_vals)
+    expected_jac = np.concatenate([df_da, df_db, df_dc, df_dd]).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_two_variables_separate_sums():
+    """Test sum(log(x) + log(y)) - two variables with separate elementwise ops."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([1.0, 2.0, 3.0, 4.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(test_values[:2]) + np.log(test_values[2:]))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    expected_jac = np.array([[1/1, 1/2, 1/3, 1/4]])
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_complex_objective_elementwise_add():
+    """Test sum(log(x + y)) + sum(exp(y + z)) + sum(log(z + x)) - elementwise after add."""
+    x = cp.Variable(3)
+    y = cp.Variable(3)
+    z = cp.Variable(3)
+    obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) + cp.sum(cp.log(z + x))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_problem(problem)
+
+    x_vals = np.array([1.0, 2.0, 3.0])
+    y_vals = np.array([0.5, 1.0, 1.5])
+    z_vals = np.array([0.2, 0.3, 0.4])
+    test_values = np.concatenate([x_vals, y_vals, z_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = (np.sum(np.log(x_vals + y_vals)) +
+                np.sum(np.exp(y_vals + z_vals)) +
+                np.sum(np.log(z_vals + x_vals)))
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for elementwise(add(...)) patterns not yet supported
+
+
+def test_complex_objective_no_add():
+    """Test sum(log(x) + exp(y) + log(z)) - multiple elementwise ops without add composition."""
+    x = cp.Variable(2)
+    y = cp.Variable(2)
+    z = cp.Variable(2)
+    obj = cp.sum(cp.log(x) + cp.exp(y) + cp.log(z))
+    problem = cp.Problem(cp.Minimize(obj))
+    c_obj, _ = convert_problem(problem)
+
+    x_vals = np.array([1.0, 2.0])
+    y_vals = np.array([0.5, 1.0])
+    z_vals = np.array([0.2, 0.3])
+    test_values = np.concatenate([x_vals, y_vals, z_vals])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(np.log(x_vals) + np.exp(y_vals) + np.log(z_vals))
+    assert np.allclose(result, expected)
+
+    # Jacobian
+    jac = get_jacobian(c_obj, test_values)
+    df_dx = 1.0 / x_vals
+    df_dy = np.exp(y_vals)
+    df_dz = 1.0 / z_vals
+    expected_jac = np.concatenate([df_dx, df_dy, df_dz]).reshape(1, -1)
+    assert np.allclose(jac.toarray(), expected_jac)
+
+
+def test_log_exp_identity():
+    """Test sum(log(exp(x))) = sum(x) identity - nested elementwise."""
+    x = cp.Variable(5)
+    problem = cp.Problem(cp.Minimize(cp.sum(cp.log(cp.exp(x)))))
+    c_obj, _ = convert_problem(problem)
+
+    test_values = np.array([-1.0, 0.0, 1.0, 2.0, 3.0])
+
+    # Forward
+    result = diffengine.forward(c_obj, test_values)
+    expected = np.sum(test_values)  # log(exp(x)) = x
+    assert np.allclose(result, expected)
+
+    # TODO: Jacobian for nested elementwise compositions not yet supported
+
+
+if __name__ == "__main__":
+    test_sum_log()
+    test_sum_exp()
+    test_two_variables_elementwise_add()
+    test_variable_reuse()
+    test_four_variables_elementwise_add()
+    test_deep_nesting()
+    test_chained_additions()
+    test_variable_used_multiple_times()
+    test_larger_variable()
+    test_matrix_variable()
+    test_mixed_sizes()
+    test_multiple_variables_log_exp()
+    test_two_variables_separate_sums()
+    test_complex_objective_elementwise_add()
+    test_complex_objective_no_add()
+    test_log_exp_identity()
+    print("All tests passed!")