Skip to content
Fetching contributors…
Cannot retrieve contributors at this time
716 lines (623 sloc) 18.3 KB
#include <Python.h>
#include <numpy/arrayobject.h>
#include "KDTree.h"
/* Must define Py_TYPE for Python 2.5 or older */
#ifndef Py_TYPE
# define Py_TYPE(o) ((o)->ob_type)
#endif
/* Must define PyVarObject_HEAD_INIT for Python 2.5 or older */
#ifndef PyVarObject_HEAD_INIT
#define PyVarObject_HEAD_INIT(type, size) \
PyObject_HEAD_INIT(type) size,
#endif
typedef struct {
PyObject_HEAD
struct Neighbor neighbor;
} PyNeighbor;
static int
PyNeighbor_init(PyNeighbor *self, PyObject *args, PyObject *kwds)
{
long int index1, index2;
float radius = 0.0;
static char *kwlist[] = {"index1", "index2", "radius", NULL};
if (! PyArg_ParseTupleAndKeywords(args, kwds, "ii|d", kwlist,
&index1, &index2, &radius))
return -1;
self->neighbor.index1 = index1;
self->neighbor.index2 = index2;
self->neighbor.radius = radius;
return 0;
}
static PyObject*
PyNeighbor_repr(PyNeighbor* self)
{
char string[64];
sprintf(string, "(%ld, %ld): %g",
self->neighbor.index1, self->neighbor.index2, self->neighbor.radius);
#if PY_MAJOR_VERSION >= 3
return PyUnicode_FromString(string);
#else
return PyString_FromString(string);
#endif
}
static char PyNeighbor_index1__doc__[] =
"index of the first neighbor";
static PyObject*
PyNeighbor_getindex1(PyNeighbor* self, void* closure)
{
#if PY_MAJOR_VERSION >= 3
return PyLong_FromLong(self->neighbor.index1);
#else
return PyInt_FromLong(self->neighbor.index1);
#endif
}
static int
PyNeighbor_setindex1(PyNeighbor* self, PyObject* value, void* closure)
{
#if PY_MAJOR_VERSION >= 3
long index1 = PyLong_AsLong(value);
#else
long index1 = PyInt_AsLong(value);
#endif
if (PyErr_Occurred()) return -1;
self->neighbor.index1 = index1;
return 0;
}
static char PyNeighbor_index2__doc__[] =
"index of the second neighbor";
static PyObject*
PyNeighbor_getindex2(PyNeighbor* self, void* closure)
{
#if PY_MAJOR_VERSION >= 3
return PyLong_FromLong(self->neighbor.index2);
#else
return PyInt_FromLong(self->neighbor.index2);
#endif
}
static int
PyNeighbor_setindex2(PyNeighbor* self, PyObject* value, void* closure)
{
#if PY_MAJOR_VERSION >= 3
long index2 = PyLong_AsLong(value);
#else
long index2 = PyInt_AsLong(value);
#endif
if (PyErr_Occurred()) return -1;
self->neighbor.index2 = index2;
return 0;
}
static PyObject*
PyNeighbor_getradius(PyNeighbor* self, void* closure)
{
float value = self->neighbor.radius;
return PyFloat_FromDouble((double)value);
}
static int
PyNeighbor_setradius(PyNeighbor* self, PyObject* value, void* closure)
{ const double radius = PyFloat_AsDouble(value);
if (PyErr_Occurred()) return -1;
self->neighbor.radius = (float)radius;
return 0;
}
static char PyNeighbor_radius__doc__[] =
"the radius\n";
static PyGetSetDef PyNeighbor_getset[] = {
{"index1", (getter)PyNeighbor_getindex1, (setter)PyNeighbor_setindex1, PyNeighbor_index1__doc__, NULL},
{"index2", (getter)PyNeighbor_getindex2, (setter)PyNeighbor_setindex2, PyNeighbor_index2__doc__, NULL},
{"radius", (getter)PyNeighbor_getradius, (setter)PyNeighbor_setradius, PyNeighbor_radius__doc__, NULL},
{NULL} /* Sentinel */
};
static char PyNeighbor_doc[] =
"A neighbor pair; members are index1, index2, and radius.\n";
static PyTypeObject PyNeighborType = {
PyVarObject_HEAD_INIT(NULL, 0)
"KDTree.Neighbor", /* tp_name*/
sizeof(PyNeighbor), /* tp_basicsize*/
0, /* tp_itemsize*/
0, /* tp_dealloc*/
0, /* tp_print*/
0, /* tp_getattr*/
0, /* tp_setattr*/
0, /* tp_compare*/
(reprfunc)PyNeighbor_repr, /* tp_repr*/
0, /* tp_as_number*/
0, /* tp_as_sequence*/
0, /* tp_as_mapping*/
0, /* tp_hash */
0, /* tp_call*/
0, /* tp_str*/
0, /* tp_getattro*/
0, /* tp_setattro*/
0, /* tp_as_buffer*/
Py_TPFLAGS_DEFAULT, /* tp_flags*/
PyNeighbor_doc, /* tp_doc */
0, /* tp_traverse */
0, /* tp_clear */
0, /* tp_richcompare */
0, /* tp_weaklistoffset */
0, /* tp_iter */
0, /* tp_iternext */
0, /* tp_methods */
0, /* tp_members */
PyNeighbor_getset, /* tp_getset */
0, /* tp_base */
0, /* tp_dict */
0, /* tp_descr_get */
0, /* tp_descr_set */
0, /* tp_dictoffset */
(initproc)PyNeighbor_init, /* tp_init */
};
typedef struct {
PyObject_HEAD
struct KDTree* tree;
} PyTree;
static void
PyTree_dealloc(PyTree* self)
{
KDTree_destroy(self->tree);
Py_TYPE(self)->tp_free((PyObject*)self);
}
static int
PyTree_init(PyTree* self, PyObject* args, PyObject* kwds)
{
int dim;
int bucket_size;
struct KDTree* tree;
if(!PyArg_ParseTuple(args, "ii:KDTree_init" ,&dim, &bucket_size)) return -1;
if (dim <= 0 || bucket_size <= 0)
{
PyErr_SetString(PyExc_ValueError, "Both arguments should be positive");
return -1;
}
tree = KDTree_init(dim, bucket_size);
if (tree==NULL)
{
PyErr_SetString(PyExc_MemoryError, "Insufficient memory for tree");
return -1;
}
self->tree = tree;
return 0;
}
static PyObject*
PyTree_get_count(PyTree* self)
{
long count;
struct KDTree* tree = self->tree;
PyObject* result;
count = KDTree_get_count(tree);
#if PY_MAJOR_VERSION >= 3
result = PyLong_FromLong(count);
#else
result = PyInt_FromLong(count);
#endif
if (!result)
{
PyErr_SetString (PyExc_MemoryError, "Failed to allocate memory for object.");
return NULL;
}
return result;
}
static PyObject*
PyTree_neighbor_get_count(PyTree* self)
{
long count;
struct KDTree* tree = self->tree;
PyObject* result;
count = KDTree_neighbor_get_count(tree);
#if PY_MAJOR_VERSION >= 3
result = PyLong_FromLong(count);
#else
result = PyInt_FromLong(count);
#endif
if (!result)
{
PyErr_SetString (PyExc_MemoryError, "Failed to allocate memory for object.");
return NULL;
}
return result;
}
static PyObject*
PyTree_set_data(PyTree* self, PyObject* args)
{
float* coords;
long int n, m, i;
PyObject *obj;
PyArrayObject *array;
struct KDTree* tree = self->tree;
int ok;
npy_intp rowstride, colstride;
const char* p;
if(!PyArg_ParseTuple(args, "O:KDTree_set_data",&obj)) return NULL;
/* Check if it is an array */
if (!PyArray_Check(obj))
{
PyErr_SetString(PyExc_TypeError, "First argument must be an array.");
return NULL;
}
array=(PyArrayObject *) obj;
if(PyArray_NDIM(array)!=2)
{
PyErr_SetString(PyExc_ValueError, "Array must be two dimensional.");
return NULL;
}
if (PyArray_TYPE(array) == NPY_DOUBLE)
{
Py_INCREF(obj);
}
else
{
/* Cast to type double */
obj = PyArray_Cast(array, NPY_DOUBLE);
if (!obj)
{
PyErr_SetString(PyExc_ValueError,
"coordinates cannot be cast to needed type.");
return NULL;
}
array = (PyArrayObject*) obj;
}
n = (long int) PyArray_DIM(array, 0);
m = (long int) PyArray_DIM(array, 1);
/* coord_data is deleted by the KDTree object */
coords= malloc(m*n*sizeof(float));
if (!coords)
{
Py_DECREF(obj);
PyErr_SetString (PyExc_MemoryError, "Failed to allocate memory for coordinates.");
return NULL;
}
rowstride = PyArray_STRIDE(array, 0);
colstride = PyArray_STRIDE(array, 1);
p = PyArray_BYTES(array);
for (i=0; i<n; i++)
{
int j;
for (j=0; j<m; j++)
{
coords[i*m+j]=*(double *) (p+i*rowstride+j*colstride);
}
}
Py_DECREF(obj);
ok = KDTree_set_data(tree, coords, n);
if (!ok)
{
PyErr_SetString (PyExc_MemoryError, "Failed to allocate memory for nodes.");
return NULL;
}
Py_INCREF(Py_None);
return Py_None;
}
static PyObject*
PyTree_search_center_radius(PyTree* self, PyObject* args)
{
PyObject *obj;
double radius;
PyArrayObject *array;
long int n, i;
float *coords;
struct KDTree* tree = self->tree;
int ok;
npy_intp stride;
const char* p;
if(!PyArg_ParseTuple(args, "Od:KDTree_search_center_radius", &obj ,&radius))
return NULL;
if(radius <= 0)
{
PyErr_SetString(PyExc_ValueError, "Radius must be positive.");
return NULL;
}
/* Check if it is an array */
if (!PyArray_Check(obj))
{
PyErr_SetString(PyExc_TypeError, "First argument must be an array.");
return NULL;
}
array=(PyArrayObject *) obj;
if(PyArray_NDIM(array)!=1)
{
PyErr_SetString(PyExc_ValueError, "Array must be one dimensional.");
return NULL;
}
if (PyArray_TYPE(array) == NPY_DOUBLE)
{
Py_INCREF(obj);
}
else
{
/* Cast to type double */
obj = PyArray_Cast(array, NPY_DOUBLE);
if (!obj)
{
PyErr_SetString(PyExc_ValueError,
"coordinates cannot be cast to needed type.");
return NULL;
}
array = (PyArrayObject*) obj;
}
n = (long int) PyArray_DIM(array, 0);
/* coord_data is deleted by the KDTree object */
coords= malloc(n*sizeof(float));
if (!coords)
{
Py_DECREF(obj);
PyErr_SetString (PyExc_MemoryError, "Failed to allocate memory for coordinates.");
return NULL;
}
stride = PyArray_STRIDE(array, 0);
p = PyArray_BYTES(array);
for (i=0; i<n; i++)
{
coords[i]=*(double *) (p+i*stride);
}
Py_DECREF(obj);
ok = KDTree_search_center_radius(tree, coords, radius);
if (!ok)
{
PyErr_SetString (PyExc_MemoryError, "Insufficient memory for calculation.");
return NULL;
}
Py_INCREF(Py_None);
return Py_None;
}
static PyObject*
PyTree_neighbor_search(PyTree* self, PyObject* args)
{
int ok;
double radius;
struct KDTree* tree = self->tree;
struct Neighbor* neighbors;
struct Neighbor* pp, *qq;
PyObject* list;
Py_ssize_t i, n;
if(!PyArg_ParseTuple(args, "d:KDTree_neighbor_search", &radius))
return NULL;
if(radius <= 0)
{
PyErr_SetString(PyExc_ValueError, "Radius must be positive.");
return NULL;
}
ok = KDTree_neighbor_search(tree, radius, &neighbors);
if (!ok)
{
PyErr_SetString(PyExc_MemoryError,
"calculation failed due to lack of memory");
return NULL;
}
pp = neighbors;
n = 0;
while (pp)
{
n+=1;
pp = pp->next;
}
list = PyList_New(n);
if (list)
{
PyNeighbor* p;
pp = neighbors;
for (i = 0; i < n; i++)
{
p = (PyNeighbor*) PyNeighborType.tp_alloc(&PyNeighborType, 0);
if(!p)
{
PyErr_SetString(PyExc_MemoryError,
"could not create node for return value");
Py_DECREF(list);
return NULL;
}
p->neighbor = *pp;
PyList_SET_ITEM(list, i, (PyObject*)p);
qq = pp->next;
free(pp);
pp = qq;
}
}
return list;
}
static PyObject*
PyTree_neighbor_simple_search(PyTree* self, PyObject* args)
{
int ok;
double radius;
struct KDTree* tree = self->tree;
struct Neighbor* neighbors;
struct Neighbor* pp, *qq;
PyObject* list;
Py_ssize_t i, n;
if(!PyArg_ParseTuple(args, "d:KDTree_neighbor_simple_search", &radius))
return NULL;
if(radius <= 0)
{
PyErr_SetString(PyExc_ValueError, "Radius must be positive.");
return NULL;
}
ok = KDTree_neighbor_simple_search(tree, radius, &neighbors);
if (!ok)
{
PyErr_SetString(PyExc_MemoryError,
"calculation failed due to lack of memory");
return NULL;
}
pp = neighbors;
n = 0;
while (pp)
{
n+=1;
pp = pp->next;
}
list = PyList_New(n);
if (list)
{
PyNeighbor* p;
pp = neighbors;
for (i = 0; i < n; i++)
{
p = (PyNeighbor*) PyNeighborType.tp_alloc(&PyNeighborType, 0);
if(!p)
{
PyErr_SetString(PyExc_MemoryError,
"could not create node for return value");
Py_DECREF(list);
return NULL;
}
p->neighbor = *pp;
PyList_SET_ITEM(list, i, (PyObject*)p);
qq = pp->next;
free(pp);
pp = qq;
}
}
return list;
}
static char PyTree_get_indices__doc__[] =
"returns indices of coordinates within radius as a Numpy array\n";
static PyObject *PyTree_get_indices(PyTree *self)
{
npy_intp length;
PyArrayObject *array;
struct KDTree* tree = self->tree;
length=KDTree_get_count(tree);
if (length==0)
{
Py_INCREF(Py_None);
return Py_None;
}
array=(PyArrayObject *) PyArray_SimpleNew(1, &length, PyArray_LONG);
if (!array)
{
PyErr_SetString(PyExc_MemoryError,
"Insufficient memory for array");
return NULL;
}
/* copy the data into the Numpy data pointer */
KDTree_copy_indices(tree, (long int *) PyArray_BYTES(array));
return PyArray_Return(array);
}
static char PyTree_get_radii__doc__[] =
"returns distances of coordinates within radius as a Numpy array.\n";
static PyObject *PyTree_get_radii(PyTree *self)
{
npy_intp length;
PyArrayObject *array;
struct KDTree* tree = self->tree;
length=KDTree_get_count(tree);
if (length==0)
{
Py_INCREF(Py_None);
return Py_None;
}
array=(PyArrayObject *) PyArray_SimpleNew(1, &length, PyArray_FLOAT);
if (!array)
{
PyErr_SetString(PyExc_MemoryError,
"Insufficient memory for array");
return NULL;
}
/* copy the data into the Numpy data pointer */
KDTree_copy_radii(tree, (float *) PyArray_BYTES(array));
return PyArray_Return(array);
}
static PyMethodDef PyTree_methods[] = {
{"get_count", (PyCFunction)PyTree_get_count, METH_NOARGS, NULL},
{"set_data", (PyCFunction)PyTree_set_data, METH_VARARGS, NULL},
{"search_center_radius", (PyCFunction)PyTree_search_center_radius, METH_VARARGS, NULL},
{"neighbor_get_count", (PyCFunction)PyTree_neighbor_get_count, METH_NOARGS, NULL},
{"neighbor_search", (PyCFunction)PyTree_neighbor_search, METH_VARARGS, NULL},
{"neighbor_simple_search", (PyCFunction)PyTree_neighbor_simple_search, METH_VARARGS, NULL},
{"get_indices", (PyCFunction)PyTree_get_indices, METH_NOARGS, PyTree_get_indices__doc__},
{"get_radii", (PyCFunction)PyTree_get_radii, METH_NOARGS, PyTree_get_radii__doc__},
{NULL} /* Sentinel */
};
static char PyTree_doc[] = "C KDTree.\n";
static PyTypeObject PyTreeType = {
PyVarObject_HEAD_INIT(NULL, 0)
"C KDTree", /*tp_name*/
sizeof(PyTree), /*tp_basicsize*/
0, /*tp_itemsize*/
(destructor)PyTree_dealloc, /*tp_dealloc*/
0, /*tp_print*/
0, /*tp_getattr*/
0, /*tp_setattr*/
0, /*tp_compare*/
0, /*tp_repr*/
0, /*tp_as_number*/
0, /*tp_as_sequence*/
0, /*tp_as_mapping*/
0, /*tp_hash */
0, /*tp_call*/
0, /*tp_str*/
0, /*tp_getattro*/
0, /*tp_setattro*/
0, /*tp_as_buffer*/
Py_TPFLAGS_DEFAULT, /*tp_flags*/
PyTree_doc, /* tp_doc */
0, /* tp_traverse */
0, /* tp_clear */
0, /* tp_richcompare */
0, /* tp_weaklistoffset */
0, /* tp_iter */
0, /* tp_iternext */
PyTree_methods, /* tp_methods */
NULL, /* tp_members */
0, /* tp_getset */
0, /* tp_base */
0, /* tp_dict */
0, /* tp_descr_get */
0, /* tp_descr_set */
0, /* tp_dictoffset */
(initproc)PyTree_init, /* tp_init */
};
/* ========================================================================== */
/* -- Initialization -------------------------------------------------------- */
/* ========================================================================== */
#if PY_MAJOR_VERSION >= 3
static struct PyModuleDef moduledef = {
PyModuleDef_HEAD_INIT,
"_CKDTree",
NULL,
-1,
NULL,
NULL,
NULL,
NULL,
NULL
};
PyObject *
PyInit__CKDTree(void)
#else
void
init_CKDTree(void)
#endif
{
PyObject *module;
import_array();
PyTreeType.tp_new = PyType_GenericNew;
PyNeighborType.tp_new = PyType_GenericNew;
if (PyType_Ready(&PyTreeType) < 0)
#if PY_MAJOR_VERSION >= 3
return NULL;
#else
return;
#endif
if (PyType_Ready(&PyNeighborType) < 0)
#if PY_MAJOR_VERSION >= 3
return NULL;
#else
return;
#endif
#if PY_MAJOR_VERSION >= 3
module = PyModule_Create(&moduledef);
if (module==NULL) return NULL;
#else
module = Py_InitModule("_CKDTree", NULL);
if (module==NULL) return;
#endif
Py_INCREF(&PyTreeType);
Py_INCREF(&PyNeighborType);
PyModule_AddObject(module, "KDTree", (PyObject*) &PyTreeType);
PyModule_AddObject(module, "Neighbor", (PyObject*) &PyNeighborType);
if (PyErr_Occurred()) Py_FatalError("can't initialize module _CKDTree");
#if PY_MAJOR_VERSION >= 3
return module;
#endif
}
Jump to Line
Something went wrong with that request. Please try again.