Expose sparse NNLS to Python (#42)

* Expose sparse NNLS to Python

* Add test for nnls

* Bump version

* ci: install scipy

.github/workflows/ci.yml

@@ -33,7 +33,7 @@ jobs:
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
-      - run: pip install numpy
+      - run: pip install numpy scipy
       - run: sudo apt-get install libsuitesparse-dev libbtbb-dev liblapack-dev libcfitsio-dev libmetis-dev libgsl-dev
       - run: cmake -DPython_EXECUTABLE=$(which python) . && make && sudo make install
       - run: python -c "import photospline"
@@ -53,7 +53,7 @@ jobs:
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
-      - run: pip install numpy
+      - run: pip install numpy scipy
       - run: brew install suite-sparse cfitsio gsl
       - run: cmake -DPython_EXECUTABLE=$(which python) . && make && make install
       - run: python -c "import photospline"
@@ -78,7 +78,7 @@ jobs:
           githubToken: ${{ github.token }}
           install: |
             apt-get update -q -y
-            apt-get install -q -y cmake clang libsuitesparse-dev libbtbb-dev liblapack-dev libcfitsio-dev libmetis-dev libgsl-dev python3-dev python3-numpy
+            apt-get install -q -y cmake clang libsuitesparse-dev libbtbb-dev liblapack-dev libcfitsio-dev libmetis-dev libgsl-dev python3-dev python3-numpy python3-scipy
           run: |
             cmake -DPython_EXECUTABLE=$(which python3) .
             make

CMakeLists.txt

@@ -1,7 +1,7 @@
 cmake_minimum_required (VERSION 3.1.0 FATAL_ERROR)
 cmake_policy(VERSION 3.1.0)
 
-project (photospline VERSION 2.2.1 LANGUAGES C CXX)
+project (photospline VERSION 2.3.0 LANGUAGES C CXX)
 
 SET(CMAKE_CXX_STANDARD 11)
 SET(CMAKE_C_STANDARD 99)
@@ -394,6 +394,14 @@ if(BUILD_SPGLAM)
      WORKING_DIRECTORY ${CMAKE_BUILD_DIR}
     )
     LIST (APPEND ALL_TESTS photospline-test-pyfit)
+    if(BUILD_SPGLAM)
+      ADD_TEST(NAME photospline-test-nnls
+        COMMAND ${PYTHON_EXECUTABLE} ${PROJECT_SOURCE_DIR}/test/test_nnls.py
+        WORKING_DIRECTORY ${CMAKE_BUILD_DIR}
+      )
+      set_property(TEST photospline-test-nnls PROPERTY ENVIRONMENT PYTHONPATH=${PROJECT_BINARY_DIR})
+      LIST (APPEND ALL_TESTS photospline-test-nnls)
+    endif()
   endif()
 endif()
 

src/python/photosplinemodule.cpp

