object k_fixed_radius_search_array(ANNkd_tree& kdtree, object qarray, double sqRad, int k, double eps)
{
BOOST_ASSERT(k <= kdtree.nPoints());
int N = len(qarray);
if( N == 0 )
return boost::python::make_tuple(numeric::array(boost::python::list()).astype("i4"),numeric::array(boost::python::list()),numeric::array(boost::python::list()));
BOOST_ASSERT(len(qarray[0])==kdtree.theDim());
ANNpointManaged annq(kdtree.theDim());
npy_intp dimsball[] = { N};
PyObject *pykball = PyArray_SimpleNew(1,dimsball, PyArray_INT);
BOOST_ASSERT(!!pykball);
int* pkball = (int*)PyArray_DATA(pykball);
if( k <= 0 ) {
for(int i = 0; i < N; ++i) {
object q = qarray[i];
for (int c = 0; c < kdtree.theDim(); ++c)
annq.pt[c] = extract<ANNcoord>(q[c]);
pkball[i] = kdtree.annkFRSearch(annq.pt, sqRad, k, NULL, NULL, eps);
}
return boost::python::make_tuple(numeric::array(boost::python::list()).astype("i4"),numeric::array(boost::python::list()),static_cast<numeric::array>(handle<>(pykball)));
}
npy_intp dims[] = { N,k};
PyObject *pydists = PyArray_SimpleNew(2,dims, sizeof(ANNdist)==8 ? PyArray_DOUBLE : PyArray_FLOAT);
if( !pydists )
Py_DECREF(pykball);
BOOST_ASSERT(!!pydists);
PyObject *pyidx = PyArray_SimpleNew(2,dims, PyArray_INT);
if( !pyidx ) {
Py_DECREF(pykball);
Py_DECREF(pydists);
}
BOOST_ASSERT(!!pyidx);
ANNdist* pdists = (ANNdist*)PyArray_DATA(pydists);
ANNidx* pidx = (ANNidx*)PyArray_DATA(pyidx);
std::vector<ANNdist> dists(k);
std::vector<ANNidx> nn_idx(k);
for(int i = 0; i < N; ++i) {
object q = qarray[i];
for (int c = 0; c < kdtree.theDim(); ++c)
annq.pt[c] = extract<ANNcoord>(q[c]);
pkball[i] = kdtree.annkFRSearch(annq.pt, sqRad, k, &nn_idx[0], &dists[0], eps);
std::copy(nn_idx.begin(),nn_idx.end(),pidx);
pidx += k;
std::copy(dists.begin(),dists.end(),pdists);
pdists += k;
}
return boost::python::make_tuple(static_cast<numeric::array>(handle<>(pyidx)), static_cast<numeric::array>(handle<>(pydists)),static_cast<numeric::array>(handle<>(pykball)));
}
object search(ANNkd_tree& kdtree, object q, int k, double eps, bool priority = false)
{
BOOST_ASSERT(k <= kdtree.nPoints() && kdtree.theDim() == len(q));
ANNpointManaged annq(kdtree.theDim());
for (int c = 0; c < kdtree.theDim(); ++c)
annq.pt[c] = extract<ANNcoord>(q[c]);
npy_intp dims[] = { k};
PyObject *pydists = PyArray_SimpleNew(1,dims, sizeof(ANNdist)==8 ? PyArray_DOUBLE : PyArray_FLOAT);
BOOST_ASSERT(!!pydists);
PyObject *pyidx = PyArray_SimpleNew(1,dims, PyArray_INT);
if( !pyidx )
Py_DECREF(pydists);
BOOST_ASSERT(!!pyidx);
ANNdist* pdists = (ANNdist*)PyArray_DATA(pydists);
ANNidx* pidx = (ANNidx*)PyArray_DATA(pyidx);
std::vector<ANNidx> nn_idx(k);
std::vector<ANNdist> dists(k);
if (priority)
kdtree.annkPriSearch(annq.pt, k, pidx, pdists, eps);
else
kdtree.annkSearch(annq.pt, k, pidx, pdists, eps);
return boost::python::make_tuple(static_cast<numeric::array>(handle<>(pyidx)), static_cast<numeric::array>(handle<>(pydists)));
}
object k_fixed_radius_search(ANNkd_tree& kdtree, object q, double sqRad, int k, double eps)
{
BOOST_ASSERT(k <= kdtree.nPoints() && kdtree.theDim() == len(q));
ANNpointManaged annq(kdtree.theDim());
for (int c = 0; c < kdtree.theDim(); ++c)
annq.pt[c] = extract<ANNcoord>(q[c]);
if( k <= 0 ) {
int kball = kdtree.annkFRSearch(annq.pt, sqRad, k, NULL, NULL, eps);
return boost::python::make_tuple(numeric::array(boost::python::list()).astype("i4"),numeric::array(boost::python::list()),kball);
}
std::vector<ANNdist> dists(k);
std::vector<ANNidx> nn_idx(k);
int kball = kdtree.annkFRSearch(annq.pt, sqRad, k, &nn_idx[0], &dists[0], eps);
if( kball <= 0 )
return boost::python::make_tuple(numeric::array(boost::python::list()).astype("i4"),numeric::array(boost::python::list()),kball);
npy_intp dims[] = { min(k,kball)};
PyObject *pydists = PyArray_SimpleNew(1,dims, sizeof(ANNdist)==8 ? PyArray_DOUBLE : PyArray_FLOAT);
BOOST_ASSERT(!!pydists);
PyObject *pyidx = PyArray_SimpleNew(1,dims, PyArray_INT);
if( !pyidx )
Py_DECREF(pydists);
BOOST_ASSERT(!!pyidx);
ANNdist* pdists = (ANNdist*)PyArray_DATA(pydists);
ANNidx* pidx = (ANNidx*)PyArray_DATA(pyidx);
int addindex=0;
for (int i = 0; i < k; ++i) {
if (nn_idx[i] != ANN_NULL_IDX) {
pdists[addindex] = dists[i];
pidx[addindex] = nn_idx[i];
addindex++;
}
}
BOOST_ASSERT(kball > k || addindex==kball);
return boost::python::make_tuple(static_cast<numeric::array>(handle<>(pyidx)), static_cast<numeric::array>(handle<>(pydists)),kball);
}
