Hi,
I'm rather new to Python, and I've just written my first Python C
module. I was wondering if some more experience Pythonista would look
over what I've written and given me some pointers (or find some bugs).
I had a few questions while writing this as well.
Also, I know that there are much better ways to compute nearest
neighbors than the brute force n^2 method, but this is plenty fast for
this application (at least the C version is).
1. How do I add docstrings to my module?
2. The Python reference counting and memory managment seemes awfully
repetetive and error prone. Is there a nicer way of doing this? I
know there are some wrappers out there such as Pyrex and SWIG that
might prove useful.
3. This module consisted of turning 4 Python sequences into C double
arrays, performing some calculations, and then turning a C int array
into a Python tuple for return. And it was a lot of work. Are there
any nice wrapper functions out there for turning Python sequences into
C arrays and vice versa?
Thanks,
Jeremy Brewer
The code is below:
#include <Python.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <float.h>
/* xmatch: Takes 2 sets of positions (ra1, dec1 of size len1 and ra2,
dec2 of
of size len2) and computes the nearest neighbors. An array
of
integers for the nearest neighbors of ra1, dec1 that
correspond to
the positions in ra2, dec2 is returned, where points further
away
than mindist (in degrees) on the sky are marked as
-(closest index + 1). */
static int *xmatch(double *ra1, double *dec1, int len1, double *ra2,
double *dec2, int len2, double mindist)
{
int i, j, closest, *neighbors;
double mindist2, dist, old_dist;
/* allocate memory for the neighbors */
neighbors = (int *) malloc(len1 * sizeof(int));
if (neighbors == NULL) {
return NULL;
}
/* compare to mindist^2 to save computational time (saves many
calls to
sqrt in inner loop) */
mindist2 = mindist * mindist;
for (i = 0; i < len1; i++) {
/* intial closest match is as far off as possible */
closest = 0;
old_dist = DBL_MAX;
for (j = 0; j < len2; j++) {
/* angular distance on the sky squared */
dist = pow((cos(dec1[i]) * (ra1[i] - ra2[j])), 2) +
pow((dec1[i] - dec2[j]), 2);
/* keep track of only the closest point */
if (dist < old_dist) {
closest = j;
old_dist = dist;
}
}
/* mark points that lie outside the specified distance */
if (old_dist > mindist2) {
neighbors[i] = -(closest + 1);
}
else {
neighbors[i] = closest;
}
}
return neighbors;
}
/* wrap_xmatch: Python wrapper function for xmatch above */
static PyObject *wrap_xmatch(PyObject *self, PyObject *args)
{
PyObject *ra1, *dec1, *ra2, *dec2, *ntup;
int i, len1, len2, len_chk1, len_chk2;
int *neighbors;
double mindist;
double *r1, *d1, *r2, *d2;
/* read 4 sequences (ra1, dec1, ra2, dec2) and a double (mindist)
*/
if (!PyArg_ParseTuple(args, "OOOOd", &ra1, &dec1, &ra2, &dec2,
&mindist)) {
return NULL;
}
/* check that mindist is > 0 */
if (mindist <= 0.0) {
PyErr_SetString(PyExc_ValueError, "mindist must be > 0");
return NULL;
}
/* convert sequences to tuples if necessary */
ra1 = PySequence_Fast(ra1, "ra1 must be iterable");
if (ra1 == NULL) {
return NULL;
}
dec1 = PySequence_Fast(dec1, "dec1 must be iterable");
if (dec1 == NULL) {
return NULL;
}
ra2 = PySequence_Fast(ra2, "ra2 must be iterable");
if (ra2 == NULL) {
return NULL;
}
dec2 = PySequence_Fast(dec2, "dec2 must be iterable");
if (dec2 == NULL) {
return NULL;
}
/* length of input sequences */
len1 = PySequence_Fast_GET_SIZE(ra1);
len_chk1 = PySequence_Fast_GET_SIZE(dec1);
if (len1 != len_chk1) {
PyErr_SetString(PyExc_ValueError, "ra1 and dec1 must be the
same size");
return NULL;
}
len2 = PySequence_Fast_GET_SIZE(ra2);
len_chk2 = PySequence_Fast_GET_SIZE(dec2);
if (len2 != len_chk2) {
PyErr_SetString(PyExc_ValueError, "ra2 and dec2 must be the
same size");
return NULL;
}
/* allocate memory for C arrays */
r1 = (double *) malloc(len1 * sizeof(double));
if (r1 == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
return PyErr_NoMemory();
}
d1 = (double *) malloc(len1 * sizeof(double));
if (d1 == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
return PyErr_NoMemory();
}
r2 = (double *) malloc(len2 * sizeof(double));
if (r2 == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
return PyErr_NoMemory();
}
d2 = (double *) malloc(len2 * sizeof(double));
if (d2 == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
return PyErr_NoMemory();
}
/* copy ra1 and dec1 */
for (i = 0; i < len1; i++) {
PyObject *fitem, *item;
/* copy ra1 */
item = PySequence_Fast_GET_ITEM(ra1, i);
if (item == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
free(r1);
free(d1);
free(r2);
free(d2);
return NULL;
}
fitem = PyNumber_Float(item);
if (fitem == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
free(r1);
free(d1);
free(r2);
free(d2);
PyErr_SetString(PyExc_TypeError, "items in ra1 must be
numbers");
return NULL;
}
r1[i] = PyFloat_AS_DOUBLE(fitem);
Py_DECREF(fitem);
/* copy dec1 */
item = PySequence_Fast_GET_ITEM(dec1, i);
if (item == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
free(r1);
free(d1);
free(r2);
free(d2);
return NULL;
}
fitem = PyNumber_Float(item);
if (fitem == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
free(r1);
free(d1);
free(r2);
free(d2);
PyErr_SetString(PyExc_TypeError, "items in dec1 must be
numbers");
return NULL;
}
d1[i] = PyFloat_AS_DOUBLE(fitem);
Py_DECREF(fitem);
}
/* copy ra2 and dec2 */
for (i = 0; i < len2; i++) {
PyObject *fitem, *item;
/* copy ra2 */
item = PySequence_Fast_GET_ITEM(ra2, i);
if (item == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
free(r1);
free(d1);
free(r2);
free(d2);
return NULL;
}
fitem = PyNumber_Float(item);
if (fitem == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
free(r1);
free(d1);
free(r2);
free(d2);
PyErr_SetString(PyExc_TypeError, "items in ra2 must be
numbers");
return NULL;
}
r2[i] = PyFloat_AS_DOUBLE(fitem);
Py_DECREF(fitem);
/* copy dec1 */
item = PySequence_Fast_GET_ITEM(dec2, i);
if (item == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
free(r1);
free(d1);
free(r2);
free(d2);
return NULL;
}
fitem = PyNumber_Float(item);
if (fitem == NULL) {
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
free(r1);
free(d1);
free(r2);
free(d2);
PyErr_SetString(PyExc_TypeError, "items in dec2 must be
numbers");
return NULL;
}
d2[i] = PyFloat_AS_DOUBLE(fitem);
Py_DECREF(fitem);
}
/* clean up Python objects */
Py_DECREF(ra1);
Py_DECREF(dec1);
Py_DECREF(ra2);
Py_DECREF(dec2);
/* now compute the cross match */
neighbors = xmatch(r1, d1, len1, r2, d2, len2, mindist);
/* clean up C arrays */
free(r1);
free(d1);
free(r2);
free(d2);
/* build a Python tuple to hold the returned neighbors */
ntup = PyTuple_New(len1);
if (ntup == NULL) {
return PyErr_NoMemory();
}
for (i = 0; i < len1; i++) {
PyObject *integer = Py_BuildValue("i", neighbors[i]);
if (integer == NULL) {
Py_DECREF(ntup);
free(neighbors);
return PyErr_NoMemory();
}
PyTuple_SET_ITEM(ntup, i, integer);
}
/* free last C array */
free(neighbors);
return ntup;
}
/* functions defined in the module */
static PyMethodDef astro_methods[] = {
{"xmatch", wrap_xmatch, METH_VARARGS, "Cross match 2 sets of
positions."},
{0} /* sentinel */
};
/* initializes the module */
void initastro(void)
{
(void) Py_InitModule("astro", astro_methods);
}
