diff --git a/.gitignore b/.gitignore index 7179491..f7de229 100644 --- a/.gitignore +++ b/.gitignore @@ -72,4 +72,7 @@ out/ *.swp *.swo *~ -.DS_Store \ No newline at end of file +.DS_Store + +.venv/ +__pycache__/ \ No newline at end of file diff --git a/include/affine.h b/include/affine.h index 12b1dab..654557b 100644 --- a/include/affine.h +++ b/include/affine.h @@ -8,9 +8,11 @@ expr *new_linear(expr *u, const CSR_Matrix *A); expr *new_add(expr *left, expr *right); +expr *new_neg(expr *child); expr *new_sum(expr *child, int axis); expr *new_hstack(expr **args, int n_args, int n_vars); +expr *new_promote(expr *child, int d1, int d2); expr *new_trace(expr *child); expr *new_constant(int d1, int d2, int n_vars, const double *values); diff --git a/include/problem.h b/include/problem.h new file mode 100644 index 0000000..3159938 --- /dev/null +++ b/include/problem.h @@ -0,0 +1,33 @@ +#ifndef PROBLEM_H +#define PROBLEM_H + +#include "expr.h" +#include "utils/CSR_Matrix.h" + +typedef struct problem +{ + expr *objective; + expr **constraints; + int n_constraints; + int n_vars; + int total_constraint_size; + + /* Allocated by new_problem */ + double *constraint_values; + double *gradient_values; + + /* Allocated by problem_init_derivatives */ + CSR_Matrix *stacked_jac; +} problem; + +/* Retains objective and constraints (shared ownership with caller) */ +problem *new_problem(expr *objective, expr **constraints, int n_constraints); +void problem_init_derivatives(problem *prob); +void free_problem(problem *prob); + +double problem_objective_forward(problem *prob, const double *u); +void problem_constraint_forward(problem *prob, const double *u); +void problem_gradient(problem *prob); +void problem_jacobian(problem *prob); + +#endif diff --git a/python/atoms/add.h b/python/atoms/add.h new file mode 100644 index 0000000..4b86b79 --- /dev/null +++ b/python/atoms/add.h @@ -0,0 +1,30 @@ +#ifndef ATOM_ADD_H +#define ATOM_ADD_H + +#include "common.h" + +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); +} + +#endif /* ATOM_ADD_H */ diff --git a/python/atoms/common.h b/python/atoms/common.h new file mode 100644 index 0000000..83101dd --- /dev/null +++ b/python/atoms/common.h @@ -0,0 +1,23 @@ +#ifndef ATOMS_COMMON_H +#define ATOMS_COMMON_H + +#define PY_SSIZE_T_CLEAN +#include +#include + +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" + +#define EXPR_CAPSULE_NAME "DNLP_EXPR" + +static void expr_capsule_destructor(PyObject *capsule) +{ + expr *node = (expr *) PyCapsule_GetPointer(capsule, EXPR_CAPSULE_NAME); + if (node) + { + free_expr(node); + } +} + +#endif /* ATOMS_COMMON_H */ diff --git a/python/atoms/constant.h b/python/atoms/constant.h new file mode 100644 index 0000000..da2af09 --- /dev/null +++ b/python/atoms/constant.h @@ -0,0 +1,34 @@ +#ifndef ATOM_CONSTANT_H +#define ATOM_CONSTANT_H + +#include "common.h" + +static PyObject *py_make_constant(PyObject *self, PyObject *args) +{ + int d1, d2, n_vars; + PyObject *values_obj; + if (!PyArg_ParseTuple(args, "iiiO", &d1, &d2, &n_vars, &values_obj)) + { + return NULL; + } + + PyArrayObject *values_array = (PyArrayObject *) PyArray_FROM_OTF( + values_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + if (!values_array) + { + return NULL; + } + + expr *node = + new_constant(d1, d2, n_vars, (const double *) PyArray_DATA(values_array)); + Py_DECREF(values_array); + + if (!node) + { + PyErr_SetString(PyExc_RuntimeError, "failed to create constant node"); + return NULL; + } + return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); +} + +#endif /* ATOM_CONSTANT_H */ diff --git a/python/atoms/exp.h b/python/atoms/exp.h new file mode 100644 index 0000000..9505daf --- /dev/null +++ b/python/atoms/exp.h @@ -0,0 +1,29 @@ +#ifndef ATOM_EXP_H +#define ATOM_EXP_H + +#include "common.h" + +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); +} + +#endif /* ATOM_EXP_H */ diff --git a/python/atoms/linear.h b/python/atoms/linear.h new file mode 100644 index 0000000..0f85b91 --- /dev/null +++ b/python/atoms/linear.h @@ -0,0 +1,60 @@ +#ifndef ATOM_LINEAR_H +#define ATOM_LINEAR_H + +#include "common.h" + +static PyObject *py_make_linear(PyObject *self, PyObject *args) +{ + PyObject *child_capsule; + PyObject *data_obj, *indices_obj, *indptr_obj; + int m, n; + if (!PyArg_ParseTuple(args, "OOOOii", &child_capsule, &data_obj, &indices_obj, + &indptr_obj, &m, &n)) + { + return NULL; + } + + expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME); + if (!child) + { + PyErr_SetString(PyExc_ValueError, "invalid child capsule"); + return NULL; + } + + PyArrayObject *data_array = + (PyArrayObject *) PyArray_FROM_OTF(data_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *indices_array = (PyArrayObject *) PyArray_FROM_OTF( + indices_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + PyArrayObject *indptr_array = (PyArrayObject *) PyArray_FROM_OTF( + indptr_obj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + + if (!data_array || !indices_array || !indptr_array) + { + Py_XDECREF(data_array); + Py_XDECREF(indices_array); + Py_XDECREF(indptr_array); + return NULL; + } + + int nnz = (int) PyArray_SIZE(data_array); + CSR_Matrix *A = new_csr_matrix(m, n, nnz); + memcpy(A->x, PyArray_DATA(data_array), nnz * sizeof(double)); + memcpy(A->i, PyArray_DATA(indices_array), nnz * sizeof(int)); + memcpy(A->p, PyArray_DATA(indptr_array), (m + 1) * sizeof(int)); + + Py_DECREF(data_array); + Py_DECREF(indices_array); + Py_DECREF(indptr_array); + + expr *node = new_linear(child, A); + free_csr_matrix(A); + + if (!node) + { + PyErr_SetString(PyExc_RuntimeError, "failed to create linear node"); + return NULL; + } + return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); +} + +#endif /* ATOM_LINEAR_H */ diff --git a/python/atoms/log.h b/python/atoms/log.h new file mode 100644 index 0000000..3d0314e --- /dev/null +++ b/python/atoms/log.h @@ -0,0 +1,29 @@ +#ifndef ATOM_LOG_H +#define ATOM_LOG_H + +#include "common.h" + +static PyObject *py_make_log(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_log(child); + if (!node) + { + PyErr_SetString(PyExc_RuntimeError, "failed to create log node"); + return NULL; + } + return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); +} + +#endif /* ATOM_LOG_H */ diff --git a/python/atoms/neg.h b/python/atoms/neg.h new file mode 100644 index 0000000..bda69d0 --- /dev/null +++ b/python/atoms/neg.h @@ -0,0 +1,29 @@ +#ifndef ATOM_NEG_H +#define ATOM_NEG_H + +#include "common.h" + +static PyObject *py_make_neg(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_neg(child); + if (!node) + { + PyErr_SetString(PyExc_RuntimeError, "failed to create neg node"); + return NULL; + } + return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); +} + +#endif /* ATOM_NEG_H */ diff --git a/python/atoms/promote.h b/python/atoms/promote.h new file mode 100644 index 0000000..fb258b4 --- /dev/null +++ b/python/atoms/promote.h @@ -0,0 +1,30 @@ +#ifndef ATOM_PROMOTE_H +#define ATOM_PROMOTE_H + +#include "common.h" + +static PyObject *py_make_promote(PyObject *self, PyObject *args) +{ + PyObject *child_capsule; + int d1, d2; + if (!PyArg_ParseTuple(args, "Oii", &child_capsule, &d1, &d2)) + { + 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_promote(child, d1, d2); + if (!node) + { + PyErr_SetString(PyExc_RuntimeError, "failed to create promote node"); + return NULL; + } + return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); +} + +#endif /* ATOM_PROMOTE_H */ diff --git a/python/atoms/sum.h b/python/atoms/sum.h new file mode 100644 index 0000000..8a8480e --- /dev/null +++ b/python/atoms/sum.h @@ -0,0 +1,30 @@ +#ifndef ATOM_SUM_H +#define ATOM_SUM_H + +#include "common.h" + +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); +} + +#endif /* ATOM_SUM_H */ diff --git a/python/atoms/variable.h b/python/atoms/variable.h new file mode 100644 index 0000000..5438c44 --- /dev/null +++ b/python/atoms/variable.h @@ -0,0 +1,23 @@ +#ifndef ATOM_VARIABLE_H +#define ATOM_VARIABLE_H + +#include "common.h" + +static PyObject *py_make_variable(PyObject *self, PyObject *args) +{ + int d1, d2, var_id, n_vars; + if (!PyArg_ParseTuple(args, "iiii", &d1, &d2, &var_id, &n_vars)) + { + return NULL; + } + + expr *node = new_variable(d1, d2, var_id, n_vars); + if (!node) + { + PyErr_SetString(PyExc_RuntimeError, "failed to create variable node"); + return NULL; + } + return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); +} + +#endif /* ATOM_VARIABLE_H */ diff --git a/python/bindings.c b/python/bindings.c index e6f4f38..197390f 100644 --- a/python/bindings.c +++ b/python/bindings.c @@ -1,14 +1,25 @@ #define PY_SSIZE_T_CLEAN #include #include -#include -#include "affine.h" -#include "elementwise_univariate.h" -#include "expr.h" - -// Capsule name for expr* pointers -#define EXPR_CAPSULE_NAME "DNLP_EXPR" +/* Include atom bindings */ +#include "atoms/add.h" +#include "atoms/constant.h" +#include "atoms/exp.h" +#include "atoms/linear.h" +#include "atoms/log.h" +#include "atoms/neg.h" +#include "atoms/promote.h" +#include "atoms/sum.h" +#include "atoms/variable.h" + +/* Include problem bindings */ +#include "problem/constraint_forward.h" +#include "problem/gradient.h" +#include "problem/init_derivatives.h" +#include "problem/jacobian.h" +#include "problem/make_problem.h" +#include "problem/objective_forward.h" static int numpy_initialized = 0; @@ -20,232 +31,28 @@ static int ensure_numpy(void) return 0; } -static void expr_capsule_destructor(PyObject *capsule) -{ - expr *node = (expr *) PyCapsule_GetPointer(capsule, EXPR_CAPSULE_NAME); - if (node) - { - free_expr(node); - } -} - -static PyObject *py_make_variable(PyObject *self, PyObject *args) -{ - int d1, d2, var_id, n_vars; - if (!PyArg_ParseTuple(args, "iiii", &d1, &d2, &var_id, &n_vars)) - { - return NULL; - } - - expr *node = new_variable(d1, d2, var_id, n_vars); - if (!node) - { - PyErr_SetString(PyExc_RuntimeError, "failed to create variable node"); - return NULL; - } - return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor); -} - -static PyObject *py_make_log(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_log(child); - if (!node) - { - PyErr_SetString(PyExc_RuntimeError, "failed to create log node"); - return NULL; - } - 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; - 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; - } - - node->forward(node, (const double *) PyArray_DATA(u_array)); - - npy_intp size = node->size; - PyObject *out = PyArray_SimpleNew(1, &size, NPY_DOUBLE); - if (!out) - { - Py_DECREF(u_array); - return NULL; - } - memcpy(PyArray_DATA((PyArrayObject *) out), node->value, size * sizeof(double)); - - Py_DECREF(u_array); - 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_constant", py_make_constant, METH_VARARGS, "Create constant node"}, + {"make_linear", py_make_linear, METH_VARARGS, "Create linear op 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"}, + {"make_neg", py_make_neg, METH_VARARGS, "Create neg node"}, + {"make_promote", py_make_promote, METH_VARARGS, "Create promote node"}, + {"make_problem", py_make_problem, METH_VARARGS, + "Create problem from objective and constraints"}, + {"problem_init_derivatives", py_problem_init_derivatives, METH_VARARGS, + "Initialize derivative structures"}, + {"problem_objective_forward", py_problem_objective_forward, METH_VARARGS, + "Evaluate objective only"}, + {"problem_constraint_forward", py_problem_constraint_forward, METH_VARARGS, + "Evaluate constraints only"}, + {"problem_gradient", py_problem_gradient, METH_VARARGS, + "Compute objective gradient"}, + {"problem_jacobian", py_problem_jacobian, METH_VARARGS, + "Compute constraint jacobian"}, {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 7a1be62..92920cb 100644 --- a/python/convert.py +++ b/python/convert.py @@ -1,15 +1,10 @@ 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] @@ -19,16 +14,25 @@ def _chain_add(children): # Mapping from CVXPY atom names to C diff engine functions +# Converters receive (expr, children) where expr is the CVXPY expression ATOM_CONVERTERS = { # Elementwise unary - "log": lambda child: diffengine.make_log(child), - "exp": lambda child: diffengine.make_exp(child), + "log": lambda expr, children: diffengine.make_log(children[0]), + "exp": lambda expr, children: diffengine.make_exp(children[0]), + + # Affine unary + "NegExpression": lambda expr, children: diffengine.make_neg(children[0]), + "Promote": lambda expr, children: diffengine.make_promote( + children[0], + expr.shape[0] if len(expr.shape) >= 1 else 1, + expr.shape[1] if len(expr.shape) >= 2 else 1, + ), # N-ary (handles 2+ args) - "AddExpression": _chain_add, + "AddExpression": lambda expr, children: _chain_add(children), # Reductions - "Sum": lambda child: diffengine.make_sum(child, -1), # axis=-1 = sum all + "Sum": lambda expr, children: diffengine.make_sum(children[0], -1), } @@ -58,31 +62,35 @@ 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 + return var_dict, n_vars -def _convert_expr(expr, var_dict: dict): +def _convert_expr(expr, var_dict: dict, n_vars: int): """Convert CVXPY expression using pre-built variable dictionary.""" # Base case: variable lookup if isinstance(expr, cp.Variable): return var_dict[expr.id] + # Base case: constant + if isinstance(expr, cp.Constant): + value = np.asarray(expr.value, dtype=np.float64).flatten() + d1 = expr.shape[0] if len(expr.shape) >= 1 else 1 + d2 = expr.shape[1] if len(expr.shape) >= 2 else 1 + return diffengine.make_constant(d1, d2, n_vars, value) + # 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]) + children = [_convert_expr(arg, var_dict, n_vars) for arg in expr.args] + return ATOM_CONVERTERS[atom_name](expr, children) raise NotImplementedError(f"Atom '{atom_name}' not supported") -def convert_problem(problem: cp.Problem) -> tuple: +def convert_expressions(problem: cp.Problem) -> tuple: """ - Convert CVXPY Problem to C expressions. + Convert CVXPY Problem to C expressions (low-level). Args: problem: CVXPY Problem object @@ -91,15 +99,63 @@ 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, n_vars = build_variable_dict(problem.variables()) # Convert objective - c_objective = _convert_expr(problem.objective.expr, var_dict) + c_objective = _convert_expr(problem.objective.expr, var_dict, n_vars) # Convert constraints (expression part only for now) c_constraints = [] for constr in problem.constraints: - c_expr = _convert_expr(constr.expr, var_dict) + c_expr = _convert_expr(constr.expr, var_dict, n_vars) c_constraints.append(c_expr) return c_objective, c_constraints + + +def convert_problem(problem: cp.Problem) -> "C_problem": + """ + Convert CVXPY Problem to C problem struct. + + Args: + problem: CVXPY Problem object + + Returns: + C_Problem wrapper around C problem struct + """ + return C_problem(problem) + + +class C_problem: + """Wrapper around C problem struct for CVXPY problems.""" + + def __init__(self, cvxpy_problem: cp.Problem): + var_dict, n_vars = build_variable_dict(cvxpy_problem.variables()) + c_obj = _convert_expr(cvxpy_problem.objective.expr, var_dict, n_vars) + c_constraints = [ + _convert_expr(c.expr, var_dict, n_vars) for c in cvxpy_problem.constraints + ] + self._capsule = diffengine.make_problem(c_obj, c_constraints) + self._allocated = False + + def init_derivatives(self): + """Initialize derivative structures. Must be called before forward/gradient/jacobian.""" + diffengine.problem_init_derivatives(self._capsule) + self._allocated = True + + def objective_forward(self, u: np.ndarray) -> float: + """Evaluate objective. Returns obj_value float.""" + return diffengine.problem_objective_forward(self._capsule, u) + + def constraint_forward(self, u: np.ndarray) -> np.ndarray: + """Evaluate constraints only. Returns constraint_values array.""" + return diffengine.problem_constraint_forward(self._capsule, u) + + def gradient(self) -> np.ndarray: + """Compute gradient of objective. Call objective_forward first. Returns gradient array.""" + return diffengine.problem_gradient(self._capsule) + + def jacobian(self) -> sparse.csr_matrix: + """Compute jacobian of constraints. Call constraint_forward first. Returns scipy CSR matrix.""" + data, indices, indptr, shape = diffengine.problem_jacobian(self._capsule) + return sparse.csr_matrix((data, indices, indptr), shape=shape) diff --git a/python/problem/common.h b/python/problem/common.h new file mode 100644 index 0000000..43e44c7 --- /dev/null +++ b/python/problem/common.h @@ -0,0 +1,24 @@ +#ifndef PROBLEM_COMMON_H +#define PROBLEM_COMMON_H + +#define PY_SSIZE_T_CLEAN +#include +#include + +#include "problem.h" + +/* Also need expr types for capsule handling */ +#include "../atoms/common.h" + +#define PROBLEM_CAPSULE_NAME "DNLP_PROBLEM" + +static void problem_capsule_destructor(PyObject *capsule) +{ + problem *prob = (problem *) PyCapsule_GetPointer(capsule, PROBLEM_CAPSULE_NAME); + if (prob) + { + free_problem(prob); + } +} + +#endif /* PROBLEM_COMMON_H */ diff --git a/python/problem/constraint_forward.h b/python/problem/constraint_forward.h new file mode 100644 index 0000000..9449db2 --- /dev/null +++ b/python/problem/constraint_forward.h @@ -0,0 +1,54 @@ +#ifndef PROBLEM_CONSTRAINT_FORWARD_H +#define PROBLEM_CONSTRAINT_FORWARD_H + +#include "common.h" + +static PyObject *py_problem_constraint_forward(PyObject *self, PyObject *args) +{ + PyObject *prob_capsule; + PyObject *u_obj; + if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj)) + { + return NULL; + } + + problem *prob = + (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME); + if (!prob) + { + PyErr_SetString(PyExc_ValueError, "invalid problem capsule"); + return NULL; + } + + PyArrayObject *u_array = + (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + if (!u_array) + { + return NULL; + } + + problem_constraint_forward(prob, (const double *) PyArray_DATA(u_array)); + Py_DECREF(u_array); + + PyObject *out = NULL; + if (prob->total_constraint_size > 0) + { + npy_intp size = prob->total_constraint_size; + out = PyArray_SimpleNew(1, &size, NPY_DOUBLE); + if (!out) + { + return NULL; + } + memcpy(PyArray_DATA((PyArrayObject *) out), prob->constraint_values, + size * sizeof(double)); + } + else + { + npy_intp size = 0; + out = PyArray_SimpleNew(1, &size, NPY_DOUBLE); + } + + return out; +} + +#endif /* PROBLEM_CONSTRAINT_FORWARD_H */ diff --git a/python/problem/gradient.h b/python/problem/gradient.h new file mode 100644 index 0000000..34b9b34 --- /dev/null +++ b/python/problem/gradient.h @@ -0,0 +1,36 @@ +#ifndef PROBLEM_GRADIENT_H +#define PROBLEM_GRADIENT_H + +#include "common.h" + +static PyObject *py_problem_gradient(PyObject *self, PyObject *args) +{ + PyObject *prob_capsule; + if (!PyArg_ParseTuple(args, "O", &prob_capsule)) + { + return NULL; + } + + problem *prob = + (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME); + if (!prob) + { + PyErr_SetString(PyExc_ValueError, "invalid problem capsule"); + return NULL; + } + + problem_gradient(prob); + + npy_intp size = prob->n_vars; + PyObject *out = PyArray_SimpleNew(1, &size, NPY_DOUBLE); + if (!out) + { + return NULL; + } + memcpy(PyArray_DATA((PyArrayObject *) out), prob->gradient_values, + size * sizeof(double)); + + return out; +} + +#endif /* PROBLEM_GRADIENT_H */ diff --git a/python/problem/init_derivatives.h b/python/problem/init_derivatives.h new file mode 100644 index 0000000..cc6298d --- /dev/null +++ b/python/problem/init_derivatives.h @@ -0,0 +1,27 @@ +#ifndef PROBLEM_INIT_DERIVATIVES_H +#define PROBLEM_INIT_DERIVATIVES_H + +#include "common.h" + +static PyObject *py_problem_init_derivatives(PyObject *self, PyObject *args) +{ + PyObject *prob_capsule; + if (!PyArg_ParseTuple(args, "O", &prob_capsule)) + { + return NULL; + } + + problem *prob = + (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME); + if (!prob) + { + PyErr_SetString(PyExc_ValueError, "invalid problem capsule"); + return NULL; + } + + problem_init_derivatives(prob); + + Py_RETURN_NONE; +} + +#endif /* PROBLEM_INIT_DERIVATIVES_H */ diff --git a/python/problem/jacobian.h b/python/problem/jacobian.h new file mode 100644 index 0000000..69f240d --- /dev/null +++ b/python/problem/jacobian.h @@ -0,0 +1,59 @@ +#ifndef PROBLEM_JACOBIAN_H +#define PROBLEM_JACOBIAN_H + +#include "common.h" + +static PyObject *py_problem_jacobian(PyObject *self, PyObject *args) +{ + PyObject *prob_capsule; + if (!PyArg_ParseTuple(args, "O", &prob_capsule)) + { + return NULL; + } + + problem *prob = + (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME); + if (!prob) + { + PyErr_SetString(PyExc_ValueError, "invalid problem capsule"); + return NULL; + } + + if (prob->n_constraints == 0) + { + /* Return empty CSR components */ + npy_intp zero = 0; + npy_intp one = 1; + PyObject *data = PyArray_SimpleNew(1, &zero, NPY_DOUBLE); + PyObject *indices = PyArray_SimpleNew(1, &zero, NPY_INT32); + PyObject *indptr = PyArray_SimpleNew(1, &one, NPY_INT32); + ((int *) PyArray_DATA((PyArrayObject *) indptr))[0] = 0; + return Py_BuildValue("(OOO(ii))", data, indices, indptr, 0, prob->n_vars); + } + + problem_jacobian(prob); + + CSR_Matrix *jac = prob->stacked_jac; + 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); + 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)); + + return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n); +} + +#endif /* PROBLEM_JACOBIAN_H */ diff --git a/python/problem/make_problem.h b/python/problem/make_problem.h new file mode 100644 index 0000000..c9a9a97 --- /dev/null +++ b/python/problem/make_problem.h @@ -0,0 +1,64 @@ +#ifndef PROBLEM_MAKE_H +#define PROBLEM_MAKE_H + +#include "common.h" + +static PyObject *py_make_problem(PyObject *self, PyObject *args) +{ + PyObject *obj_capsule; + PyObject *constraints_list; + if (!PyArg_ParseTuple(args, "OO", &obj_capsule, &constraints_list)) + { + return NULL; + } + + expr *objective = (expr *) PyCapsule_GetPointer(obj_capsule, EXPR_CAPSULE_NAME); + if (!objective) + { + PyErr_SetString(PyExc_ValueError, "invalid objective capsule"); + return NULL; + } + + if (!PyList_Check(constraints_list)) + { + PyErr_SetString(PyExc_TypeError, "constraints must be a list"); + return NULL; + } + + Py_ssize_t n_constraints = PyList_Size(constraints_list); + expr **constraints = NULL; + if (n_constraints > 0) + { + constraints = (expr **) malloc(n_constraints * sizeof(expr *)); + if (!constraints) + { + PyErr_NoMemory(); + return NULL; + } + for (Py_ssize_t i = 0; i < n_constraints; i++) + { + PyObject *c_capsule = PyList_GetItem(constraints_list, i); + constraints[i] = + (expr *) PyCapsule_GetPointer(c_capsule, EXPR_CAPSULE_NAME); + if (!constraints[i]) + { + free(constraints); + PyErr_SetString(PyExc_ValueError, "invalid constraint capsule"); + return NULL; + } + } + } + + problem *prob = new_problem(objective, constraints, (int) n_constraints); + free(constraints); + + if (!prob) + { + PyErr_SetString(PyExc_RuntimeError, "failed to create problem"); + return NULL; + } + + return PyCapsule_New(prob, PROBLEM_CAPSULE_NAME, problem_capsule_destructor); +} + +#endif /* PROBLEM_MAKE_H */ diff --git a/python/problem/objective_forward.h b/python/problem/objective_forward.h new file mode 100644 index 0000000..51b1060 --- /dev/null +++ b/python/problem/objective_forward.h @@ -0,0 +1,37 @@ +#ifndef PROBLEM_OBJECTIVE_FORWARD_H +#define PROBLEM_OBJECTIVE_FORWARD_H + +#include "common.h" + +static PyObject *py_problem_objective_forward(PyObject *self, PyObject *args) +{ + PyObject *prob_capsule; + PyObject *u_obj; + if (!PyArg_ParseTuple(args, "OO", &prob_capsule, &u_obj)) + { + return NULL; + } + + problem *prob = + (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME); + if (!prob) + { + PyErr_SetString(PyExc_ValueError, "invalid problem capsule"); + return NULL; + } + + PyArrayObject *u_array = + (PyArrayObject *) PyArray_FROM_OTF(u_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + if (!u_array) + { + return NULL; + } + + double obj_val = + problem_objective_forward(prob, (const double *) PyArray_DATA(u_array)); + + Py_DECREF(u_array); + return Py_BuildValue("d", obj_val); +} + +#endif /* PROBLEM_OBJECTIVE_FORWARD_H */ diff --git a/python/tests/convert_tests.py b/python/tests/convert_tests.py deleted file mode 100644 index 0bed90a..0000000 --- a/python/tests/convert_tests.py +++ /dev/null @@ -1,358 +0,0 @@ -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!") diff --git a/python/tests/test_constrained.py b/python/tests/test_constrained.py new file mode 100644 index 0000000..9574d3d --- /dev/null +++ b/python/tests/test_constrained.py @@ -0,0 +1,202 @@ +import cvxpy as cp +import numpy as np +from convert import C_problem + +# Note: CVXPY converts constraints A >= B to B - A <= 0 +# So constr.expr for "log(x) >= 0" is "0 - log(x)" = -log(x) +# All constraint values and jacobians are negated compared to the LHS + + +def test_single_constraint(): + """Test C_problem with single constraint.""" + x = cp.Variable(3) + obj = cp.sum(cp.log(x)) + constraints = [cp.log(x) >= 0] # becomes 0 - log(x) <= 0 + + problem = cp.Problem(cp.Minimize(obj), constraints) + prob = C_problem(problem) + + u = np.array([1.0, 2.0, 3.0]) + prob.init_derivatives() + + # Test constraint_forward: constr.expr = -log(x) + constraint_vals = prob.constraint_forward(u) + assert np.allclose(constraint_vals, -np.log(u)) + + # Test jacobian: d/dx(-log(x)) = -1/x + jac = prob.jacobian() + expected_jac = np.diag(-1.0 / u) + assert np.allclose(jac.toarray(), expected_jac) + + +def test_two_constraints(): + """Test C_problem with two constraints.""" + x = cp.Variable(2) + obj = cp.sum(cp.log(x)) + constraints = [ + cp.log(x) >= 0, # becomes -log(x) <= 0 + cp.exp(x) >= 0, # becomes -exp(x) <= 0 + ] + + problem = cp.Problem(cp.Minimize(obj), constraints) + prob = C_problem(problem) + + u = np.array([1.0, 2.0]) + prob.init_derivatives() + + # Test constraint_forward: [-log(u), -exp(u)] + expected_constraint_vals = np.concatenate([-np.log(u), -np.exp(u)]) + constraint_vals = prob.constraint_forward(u) + assert np.allclose(constraint_vals, expected_constraint_vals) + + # Test jacobian - stacked vertically + jac = prob.jacobian() + assert jac.shape == (4, 2) + expected_jac = np.vstack([np.diag(-1.0 / u), np.diag(-np.exp(u))]) + assert np.allclose(jac.toarray(), expected_jac) + + +def test_three_constraints_different_sizes(): + """Test C_problem with three constraints of different types.""" + x = cp.Variable(3) + obj = cp.sum(cp.exp(x)) + constraints = [ + cp.log(x) >= 0, # 3 constraints, becomes -log(x) <= 0 + cp.exp(x) >= 0, # 3 constraints, becomes -exp(x) <= 0 + cp.sum(cp.log(x)) >= 0, # 1 constraint, becomes -sum(log(x)) <= 0 + ] + + problem = cp.Problem(cp.Minimize(obj), constraints) + prob = C_problem(problem) + + u = np.array([1.0, 2.0, 3.0]) + prob.init_derivatives() + + # Test constraint_forward + expected_constraint_vals = np.concatenate([ + -np.log(u), + -np.exp(u), + [-np.sum(np.log(u))] + ]) + constraint_vals = prob.constraint_forward(u) + assert np.allclose(constraint_vals, expected_constraint_vals) + + # Test jacobian shape and values + jac = prob.jacobian() + assert jac.shape == (7, 3) + # First 3 rows: -diag(1/u), next 3 rows: -diag(exp(u)), last row: -1/u + expected_jac = np.zeros((7, 3)) + expected_jac[:3, :] = np.diag(-1.0 / u) + expected_jac[3:6, :] = np.diag(-np.exp(u)) + expected_jac[6, :] = -1.0 / u + assert np.allclose(jac.toarray(), expected_jac) + + +def test_multiple_variables(): + """Test C_problem with multiple CVXPY variables.""" + x = cp.Variable(2) + y = cp.Variable(2) + obj = cp.sum(cp.log(x)) + cp.sum(cp.exp(y)) + constraints = [ + cp.log(x) >= 0, # becomes -log(x) <= 0 + cp.exp(y) >= 0, # becomes -exp(y) <= 0 + ] + + problem = cp.Problem(cp.Minimize(obj), constraints) + prob = C_problem(problem) + + x_vals = np.array([1.0, 2.0]) + y_vals = np.array([0.5, 1.0]) + u = np.concatenate([x_vals, y_vals]) + prob.init_derivatives() + + # Test constraint_forward + expected_constraint_vals = np.concatenate([-np.log(x_vals), -np.exp(y_vals)]) + constraint_vals = prob.constraint_forward(u) + assert np.allclose(constraint_vals, expected_constraint_vals) + + # Test jacobian + jac = prob.jacobian() + assert jac.shape == (4, 4) + expected_jac = np.zeros((4, 4)) + expected_jac[0, 0] = -1.0 / x_vals[0] + expected_jac[1, 1] = -1.0 / x_vals[1] + expected_jac[2, 2] = -np.exp(y_vals[0]) + expected_jac[3, 3] = -np.exp(y_vals[1]) + assert np.allclose(jac.toarray(), expected_jac) + + +def test_larger_scale(): + """Test C_problem with larger variables and multiple constraints.""" + n = 50 + x = cp.Variable(n) + obj = cp.sum(cp.log(x)) + constraints = [ + cp.log(x) >= 0, # n constraints + cp.exp(x) >= 0, # n constraints + cp.sum(cp.log(x)) >= 0, # 1 constraint + cp.sum(cp.exp(x)) >= 0, # 1 constraint + ] + + problem = cp.Problem(cp.Minimize(obj), constraints) + prob = C_problem(problem) + + u = np.linspace(1.0, 5.0, n) + prob.init_derivatives() + + # Test constraint_forward + expected_constraint_vals = np.concatenate([ + -np.log(u), + -np.exp(u), + [-np.sum(np.log(u))], + [-np.sum(np.exp(u))], + ]) + constraint_vals = prob.constraint_forward(u) + assert np.allclose(constraint_vals, expected_constraint_vals) + + # Test jacobian shape + jac = prob.jacobian() + assert jac.shape == (n + n + 1 + 1, n) + + +def test_repeated_evaluations(): + """Test C_problem with repeated evaluations at different points.""" + x = cp.Variable(3) + obj = cp.sum(cp.log(x)) + constraints = [cp.exp(x) >= 0] # becomes -exp(x) <= 0 + + problem = cp.Problem(cp.Minimize(obj), constraints) + prob = C_problem(problem) + + u1 = np.array([1.0, 2.0, 3.0]) + prob.init_derivatives() + + # First evaluation + constraint_vals1 = prob.constraint_forward(u1) + jac1 = prob.jacobian() + + # Second evaluation at different point + u2 = np.array([2.0, 3.0, 4.0]) + constraint_vals2 = prob.constraint_forward(u2) + jac2 = prob.jacobian() + + assert np.allclose(constraint_vals1, -np.exp(u1)) + assert np.allclose(constraint_vals2, -np.exp(u2)) + assert np.allclose(jac1.toarray(), np.diag(-np.exp(u1))) + assert np.allclose(jac2.toarray(), np.diag(-np.exp(u2))) + + +if __name__ == "__main__": + test_single_constraint() + print("test_single_constraint passed!") + test_two_constraints() + print("test_two_constraints passed!") + test_three_constraints_different_sizes() + print("test_three_constraints_different_sizes passed!") + test_multiple_variables() + print("test_multiple_variables passed!") + test_larger_scale() + print("test_larger_scale passed!") + test_repeated_evaluations() + print("test_repeated_evaluations passed!") + print("\nAll constrained tests passed!") diff --git a/python/tests/test_problem_native.py b/python/tests/test_problem_native.py new file mode 100644 index 0000000..2e3a310 --- /dev/null +++ b/python/tests/test_problem_native.py @@ -0,0 +1,168 @@ +import numpy as np +from scipy import sparse +import DNLP_diff_engine as diffengine + + +def test_problem_objective_forward(): + """Test problem_objective_forward and problem_constraint_forward (low-level).""" + n_vars = 3 + x = diffengine.make_variable(n_vars, 1, 0, n_vars) + log_x = diffengine.make_log(x) + objective = diffengine.make_sum(log_x, -1) + constraints = [log_x] + + prob = diffengine.make_problem(objective, constraints) + u = np.array([1.0, 2.0, 3.0]) + diffengine.problem_init_derivatives(prob) + + obj_val = diffengine.problem_objective_forward(prob, u) + constraint_vals = diffengine.problem_constraint_forward(prob, u) + + expected_obj = np.sum(np.log(u)) + assert np.allclose(obj_val, expected_obj) + assert np.allclose(constraint_vals, np.log(u)) + + +def test_problem_constraint_forward(): + """Test problem_constraint_forward for constraint values only (low-level).""" + n_vars = 2 + x = diffengine.make_variable(n_vars, 1, 0, n_vars) + + log_obj = diffengine.make_log(x) + objective = diffengine.make_sum(log_obj, -1) + + log_c = diffengine.make_log(x) + exp_c = diffengine.make_exp(x) + constraints = [log_c, exp_c] + + prob = diffengine.make_problem(objective, constraints) + u = np.array([2.0, 4.0]) + diffengine.problem_init_derivatives(prob) + + constraint_vals = diffengine.problem_constraint_forward(prob, u) + + # Expected: [log(2), log(4), exp(2), exp(4)] + expected = np.concatenate([np.log(u), np.exp(u)]) + assert np.allclose(constraint_vals, expected) + + +def test_problem_gradient(): + """Test problem_gradient for objective gradient (low-level).""" + n_vars = 3 + x = diffengine.make_variable(n_vars, 1, 0, n_vars) + log_x = diffengine.make_log(x) + objective = diffengine.make_sum(log_x, -1) + + prob = diffengine.make_problem(objective, []) + u = np.array([1.0, 2.0, 4.0]) + diffengine.problem_init_derivatives(prob) + + diffengine.problem_objective_forward(prob, u) + grad = diffengine.problem_gradient(prob) + expected_grad = 1.0 / u + assert np.allclose(grad, expected_grad) + + +def test_problem_jacobian(): + """Test problem_jacobian for constraint jacobian (low-level).""" + n_vars = 2 + x = diffengine.make_variable(n_vars, 1, 0, n_vars) + log_x = diffengine.make_log(x) + objective = diffengine.make_sum(log_x, -1) + constraints = [log_x] + + prob = diffengine.make_problem(objective, constraints) + u = np.array([2.0, 4.0]) + diffengine.problem_init_derivatives(prob) + + diffengine.problem_constraint_forward(prob, u) + data, indices, indptr, shape = diffengine.problem_jacobian(prob) + jac = sparse.csr_matrix((data, indices, indptr), shape=shape) + + expected_jac = np.diag(1.0 / u) + assert np.allclose(jac.toarray(), expected_jac) + + +def test_problem_no_constraints(): + """Test Problem with no constraints (low-level).""" + n_vars = 3 + x = diffengine.make_variable(n_vars, 1, 0, n_vars) + log_x = diffengine.make_log(x) + objective = diffengine.make_sum(log_x, -1) + + prob = diffengine.make_problem(objective, []) + u = np.array([1.0, 2.0, 3.0]) + diffengine.problem_init_derivatives(prob) + + obj_val = diffengine.problem_objective_forward(prob, u) + constraint_vals = diffengine.problem_constraint_forward(prob, u) + assert np.allclose(obj_val, np.sum(np.log(u))) + assert len(constraint_vals) == 0 + + diffengine.problem_objective_forward(prob, u) + grad = diffengine.problem_gradient(prob) + assert np.allclose(grad, 1.0 / u) + + diffengine.problem_constraint_forward(prob, u) + data, indices, indptr, shape = diffengine.problem_jacobian(prob) + jac = sparse.csr_matrix((data, indices, indptr), shape=shape) + assert jac.shape == (0, 3) + + +def test_problem_multiple_constraints(): + """Test problem with multiple different constraints (low-level).""" + n_vars = 3 + x = diffengine.make_variable(n_vars, 1, 0, n_vars) + + # Objective: sum(log(x)) + log_x = diffengine.make_log(x) + objective = diffengine.make_sum(log_x, -1) + + # Multiple constraints using the same variable: + # Constraint 1: log(x) - reused from objective + # Constraint 2: exp(x) + exp_x = diffengine.make_exp(x) + constraints = [log_x, exp_x] + + prob = diffengine.make_problem(objective, constraints) + u = np.array([1.0, 2.0, 3.0]) + diffengine.problem_init_derivatives(prob) + + # Test forward pass + obj_val = diffengine.problem_objective_forward(prob, u) + constraint_vals = diffengine.problem_constraint_forward(prob, u) + expected_obj = np.sum(np.log(u)) + expected_constraints = np.concatenate([np.log(u), np.exp(u)]) + assert np.allclose(obj_val, expected_obj) + assert np.allclose(constraint_vals, expected_constraints) + + # Test gradient + diffengine.problem_objective_forward(prob, u) + grad = diffengine.problem_gradient(prob) + expected_grad = 1.0 / u + assert np.allclose(grad, expected_grad) + + # Test Jacobian + diffengine.problem_constraint_forward(prob, u) + data, indices, indptr, shape = diffengine.problem_jacobian(prob) + jac = sparse.csr_matrix((data, indices, indptr), shape=shape) + + # Expected Jacobian: + # log(x): diag(1/u) + # exp(x): diag(exp(u)) + expected_jac = np.vstack([ + np.diag(1.0 / u), + np.diag(np.exp(u)) + ]) + assert jac.shape == (6, 3) + assert np.allclose(jac.toarray(), expected_jac) + + +if __name__ == "__main__": + test_problem_objective_forward() + test_problem_constraint_forward() + test_problem_gradient() + test_problem_jacobian() + test_problem_no_constraints() + test_problem_multiple_constraints() + print("All native problem tests passed!") diff --git a/python/tests/test_unconstrained.py b/python/tests/test_unconstrained.py new file mode 100644 index 0000000..99994a6 --- /dev/null +++ b/python/tests/test_unconstrained.py @@ -0,0 +1,433 @@ +import cvxpy as cp +import numpy as np +from convert import C_problem + +# Note: Current implementation supports: +# - Elementwise ops on variables: log(x), exp(x) +# - Add of elementwise results: log(x) + exp(y) +# - Sum reductions: sum(log(x)) +# - Multiple variables with separate operations +# +# NOT YET SUPPORTED (see unsupported section at bottom of file): +# - Elementwise ops on add results: log(x + y) + + +def test_sum_log(): + """Test sum(log(x)) objective and gradient.""" + x = cp.Variable(4) + problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x)))) + prob = C_problem(problem) + + u = np.array([1.0, 2.0, 3.0, 4.0]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.log(u)) + assert np.allclose(obj_val, expected) + + # Gradient: d/dx sum(log(x)) = 1/x + grad = prob.gradient() + expected_grad = 1.0 / u + assert np.allclose(grad, expected_grad) + + +def test_sum_exp(): + """Test sum(exp(x)) objective and gradient.""" + x = cp.Variable(3) + problem = cp.Problem(cp.Minimize(cp.sum(cp.exp(x)))) + prob = C_problem(problem) + + u = np.array([0.0, 1.0, 2.0]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.exp(u)) + assert np.allclose(obj_val, expected) + + # Gradient: d/dx sum(exp(x)) = exp(x) + grad = prob.gradient() + expected_grad = np.exp(u) + assert np.allclose(grad, expected_grad) + + +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)))) + prob = C_problem(problem) + + u = np.array([1.0, 2.0]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.log(u) + np.exp(u)) + assert np.allclose(obj_val, expected) + + # Gradient: d/dx_i = 1/x_i + exp(x_i) + grad = prob.gradient() + expected_grad = 1.0 / u + np.exp(u) + assert np.allclose(grad, expected_grad) + + +def test_variable_used_multiple_times(): + """Test sum(log(x) + exp(x) + log(x)) - variable used 3 times.""" + x = cp.Variable(3) + problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.exp(x) + cp.log(x)))) + prob = C_problem(problem) + + u = np.array([1.0, 2.0, 3.0]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(2 * np.log(u) + np.exp(u)) + assert np.allclose(obj_val, expected) + + # Gradient: d/dx_i = 2/x_i + exp(x_i) + grad = prob.gradient() + expected_grad = 2.0 / u + np.exp(u) + assert np.allclose(grad, expected_grad) + + +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)))) + prob = C_problem(problem) + + u = np.linspace(1.0, 10.0, 100) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.log(u)) + assert np.allclose(obj_val, expected) + + # Gradient + grad = prob.gradient() + expected_grad = 1.0 / u + assert np.allclose(grad, expected_grad) + + +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)))) + prob = C_problem(problem) + + u = np.arange(1.0, 13.0) # 12 elements + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.log(u)) + assert np.allclose(obj_val, expected) + + # Gradient + grad = prob.gradient() + expected_grad = 1.0 / u + assert np.allclose(grad, expected_grad) + + +def test_two_variables_separate_ops(): + """Test sum(log(x)) + sum(exp(y)) - two variables with separate ops.""" + x = cp.Variable(2) + y = cp.Variable(2) + problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x)) + cp.sum(cp.exp(y)))) + prob = C_problem(problem) + + x_vals = np.array([1.0, 2.0]) + y_vals = np.array([0.5, 1.0]) + u = np.concatenate([x_vals, y_vals]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.log(x_vals)) + np.sum(np.exp(y_vals)) + assert np.allclose(obj_val, expected) + + # Gradient + grad = prob.gradient() + expected_grad = np.concatenate([1.0 / x_vals, np.exp(y_vals)]) + assert np.allclose(grad, expected_grad) + + +def test_two_variables_same_sum(): + """Test sum(log(x) + log(y)) - two variables added before sum.""" + x = cp.Variable(2) + y = cp.Variable(2) + problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x) + cp.log(y)))) + prob = C_problem(problem) + + x_vals = np.array([1.0, 2.0]) + y_vals = np.array([3.0, 4.0]) + u = np.concatenate([x_vals, y_vals]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.log(x_vals) + np.log(y_vals)) + assert np.allclose(obj_val, expected) + + # Gradient + grad = prob.gradient() + expected_grad = np.concatenate([1.0 / x_vals, 1.0 / y_vals]) + assert np.allclose(grad, expected_grad) + + +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)) + )) + prob = C_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]) + u = np.concatenate([a_vals, b_vals, c_vals]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.log(a_vals)) + np.sum(np.log(b_vals)) + np.sum(np.log(c_vals)) + assert np.allclose(obj_val, expected) + + # Gradient + grad = prob.gradient() + expected_grad = 1.0 / u + assert np.allclose(grad, expected_grad) + + +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)) + prob = C_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]) + u = np.concatenate([a_vals, b_vals, c_vals, d_vals]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + 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(obj_val, expected) + + # Gradient + grad = prob.gradient() + expected_grad = np.concatenate([ + 1.0 / a_vals, + 1.0 / b_vals, + np.exp(c_vals), + np.exp(d_vals) + ]) + assert np.allclose(grad, expected_grad) + + +def test_three_variables_mixed(): + """Test sum(log(x) + exp(y) + log(z)) - three variables mixed.""" + 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)) + prob = C_problem(problem) + + x_vals = np.array([1.0, 2.0]) + y_vals = np.array([0.5, 1.0]) + z_vals = np.array([2.0, 3.0]) + u = np.concatenate([x_vals, y_vals, z_vals]) + prob.init_derivatives() + + # Objective + obj_val = prob.objective_forward(u) + expected = np.sum(np.log(x_vals) + np.exp(y_vals) + np.log(z_vals)) + assert np.allclose(obj_val, expected) + + # Gradient + grad = prob.gradient() + expected_grad = np.concatenate([1.0 / x_vals, np.exp(y_vals), 1.0 / z_vals]) + assert np.allclose(grad, expected_grad) + + +def test_repeated_evaluations(): + """Test repeated evaluations at different points.""" + x = cp.Variable(3) + problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x)))) + prob = C_problem(problem) + + u1 = np.array([1.0, 2.0, 3.0]) + prob.init_derivatives() + + # First evaluation + obj1 = prob.objective_forward(u1) + grad1 = prob.gradient() + + # Second evaluation + u2 = np.array([2.0, 3.0, 4.0]) + obj2 = prob.objective_forward(u2) + grad2 = prob.gradient() + + assert np.allclose(obj1, np.sum(np.log(u1))) + assert np.allclose(obj2, np.sum(np.log(u2))) + assert np.allclose(grad1, 1.0 / u1) + assert np.allclose(grad2, 1.0 / u2) + + +# ============================================================================= +# UNSUPPORTED PATTERNS +# ============================================================================= +# The following patterns are NOT YET SUPPORTED because the current +# implementation of elementwise univariate ops (log, exp, etc.) in +# src/elementwise_univariate/common.c assumes that the child expression +# is either a Variable or a LinearOp (which includes negation as -I). +# +# When the child is an Add expression (like x + y), the code incorrectly +# casts it to linear_op_expr and reads invalid memory, causing crashes. +# +# To support these patterns, we would need to generalize the jacobian +# computation in elementwise ops to handle arbitrary child expressions +# by using the child's jacobian directly rather than assuming specific +# structure. +# ============================================================================= + +# def test_elementwise_on_add(): +# """NOT SUPPORTED: log(x + y) - elementwise op on add result. +# +# This fails because log's jacobian computation assumes its child +# is a Variable or LinearOp, but here it's an Add expression. +# """ +# x = cp.Variable(2) +# y = cp.Variable(2) +# problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y)))) +# prob = C_problem(problem) +# +# x_vals = np.array([1.0, 2.0]) +# y_vals = np.array([0.5, 1.0]) +# u = np.concatenate([x_vals, y_vals]) +# prob.init_derivatives() +# +# # Expected objective: sum(log(x + y)) +# obj_val = prob.objective_forward(u) +# expected = np.sum(np.log(x_vals + y_vals)) +# assert np.allclose(obj_val, expected) +# +# # Expected gradient: d/dx_i = 1/(x_i + y_i), d/dy_i = 1/(x_i + y_i) +# grad = prob.gradient() +# grad_xy = 1.0 / (x_vals + y_vals) +# expected_grad = np.concatenate([grad_xy, grad_xy]) +# assert np.allclose(grad, expected_grad) + + +# def test_log_of_sum(): +# """NOT SUPPORTED: log(sum(x)) - elementwise op on sum result. +# +# This fails because log's jacobian computation assumes its child +# is a Variable or LinearOp, but here it's a Sum expression. +# """ +# x = cp.Variable(4) +# problem = cp.Problem(cp.Minimize(cp.log(cp.sum(x)))) +# prob = C_problem(problem) +# +# u = np.array([1.0, 2.0, 3.0, 4.0]) +# prob.init_derivatives() +# +# # Expected objective: log(sum(x)) +# obj_val = prob.objective_forward(u) +# expected = np.log(np.sum(u)) +# assert np.allclose(obj_val, expected) +# +# # Expected gradient: d/dx_i log(sum(x)) = 1/sum(x) +# grad = prob.gradient() +# expected_grad = np.ones_like(u) / np.sum(u) +# assert np.allclose(grad, expected_grad) + + +# def test_mixed_elementwise_on_add(): +# """NOT SUPPORTED: log(x + y) + exp(y + z) - multiple elementwise on adds. +# +# Same underlying issue - elementwise ops don't support Add children. +# """ +# x = cp.Variable(2) +# y = cp.Variable(2) +# z = cp.Variable(2) +# obj = cp.sum(cp.log(x + y)) + cp.sum(cp.exp(y + z)) +# problem = cp.Problem(cp.Minimize(obj)) +# prob = C_problem(problem) +# +# x_vals = np.array([1.0, 2.0]) +# y_vals = np.array([0.5, 1.0]) +# z_vals = np.array([0.1, 0.2]) +# u = np.concatenate([x_vals, y_vals, z_vals]) +# prob.init_derivatives() +# +# # Expected objective +# obj_val = prob.objective_forward(u) +# expected = np.sum(np.log(x_vals + y_vals)) + np.sum(np.exp(y_vals + z_vals)) +# assert np.allclose(obj_val, expected) + + +# def test_nested_add_in_elementwise(): +# """NOT SUPPORTED: log(x + y + z) - nested adds inside elementwise. +# +# Same underlying issue. +# """ +# x = cp.Variable(2) +# y = cp.Variable(2) +# z = cp.Variable(2) +# problem = cp.Problem(cp.Minimize(cp.sum(cp.log(x + y + z)))) +# prob = C_problem(problem) +# +# x_vals = np.array([1.0, 2.0]) +# y_vals = np.array([0.5, 1.0]) +# z_vals = np.array([0.5, 0.5]) +# u = np.concatenate([x_vals, y_vals, z_vals]) +# prob.init_derivatives() +# +# # Expected objective +# obj_val = prob.objective_forward(u) +# expected = np.sum(np.log(x_vals + y_vals + z_vals)) +# assert np.allclose(obj_val, expected) + + +if __name__ == "__main__": + test_sum_log() + print("test_sum_log passed!") + test_sum_exp() + print("test_sum_exp passed!") + test_variable_reuse() + print("test_variable_reuse passed!") + test_variable_used_multiple_times() + print("test_variable_used_multiple_times passed!") + test_larger_variable() + print("test_larger_variable passed!") + test_matrix_variable() + print("test_matrix_variable passed!") + test_two_variables_separate_ops() + print("test_two_variables_separate_ops passed!") + test_two_variables_same_sum() + print("test_two_variables_same_sum passed!") + test_mixed_sizes() + print("test_mixed_sizes passed!") + test_multiple_variables_log_exp() + print("test_multiple_variables_log_exp passed!") + test_three_variables_mixed() + print("test_three_variables_mixed passed!") + test_repeated_evaluations() + print("test_repeated_evaluations passed!") + print("\nAll unconstrained tests passed!") diff --git a/src/affine/constant.c b/src/affine/constant.c index d427ec9..577029b 100644 --- a/src/affine/constant.c +++ b/src/affine/constant.c @@ -8,6 +8,19 @@ static void forward(expr *node, const double *u) (void) u; } +static void jacobian_init(expr *node) +{ + /* Constant jacobian is all zeros: size x n_vars with 0 nonzeros. + * new_csr_matrix uses calloc for row pointers, so they're already 0. */ + node->jacobian = new_csr_matrix(node->size, node->n_vars, 0); +} + +static void eval_jacobian(expr *node) +{ + /* Constant jacobian never changes - nothing to evaluate */ + (void) node; +} + static bool is_affine(const expr *node) { (void) node; @@ -20,6 +33,8 @@ expr *new_constant(int d1, int d2, int n_vars, const double *values) memcpy(node->value, values, d1 * d2 * sizeof(double)); node->forward = forward; node->is_affine = is_affine; + node->jacobian_init = jacobian_init; + node->eval_jacobian = eval_jacobian; return node; } diff --git a/src/affine/neg.c b/src/affine/neg.c new file mode 100644 index 0000000..632372b --- /dev/null +++ b/src/affine/neg.c @@ -0,0 +1,90 @@ +#include "affine.h" +#include + +static void forward(expr *node, const double *u) +{ + /* child's forward pass */ + node->left->forward(node->left, u); + + /* negate values */ + for (int i = 0; i < node->size; i++) + { + node->value[i] = -node->left->value[i]; + } +} + +static void jacobian_init(expr *node) +{ + /* initialize child's jacobian */ + node->left->jacobian_init(node->left); + + /* same sparsity pattern as child */ + CSR_Matrix *child_jac = node->left->jacobian; + node->jacobian = new_csr_matrix(child_jac->m, child_jac->n, child_jac->nnz); + + /* copy row pointers and column indices (sparsity pattern is constant) */ + memcpy(node->jacobian->p, child_jac->p, (child_jac->m + 1) * sizeof(int)); + memcpy(node->jacobian->i, child_jac->i, child_jac->nnz * sizeof(int)); + node->jacobian->nnz = child_jac->nnz; +} + +static void eval_jacobian(expr *node) +{ + /* evaluate child's jacobian */ + node->left->eval_jacobian(node->left); + + /* negate values only (sparsity pattern set in jacobian_init) */ + CSR_Matrix *child_jac = node->left->jacobian; + for (int k = 0; k < child_jac->nnz; k++) + { + node->jacobian->x[k] = -child_jac->x[k]; + } +} + +static void wsum_hess_init(expr *node) +{ + /* initialize child's wsum_hess */ + node->left->wsum_hess_init(node->left); + + /* same sparsity pattern as child */ + CSR_Matrix *child_hess = node->left->wsum_hess; + node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz); + + /* copy row pointers and column indices (sparsity pattern is constant) */ + memcpy(node->wsum_hess->p, child_hess->p, (child_hess->m + 1) * sizeof(int)); + memcpy(node->wsum_hess->i, child_hess->i, child_hess->nnz * sizeof(int)); + node->wsum_hess->nnz = child_hess->nnz; +} + +static void eval_wsum_hess(expr *node, const double *w) +{ + /* For f(x) = -g(x): d²f/dx² = -d²g/dx² */ + node->left->eval_wsum_hess(node->left, w); + + /* negate values (sparsity pattern set in wsum_hess_init) */ + CSR_Matrix *child_hess = node->left->wsum_hess; + for (int k = 0; k < child_hess->nnz; k++) + { + node->wsum_hess->x[k] = -child_hess->x[k]; + } +} + +static bool is_affine(const expr *node) +{ + return node->left->is_affine(node->left); +} + +expr *new_neg(expr *child) +{ + expr *node = new_expr(child->d1, child->d2, child->n_vars); + node->left = child; + expr_retain(child); + node->forward = forward; + node->is_affine = is_affine; + node->jacobian_init = jacobian_init; + node->eval_jacobian = eval_jacobian; + node->wsum_hess_init = wsum_hess_init; + node->eval_wsum_hess = eval_wsum_hess; + + return node; +} diff --git a/src/affine/promote.c b/src/affine/promote.c new file mode 100644 index 0000000..fca5c80 --- /dev/null +++ b/src/affine/promote.c @@ -0,0 +1,110 @@ +#include "affine.h" +#include +#include + +/* Promote broadcasts a scalar expression to a vector/matrix shape. + * This matches CVXPY's promote atom which only handles scalars. */ + +static void forward(expr *node, const double *u) +{ + node->left->forward(node->left, u); + + /* broadcast scalar value to all output elements */ + double val = node->left->value[0]; + for (int i = 0; i < node->size; i++) + { + node->value[i] = val; + } +} + +static void jacobian_init(expr *node) +{ + node->left->jacobian_init(node->left); + + /* Each output row copies the single row from child's jacobian */ + CSR_Matrix *child_jac = node->left->jacobian; + int nnz = node->size * child_jac->nnz; + + node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz); + CSR_Matrix *jac = node->jacobian; + + /* Build sparsity pattern by replicating child's single row */ + int child_nnz = child_jac->p[1] - child_jac->p[0]; + jac->nnz = 0; + for (int row = 0; row < node->size; row++) + { + jac->p[row] = jac->nnz; + memcpy(jac->i + jac->nnz, child_jac->i + child_jac->p[0], + child_nnz * sizeof(int)); + jac->nnz += child_nnz; + } + jac->p[node->size] = jac->nnz; +} + +static void eval_jacobian(expr *node) +{ + node->left->eval_jacobian(node->left); + + CSR_Matrix *child_jac = node->left->jacobian; + CSR_Matrix *jac = node->jacobian; + int child_nnz = child_jac->p[1] - child_jac->p[0]; + + /* Copy child's row values to each output row */ + for (int row = 0; row < node->size; row++) + { + memcpy(jac->x + row * child_nnz, child_jac->x + child_jac->p[0], + child_nnz * sizeof(double)); + } +} + +static void wsum_hess_init(expr *node) +{ + node->left->wsum_hess_init(node->left); + + /* same sparsity as child since we're summing weights */ + CSR_Matrix *child_hess = node->left->wsum_hess; + node->wsum_hess = new_csr_matrix(child_hess->m, child_hess->n, child_hess->nnz); + + /* copy sparsity pattern */ + memcpy(node->wsum_hess->p, child_hess->p, (child_hess->m + 1) * sizeof(int)); + memcpy(node->wsum_hess->i, child_hess->i, child_hess->nnz * sizeof(int)); + node->wsum_hess->nnz = child_hess->nnz; +} + +static void eval_wsum_hess(expr *node, const double *w) +{ + /* Sum all weights (they all correspond to the same scalar child) */ + double sum_w = 0.0; + for (int i = 0; i < node->size; i++) + { + sum_w += w[i]; + } + + /* evaluate child's wsum_hess with summed weight */ + node->left->eval_wsum_hess(node->left, &sum_w); + + /* copy values */ + CSR_Matrix *child_hess = node->left->wsum_hess; + memcpy(node->wsum_hess->x, child_hess->x, child_hess->nnz * sizeof(double)); +} + +static bool is_affine(const expr *node) +{ + return node->left->is_affine(node->left); +} + +expr *new_promote(expr *child, int d1, int d2) +{ + expr *node = new_expr(d1, d2, child->n_vars); + node->forward = forward; + node->jacobian_init = jacobian_init; + node->eval_jacobian = eval_jacobian; + node->is_affine = is_affine; + node->wsum_hess_init = wsum_hess_init; + node->eval_wsum_hess = eval_wsum_hess; + + node->left = child; + expr_retain(child); + + return node; +} diff --git a/src/affine/variable.c b/src/affine/variable.c index 43b706f..8ce8c6e 100644 --- a/src/affine/variable.c +++ b/src/affine/variable.c @@ -21,7 +21,7 @@ static void jacobian_init(expr *node) static void eval_jacobian(expr *node) { - /* Jacobian is already initialized with correct values */ + /* Variable jacobian never changes - nothing to evaluate */ (void) node; } diff --git a/src/problem.c b/src/problem.c new file mode 100644 index 0000000..945d0a2 --- /dev/null +++ b/src/problem.c @@ -0,0 +1,178 @@ +#include "problem.h" +#include +#include + +problem *new_problem(expr *objective, expr **constraints, int n_constraints) +{ + problem *prob = (problem *) calloc(1, sizeof(problem)); + if (!prob) return NULL; + + /* Retain objective (shared ownership with caller) */ + prob->objective = objective; + expr_retain(objective); + + /* Copy and retain constraints array */ + prob->n_constraints = n_constraints; + if (n_constraints > 0) + { + prob->constraints = (expr **) malloc(n_constraints * sizeof(expr *)); + for (int i = 0; i < n_constraints; i++) + { + prob->constraints[i] = constraints[i]; + expr_retain(constraints[i]); + } + } + else + { + prob->constraints = NULL; + } + + /* Compute total constraint size */ + prob->total_constraint_size = 0; + for (int i = 0; i < n_constraints; i++) + { + prob->total_constraint_size += constraints[i]->size; + } + + prob->n_vars = objective->n_vars; + + /* Allocate value arrays */ + if (prob->total_constraint_size > 0) + { + prob->constraint_values = + (double *) calloc(prob->total_constraint_size, sizeof(double)); + } + else + { + prob->constraint_values = NULL; + } + prob->gradient_values = (double *) calloc(prob->n_vars, sizeof(double)); + + /* Derivative structures allocated by problem_init_derivatives */ + prob->stacked_jac = NULL; + + return prob; +} + +void problem_init_derivatives(problem *prob) +{ + /* 1. Initialize objective jacobian */ + prob->objective->jacobian_init(prob->objective); + + /* 2. Initialize constraint jacobians and count total nnz */ + int total_nnz = 0; + for (int i = 0; i < prob->n_constraints; i++) + { + expr *c = prob->constraints[i]; + c->jacobian_init(c); + total_nnz += c->jacobian->nnz; + } + + /* 3. Allocate stacked jacobian */ + if (prob->total_constraint_size > 0) + { + prob->stacked_jac = + new_csr_matrix(prob->total_constraint_size, prob->n_vars, total_nnz); + } + + /* TODO: 4. Initialize objective wsum_hess */ + + /* TODO: 5. Initialize constraint wsum_hess */ +} + +void free_problem(problem *prob) +{ + if (prob == NULL) return; + + /* Free allocated arrays */ + free(prob->constraint_values); + free(prob->gradient_values); + free_csr_matrix(prob->stacked_jac); + + /* Release expression references (decrements refcount) */ + free_expr(prob->objective); + for (int i = 0; i < prob->n_constraints; i++) + { + free_expr(prob->constraints[i]); + } + free(prob->constraints); + + /* Free problem struct */ + free(prob); +} + +double problem_objective_forward(problem *prob, const double *u) +{ + /* Evaluate objective only */ + prob->objective->forward(prob->objective, u); + return prob->objective->value[0]; +} + +void problem_constraint_forward(problem *prob, const double *u) +{ + /* Evaluate constraints only and copy values */ + 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; + } +} + +void problem_gradient(problem *prob) +{ + /* Jacobian on objective */ + prob->objective->eval_jacobian(prob->objective); + + /* Zero gradient array */ + memset(prob->gradient_values, 0, prob->n_vars * sizeof(double)); + + /* Copy sparse jacobian row to dense gradient + * Objective jacobian is 1 x n_vars */ + CSR_Matrix *jac = prob->objective->jacobian; + for (int k = jac->p[0]; k < jac->p[1]; k++) + { + int col = jac->i[k]; + prob->gradient_values[col] = jac->x[k]; + } +} + +void problem_jacobian(problem *prob) +{ + CSR_Matrix *stacked = prob->stacked_jac; + + /* Initialize row pointers */ + stacked->p[0] = 0; + + int row_offset = 0; + int nnz_offset = 0; + + for (int i = 0; i < prob->n_constraints; i++) + { + expr *c = prob->constraints[i]; + + /* Evaluate jacobian */ + c->eval_jacobian(c); + + CSR_Matrix *cjac = c->jacobian; + + /* Copy row pointers with offset */ + for (int r = 1; r <= cjac->m; r++) + { + stacked->p[row_offset + r] = nnz_offset + cjac->p[r]; + } + + /* Copy column indices and values */ + int constraint_nnz = cjac->p[cjac->m]; + memcpy(stacked->i + nnz_offset, cjac->i, constraint_nnz * sizeof(int)); + memcpy(stacked->x + nnz_offset, cjac->x, constraint_nnz * sizeof(double)); + + row_offset += cjac->m; + nnz_offset += constraint_nnz; + } + + /* Update actual nnz (may be less than allocated) */ + stacked->nnz = nnz_offset; +} diff --git a/tests/all_tests.c b/tests/all_tests.c index 768ba86..502116e 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -6,6 +6,8 @@ #include "forward_pass/affine/test_add.h" #include "forward_pass/affine/test_hstack.h" #include "forward_pass/affine/test_linear_op.h" +#include "forward_pass/affine/test_neg.h" +#include "forward_pass/affine/test_promote.h" #include "forward_pass/affine/test_sum.h" #include "forward_pass/affine/test_variable_constant.h" #include "forward_pass/composite/test_composite.h" @@ -13,13 +15,16 @@ #include "forward_pass/elementwise/test_log.h" #include "jacobian_tests/test_composite.h" #include "jacobian_tests/test_elementwise_mult.h" +#include "jacobian_tests/test_neg.h" #include "jacobian_tests/test_hstack.h" #include "jacobian_tests/test_log.h" +#include "jacobian_tests/test_promote.h" #include "jacobian_tests/test_prod.h" #include "jacobian_tests/test_quad_form.h" #include "jacobian_tests/test_quad_over_lin.h" #include "jacobian_tests/test_rel_entr.h" #include "jacobian_tests/test_sum.h" +#include "problem/test_problem.h" #include "jacobian_tests/test_trace.h" #include "utils/test_csc_matrix.h" #include "utils/test_csr_matrix.h" @@ -51,6 +56,8 @@ int main(void) mu_run_test(test_constant, tests_run); mu_run_test(test_addition, tests_run); mu_run_test(test_linear_op, tests_run); + mu_run_test(test_neg_forward, tests_run); + mu_run_test(test_promote_scalar_to_vector, tests_run); mu_run_test(test_exp, tests_run); mu_run_test(test_log, tests_run); mu_run_test(test_composite, tests_run); @@ -61,6 +68,8 @@ int main(void) mu_run_test(test_hstack_forward_matrix, tests_run); printf("\n--- Jacobian Tests ---\n"); + mu_run_test(test_neg_jacobian, tests_run); + mu_run_test(test_neg_chain, tests_run); mu_run_test(test_jacobian_log, tests_run); mu_run_test(test_jacobian_log_matrix, tests_run); mu_run_test(test_jacobian_composite_log, tests_run); @@ -89,6 +98,10 @@ int main(void) mu_run_test(test_jacobian_sum_log_axis_1, tests_run); mu_run_test(test_jacobian_hstack_vectors, tests_run); mu_run_test(test_jacobian_hstack_matrix, tests_run); + mu_run_test(test_promote_scalar_jacobian, tests_run); + mu_run_test(test_promote_scalar_to_matrix_jacobian, tests_run); + mu_run_test(test_wsum_hess_multiply_1, tests_run); + mu_run_test(test_wsum_hess_multiply_2, tests_run); mu_run_test(test_jacobian_trace_variable, tests_run); mu_run_test(test_jacobian_trace_composite, tests_run); @@ -150,6 +163,14 @@ int main(void) mu_run_test(test_ATA_alloc_random2, tests_run); mu_run_test(test_BTA_alloc_and_BTDA_fill, tests_run); + printf("\n--- Problem Struct Tests ---\n"); + mu_run_test(test_problem_new_free, tests_run); + mu_run_test(test_problem_objective_forward, tests_run); + mu_run_test(test_problem_gradient, tests_run); + mu_run_test(test_problem_jacobian, tests_run); + mu_run_test(test_problem_jacobian_multi, tests_run); + mu_run_test(test_problem_constraint_forward, tests_run); + printf("\n=== All %d tests passed ===\n", tests_run); return 0; diff --git a/tests/forward_pass/affine/test_neg.h b/tests/forward_pass/affine/test_neg.h new file mode 100644 index 0000000..0fad2b1 --- /dev/null +++ b/tests/forward_pass/affine/test_neg.h @@ -0,0 +1,20 @@ +#include +#include +#include + +#include "affine.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_neg_forward(void) +{ + double u[3] = {1.0, 2.0, 3.0}; + expr *var = new_variable(3, 1, 0, 3); + expr *neg_node = new_neg(var); + neg_node->forward(neg_node, u); + double expected[3] = {-1.0, -2.0, -3.0}; + mu_assert("neg forward failed", cmp_double_array(neg_node->value, expected, 3)); + free_expr(neg_node); + return 0; +} diff --git a/tests/forward_pass/affine/test_promote.h b/tests/forward_pass/affine/test_promote.h new file mode 100644 index 0000000..dcc9060 --- /dev/null +++ b/tests/forward_pass/affine/test_promote.h @@ -0,0 +1,24 @@ +#include +#include +#include + +#include "affine.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_promote_scalar_to_vector(void) +{ + /* Promote scalar to 4-element vector */ + double u[1] = {5.0}; + expr *var = new_variable(1, 1, 0, 1); + expr *promote_node = new_promote(var, 4, 1); + promote_node->forward(promote_node, u); + + double expected[4] = {5.0, 5.0, 5.0, 5.0}; + mu_assert("promote scalar->vector forward failed", + cmp_double_array(promote_node->value, expected, 4)); + + free_expr(promote_node); + return 0; +} diff --git a/tests/jacobian_tests/test_neg.h b/tests/jacobian_tests/test_neg.h new file mode 100644 index 0000000..f5377dc --- /dev/null +++ b/tests/jacobian_tests/test_neg.h @@ -0,0 +1,62 @@ +#include + +#include "affine.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_neg_jacobian(void) +{ + double u[3] = {1.0, 2.0, 3.0}; + expr *var = new_variable(3, 1, 0, 3); + expr *neg_node = new_neg(var); + neg_node->forward(neg_node, u); + neg_node->jacobian_init(neg_node); + neg_node->eval_jacobian(neg_node); + + /* Jacobian of neg(x) is -I (diagonal with -1) */ + double expected_x[3] = {-1.0, -1.0, -1.0}; + int expected_p[4] = {0, 1, 2, 3}; + int expected_i[3] = {0, 1, 2}; + + mu_assert("neg jacobian vals fail", + cmp_double_array(neg_node->jacobian->x, expected_x, 3)); + mu_assert("neg jacobian rows fail", + cmp_int_array(neg_node->jacobian->p, expected_p, 4)); + mu_assert("neg jacobian cols fail", + cmp_int_array(neg_node->jacobian->i, expected_i, 3)); + + free_expr(neg_node); + return 0; +} + +const char *test_neg_chain(void) +{ + /* Test neg(neg(x)) = x */ + double u[3] = {1.0, 2.0, 3.0}; + expr *var = new_variable(3, 1, 0, 3); + expr *neg1 = new_neg(var); + expr *neg2 = new_neg(neg1); + neg2->forward(neg2, u); + + /* neg(neg(x)) should equal x */ + mu_assert("neg chain forward failed", cmp_double_array(neg2->value, u, 3)); + + neg2->jacobian_init(neg2); + neg2->eval_jacobian(neg2); + + /* Jacobian of neg(neg(x)) is (-1)*(-1)*I = I */ + double expected_x[3] = {1.0, 1.0, 1.0}; + int expected_p[4] = {0, 1, 2, 3}; + int expected_i[3] = {0, 1, 2}; + + mu_assert("neg chain jacobian vals fail", + cmp_double_array(neg2->jacobian->x, expected_x, 3)); + mu_assert("neg chain jacobian rows fail", + cmp_int_array(neg2->jacobian->p, expected_p, 4)); + mu_assert("neg chain jacobian cols fail", + cmp_int_array(neg2->jacobian->i, expected_i, 3)); + + free_expr(neg2); + return 0; +} diff --git a/tests/jacobian_tests/test_promote.h b/tests/jacobian_tests/test_promote.h new file mode 100644 index 0000000..978d31b --- /dev/null +++ b/tests/jacobian_tests/test_promote.h @@ -0,0 +1,66 @@ +#include +#include +#include + +#include "affine.h" +#include "expr.h" +#include "minunit.h" +#include "test_helpers.h" + +const char *test_promote_scalar_jacobian(void) +{ + /* Promote scalar to 3-element vector, check jacobian */ + double u[1] = {2.0}; + expr *var = new_variable(1, 1, 0, 1); + expr *promote_node = new_promote(var, 3, 1); + promote_node->forward(promote_node, u); + promote_node->jacobian_init(promote_node); + promote_node->eval_jacobian(promote_node); + + /* Jacobian is 3x1 with all 1s (each output depends on same input) */ + double expected_x[3] = {1.0, 1.0, 1.0}; + int expected_p[4] = {0, 1, 2, 3}; + int expected_i[3] = {0, 0, 0}; + + mu_assert("promote jacobian vals fail", + cmp_double_array(promote_node->jacobian->x, expected_x, 3)); + mu_assert("promote jacobian rows fail", + cmp_int_array(promote_node->jacobian->p, expected_p, 4)); + mu_assert("promote jacobian cols fail", + cmp_int_array(promote_node->jacobian->i, expected_i, 3)); + + free_expr(promote_node); + return 0; +} + +const char *test_promote_scalar_to_matrix_jacobian(void) +{ + /* Promote scalar to 2x3 matrix, check jacobian */ + double u[1] = {7.0}; + expr *var = new_variable(1, 1, 0, 1); + expr *promote_node = new_promote(var, 2, 3); + promote_node->forward(promote_node, u); + + /* Forward: all 6 elements should be 7.0 */ + double expected_val[6] = {7.0, 7.0, 7.0, 7.0, 7.0, 7.0}; + mu_assert("promote scalar->matrix forward failed", + cmp_double_array(promote_node->value, expected_val, 6)); + + promote_node->jacobian_init(promote_node); + promote_node->eval_jacobian(promote_node); + + /* Jacobian is 6x1 with all 1s (each output depends on same scalar input) */ + double expected_x[6] = {1.0, 1.0, 1.0, 1.0, 1.0, 1.0}; + int expected_p[7] = {0, 1, 2, 3, 4, 5, 6}; + int expected_i[6] = {0, 0, 0, 0, 0, 0}; + + mu_assert("promote matrix jacobian vals fail", + cmp_double_array(promote_node->jacobian->x, expected_x, 6)); + mu_assert("promote matrix jacobian rows fail", + cmp_int_array(promote_node->jacobian->p, expected_p, 7)); + mu_assert("promote matrix jacobian cols fail", + cmp_int_array(promote_node->jacobian->i, expected_i, 6)); + + free_expr(promote_node); + return 0; +} diff --git a/tests/problem/test_problem.h b/tests/problem/test_problem.h new file mode 100644 index 0000000..7ea502b --- /dev/null +++ b/tests/problem/test_problem.h @@ -0,0 +1,270 @@ +#ifndef TEST_PROBLEM_H +#define TEST_PROBLEM_H + +#include +#include + +#include "affine.h" +#include "elementwise_univariate.h" +#include "expr.h" +#include "minunit.h" +#include "problem.h" +#include "test_helpers.h" + +/* + * Test problem: minimize sum(log(x)) + * subject to x >= 1 (as x - 1 >= 0) + * + * With x of size 3, n_vars = 3 + */ +const char *test_problem_new_free(void) +{ + /* Create expressions */ + expr *x = new_variable(3, 1, 0, 3); + expr *log_x = new_log(x); + expr *objective = new_sum(log_x, -1); + + /* Create constraint: x - 1 (represented as just x for simplicity) */ + expr *x_constraint = new_variable(3, 1, 0, 3); + + expr *constraints[1] = {x_constraint}; + + /* Create problem */ + problem *prob = new_problem(objective, constraints, 1); + + mu_assert("new_problem failed", prob != NULL); + mu_assert("n_vars wrong", prob->n_vars == 3); + mu_assert("n_constraints wrong", prob->n_constraints == 1); + mu_assert("total_constraint_size wrong", prob->total_constraint_size == 3); + + /* Free problem (recursively frees expressions) */ + free_problem(prob); + + return 0; +} + +/* + * Test problem_objective_forward: minimize sum(log(x)) + * subject to x (as constraint) + */ +const char *test_problem_objective_forward(void) +{ + expr *x = new_variable(3, 1, 0, 3); + expr *log_x = new_log(x); + expr *objective = new_sum(log_x, -1); + + expr *x_constraint = new_variable(3, 1, 0, 3); + expr *constraints[1] = {x_constraint}; + + problem *prob = new_problem(objective, constraints, 1); + + double u[3] = {1.0, 2.0, 3.0}; + problem_init_derivatives(prob); + + double obj_val = problem_objective_forward(prob, u); + + /* Expected: sum(log([1, 2, 3])) = 0 + log(2) + log(3) */ + double expected_obj = log(1.0) + log(2.0) + log(3.0); + mu_assert("objective value wrong", fabs(obj_val - expected_obj) < 1e-10); + + /* Now evaluate constraints separately */ + problem_constraint_forward(prob, u); + + /* Constraint values should be [1, 2, 3] */ + double expected_constraints[3] = {1.0, 2.0, 3.0}; + mu_assert("constraint values wrong", + cmp_double_array(prob->constraint_values, expected_constraints, 3)); + + free_problem(prob); + + return 0; +} + +/* + * Test problem_constraint_forward: evaluate constraints only + * Constraint 1: log(x) -> [log(2), log(4)] + * Constraint 2: exp(x) -> [exp(2), exp(4)] + */ +const char *test_problem_constraint_forward(void) +{ + int n_vars = 2; + + /* Shared variable */ + expr *x = new_variable(2, 1, 0, n_vars); + + /* Objective: sum(log(x)) */ + expr *log_obj = new_log(x); + expr *objective = new_sum(log_obj, -1); + + /* Constraint 1: log(x) */ + expr *log_c1 = new_log(x); + + /* Constraint 2: exp(x) */ + expr *exp_c2 = new_exp(x); + + expr *constraints[2] = {log_c1, exp_c2}; + + problem *prob = new_problem(objective, constraints, 2); + + double u[2] = {2.0, 4.0}; + problem_init_derivatives(prob); + + problem_constraint_forward(prob, u); + + /* Check constraint values: [log(2), log(4), exp(2), exp(4)] */ + double expected_constraints[4] = {log(2.0), log(4.0), exp(2.0), exp(4.0)}; + mu_assert("constraint values wrong", + cmp_double_array(prob->constraint_values, expected_constraints, 4)); + + free_problem(prob); + + return 0; +} + +/* + * Test problem_gradient: gradient of sum(log(x)) = [1/x1, 1/x2, 1/x3] + */ +const char *test_problem_gradient(void) +{ + expr *x = new_variable(3, 1, 0, 3); + expr *log_x = new_log(x); + expr *objective = new_sum(log_x, -1); + + problem *prob = new_problem(objective, NULL, 0); + + double u[3] = {1.0, 2.0, 4.0}; + problem_init_derivatives(prob); + + problem_objective_forward(prob, u); + problem_gradient(prob); + + /* Expected gradient: [1/1, 1/2, 1/4] = [1.0, 0.5, 0.25] */ + double expected_grad[3] = {1.0, 0.5, 0.25}; + mu_assert("gradient wrong", + cmp_double_array(prob->gradient_values, expected_grad, 3)); + + free_problem(prob); + + return 0; +} + +/* + * Test problem_jacobian: one constraint log(x) + * Jacobian of log(x): diag([1/x1, 1/x2]) + */ +const char *test_problem_jacobian(void) +{ + int n_vars = 2; + + /* Create separate expression trees */ + expr *x_obj = new_variable(2, 1, 0, n_vars); + expr *log_obj = new_log(x_obj); + expr *objective = new_sum(log_obj, -1); + + expr *x_c1 = new_variable(2, 1, 0, n_vars); + expr *log_c1 = new_log(x_c1); + + expr *constraints[1] = {log_c1}; + + problem *prob = new_problem(objective, constraints, 1); + + double u[2] = {2.0, 4.0}; + problem_init_derivatives(prob); + + problem_constraint_forward(prob, u); + problem_jacobian(prob); + + CSR_Matrix *jac = prob->stacked_jac; + + /* Check dimensions */ + mu_assert("jac rows wrong", jac->m == 2); + mu_assert("jac cols wrong", jac->n == 2); + + /* Check row pointers: each row has 1 element */ + int expected_p[3] = {0, 1, 2}; + mu_assert("jac->p wrong", cmp_int_array(jac->p, expected_p, 3)); + + /* Check column indices */ + int expected_i[2] = {0, 1}; + mu_assert("jac->i wrong", cmp_int_array(jac->i, expected_i, 2)); + + /* Check values: [1/2, 1/4] */ + double expected_x[2] = {0.5, 0.25}; + mu_assert("jac->x wrong", cmp_double_array(jac->x, expected_x, 2)); + + free_problem(prob); + + return 0; +} + +/* + * Test problem_jacobian with multiple constraints and SHARED variable: + * Constraint 1: log(x) -> Jacobian diag([1/x1, 1/x2]) + * Constraint 2: exp(x) -> Jacobian diag([exp(x1), exp(x2)]) + * + * Stacked jacobian (4x2): + * [ 1/x1 0 ] + * [ 0 1/x2 ] + * [exp(x1) 0 ] + * [ 0 exp(x2)] + * + * Note: All expressions share the same variable node x, testing that + * free_problem correctly handles shared nodes without double-free. + */ +const char *test_problem_jacobian_multi(void) +{ + int n_vars = 2; + + /* Single shared variable used across all expressions */ + expr *x = new_variable(2, 1, 0, n_vars); + + /* Objective: sum(log(x)) */ + expr *log_obj = new_log(x); + expr *objective = new_sum(log_obj, -1); + + /* Constraint 1: log(x) - shares x */ + expr *log_c1 = new_log(x); + + /* Constraint 2: exp(x) - shares x */ + expr *exp_c2 = new_exp(x); + + expr *constraints[2] = {log_c1, exp_c2}; + + problem *prob = new_problem(objective, constraints, 2); + + double u[2] = {2.0, 4.0}; + problem_init_derivatives(prob); + + problem_constraint_forward(prob, u); + problem_jacobian(prob); + + CSR_Matrix *jac = prob->stacked_jac; + + /* Check dimensions: 4 rows (2 + 2), 2 cols */ + mu_assert("jac rows wrong", jac->m == 4); + mu_assert("jac cols wrong", jac->n == 2); + mu_assert("jac nnz wrong", jac->nnz == 4); + + /* Check row pointers: each row has 1 element */ + int expected_p[5] = {0, 1, 2, 3, 4}; + mu_assert("jac->p wrong", cmp_int_array(jac->p, expected_p, 5)); + + /* Check column indices: diagonal pattern */ + int expected_i[4] = {0, 1, 0, 1}; + mu_assert("jac->i wrong", cmp_int_array(jac->i, expected_i, 4)); + + /* Check values: + * Row 0: 1/2 = 0.5 + * Row 1: 1/4 = 0.25 + * Row 2: exp(2) ≈ 7.389 + * Row 3: exp(4) ≈ 54.598 + */ + double expected_x[4] = {0.5, 0.25, exp(2.0), exp(4.0)}; + mu_assert("jac->x wrong", cmp_double_array(jac->x, expected_x, 4)); + + free_problem(prob); + + return 0; +} + +#endif /* TEST_PROBLEM_H */