static anbool rs_within_range(void* vparams, kdtree_t* xtree, int xnode, kdtree_t* ytree, int ynode) { rs_params* p = (rs_params*)vparams; double maxd2; // count-in-range is actually more like rangesearch... if (p->count_in_range) { if (kdtree_node_node_mindist2_exceeds(xtree, xnode, ytree, ynode, p->d2)) return FALSE; return TRUE; } if (kdtree_node_node_mindist2_exceeds(xtree, xnode, ytree, ynode, p->node_nearest_d2[ynode])) return FALSE; maxd2 = kdtree_node_node_maxdist2(xtree, xnode, ytree, ynode); if (maxd2 < p->node_nearest_d2[ynode]) { // update this node and its children. p->node_nearest_d2[ynode] = maxd2; if (!KD_IS_LEAF(ytree, ynode)) { int child = KD_CHILD_LEFT(ynode); p->node_nearest_d2[child] = MIN(p->node_nearest_d2[child], maxd2); child = KD_CHILD_RIGHT(ynode); p->node_nearest_d2[child] = MIN(p->node_nearest_d2[child], maxd2); } } return TRUE; }
static PyObject* spherematch_nn2(PyObject* self, PyObject* args) { int i, j, NY, N; long p1, p2; kdtree_t *kd1, *kd2; npy_intp dims[1]; PyObject* I; PyObject* J; PyObject* dist2s; PyObject* counts = NULL; int *pi; int *pj; int *pc = NULL; double *pd; double rad; anbool notself; anbool docount; int* tempinds; int* tempcount = NULL; int** ptempcount = NULL; double* tempd2; PyObject* rtn; // So that ParseTuple("b") with a C "anbool" works assert(sizeof(anbool) == sizeof(unsigned char)); if (!PyArg_ParseTuple(args, "lldbb", &p1, &p2, &rad, ¬self, &docount)) { PyErr_SetString(PyExc_ValueError, "need five args: two kdtree identifiers (ints), search radius, notself (bool) and docount (bool)"); return NULL; } // Nasty! kd1 = (kdtree_t*)p1; kd2 = (kdtree_t*)p2; // quick check for no-overlap case if (kdtree_node_node_mindist2_exceeds(kd1, 0, kd2, 0, rad*rad)) { // allocate empty return arrays dims[0] = 0; I = PyArray_SimpleNew(1, dims, NPY_INT); J = PyArray_SimpleNew(1, dims, NPY_INT); dist2s = PyArray_SimpleNew(1, dims, NPY_DOUBLE); if (docount) { counts = PyArray_SimpleNew(1, dims, NPY_INT); rtn = Py_BuildValue("(OOOO)", I, J, dist2s, counts); Py_DECREF(counts); } else { rtn = Py_BuildValue("(OOO)", I, J, dist2s); } Py_DECREF(I); Py_DECREF(J); Py_DECREF(dist2s); return rtn; } NY = kdtree_n(kd2); tempinds = (int*)malloc(NY * sizeof(int)); tempd2 = (double*)malloc(NY * sizeof(double)); if (docount) { tempcount = (int*)calloc(NY, sizeof(int)); ptempcount = &tempcount; } dualtree_nearestneighbour(kd1, kd2, rad*rad, &tempd2, &tempinds, ptempcount, notself); // count number of matches N = 0; for (i=0; i<NY; i++) if (tempinds[i] != -1) N++; // allocate return arrays dims[0] = N; I = PyArray_SimpleNew(1, dims, NPY_INT); J = PyArray_SimpleNew(1, dims, NPY_INT); dist2s = PyArray_SimpleNew(1, dims, NPY_DOUBLE); pi = PyArray_DATA(I); pj = PyArray_DATA(J); pd = PyArray_DATA(dist2s); if (docount) { counts = PyArray_SimpleNew(1, dims, NPY_INT); pc = PyArray_DATA(counts); } j = 0; for (i=0; i<NY; i++) { if (tempinds[i] == -1) continue; pi[j] = kdtree_permute(kd1, tempinds[i]); pj[j] = kdtree_permute(kd2, i); pd[j] = tempd2[i]; if (docount) pc[j] = tempcount[i]; j++; } free(tempinds); free(tempd2); free(tempcount); if (docount) { rtn = Py_BuildValue("(OOOO)", I, J, dist2s, counts); Py_DECREF(counts); } else { rtn = Py_BuildValue("(OOO)", I, J, dist2s); } Py_DECREF(I); Py_DECREF(J); Py_DECREF(dist2s); return rtn; }