Skip to content

Commit 91b1c48

Browse files
dance858dance858
authored
Merge pull request #2 from SparseDifferentiation/coo-support
[WIP] COO support
2 parents ba3a5a2 + a456557 commit 91b1c48

8 files changed

Lines changed: 229 additions & 3 deletions

File tree

‎SparseDiffEngine‎

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
44

55
[project]
66
name = "sparsediffpy"
7-
version = "0.1.2"
7+
version = "0.1.3"
88
description = "Python bindings for SparseDiffEngine automatic differentiation"
99
requires-python = ">=3.11"
1010
dependencies = ["numpy >= 2.0.0"]

‎setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ def finalize_options(self) -> None:
6161
"-std=c99",
6262
"-Wall",
6363
not_on_windows("-Wextra"),
64-
'-DDIFF_ENGINE_VERSION="0.1.0"',
64+
'-DDIFF_ENGINE_VERSION="0.1.3"',
6565
],
6666
extra_link_args=["-lm"] if platform.system().lower() != "windows" else [],
6767
)

‎sparsediffpy/_bindings/bindings.c‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,19 @@ static PyMethodDef DNLPMethods[] = {
150150
"Compute Lagrangian Hessian"},
151151
{"get_hessian", py_get_hessian, METH_VARARGS,
152152
"Get Lagrangian Hessian without recomputing"},
153+
{"problem_init_jacobian_coo", py_problem_init_jacobian_coo, METH_VARARGS,
154+
"Initialize Jacobian COO structures"},
155+
{"get_jacobian_sparsity_coo", py_get_jacobian_sparsity_coo, METH_VARARGS,
156+
"Get Jacobian sparsity in COO format"},
157+
{"problem_eval_jacobian_vals", py_problem_eval_jacobian_vals, METH_VARARGS,
158+
"Evaluate Jacobian and return values array"},
159+
{"problem_init_hessian_coo_lower_triangular",
160+
py_problem_init_hessian_coo_lower_triangular, METH_VARARGS,
161+
"Initialize lower-triangular Hessian COO structures"},
162+
{"get_problem_hessian_sparsity_coo", py_get_problem_hessian_sparsity_coo, METH_VARARGS,
163+
"Get Hessian sparsity in COO format (lower triangular)"},
164+
{"problem_eval_hessian_vals_coo", py_problem_eval_hessian_vals_coo, METH_VARARGS,
165+
"Evaluate Hessian and return COO values array"},
153166
{NULL, NULL, 0, NULL}};
154167

155168
static struct PyModuleDef sparsediffpy_module = {

‎sparsediffpy/_bindings/problem/hessian.h‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,4 +120,99 @@ static PyObject *py_get_hessian(PyObject *self, PyObject *args)
120120
return Py_BuildValue("(OOO(ii))", data, indices, indptr, H->m, H->n);
121121
}
122122

