New dimension convention and bug fixes

TODO.md

@@ -1,9 +1,10 @@
 3. more tests for chain rule elementwise univariate hessian
 4. in the refactor, add consts
-5. multiply with one constant vector/scalar argument
 6. AX where X is a matrix. Can that happen? How is that canonicalized? Maybe it can't happen.
 7. Must be able to compute jacobian and hessian of A @ phi(x), so linear operator needs other code! This requires new infrastructure, I think.
 8. Shortcut hessians of affine stuff?
+10. For performance reasons, is it useful to have a dense matmul with A and B as dense matrices?
+11. right matmul, add broadcasting logic as in left matmul
 
 Going through all atoms to see that sparsity pattern is computed in initialization of jacobian:
 2. trace - not ok

include/affine.h

@@ -18,7 +18,7 @@ expr *new_trace(expr *child);
 expr *new_constant(int d1, int d2, int n_vars, const double *values);
 expr *new_variable(int d1, int d2, int var_id, int n_vars);
 
-expr *new_index(expr *child, const int *indices, int n_idxs);
+expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs);
 expr *new_reshape(expr *child, int d1, int d2);
 expr *new_broadcast(expr *child, int target_d1, int target_d2);
 

include/subexpr.h

@@ -130,8 +130,6 @@ typedef struct broadcast_expr
 {
     expr base;
     broadcast_type type;
-    int m; /* target rows */
-    int n; /* target cols */
 } broadcast_expr;
 
 #endif /* SUBEXPR_H */

python/atoms/index.h

@@ -11,9 +11,10 @@
 static PyObject *py_make_index(PyObject *self, PyObject *args)
 {
     PyObject *child_capsule;
+    int d1, d2;
     PyObject *indices_obj;
 
-    if (!PyArg_ParseTuple(args, "OO", &child_capsule, &indices_obj))
+    if (!PyArg_ParseTuple(args, "OiiO", &child_capsule, &d1, &d2, &indices_obj))
     {
         return NULL;
     }
@@ -37,7 +38,7 @@ static PyObject *py_make_index(PyObject *self, PyObject *args)
     int n_idxs = (int) PyArray_SIZE(indices_array);
     int *indices_data = (int *) PyArray_DATA(indices_array);
 
-    expr *node = new_index(child, indices_data, n_idxs);
+    expr *node = new_index(child, d1, d2, indices_data, n_idxs);
 
     Py_DECREF(indices_array);
 

python/atoms/rel_entr_scalar_vector.h

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef ATOM_REL_ENTR_SCALAR_VECTOR_H
+#define ATOM_REL_ENTR_SCALAR_VECTOR_H
+
+#include "bivariate.h"
+#include "common.h"
+
+/* rel_entr_scalar_vector: rel_entr(x, y) where x is scalar, y is vector */
+static PyObject *py_make_rel_entr_scalar_vector(PyObject *self, PyObject *args)
+{
+    (void) self;
+    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_rel_entr_first_arg_scalar(left, right);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "failed to create rel_entr_scalar_vector node");
+        return NULL;
+    }
+    expr_retain(node); /* Capsule owns a reference */
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_REL_ENTR_SCALAR_VECTOR_H */

python/atoms/rel_entr_vector_scalar.h

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef ATOM_REL_ENTR_VECTOR_SCALAR_H
+#define ATOM_REL_ENTR_VECTOR_SCALAR_H
+
+#include "bivariate.h"
+#include "common.h"
+
+/* rel_entr_vector_scalar: rel_entr(x, y) where x is vector, y is scalar */
+static PyObject *py_make_rel_entr_vector_scalar(PyObject *self, PyObject *args)
+{
+    (void) self;
+    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_rel_entr_second_arg_scalar(left, right);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "failed to create rel_entr_vector_scalar node");
+        return NULL;
+    }
+    expr_retain(node); /* Capsule owns a reference */
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_REL_ENTR_VECTOR_SCALAR_H */

python/bindings.c

@@ -27,6 +27,8 @@
 #include "atoms/quad_form.h"
 #include "atoms/quad_over_lin.h"
 #include "atoms/rel_entr.h"
+#include "atoms/rel_entr_scalar_vector.h"
+#include "atoms/rel_entr_vector_scalar.h"
 #include "atoms/reshape.h"
 #include "atoms/right_matmul.h"
 #include "atoms/sin.h"
