Add pickle support

.github/workflows/ci.yml

@@ -45,8 +45,8 @@ jobs:
       fail-fast: false
       matrix:
         python-version:
-          - 3.8
           - 3.9
+          - 3.12
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python ${{ matrix.python-version }}
@@ -55,7 +55,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - run: pip install numpy scipy
       - run: brew install suite-sparse cfitsio gsl
-      - run: cmake -DPython_EXECUTABLE=$(which python) . && make && make install
+      - run: cmake -DPython_EXECUTABLE=$(which python) -DCMAKE_INSTALL_PREFIX=install . && make install
       - run: python -c "import photospline"
       - run: ctest --output-on-failure
   architecture-zoo:

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.3.1 LANGUAGES C CXX)
+project (photospline VERSION 2.4.0 LANGUAGES C CXX)
 
 SET(CMAKE_CXX_STANDARD 11)
 SET(CMAKE_C_STANDARD 99)
@@ -374,6 +374,9 @@ if(PYTHON_FOUND AND NUMPY_FOUND)
   ADD_TEST(photospline-test-pyeval ${PYTHON_EXECUTABLE} ${PROJECT_SOURCE_DIR}/test/test_eval.py)
   set_property(TEST photospline-test-pyeval PROPERTY ENVIRONMENT PYTHONPATH=${PROJECT_BINARY_DIR})
   LIST (APPEND ALL_TESTS photospline-test-pyeval)
+  ADD_TEST(photospline-test-pickle ${PYTHON_EXECUTABLE} ${PROJECT_SOURCE_DIR}/test/test_pickle.py)
+  set_property(TEST photospline-test-pickle PROPERTY ENVIRONMENT PYTHONPATH=${PROJECT_BINARY_DIR})
+  LIST (APPEND ALL_TESTS photospline-test-pickle)
 endif()
 
 if(BUILD_SPGLAM)

src/python/photosplinemodule.cpp

@@ -318,21 +318,30 @@ static inline handle new_reference(PyObject *ptr)
 	return handle(ptr, deleter);
 }
 
+static PyObject*
+pysplinetable_init_empty(pysplinetable *self){
+	try{
+		self->table=new photospline::splinetable<>();
+	}catch(std::exception& ex){
+		PyErr_SetString(PyExc_Exception,
+						(std::string("Unable to allocate spline table: ")+ex.what()).c_str());
+		return(NULL);
+	}catch(...){
+		PyErr_SetString(PyExc_Exception, "Unable to allocate spline table");
+		return(NULL);
+	}
+
+	return (PyObject*)self;
+}
+
 static PyObject*
 pysplinetable_new(PyTypeObject* type, PyObject* args, PyObject* kwds){
 	pysplinetable* self;
 
 	self = (pysplinetable*)type->tp_alloc(type, 0);
 	if(self){
-		try{
-			self->table=new photospline::splinetable<>();
-		}catch(std::exception& ex){
-			PyErr_SetString(PyExc_Exception,
-			                (std::string("Unable to allocate spline table: ")+ex.what()).c_str());
-			return(NULL);
-		}catch(...){
-			PyErr_SetString(PyExc_Exception, "Unable to allocate spline table");
-			return(NULL);
+		if (pysplinetable_init_empty(self) == NULL){
+			return NULL;
 		}
 	}
 
@@ -1217,6 +1226,35 @@ pyphotospline_bspline(pysplinetable* self, PyObject* args, PyObject* kwds){
 	return PyFloat_FromDouble(photospline::bspline(knots, x, i, order));
 }
 
+static PyObject*
+pysplinetable_getstate(pysplinetable* self, PyObject *Py_UNUSED(ignored)){
+	auto buffer=self->table->write_fits_mem();
+	std::unique_ptr<void,void(*)(void*)> data(buffer.first,&free);
+	return PyBytes_FromStringAndSize((char*)buffer.first, buffer.second);
+}
+
+static PyObject*
+pysplinetable_setstate(pysplinetable* self, PyObject* state){
+    if (!PyBytes_CheckExact(state)) {
+        PyErr_SetString(PyExc_ValueError, "Pickled object is not bytes.");
+        return NULL;
+    }
+	char *buffer = PyBytes_AsString(state);
+	if (!buffer) {
+		return NULL;
+	}
+	if (!pysplinetable_init_empty(self)) {
+		return NULL;
+	}
+	try{
+		self->table->read_fits_mem(buffer, PyBytes_Size(state));
+	}catch(std::exception& ex){
+		PyErr_SetString(PyExc_Exception,ex.what());
+		return(NULL);
+	}
+
+	Py_RETURN_NONE;
+}
 
 static PyGetSetDef pysplinetable_properties[] = {
 	{(char*)"order", (getter)pysplinetable_getorder, NULL, (char*)"Order of spline in each dimension", NULL},
@@ -1258,12 +1296,14 @@ static PyMethodDef pysplinetable_methods[] = {
 	 ":returns: an array of spline evaluates with size `len(coord[dim])` in each dimension"},
 #endif
 #endif
+	{"__getstate__", (PyCFunction)pysplinetable_getstate, METH_NOARGS, "Pickle the spline"},
+	{"__setstate__", (PyCFunction)pysplinetable_setstate, METH_O, "Unpickle the spline"},
 	{NULL}  /* Sentinel */
 };
 
 static PyTypeObject pysplinetableType = {
 	PyVarObject_HEAD_INIT(NULL, 0)
-	"pyphotospline.Splinetable", /*tp_name*/
+	"photospline.SplineTable", /*tp_name*/
 	sizeof(pysplinetable),     /*tp_basicsize*/
 	0,                         /*tp_itemsize*/
 	(destructor)pysplinetable_dealloc, /*tp_dealloc*/

test/test_pickle.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+
+import pickle
+import unittest
+import photospline
+
+import numpy as np
+
+from pathlib import Path
+
+
+class TestEvaluation(unittest.TestCase):
+    def setUp(self):
+        self.testdata = Path(__file__).parent / "test_data"
+        self.spline = photospline.SplineTable(self.testdata / "test_spline_4d.fits")
+        extents = np.array(self.spline.extents)
+        loc = extents[:, :1]
+        scale = np.diff(extents, axis=1)
+        self.x = (np.random.uniform(0, 1, size=(self.spline.ndim, 10)) + loc) * scale
+
+    def test_pickle(self):
+        restored = pickle.loads(pickle.dumps(self.spline))
+        np.testing.assert_allclose(self.spline.evaluate_simple(self.x), restored.evaluate_simple(self.x), rtol=0, atol=0)
+
+if __name__ == "__main__":
+    unittest.main()