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); }