@@ -1098,6 +1098,9 @@ pysplinetable_deriv(pysplinetable* self, PyObject* args, PyObject* kwds){
 static PyObject*
 pyphotospline_glam_fit(PyObject* self, PyObject* args, PyObject* kwds);
 
+static PyObject*
+pyphotospline_nnls(PyObject* self, PyObject* args, PyObject* kwds);
+
 #ifdef HAVE_NUMPY
 static PyObject*
 pysplinetable_grideval(pysplinetable* self, PyObject* args, PyObject* kwds);
@@ -1313,6 +1316,19 @@ static PyMethodDef photospline_methods[] = {
 	":param monodim: if specified, the fit function will be monotonic (strictly non-decreasing) along the given axis\n"
 	":param verbose: if True, will print status information to standard output while running\n"
 	":returns: a spline table object\n"},
+	{"nnls", (PyCFunction)pyphotospline_nnls, METH_VARARGS | METH_KEYWORDS,
+	 "Solve argmin_x || Ax - b ||_2 for x>=0, where A is a sparse matrix\n\n"
+	 ":param A: model matrix\n"
+	 ":param b: right-hand side vector\n"
+	 ":param tolerance: terminate once the largest element of the gradient is smaller than the tolerance\n"
+	 ":param min_iterations: minimum number of iterations (least-squares solutions in a constrained subspace) before terminating\n"
+	 ":param max_iterations: minimum number of iterations (least-squares solutions in a constrained subspace) before terminating\n"
+	 ":param npos: maximum number of positive coefficients to return\n"
+	 ":param normaleq: if True, the inputs are the normal equations At*A and At*b instead of A and b\n"
+     ":param verbose: print debugging info\n"
+	 ":returns: x, the solution vector\n\n"
+	 "This is a wrapper around a sparse implementation of Algorithm NNLS from \"Solving Least Squares Problems\", Charles Lawson and Richard Hanson. Prentice-Hall, 1974.\n"
+	 },
 #endif
 	{"bspline", (PyCFunction)pyphotospline_bspline, METH_VARARGS | METH_KEYWORDS,
 	 "Evaluate the `i`th B-spline on knot vector `knots` at `x`\n\n"
@@ -1768,6 +1784,139 @@ pysplinetable_grideval(pysplinetable* self, PyObject* args, PyObject* kwds){
 }
 #endif
 
+namespace {
+namespace detail {
+
+// A simple self-destructing wrapper around PyObject*
+struct ref {
+	static inline void decref(PyObject *x) { Py_DECREF(x); };
+	std::unique_ptr<PyObject, decltype(&decref)> px;
+	ref(PyObject *ptr) : px(ptr, &decref) {};
+	operator bool() { return bool(px); };
+	operator PyObject*() { return px.get(); };
+	PyObject* ptr() { return px.get(); }
+};
+
+}
+}
+
+static PyObject*
+pyphotospline_nnls(PyObject* self, PyObject* args, PyObject* kwds){
+	static const char* kwlist[] = {"A", "b", "tolerance", "min_iterations", "max_iterations", "npos", "normaleq", "verbose", NULL};
+
+	PyObject* py_A=NULL;
+	PyObject* py_b=NULL;
+	const char *method;
+	double tolerance=0;
+	int min_iterations=0;
+	int max_iterations=INT_MAX;
+	unsigned npos=0;
+	bool normaleq=false;
+	bool verbose=false;
+
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|diiIpp", (char**)kwlist,
+									 &py_A, &py_b, &tolerance,
+									 &min_iterations, &max_iterations, &npos, &normaleq, &verbose))
+		return(NULL);
+
+	cholmod_common cholmod_state;
+	cholmod_l_start(&cholmod_state);
+
+	// create a cholmod_sparse view into csc matrix A
+	cholmod_sparse sparse_A;
+	{
+		detail::ref sparse(PyImport_ImportModule("scipy.sparse"));
+		if (!sparse) {
+			return(NULL);
+		}
+		detail::ref isspmatrix_csc(PyObject_GetAttrString(sparse, "isspmatrix_csc"));
+		if (!isspmatrix_csc) {
+			return(NULL);
+		}
+		if (PyObject_CallFunctionObjArgs(isspmatrix_csc, py_A, NULL) != Py_True) {
+			detail::ref type(PyObject_Type(py_A));
+			detail::ref name(PyObject_GetAttrString(type, "__name__"));
+			PyErr_Format(PyExc_TypeError, "Expected scipy.sparse.csc_matrix, got %S", name.ptr());
+			return(NULL);
+		}
+	}
+	detail::ref data(
+		PyArray_FromAny(
+			detail::ref(PyObject_GetAttrString(py_A, "data")),
+			PyArray_DescrFromType(NPY_DOUBLE), 1, 1, NPY_ARRAY_C_CONTIGUOUS, NULL
+		));
+	if (!data) {
+		return(NULL);
+	}
+	detail::ref indices(
+		PyArray_FromAny(
+			detail::ref(PyObject_GetAttrString(py_A, "indices")),
+			PyArray_DescrFromType(NPY_LONGLONG), 1, 1, NPY_ARRAY_C_CONTIGUOUS, NULL
+		));
+	if (!indices) {
+		return(NULL);
+	}
+	detail::ref indptr(
+		PyArray_FromAny(
+			detail::ref(PyObject_GetAttrString(py_A, "indptr")),
+			PyArray_DescrFromType(NPY_LONGLONG), 1, 1, NPY_ARRAY_C_CONTIGUOUS, NULL
+		));
+	if (!indptr) {
+		return(NULL);
+	}
+
+	sparse_A.itype = cholmod_state.itype;
+	sparse_A.xtype = CHOLMOD_REAL;
+	sparse_A.dtype = CHOLMOD_DOUBLE;
+	sparse_A.stype = 0; // unsymmetric; all nonzero entries stored
+	sparse_A.nz = NULL; // packed storage
+	sparse_A.z = NULL; // real matrix
+	sparse_A.sorted = true;
+	sparse_A.packed = true;
+	{
+		detail::ref shape = PyObject_GetAttrString(py_A, "shape");
+		sparse_A.nrow = PyLong_AsLong(PyTuple_GetItem(shape, 0));
+		sparse_A.ncol = PyLong_AsLong(PyTuple_GetItem(shape, 1));
+	}
+	{
+		detail::ref nnz(PyObject_GetAttrString(py_A, "nnz"));
+		sparse_A.nzmax = PyLong_AsLong(nnz);
+	}
+	sparse_A.p = ExtractPyArrayDataPtr(indptr.ptr());
+	sparse_A.i = ExtractPyArrayDataPtr(indices.ptr());
+	sparse_A.x = ExtractPyArrayDataPtr(data.ptr());
+
+	// cholmod_dense view into dense matrix b
+	detail::ref array_b(PyArray_FromAny(py_b, PyArray_DescrFromType(NPY_DOUBLE), 1, 1, NPY_ARRAY_C_CONTIGUOUS, NULL));
+	cholmod_dense dense_b;
+	dense_b.xtype = CHOLMOD_REAL;
+	dense_b.dtype = CHOLMOD_DOUBLE;
+	dense_b.nzmax = PyArray_SIZE((PyArrayObject*)array_b.ptr());
+	dense_b.nrow = dense_b.nzmax;
+	dense_b.ncol = 1;
+	dense_b.d = dense_b.nrow;
+	dense_b.x = ExtractPyArrayDataPtr(array_b.ptr());
+	dense_b.z = NULL;
+
+	cholmod_dense *dense_x = nnls_lawson_hanson(&sparse_A, &dense_b,
+		tolerance,/* double tolerance */
+		min_iterations, /* unsigned min_iterations */
+		max_iterations, /* unsigned max_iterations */
+		npos, /* unsigned npos */
+		normaleq, /* int normaleq */
+		verbose, /* verbose */
+		&cholmod_state
+	);
+	npy_intp dims = dense_x->nrow;
+	PyObject *x = PyArray_SimpleNew(1, &dims, NPY_DOUBLE);
+	std::copy((double*)(dense_x->x), (double*)(dense_x->x)+dense_x->nrow, (double*)ExtractPyArrayDataPtr(x));
+
+	cholmod_l_free_dense(&dense_x, &cholmod_state);
+	cholmod_l_finish(&cholmod_state);
+
+	return x;
+}
+
 #endif //PHOTOSPLINE_INCLUDES_SPGLAM
 
 static int

