Skip to content

Commit

Permalink
Add pickle support (#44)
Browse files Browse the repository at this point in the history
* Add pickle support

* Drop py38, add py312

* skip install on macos

* install in local dir on osx
  • Loading branch information
jvansanten committed Sep 11, 2024
1 parent 9610bac commit c6fb3ea
Show file tree
Hide file tree
Showing 4 changed files with 82 additions and 13 deletions.
4 changes: 2 additions & 2 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
Expand All @@ -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:
Expand Down
5 changes: 4 additions & 1 deletion CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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)
Expand Down Expand Up @@ -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)
Expand Down
60 changes: 50 additions & 10 deletions src/python/photosplinemodule.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}

Expand Down Expand Up @@ -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},
Expand Down Expand Up @@ -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*/
Expand Down
26 changes: 26 additions & 0 deletions test/test_pickle.py
Original file line number Diff line number Diff line change
@@ -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()

0 comments on commit c6fb3ea

Please sign in to comment.