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;
}
Example #2
0
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, &notself, &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;
}