[Ready for review] Adds problem struct, neg and promote atoms and a lot more tests

.gitignore

@@ -72,4 +72,7 @@ out/
 *.swp
 *.swo
 *~
-.DS_Store
+.DS_Store
+
+.venv/
+__pycache__/

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);

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

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 */

python/atoms/common.h

@@ -0,0 +1,23 @@
+#ifndef ATOMS_COMMON_H
+#define ATOMS_COMMON_H
+
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include <numpy/arrayobject.h>
+
+#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 */

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 */

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 */

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 */

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 */

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 */

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 */

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 */

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 */