test/test_nnls.py

@@ -0,0 +1,31 @@
+import unittest
+import numpy as np
+import photospline
+
+
+class TestNNLS(unittest.TestCase):
+    def setUp(self):
+        try:
+            from scipy import sparse  # type: ignore[import]
+        except ImportError:
+            raise unittest.SkipTest("test requires scipy")
+
+        self.A = np.array([[1, 0], [1, 0], [0, 1]])
+        self.Asp = sparse.csc_matrix(self.A)
+        self.b = np.array([2, 1, 1])
+
+    def testSparseIsSparse(self):
+        self.assertEqual(len(self.Asp.data), 3)
+
+    def testImplementationMatchesScipy(self):
+        from scipy import optimize  # type: ignore[import]
+
+        np.testing.assert_allclose(
+            photospline.nnls(self.Asp, self.b),
+            optimize.nnls(self.A, self.b)[0],
+            err_msg="sparse result does not aggree with dense",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

typings/photospline-stubs/__init__.pyi

@@ -1,6 +1,7 @@
-from typing import Sequence, Union, Any, overload
+from typing import Annotated, Sequence, Union, Any, overload
 from os import PathLike
 import numpy as np
+from scipy import sparse
 import numpy.typing as npt
 
 class SplineTable:
@@ -79,3 +80,14 @@ def glam_fit(
     monodim: None | int = None,
     verbose: int = 1,
 ) -> SplineTable: ...
+
+def nnls(
+    A: sparse.csc_matrix,
+    b: npt.NDArray[np.float64],
+    tolerance: float=0.,
+    min_iterations: int=0,
+    max_iterations: int=np.iinfo(np.int32).max,
+    npos: int=0,
+    normaleq: bool=False,
+    verbose: bool=False,
+) -> npt.NDArray[np.float64]: ...