@@ -96,6 +98,10 @@ static PyMethodDef DNLPMethods[] = {
      "Create quad_over_lin node (sum(x^2) / y)"},
     {"make_rel_entr", py_make_rel_entr, METH_VARARGS,
      "Create rel_entr node: x * log(x/y) elementwise"},
+    {"make_rel_entr_vector_scalar", py_make_rel_entr_vector_scalar, METH_VARARGS,
+     "Create rel_entr node with vector first arg, scalar second arg"},
+    {"make_rel_entr_scalar_vector", py_make_rel_entr_scalar_vector, METH_VARARGS,
+     "Create rel_entr node with scalar first arg, vector second arg"},
     {"get_expr_dimensions", py_get_expr_dimensions, METH_VARARGS,
      "Get the dimensions (d1, d2) of an expression"},
     {"get_expr_size", py_get_expr_size, METH_VARARGS,
@@ -113,8 +119,12 @@ static PyMethodDef DNLPMethods[] = {
      "Compute objective gradient"},
     {"problem_jacobian", py_problem_jacobian, METH_VARARGS,
      "Compute constraint jacobian"},
+    {"get_jacobian", py_get_jacobian, METH_VARARGS,
+     "Get constraint jacobian without recomputing"},
     {"problem_hessian", py_problem_hessian, METH_VARARGS,
      "Compute Lagrangian Hessian"},
+    {"get_hessian", py_get_hessian, METH_VARARGS,
+     "Get Lagrangian Hessian without recomputing"},
     {NULL, NULL, 0, NULL}};
 
 static struct PyModuleDef dnlp_module = {

python/problem/hessian.h

@@ -9,7 +9,8 @@
  * Args:
  *   prob_capsule: PyCapsule containing problem pointer
  *   obj_factor: Scaling factor for objective Hessian (double)
- *   lagrange: Array of Lagrange multipliers (numpy array, length = total_constraint_size)
+ *   lagrange: Array of Lagrange multipliers (numpy array, length =
+ * total_constraint_size)
  *
  * Returns:
  *   Tuple of (data, indices, indptr, (m, n)) for scipy.sparse.csr_matrix
@@ -73,4 +74,50 @@ static PyObject *py_problem_hessian(PyObject *self, PyObject *args)
     return Py_BuildValue("(OOO(ii))", data, indices, indptr, H->m, H->n);
 }
 
+static PyObject *py_get_hessian(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->lagrange_hessian)
+    {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "hessian not initialized - call problem_hessian first");
+        return NULL;
+    }
+
+    CSR_Matrix *H = prob->lagrange_hessian;
+    npy_intp nnz = H->nnz;
+    npy_intp n_plus_1 = H->n + 1;
+
+    PyObject *data = PyArray_SimpleNew(1, &nnz, NPY_DOUBLE);
+    PyObject *indices = PyArray_SimpleNew(1, &nnz, NPY_INT32);
+    PyObject *indptr = PyArray_SimpleNew(1, &n_plus_1, NPY_INT32);
+
+    if (!data || !indices || !indptr)
+    {
+        Py_XDECREF(data);
+        Py_XDECREF(indices);
+        Py_XDECREF(indptr);
+        return NULL;
+    }
+
+    memcpy(PyArray_DATA((PyArrayObject *) data), H->x, nnz * sizeof(double));
+    memcpy(PyArray_DATA((PyArrayObject *) indices), H->i, nnz * sizeof(int));
+    memcpy(PyArray_DATA((PyArrayObject *) indptr), H->p, n_plus_1 * sizeof(int));
+
+    return Py_BuildValue("(OOO(ii))", data, indices, indptr, H->m, H->n);
+}
+
 #endif /* PROBLEM_HESSIAN_H */

python/problem/jacobian.h

@@ -56,4 +56,50 @@ static PyObject *py_problem_jacobian(PyObject *self, PyObject *args)
     return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
 }
 
+static PyObject *py_get_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->jacobian)
+    {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "jacobian not initialized - call problem_jacobian first");
+        return NULL;
+    }
+
+    CSR_Matrix *jac = prob->jacobian;
+    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 */

src/affine/add.c

@@ -1,5 +1,6 @@
 #include "affine.h"
 #include <assert.h>
+#include <stdio.h>
 
 static void forward(expr *node, const double *u)
 {