diff --git a/docs/problem_struct_design.md b/docs/problem_struct_design.md new file mode 100644 index 0000000..4c9144e --- /dev/null +++ b/docs/problem_struct_design.md @@ -0,0 +1,241 @@ +# 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 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); +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)` +- 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) +- 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 +- **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/bindings.c b/python/bindings.c index b4bcd38..68e4b76 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; @@ -103,10 +178,73 @@ 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"}, + {"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"}, + {"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 new file mode 100644 index 0000000..7a1be62 --- /dev/null +++ b/python/convert.py @@ -0,0 +1,105 @@ +import cvxpy as cp +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] + 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, _ = 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 + + +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 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) 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!")