[WIP] COO support

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sparsediffpy"
-version = "0.1.2"
+version = "0.1.3"
 description = "Python bindings for SparseDiffEngine automatic differentiation"
 requires-python = ">=3.11"
 dependencies = ["numpy >= 2.0.0"]

setup.py

@@ -61,7 +61,7 @@ def finalize_options(self) -> None:
         "-std=c99",
         "-Wall",
         not_on_windows("-Wextra"),
-        '-DDIFF_ENGINE_VERSION="0.1.0"',
+        '-DDIFF_ENGINE_VERSION="0.1.3"',
     ],
     extra_link_args=["-lm"] if platform.system().lower() != "windows" else [],
 )

sparsediffpy/_bindings/bindings.c

@@ -150,6 +150,19 @@ static PyMethodDef DNLPMethods[] = {
      "Compute Lagrangian Hessian"},
     {"get_hessian", py_get_hessian, METH_VARARGS,
      "Get Lagrangian Hessian without recomputing"},
+    {"problem_init_jacobian_coo", py_problem_init_jacobian_coo, METH_VARARGS,
+     "Initialize Jacobian COO structures"},
+    {"get_jacobian_sparsity_coo", py_get_jacobian_sparsity_coo, METH_VARARGS,
+     "Get Jacobian sparsity in COO format"},
+    {"problem_eval_jacobian_vals", py_problem_eval_jacobian_vals, METH_VARARGS,
+     "Evaluate Jacobian and return values array"},
+    {"problem_init_hessian_coo_lower_triangular",
+     py_problem_init_hessian_coo_lower_triangular, METH_VARARGS,
+     "Initialize lower-triangular Hessian COO structures"},
+    {"get_problem_hessian_sparsity_coo", py_get_problem_hessian_sparsity_coo, METH_VARARGS,
+     "Get Hessian sparsity in COO format (lower triangular)"},
+    {"problem_eval_hessian_vals_coo", py_problem_eval_hessian_vals_coo, METH_VARARGS,
+     "Evaluate Hessian and return COO values array"},
     {NULL, NULL, 0, NULL}};
 
 static struct PyModuleDef sparsediffpy_module = {

sparsediffpy/_bindings/problem/hessian.h

@@ -120,4 +120,99 @@ static PyObject *py_get_hessian(PyObject *self, PyObject *args)
     return Py_BuildValue("(OOO(ii))", data, indices, indptr, H->m, H->n);
 }
 
+static PyObject *py_get_problem_hessian_sparsity_coo(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_coo)
+    {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "hessian COO not initialized - call "
+                        "problem_init_hessian_coo_lower_triangular first");
+        return NULL;
+    }
+
+    COO_Matrix *coo = prob->lagrange_hessian_coo;
+    npy_intp nnz = coo->nnz;
+
+    PyObject *rows = PyArray_SimpleNew(1, &nnz, NPY_INT32);
+    PyObject *cols = PyArray_SimpleNew(1, &nnz, NPY_INT32);
+
+    if (!rows || !cols)
+    {
+        Py_XDECREF(rows);
+        Py_XDECREF(cols);
+        return NULL;
+    }
+
+    memcpy(PyArray_DATA((PyArrayObject *) rows), coo->rows, nnz * sizeof(int));
+    memcpy(PyArray_DATA((PyArrayObject *) cols), coo->cols, nnz * sizeof(int));
+
+    return Py_BuildValue("(OO(ii))", rows, cols, coo->m, coo->n);
+}
+
+static PyObject *py_problem_eval_hessian_vals_coo(PyObject *self, PyObject *args)
+{
+    PyObject *prob_capsule;
+    double obj_factor;
+    PyObject *lagrange_obj;
+
+    if (!PyArg_ParseTuple(args, "OdO", &prob_capsule, &obj_factor, &lagrange_obj))
+    {
+        return NULL;
+    }
+
+    problem *prob =
+        (problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
+    if (!prob)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
+        return NULL;
+    }
+
+    /* Convert lagrange to contiguous C array */
+    PyArrayObject *lagrange_arr = (PyArrayObject *) PyArray_FROM_OTF(
+        lagrange_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
+    if (!lagrange_arr)
+    {
+        return NULL;
+    }
+
+    double *lagrange = (double *) PyArray_DATA(lagrange_arr);
+
+    /* Compute Hessian */
+    problem_hessian(prob, obj_factor, lagrange);
+
+    Py_DECREF(lagrange_arr);
+
+    /* Refresh COO values from the CSR Hessian */
+    refresh_lower_triangular_coo(prob->lagrange_hessian_coo,
+                                 prob->lagrange_hessian->x);
+
+    COO_Matrix *coo = prob->lagrange_hessian_coo;
+    npy_intp nnz = coo->nnz;
+
+    PyObject *data = PyArray_SimpleNew(1, &nnz, NPY_DOUBLE);
+    if (!data)
+    {
+        return NULL;
+    }
+
+    memcpy(PyArray_DATA((PyArrayObject *) data), coo->x, nnz * sizeof(double));
+
+    return data;
+}
+
 #endif /* PROBLEM_HESSIAN_H */

sparsediffpy/_bindings/problem/init_hessian.h

@@ -39,4 +39,26 @@ static PyObject *py_problem_init_hessian(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+static PyObject *py_problem_init_hessian_coo_lower_triangular(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_hessian_coo_lower_triangular(prob);
+
+    Py_RETURN_NONE;
+}
+
 #endif /* PROBLEM_INIT_HESSIAN_H */

sparsediffpy/_bindings/problem/init_jacobian.h

@@ -39,4 +39,25 @@ static PyObject *py_problem_init_jacobian(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+static PyObject *py_problem_init_jacobian_coo(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_jacobian_coo(prob);
+
+    Py_RETURN_NONE;
+}
+
 #endif /* PROBLEM_INIT_JACOBIAN_H */

sparsediffpy/_bindings/problem/jacobian.h

@@ -102,4 +102,79 @@ static PyObject *py_get_jacobian(PyObject *self, PyObject *args)
     return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
 }
 
+static PyObject *py_get_jacobian_sparsity_coo(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_coo)
+    {
+        PyErr_SetString(
+            PyExc_RuntimeError,
+            "jacobian COO not initialized - call problem_init_jacobian_coo first");
+        return NULL;
+    }
+
+    COO_Matrix *coo = prob->jacobian_coo;
+    npy_intp nnz = coo->nnz;
+
+    PyObject *rows = PyArray_SimpleNew(1, &nnz, NPY_INT32);
+    PyObject *cols = PyArray_SimpleNew(1, &nnz, NPY_INT32);
+
+    if (!rows || !cols)
+    {
+        Py_XDECREF(rows);
+        Py_XDECREF(cols);
+        return NULL;
+    }
+
+    memcpy(PyArray_DATA((PyArrayObject *) rows), coo->rows, nnz * sizeof(int));
+    memcpy(PyArray_DATA((PyArrayObject *) cols), coo->cols, nnz * sizeof(int));
+
+    return Py_BuildValue("(OO(ii))", rows, cols, coo->m, coo->n);
+}
+
+static PyObject *py_problem_eval_jacobian_vals(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_jacobian(prob);
+
+    CSR_Matrix *jac = prob->jacobian;
+    npy_intp nnz = jac->nnz;
+
+    PyObject *data = PyArray_SimpleNew(1, &nnz, NPY_DOUBLE);
+    if (!data)
+    {
+        return NULL;
+    }
+
+    memcpy(PyArray_DATA((PyArrayObject *) data), jac->x, nnz * sizeof(double));
+
+    return data;
+}
+
 #endif /* PROBLEM_JACOBIAN_H */