void apply_polyaffine(PyArrayObject* XYZ,
const PyArrayObject* Centers,
const PyArrayObject* Affines,
const PyArrayObject* Sigma)
{
PyArrayIterObject *iter_xyz, *iter_centers, *iter_affines;
int axis = 1;
double *xyz, *center, *affine, *sigma;
double w, W;
double mat[12], t_xyz[3];
size_t bytes_mat = 12*sizeof(double);
size_t bytes_xyz = 3*sizeof(double);
/* Initialize arrays and iterators */
sigma = PyArray_DATA(Sigma);
iter_xyz = (PyArrayIterObject*)PyArray_IterAllButAxis((PyObject*)XYZ, &axis);
iter_centers = (PyArrayIterObject*)PyArray_IterAllButAxis((PyObject*)Centers, &axis);
iter_affines = (PyArrayIterObject*)PyArray_IterAllButAxis((PyObject*)Affines, &axis);
/* Loop over input points */
while(iter_xyz->index < iter_xyz->size) {
xyz = PyArray_ITER_DATA(iter_xyz);
PyArray_ITER_RESET(iter_centers);
PyArray_ITER_RESET(iter_affines);
memset((void*)mat, 0, bytes_mat);
W = 0.0;
/* Loop over centers */
while(iter_centers->index < iter_centers->size) {
center = PyArray_ITER_DATA(iter_centers);
affine = PyArray_ITER_DATA(iter_affines);
w = _gaussian(xyz, center, sigma);
W += w;
_add_weighted_affine(mat, affine, w);
PyArray_ITER_NEXT(iter_centers);
PyArray_ITER_NEXT(iter_affines);
}
/* Apply matrix */
_apply_affine(t_xyz, mat, xyz, W);
memcpy((void*)xyz, (void*)t_xyz, bytes_xyz);
/* Update xyz iterator */
PyArray_ITER_NEXT(iter_xyz);
}
/* Free memory */
Py_XDECREF(iter_xyz);
Py_XDECREF(iter_centers);
Py_XDECREF(iter_affines);
return;
}
/*
Compute the interaction energy:
sum_i,j qi^T U qj
= sum_i qi^T sum_j U qj
*/
double interaction_energy(PyArrayObject* ppm,
const PyArrayObject* XYZ,
const PyArrayObject* U,
int ngb_size)
{
npy_intp k, x, y, z, pos;
double *p, *buf;
double res = 0.0, tmp;
PyArrayIterObject* iter;
int axis = 1;
double* ppm_data;
npy_intp K = ppm->dimensions[3];
npy_intp u2 = ppm->dimensions[2]*K;
npy_intp u1 = ppm->dimensions[1]*u2;
npy_intp* xyz;
const double* U_data = (double*)U->data;
int* ngb;
/* Neighborhood system */
ngb = _select_neighborhood_system(ngb_size);
/* Pointer to ppm array */
ppm_data = (double*)ppm->data;
/* Allocate auxiliary vector */
p = (double*)calloc(K, sizeof(double));
/* Loop over points */
iter = (PyArrayIterObject*)PyArray_IterAllButAxis((PyObject*)XYZ, &axis);
while(iter->index < iter->size) {
/* Compute the average ppm in the neighborhood */
xyz = PyArray_ITER_DATA(iter);
x = xyz[0];
y = xyz[1];
z = xyz[2];
_ngb_integrate(p, ppm, x, y, z, U_data, (const int*)ngb, ngb_size);
/* Calculate the dot product qi^T p where qi is the local
posterior */
tmp = 0.0;
pos = x*u1 + y*u2 + z*K;
for (k=0, buf=p; k<K; k++, pos++, buf++)
tmp += ppm_data[pos]*(*buf);
/* Update overall energy */
res += tmp;
/* Update iterator */
PyArray_ITER_NEXT(iter);
}
/* Free memory */
free(p);
Py_XDECREF(iter);
return res;
}
static void _cubic_spline_transform(PyArrayObject* res, int axis, double* work)
{
PyArrayIterObject* iter;
unsigned int dim, stride;
/* Instantiate iterator and views */
iter = (PyArrayIterObject*)PyArray_IterAllButAxis((PyObject*)res, &axis);
dim = PyArray_DIM((PyArrayObject*)iter->ao, axis);
stride = PyArray_STRIDE((PyArrayObject*)iter->ao, axis)/sizeof(double);
/* Apply the cubic spline transform along given axis */
while(iter->index < iter->size) {
_copy_double_buffer(work, PyArray_ITER_DATA(iter), dim, stride);
_cubic_spline_transform1d(PyArray_ITER_DATA(iter), work, dim, stride, 1);
PyArray_ITER_NEXT(iter);
}
/* Free local structures */
Py_DECREF(iter);
return;
}
void ve_step(PyArrayObject* ppm,
const PyArrayObject* ref,
const PyArrayObject* XYZ,
const PyArrayObject* U,
int ngb_size,
double beta)
{
npy_intp k, x, y, z, pos;
double *p, *buf, *ppm_data;
double psum, tmp;
PyArrayIterObject* iter;
int axis = 1;
npy_intp K = ppm->dimensions[3];
npy_intp u2 = ppm->dimensions[2]*K;
npy_intp u1 = ppm->dimensions[1]*u2;
const double* ref_data = (double*)ref->data;
const double* U_data = (double*)U->data;
npy_intp* xyz;
int* ngb;
/* Neighborhood system */
ngb = _select_neighborhood_system(ngb_size);
/* Pointer to the data array */
ppm_data = (double*)ppm->data;
/* Allocate auxiliary vectors */
p = (double*)calloc(K, sizeof(double));
/* Loop over points */
iter = (PyArrayIterObject*)PyArray_IterAllButAxis((PyObject*)XYZ, &axis);
while(iter->index < iter->size) {
/* Integrate the energy over the neighborhood */
xyz = PyArray_ITER_DATA(iter);
x = xyz[0];
y = xyz[1];
z = xyz[2];
_ngb_integrate(p, ppm, x, y, z, U_data, (const int*)ngb, ngb_size);
/* Apply exponential transform, multiply with reference and
compute normalization constant */
psum = 0.0;
for (k=0, pos=(iter->index)*K, buf=p; k<K; k++, pos++, buf++) {
tmp = exp(-2 * beta * (*buf)) * ref_data[pos];
psum += tmp;
*buf = tmp;
}
/* Normalize to unitary sum */
pos = x*u1 + y*u2 + z*K;
if (psum > TINY)
for (k=0, buf=p; k<K; k++, pos++, buf++)
ppm_data[pos] = *buf/psum;
else
for (k=0, buf=p; k<K; k++, pos++, buf++)
ppm_data[pos] = (*buf+TINY/(double)K)/(psum+TINY);
/* Update iterator */
PyArray_ITER_NEXT(iter);
}
/* Free memory */
free(p);
Py_XDECREF(iter);
return;
}
static PyObject* py_fitexpsin(PyObject *obj, PyObject *args, PyObject *kwds)
{
PyArrayObject *data = NULL;
PyArrayObject *fitt = NULL;
PyArrayObject *rslt = NULL;
PyArrayIterObject *data_it = NULL;
PyArrayIterObject *fitt_it = NULL;
PyArrayIterObject *rslt_it = NULL;
Py_ssize_t newshape[NPY_MAXDIMS];
double *poly = NULL;
double *coef = NULL;
double *buff = NULL;
int i, j, error, lastaxis, numdata;
int startcoef = -1;
int numcoef = MAXCOEF;
int axis = NPY_MAXDIMS;
double deltat = 1.0;
static char *kwlist[] = {"data", "numcoef",
"deltat", "axis", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|idO&", kwlist,
PyConverter_AnyDoubleArray, &data, &numcoef, &deltat,
PyArray_AxisConverter, &axis)) return NULL;
if (axis < 0) {
axis += PyArray_NDIM(data);
}
if ((axis < 0) || (axis >= NPY_MAXDIMS)) {
PyErr_Format(PyExc_ValueError, "invalid axis");
goto _fail;
}
lastaxis = PyArray_NDIM(data) - 1;
if ((numcoef < 1) || (numcoef > MAXCOEF)) {
PyErr_Format(PyExc_ValueError, "numcoef out of bounds");
goto _fail;
}
if (startcoef < 0) { /* start regression away from zero coefficients */
startcoef = 4;
}
if (startcoef > numcoef - 2) {
PyErr_Format(PyExc_ValueError, "startcoef out of bounds");
goto _fail;
}
numdata = (int)PyArray_DIM(data, axis);
if (numcoef > numdata)
numcoef = numdata;
if ((numcoef - startcoef - 1) < 3) {
PyErr_Format(PyExc_ValueError,
"number of coefficients insufficient to fit data");
goto _fail;
}
/* fitted data */
fitt = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(data),
PyArray_DIMS(data), NPY_DOUBLE);
if (fitt == NULL) {
PyErr_Format(PyExc_MemoryError, "unable to allocate fitt array");
goto _fail;
}
/* fitted parameters */
j = 0;
for (i = 0; i < PyArray_NDIM(data); i++) {
if (i != axis)
newshape[j++] = PyArray_DIM(data, i);
}
newshape[j] = 5;
rslt = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(data),
newshape, NPY_DOUBLE);
if (rslt == NULL) {
PyErr_Format(PyExc_MemoryError, "unable to allocate rslt array");
goto _fail;
}
/* working buffer */
buff = (double *)PyMem_Malloc(3*numdata * sizeof(double));
if (buff == NULL) {
PyErr_Format(PyExc_MemoryError, "unable to allocate buff array");
goto _fail;
}
/* buffer for differential coefficients */
coef = (double *)PyMem_Malloc((3+1)*(numcoef+1) * sizeof(double));
if (coef == NULL) {
PyErr_Format(PyExc_MemoryError, "unable to allocate coef array");
goto _fail;
}
/* precalculate normalized Chebyshev polynomial */
poly = (double *)PyMem_Malloc(numdata * (numcoef+1) * sizeof(double));
if (poly == NULL) {
PyErr_Format(PyExc_MemoryError, "unable to allocate poly");
goto _fail;
}
error = chebypoly(numdata, numcoef, poly, 1);
if (error != 0) {
PyErr_Format(PyExc_ValueError,
"chebypoly() failed with error code %i", error);
goto _fail;
}
/* iterate over all but specified axis */
data_it = (PyArrayIterObject *)PyArray_IterAllButAxis(
(PyObject *)data, &axis);
fitt_it = (PyArrayIterObject *)PyArray_IterAllButAxis(
(PyObject *)fitt, &axis);
rslt_it = (PyArrayIterObject *)PyArray_IterAllButAxis(
(PyObject *)rslt, &lastaxis);
while (data_it->index < data_it->size) {
error = fitexpsin(
(char *)data_it->dataptr,
(int)PyArray_STRIDE(data, axis),
numdata,
poly,
coef,
numcoef,
deltat,
startcoef,
buff,
(double *)rslt_it->dataptr,
(char *)fitt_it->dataptr,
(int)PyArray_STRIDE(fitt, axis));
if (error != 0) {
PyErr_Format(PyExc_ValueError,
"fitexpsin() failed with error code %i", error);
goto _fail;
}
PyArray_ITER_NEXT(data_it);
PyArray_ITER_NEXT(fitt_it);
PyArray_ITER_NEXT(rslt_it);
}
Py_XDECREF(data_it);
Py_XDECREF(fitt_it);
Py_XDECREF(rslt_it);
Py_XDECREF(data);
PyMem_Free(poly);
PyMem_Free(coef);
PyMem_Free(buff);
return Py_BuildValue("(N, N)", rslt, fitt);
_fail:
Py_XDECREF(data_it);
Py_XDECREF(fitt_it);
Py_XDECREF(rslt_it);
Py_XDECREF(data);
Py_XDECREF(fitt);
Py_XDECREF(rslt);
PyMem_Free(poly);
PyMem_Free(coef);
PyMem_Free(buff);
return NULL;
}
//Wrapper for Brute-Force neighbor search
static PyObject*
BallTree_knn_brute(PyObject *self, PyObject *args, PyObject *kwds){
long int k = 1;
std::vector<BallTree_Point*> Points;
PyObject *arg1 = NULL;
PyObject *arg2 = NULL;
PyObject *arr1 = NULL;
PyObject *arr2 = NULL;
PyObject *nbrs = NULL;
long int* nbrs_data;
PyArrayIterObject *arr2_iter = NULL;
PyArrayIterObject *nbrs_iter = NULL;
static char *kwlist[] = {"x", "pt", "k", NULL};
npy_intp* dim;
int nd, pt_size, pt_inc;
long int N;
long int D;
//parse arguments. If k is not provided, the default is 1
if(!PyArg_ParseTupleAndKeywords(args,kwds,"OO|l",kwlist,
&arg1,&arg2,&k))
goto fail;
//First array should be a 2D array of doubles
arr1 = PyArray_FROM_OTF(arg1,NPY_DOUBLE,0);
if(arr1==NULL)
goto fail;
if( PyArray_NDIM(arr1) != 2){
PyErr_SetString(PyExc_ValueError,
"x must be two dimensions");
goto fail;
}
//Second array should be a 1D array of doubles
arr2 = PyArray_FROM_OTF(arg2,NPY_DOUBLE,0);
if(arr2==NULL)
goto fail;
nd = PyArray_NDIM(arr2);
if(nd == 0){
PyErr_SetString(PyExc_ValueError,
"pt cannot be zero-sized array");
goto fail;
}
pt_size = PyArray_DIM(arr2,nd-1);
//Check that dimensions match
N = PyArray_DIMS(arr1)[0];
D = PyArray_DIMS(arr1)[1];
if( pt_size != D ){
PyErr_SetString(PyExc_ValueError,
"pt must be same dimension as x");
goto fail;
}
//check the value of k
if(k<1){
PyErr_SetString(PyExc_ValueError,
"k must be a positive integer");
goto fail;
}
if(k>N){
PyErr_SetString(PyExc_ValueError,
"k must be less than the number of points");
goto fail;
}
//create a neighbors array and distance array
dim = new npy_intp[nd];
for(int i=0; i<nd-1;i++)
dim[i] = PyArray_DIM(arr2,i);
dim[nd-1] = k;
nbrs = (PyObject*)PyArray_SimpleNew(nd,dim,PyArray_LONG);
delete[] dim;
if(nbrs==NULL)
goto fail;
//create iterators to cycle through points
nd-=1;
arr2_iter = (PyArrayIterObject*)PyArray_IterAllButAxis(arr2,&nd);
nbrs_iter = (PyArrayIterObject*)PyArray_IterAllButAxis(nbrs,&nd);
nd+=1;
if( arr2_iter==NULL ||
nbrs_iter==NULL ||
(arr2_iter->size != nbrs_iter->size) ){
PyErr_SetString(PyExc_ValueError,
"failure constructing iterators");
goto fail;
}
pt_inc = PyArray_STRIDES(arr2)[nd-1] / PyArray_DESCR(arr2)->elsize;
if(PyArray_STRIDES(nbrs)[nd-1] != PyArray_DESCR(nbrs)->elsize ){
PyErr_SetString(PyExc_ValueError,
"nbrs not allocated as a C-array");
goto fail;
}
//create the list of points
pt_inc = PyArray_STRIDES(arr1)[1]/PyArray_DESCR(arr1)->elsize;
Points.resize(N);
for(int i=0;i<N;i++)
Points[i] = new BallTree_Point(arr1,
(double*)PyArray_GETPTR2(arr1,i,0),
pt_inc, PyArray_DIM(arr1,1));
//iterate through points and determine neighbors
//warning: if nbrs is not a C-array, or if we're not iterating
// over the last dimension, this may cause a seg fault.
while(arr2_iter->index < arr2_iter->size){
BallTree_Point Query_Point(arr2,
(double*)PyArray_ITER_DATA(arr2_iter),
pt_inc,pt_size);
nbrs_data = (long int*)(PyArray_ITER_DATA(nbrs_iter));
BruteForceNeighbors(Points, Query_Point,
k, nbrs_data );
PyArray_ITER_NEXT(arr2_iter);
PyArray_ITER_NEXT(nbrs_iter);
}
for(int i=0;i<N;i++)
delete Points[i];
//if only one neighbor is requested, then resize the neighbors array
if(k==1){
PyArray_Dims dims;
dims.ptr = PyArray_DIMS(arr2);
dims.len = PyArray_NDIM(arr2)-1;
//PyArray_Resize returns None - this needs to be picked
// up and dereferenced.
PyObject *NoneObj = PyArray_Resize( (PyArrayObject*)nbrs, &dims,
0, NPY_ANYORDER );
if (NoneObj == NULL){
goto fail;
}
Py_DECREF(NoneObj);
}
return nbrs;
fail:
Py_XDECREF(arr1);
Py_XDECREF(arr2);
Py_XDECREF(nbrs);
Py_XDECREF(arr2_iter);
Py_XDECREF(nbrs_iter);
return NULL;
}
//query the ball tree. Arguments are the array of search points
// and the radius around each point to search
static PyObject *
BallTree_queryball(BallTreeObject *self, PyObject *args, PyObject *kwds){
//we use goto statements : all variables should be declared up front
int count_only = 0;
double r;
PyObject *arg = NULL;
PyObject *arr = NULL;
PyObject *nbrs = NULL;
PyArrayIterObject* arr_iter = NULL;
PyArrayIterObject* nbrs_iter = NULL;
int nd, pt_size, pt_inc;
static char *kwlist[] = {"x", "r", "count_only", NULL};
//parse arguments. If kmax is not provided, the default is 20
if(!PyArg_ParseTupleAndKeywords(args,kwds,"Od|i",
kwlist,&arg,&r,&count_only)){
goto fail;
}
//check value of r
if(r < 0){
PyErr_SetString(PyExc_ValueError,
"r must not be negative");
goto fail;
}
//get the array object from the first argument
arr = PyArray_FROM_OTF(arg,NPY_DOUBLE,NPY_ALIGNED);
//check that the array was properly constructed
if(arr==NULL){
PyErr_SetString(PyExc_ValueError,
"pt must be convertable to array");
goto fail;
}
nd = PyArray_NDIM(arr);
if(nd == 0){
PyErr_SetString(PyExc_ValueError,
"pt cannot be zero-sized array");
goto fail;
}
pt_size = PyArray_DIM(arr,nd-1);
if( pt_size != self->tree->PointSize() ){
PyErr_SetString(PyExc_ValueError,
"points are incorrect dimension");
goto fail;
}
// Case 1: return arrays of all neighbors for each point
//
if(!count_only){
//create a neighbors array. This is an array of python objects.
// each of which will be a numpy array of neighbors
nbrs = (PyObject*)PyArray_SimpleNew(nd-1,PyArray_DIMS(arr),
PyArray_OBJECT);
if(nbrs==NULL){
goto fail;
}
//create iterators to cycle through points
--nd;
arr_iter = (PyArrayIterObject*)PyArray_IterAllButAxis(arr,&nd);
nbrs_iter = (PyArrayIterObject*)PyArray_IterNew(nbrs);
++nd;
if( arr_iter==NULL || nbrs_iter==NULL ||
(arr_iter->size != nbrs_iter->size)){
PyErr_SetString(PyExc_ValueError,
"unable to construct iterator");
goto fail;
}
pt_inc = PyArray_STRIDES(arr)[nd-1] / PyArray_DESCR(arr)->elsize;
if(PyArray_NDIM(nbrs)==0){
BallTree_Point pt(arr,
(double*)PyArray_ITER_DATA(arr_iter),
pt_inc,pt_size);
std::vector<long int> nbrs_vec;
self->tree->query_ball(pt,r,nbrs_vec);
npy_intp N_nbrs = nbrs_vec.size();
PyObject* nbrs_obj = PyArray_SimpleNew(1, &N_nbrs, PyArray_LONG);
long int* data = (long int*)PyArray_DATA(nbrs_obj);
for(int i=0; i<N_nbrs; i++)
data[i] = nbrs_vec[i];
PyObject* tmp = nbrs;
nbrs = nbrs_obj;
Py_DECREF(tmp);
}else{
while(arr_iter->index < arr_iter->size){
BallTree_Point pt(arr,
(double*)PyArray_ITER_DATA(arr_iter),
pt_inc,pt_size);
std::vector<long int> nbrs_vec;
self->tree->query_ball(pt,r,nbrs_vec);
npy_intp N_nbrs = nbrs_vec.size();
PyObject* nbrs_obj = PyArray_SimpleNew(1, &N_nbrs,
PyArray_LONG);
long int* data = (long int*)PyArray_DATA(nbrs_obj);
for(int i=0; i<N_nbrs; i++)
data[i] = nbrs_vec[i];
PyObject** nbrs_data = (PyObject**)PyArray_ITER_DATA(nbrs_iter);
PyObject* tmp = nbrs_data[0];
nbrs_data[0] = nbrs_obj;
Py_XDECREF(tmp);
PyArray_ITER_NEXT(arr_iter);
PyArray_ITER_NEXT(nbrs_iter);
}
}
}
// Case 2 : return number of neighbors for each point
else{
//create an array to keep track of the count
nbrs = (PyObject*)PyArray_SimpleNew(nd-1,PyArray_DIMS(arr),
PyArray_LONG);
if(nbrs==NULL){
goto fail;
}
//create iterators to cycle through points
--nd;
arr_iter = (PyArrayIterObject*)PyArray_IterAllButAxis(arr,&nd);
nbrs_iter = (PyArrayIterObject*)PyArray_IterNew(nbrs);
++nd;
if( arr_iter==NULL || nbrs_iter==NULL ||
(arr_iter->size != nbrs_iter->size)){
PyErr_SetString(PyExc_ValueError,
"unable to construct iterator");
goto fail;
}
//go through points and call BallTree::query_ball to count neighbors
pt_inc = PyArray_STRIDES(arr)[nd-1] / PyArray_DESCR(arr)->elsize;
while(arr_iter->index < arr_iter->size){
BallTree_Point pt(arr,
(double*)PyArray_ITER_DATA(arr_iter),
pt_inc,pt_size);
long int* nbrs_count = (long int*)PyArray_ITER_DATA(nbrs_iter);
*nbrs_count = self->tree->query_ball(pt,r);
PyArray_ITER_NEXT(arr_iter);
PyArray_ITER_NEXT(nbrs_iter);
}
}
Py_DECREF(nbrs_iter);
Py_DECREF(arr_iter);
Py_DECREF(arr);
return nbrs;
fail:
Py_XDECREF(nbrs_iter);
Py_XDECREF(arr_iter);
Py_XDECREF(arr);
Py_XDECREF(nbrs);
return NULL;
}
//query the ball tree. Arguments are the array of search points
// and the number of nearest neighbors, k (optional)
static PyObject *
BallTree_query(BallTreeObject *self, PyObject *args, PyObject *kwds){
//we use goto statements : all variables should be declared up front
int return_distance = 1;
long int k = 1;
PyObject *arg = NULL;
PyObject *arr = NULL;
PyObject *nbrs = NULL;
PyObject *dist = NULL;
PyArrayIterObject *arr_iter = NULL;
PyArrayIterObject *nbrs_iter = NULL;
PyArrayIterObject *dist_iter = NULL;
long int* nbrs_data;
double* dist_data;
static char *kwlist[] = {"x", "k", "return_distance", NULL};
int nd, pt_size, pt_inc;
npy_intp* dim;
//parse arguments. If k is not provided, the default is 1
if(!PyArg_ParseTupleAndKeywords(args,kwds,"O|li",kwlist,
&arg,&k,&return_distance)){
goto fail;
}
//check value of k
if(k < 1){
PyErr_SetString(PyExc_ValueError,
"k must be positive");
goto fail;
}
if(k > self->size){
PyErr_SetString(PyExc_ValueError,
"k must not be greater than number of points");
goto fail;
}
//get the array object from the first argument
arr = PyArray_FROM_OTF(arg,NPY_DOUBLE,NPY_ALIGNED);
//check that the array was properly constructed
if(arr==NULL){
PyErr_SetString(PyExc_ValueError,
"pt must be convertable to array");
goto fail;
}
nd = PyArray_NDIM(arr);
if(nd == 0){
PyErr_SetString(PyExc_ValueError,
"pt cannot be zero-sized array");
goto fail;
}
pt_size = PyArray_DIM(arr,nd-1);
if( pt_size != self->tree->PointSize() ){
PyErr_SetString(PyExc_ValueError,
"points are incorrect dimension");
goto fail;
}
//create a neighbors array and distance array
dim = new npy_intp[nd];
for(int i=0; i<nd-1;i++)
dim[i] = PyArray_DIM(arr,i);
dim[nd-1] = k;
nbrs = (PyObject*)PyArray_SimpleNew(nd,dim,PyArray_LONG);
if(return_distance)
dist = (PyObject*)PyArray_SimpleNew(nd,dim,PyArray_DOUBLE);
delete[] dim;
if(nbrs==NULL)
goto fail;
//create iterators to cycle through points
nd-=1;
arr_iter = (PyArrayIterObject*)PyArray_IterAllButAxis(arr,&nd);
nbrs_iter = (PyArrayIterObject*)PyArray_IterAllButAxis(nbrs,&nd);
if(return_distance)
dist_iter = (PyArrayIterObject*)PyArray_IterAllButAxis(dist,&nd);
nd+=1;
if( arr_iter==NULL || nbrs_iter==NULL ||
(arr_iter->size != nbrs_iter->size) ||
(return_distance &&
(dist_iter==NULL || (arr_iter->size != dist_iter->size))) ){
PyErr_SetString(PyExc_ValueError,
"failure constructing iterators");
goto fail;
}
pt_inc = PyArray_STRIDES(arr)[nd-1] / PyArray_DESCR(arr)->elsize;
if( (PyArray_STRIDES(nbrs)[nd-1] != PyArray_DESCR(nbrs)->elsize ) ||
(return_distance &&
(PyArray_STRIDES(dist)[nd-1] != PyArray_DESCR(dist)->elsize )) ){
PyErr_SetString(PyExc_ValueError,
"nbrs & dist not allocated as a C-array");
goto fail;
}
//iterate through points and determine neighbors
//warning: if nbrs is not a C-array, or if we're not iterating
// over the last dimension, this may cause a seg fault.
if(return_distance){
while(arr_iter->index < arr_iter->size){
BallTree_Point pt(arr,
(double*)PyArray_ITER_DATA(arr_iter),
pt_inc,pt_size);
nbrs_data = (long int*)(PyArray_ITER_DATA(nbrs_iter));
dist_data = (double*)(PyArray_ITER_DATA(dist_iter));
self->tree->query(pt,k,nbrs_data,dist_data);
PyArray_ITER_NEXT(arr_iter);
PyArray_ITER_NEXT(nbrs_iter);
PyArray_ITER_NEXT(dist_iter);
}
}else{
while(arr_iter->index < arr_iter->size){
BallTree_Point pt(arr,
(double*)PyArray_ITER_DATA(arr_iter),
pt_inc,pt_size);
nbrs_data = (long int*)(PyArray_ITER_DATA(nbrs_iter));
self->tree->query(pt,k,nbrs_data);
PyArray_ITER_NEXT(arr_iter);
PyArray_ITER_NEXT(nbrs_iter);
}
}
//if only one neighbor is requested, then resize the neighbors array
if(k==1){
PyArray_Dims dims;
dims.ptr = PyArray_DIMS(arr);
dims.len = PyArray_NDIM(arr)-1;
//PyArray_Resize returns None - this needs to be picked
// up and dereferenced.
PyObject *NoneObj = PyArray_Resize( (PyArrayObject*)nbrs, &dims,
0, NPY_ANYORDER );
if (NoneObj == NULL){
goto fail;
}
Py_DECREF(NoneObj);
if(return_distance){
NoneObj = PyArray_Resize( (PyArrayObject*)dist, &dims,
0, NPY_ANYORDER );
if (NoneObj == NULL){
goto fail;
}
Py_DECREF(NoneObj);
}
}
if(return_distance){
Py_DECREF(arr_iter);
Py_DECREF(nbrs_iter);
Py_DECREF(dist_iter);
Py_DECREF(arr);
arr = Py_BuildValue("(OO)",dist,nbrs);
Py_DECREF(nbrs);
Py_DECREF(dist);
return arr;
}else{
Py_DECREF(arr_iter);
Py_DECREF(nbrs_iter);
Py_DECREF(arr);
return nbrs;
}
fail:
Py_XDECREF(arr);
Py_XDECREF(nbrs);
Py_XDECREF(dist);
Py_XDECREF(arr_iter);
Py_XDECREF(nbrs_iter);
Py_XDECREF(dist_iter);
return NULL;
}
static PyObject *
py_interpolate(
PyObject *obj,
PyObject *args,
PyObject *kwds)
{
PyArrayObject *xdata = NULL;
PyArrayObject *data = NULL;
PyArrayObject *xout = NULL;
PyArrayObject *out = NULL;
PyArrayObject *oout = NULL;
PyArrayIterObject *dit = NULL;
PyArrayIterObject *oit = NULL;
npy_intp dstride, ostride, xdstride, xostride, size, outsize;
Py_ssize_t newshape[NPY_MAXDIMS];
int axis = NPY_MAXDIMS;
int i, ndim, error;
double *buffer = NULL;
static char *kwlist[] = {"x", "y", "x_new", "axis", "out", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&O&O&|O&O&", kwlist,
PyConverter_AnyDoubleArray, &xdata,
PyConverter_AnyDoubleArray, &data,
PyConverter_AnyDoubleArray, &xout,
PyArray_AxisConverter, &axis,
PyOutputConverter_AnyDoubleArrayOrNone, &oout))
goto _fail;
/* check axis */
ndim = PyArray_NDIM(data);
if ((axis == NPY_MAXDIMS) || (axis == -1)) {
axis = ndim - 1;
} else if ((axis < 0) || (axis > NPY_MAXDIMS)) {
PyErr_Format(PyExc_ValueError, "invalid axis");
goto _fail;
}
if ((PyArray_NDIM(xdata) != 1) || (PyArray_NDIM(xout) != 1)) {
PyErr_Format(PyExc_ValueError,
"x-arrays must be one dimensional");
goto _fail;
}
size = PyArray_DIM(data, axis);
outsize = PyArray_DIM(xout, 0);
if (size < 3) {
PyErr_Format(PyExc_ValueError, "size along axis is too small");
goto _fail;
}
if (size != PyArray_DIM(xdata, 0)) {
PyErr_Format(PyExc_ValueError,
"size of x-array must match data shape at axis");
goto _fail;
}
for (i = 0; i < ndim; i++) {
newshape[i] = (i == axis) ? outsize : PyArray_DIM(data, i);
}
if (oout == NULL) {
/* create a new output array */
out = (PyArrayObject*)PyArray_SimpleNew(ndim, newshape, NPY_DOUBLE);
if (out == NULL) {
PyErr_Format(PyExc_ValueError, "failed to allocate output array");
goto _fail;
}
} else if (ndim != PyArray_NDIM(oout)) {
PyErr_Format(PyExc_ValueError,
"output and data array dimension mismatch");
goto _fail;
} else {
for (i = 0; i < ndim; i++) {
if (newshape[i] != PyArray_DIM(oout, i)) {
PyErr_Format(PyExc_ValueError, "wrong output shape");
goto _fail;
}
}
out = oout;
}
/* iterate over all but specified axis */
dit = (PyArrayIterObject *)PyArray_IterAllButAxis((PyObject *)data, &axis);
oit = (PyArrayIterObject *)PyArray_IterAllButAxis((PyObject *)out, &axis);
dstride = PyArray_STRIDE(data, axis);
ostride = PyArray_STRIDE(out, axis);
xdstride = PyArray_STRIDE(xdata, 0);
xostride = PyArray_STRIDE(xout, 0);
buffer = (double *)PyMem_Malloc((size * 4 + 4) * sizeof(double));
if (buffer == NULL) {
PyErr_Format(PyExc_ValueError, "failed to allocate output buffer");
goto _fail;
}
while (dit->index < dit->size) {
error = interpolate(
size,
PyArray_DATA(xdata), xdstride,
dit->dataptr, dstride,
outsize,
PyArray_DATA(xout), xostride,
oit->dataptr, ostride,
buffer);
if (error != 0) {
PyErr_Format(PyExc_ValueError, "interpolate() failed");
goto _fail;
}
PyArray_ITER_NEXT(oit);
PyArray_ITER_NEXT(dit);
}
PyMem_Free(buffer);
Py_DECREF(oit);
Py_DECREF(dit);
Py_DECREF(data);
Py_DECREF(xout);
Py_DECREF(xdata);
/* Return output vector if not provided as argument */
if (oout == NULL) {
return PyArray_Return(out);
} else {
Py_INCREF(Py_None);
return Py_None;
}
_fail:
Py_XDECREF(xdata);
Py_XDECREF(xout);
Py_XDECREF(data);
Py_XDECREF(oit);
Py_XDECREF(dit);
if (buffer != NULL)
PyMem_Free(buffer);
if (oout == NULL)
Py_XDECREF(out);
else
Py_XDECREF(oout);
return NULL;
}
fffpy_multi_iterator* fffpy_multi_iterator_new(int narr, int axis, ...)
{
fffpy_multi_iterator* thisone;
va_list va;
fff_vector** vector;
PyArrayMultiIterObject *multi;
PyObject *current, *arr;
int i, err=0;
/* Create new instance */
thisone = (fffpy_multi_iterator*)malloc(sizeof(fffpy_multi_iterator));
multi = PyArray_malloc(sizeof(PyArrayMultiIterObject));
vector = (fff_vector**)malloc(narr*sizeof(fff_vector*));
/* Initialize the PyArrayMultiIterObject instance from the variadic arguments */
PyObject_Init((PyObject *)multi, &PyArrayMultiIter_Type);
for (i=0; i<narr; i++)
multi->iters[i] = NULL;
multi->numiter = narr;
multi->index = 0;
va_start(va, axis);
for (i=0; i<narr; i++) {
current = va_arg(va, PyObject *);
arr = PyArray_FROM_O(current);
if (arr==NULL) {
err=1; break;
}
else {
multi->iters[i] = (PyArrayIterObject *)PyArray_IterAllButAxis(arr, &axis);
Py_DECREF(arr);
}
}
va_end(va);
/* Test */
if (!err && _PyArray_BroadcastAllButAxis(multi, axis) < 0)
err=1;
if (err) {
FFF_ERROR("Cannot create broadcast object", ENOMEM);
free(thisone);
free(vector);
Py_DECREF(multi);
return NULL;
}
/* Initialize the multi iterator */
PyArray_MultiIter_RESET(multi);
/* Create the fff vectors (views or copies) */
for(i=0; i<narr; i++)
vector[i] = _fff_vector_new_from_PyArrayIter((const PyArrayIterObject*)multi->iters[i], axis);
/* Instantiate fiels */
thisone->narr = narr;
thisone->axis = axis;
thisone->vector = vector;
thisone->multi = multi;
thisone->index = thisone->multi->index;
thisone->size = thisone->multi->size;
return thisone;
}