123+
static PyObject *py_get_problem_hessian_sparsity_coo(PyObject *self, PyObject *args)
124+
{
125+
PyObject *prob_capsule;
126+
if (!PyArg_ParseTuple(args, "O", &prob_capsule))
127+
{
128+
return NULL;
129+
}
130+
131+
problem *prob =
132+
(problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
133+
if (!prob)
134+
{
135+
PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
136+
return NULL;
137+
}
138+
139+
if (!prob->lagrange_hessian_coo)
140+
{
141+
PyErr_SetString(PyExc_RuntimeError,
142+
"hessian COO not initialized - call "
143+
"problem_init_hessian_coo_lower_triangular first");
144+
return NULL;
145+
}
146+
147+
COO_Matrix *coo = prob->lagrange_hessian_coo;
148+
npy_intp nnz = coo->nnz;
149+
150+
PyObject *rows = PyArray_SimpleNew(1, &nnz, NPY_INT32);
151+
PyObject *cols = PyArray_SimpleNew(1, &nnz, NPY_INT32);
152+
153+
if (!rows || !cols)
154+
{
155+
Py_XDECREF(rows);
156+
Py_XDECREF(cols);
157+
return NULL;
158+
}
159+
160+
memcpy(PyArray_DATA((PyArrayObject *) rows), coo->rows, nnz * sizeof(int));
161+
memcpy(PyArray_DATA((PyArrayObject *) cols), coo->cols, nnz * sizeof(int));
162+
163+
return Py_BuildValue("(OO(ii))", rows, cols, coo->m, coo->n);
164+
}
165+
166+
static PyObject *py_problem_eval_hessian_vals_coo(PyObject *self, PyObject *args)
167+
{
168+
PyObject *prob_capsule;
169+
double obj_factor;
170+
PyObject *lagrange_obj;
171+
172+
if (!PyArg_ParseTuple(args, "OdO", &prob_capsule, &obj_factor, &lagrange_obj))
173+
{
174+
return NULL;
175+
}
176+
177+
problem *prob =
178+
(problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
179+
if (!prob)
180+
{
181+
PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
182+
return NULL;
183+
}
184+
185+
/* Convert lagrange to contiguous C array */
186+
PyArrayObject *lagrange_arr = (PyArrayObject *) PyArray_FROM_OTF(
187+
lagrange_obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
188+
if (!lagrange_arr)
189+
{
190+
return NULL;
191+
}
192+
193+
double *lagrange = (double *) PyArray_DATA(lagrange_arr);
194+
195+
/* Compute Hessian */
196+
problem_hessian(prob, obj_factor, lagrange);
197+
198+
Py_DECREF(lagrange_arr);
199+
200+
/* Refresh COO values from the CSR Hessian */
201+
refresh_lower_triangular_coo(prob->lagrange_hessian_coo,
202+
prob->lagrange_hessian->x);
203+
204+
COO_Matrix *coo = prob->lagrange_hessian_coo;
205+
npy_intp nnz = coo->nnz;
206+
207+
PyObject *data = PyArray_SimpleNew(1, &nnz, NPY_DOUBLE);
208+
if (!data)
209+
{
210+
return NULL;
211+
}
212+
213+
memcpy(PyArray_DATA((PyArrayObject *) data), coo->x, nnz * sizeof(double));
214+
215+
return data;
216+
}
217+
123218
#endif /* PROBLEM_HESSIAN_H */

‎sparsediffpy/_bindings/problem/init_hessian.h‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,4 +39,26 @@ static PyObject *py_problem_init_hessian(PyObject *self, PyObject *args)
3939
Py_RETURN_NONE;
4040
}
4141

42+
static PyObject *py_problem_init_hessian_coo_lower_triangular(PyObject *self,
43+
PyObject *args)
44+
{
45+
PyObject *prob_capsule;
46+
if (!PyArg_ParseTuple(args, "O", &prob_capsule))
47+
{
48+
return NULL;
49+
}
50+
51+
problem *prob =
52+
(problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
53+
if (!prob)
54+
{
55+
PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
56+
return NULL;
57+
}
58+
59+
problem_init_hessian_coo_lower_triangular(prob);
60+
61+
Py_RETURN_NONE;
62+
}
63+
4264
#endif /* PROBLEM_INIT_HESSIAN_H */

‎sparsediffpy/_bindings/problem/init_jacobian.h‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,4 +39,25 @@ static PyObject *py_problem_init_jacobian(PyObject *self, PyObject *args)
3939
Py_RETURN_NONE;
4040
}
4141

42+
static PyObject *py_problem_init_jacobian_coo(PyObject *self, PyObject *args)
43+
{
44+
PyObject *prob_capsule;
45+
if (!PyArg_ParseTuple(args, "O", &prob_capsule))
46+
{
47+
return NULL;
48+
}
49+
50+
problem *prob =
51+
(problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
52+
if (!prob)
53+
{
54+
PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
55+
return NULL;
56+
}
57+
58+
problem_init_jacobian_coo(prob);
59+
60+
Py_RETURN_NONE;
61+
}
62+
4263
#endif /* PROBLEM_INIT_JACOBIAN_H */

‎sparsediffpy/_bindings/problem/jacobian.h‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,4 +102,79 @@ static PyObject *py_get_jacobian(PyObject *self, PyObject *args)
102102
return Py_BuildValue("(OOO(ii))", data, indices, indptr, jac->m, jac->n);
103103
}
104104

105+
static PyObject *py_get_jacobian_sparsity_coo(PyObject *self, PyObject *args)
106+
{
107+
PyObject *prob_capsule;
108+
if (!PyArg_ParseTuple(args, "O", &prob_capsule))
109+
{
110+
return NULL;
111+
}
112+
113+
problem *prob =
114+
(problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
115+
if (!prob)
116+
{
117+
PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
118+
return NULL;
119+
}
120+
121+
if (!prob->jacobian_coo)
122+
{
123+
PyErr_SetString(
124+
PyExc_RuntimeError,
125+
"jacobian COO not initialized - call problem_init_jacobian_coo first");
126+
return NULL;
127+
}
128+
129+
COO_Matrix *coo = prob->jacobian_coo;
130+
npy_intp nnz = coo->nnz;
131+
132+
PyObject *rows = PyArray_SimpleNew(1, &nnz, NPY_INT32);
133+
PyObject *cols = PyArray_SimpleNew(1, &nnz, NPY_INT32);
134+
135+
if (!rows || !cols)
136+
{
137+
Py_XDECREF(rows);
138+
Py_XDECREF(cols);
139+
return NULL;
140+
}
141+
142+
memcpy(PyArray_DATA((PyArrayObject *) rows), coo->rows, nnz * sizeof(int));
143+
memcpy(PyArray_DATA((PyArrayObject *) cols), coo->cols, nnz * sizeof(int));
144+
145+
return Py_BuildValue("(OO(ii))", rows, cols, coo->m, coo->n);
146+
}
147+
148+
static PyObject *py_problem_eval_jacobian_vals(PyObject *self, PyObject *args)
149+
{
150+
PyObject *prob_capsule;
151+
if (!PyArg_ParseTuple(args, "O", &prob_capsule))
152+
{
153+
return NULL;
154+
}
155+
156+
problem *prob =
157+
(problem *) PyCapsule_GetPointer(prob_capsule, PROBLEM_CAPSULE_NAME);
158+
if (!prob)
159+
{
160+
PyErr_SetString(PyExc_ValueError, "invalid problem capsule");
161+
return NULL;
162+
}
163+
164+
problem_jacobian(prob);
165+
166+
CSR_Matrix *jac = prob->jacobian;
167+
npy_intp nnz = jac->nnz;
168+
169+
PyObject *data = PyArray_SimpleNew(1, &nnz, NPY_DOUBLE);
170+
if (!data)
171+
{
172+
return NULL;
173+
}
174+
175+
memcpy(PyArray_DATA((PyArrayObject *) data), jac->x, nnz * sizeof(double));
176+
177+
return data;
178+
}
179+
105180
#endif /* PROBLEM_JACOBIAN_H */

0 commit comments

Comments
 (0)