static PyObject* spherematch_kdtree_n(PyObject* self, PyObject* args) {
long i;
kdtree_t* kd;
if (!PyArg_ParseTuple(args, "l", &i)) {
PyErr_SetString(PyExc_ValueError, "need one arg: kdtree identifier (int)");
return NULL;
}
// Nasty!
kd = (kdtree_t*)i;
return PyInt_FromLong(kdtree_n(kd));
}
static PyObject* spherematch_match2(PyObject* self, PyObject* args) {
int i, N;
long p1, p2;
struct dualtree_results2 dtresults;
kdtree_t *kd1, *kd2;
double rad;
PyObject* indlist;
anbool notself;
anbool permute;
// So that ParseTuple("b") with a C "anbool" works
assert(sizeof(anbool) == sizeof(unsigned char));
if (!PyArg_ParseTuple(args, "lldbb", &p1, &p2, &rad, ¬self, &permute)) {
PyErr_SetString(PyExc_ValueError, "spherematch_c.match: need five args: two kdtree identifiers (ints), search radius (float), notself (boolean), permuted (boolean)");
return NULL;
}
// Nasty!
kd1 = (kdtree_t*)p1;
kd2 = (kdtree_t*)p2;
N = kdtree_n(kd1);
indlist = PyList_New(N);
assert(indlist);
dtresults.kd1 = kd1;
dtresults.kd2 = kd2;
dtresults.indlist = indlist;
dtresults.permute = permute;
dualtree_rangesearch(kd1, kd2, 0.0, rad, notself, NULL,
callback_dualtree2, &dtresults,
NULL, NULL);
// set empty slots to None, not NULL.
for (i=0; i<N; i++) {
if (PyList_GET_ITEM(indlist, i))
continue;
Py_INCREF(Py_None);
PyList_SET_ITEM(indlist, i, Py_None);
}
return indlist;
}
void dualtree_nearestneighbour(kdtree_t* xtree, kdtree_t* ytree, double maxdist2,
double** nearest_d2, int** nearest_ind,
int** count_in_range,
int notself) {
int i, NY, NNY;
// dual-tree search callback functions
dualtree_callbacks callbacks;
rs_params params;
// These two inputs must be non-NULL (they are essential return values);
// but they may point to pointers that are NULL (indicating that the caller wants us to
// allocate and return new arrays).
assert(nearest_d2);
assert(nearest_ind);
memset(&callbacks, 0, sizeof(dualtree_callbacks));
callbacks.decision = rs_within_range;
callbacks.decision_extra = ¶ms;
callbacks.result = rs_handle_result;
callbacks.result_extra = ¶ms;
// set search params
NY = kdtree_n(ytree);
memset(¶ms, 0, sizeof(params));
params.xtree = xtree;
params.ytree = ytree;
params.notself = notself;
params.d2 = maxdist2;
params.count_in_range = NULL;
if (count_in_range) {
if (!(*count_in_range)) {
*count_in_range = (int*)calloc(NY, sizeof(int));
}
params.count_in_range = *count_in_range;
}
// were we given a d2 array?
if (*nearest_d2)
params.nearest_d2 = *nearest_d2;
else
params.nearest_d2 = malloc(NY * sizeof(double));
if (maxdist2 == 0.0)
maxdist2 = HUGE_VAL;
for (i=0; i<NY; i++)
params.nearest_d2[i] = maxdist2;
// were we given an ind array?
if (*nearest_ind)
params.nearest_ind = *nearest_ind;
else
params.nearest_ind = malloc(NY * sizeof(int));
for (i=0; i<NY; i++)
params.nearest_ind[i] = -1;
NNY = kdtree_nnodes(ytree);
params.node_nearest_d2 = malloc(NNY * sizeof(double));
for (i=0; i<NNY; i++)
params.node_nearest_d2[i] = maxdist2;
dualtree_search(xtree, ytree, &callbacks);
// Return array addresses
*nearest_d2 = params.nearest_d2;
*nearest_ind = params.nearest_ind;
free(params.node_nearest_d2);
}
int henry_draper_n(const hd_catalog_t* hd) {
assert(hd);
assert(hd->kd);
return kdtree_n(hd->kd);
}
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;
}
static PyObject* spherematch_nn(PyObject* self, PyObject* args) {
int i, NY;
long p1, p2;
kdtree_t *kd1, *kd2;
npy_intp dims[1];
PyArrayObject* inds;
PyArrayObject* dist2s;
int *pinds;
double *pdist2s;
double rad;
anbool notself;
int* tempinds;
PyObject* rtn;
// So that ParseTuple("b") with a C "anbool" works
assert(sizeof(anbool) == sizeof(unsigned char));
if (!PyArg_ParseTuple(args, "lldb", &p1, &p2, &rad, ¬self)) {
PyErr_SetString(PyExc_ValueError, "need three args: two kdtree identifiers (ints), and search radius");
return NULL;
}
// Nasty!
kd1 = (kdtree_t*)p1;
kd2 = (kdtree_t*)p2;
NY = kdtree_n(kd2);
dims[0] = NY;
inds = (PyArrayObject*)PyArray_SimpleNew(1, dims, PyArray_INT);
dist2s = (PyArrayObject*)PyArray_SimpleNew(1, dims, PyArray_DOUBLE);
assert(PyArray_ITEMSIZE(inds) == sizeof(int));
assert(PyArray_ITEMSIZE(dist2s) == sizeof(double));
// YUCK!
tempinds = (int*)malloc(NY * sizeof(int));
double* tempdists = (double*)malloc(NY * sizeof(double));
pinds = tempinds; //PyArray_DATA(inds);
//pdist2s = PyArray_DATA(dist2s);
pdist2s = tempdists;
dualtree_nearestneighbour(kd1, kd2, rad*rad, &pdist2s, &pinds, NULL, notself);
// now we have to apply kd1's permutation array!
for (i=0; i<NY; i++)
if (pinds[i] != -1)
pinds[i] = kdtree_permute(kd1, pinds[i]);
pinds = PyArray_DATA(inds);
pdist2s = PyArray_DATA(dist2s);
for (i=0; i<NY; i++) {
pinds[i] = -1;
pdist2s[i] = HUGE_VAL;
}
// and apply kd2's permutation array!
for (i=0; i<NY; i++) {
if (tempinds[i] != -1) {
int j = kdtree_permute(kd2, i);
pinds[j] = tempinds[i];
pdist2s[j] = tempdists[i];
}
}
free(tempinds);
free(tempdists);
rtn = Py_BuildValue("(OO)", inds, dist2s);
Py_DECREF(inds);
Py_DECREF(dist2s);
return rtn;
}
