Skip to content

Commit

Permalink
Rewrite core in C using Numpy ufuncs
Browse files Browse the repository at this point in the history 
As a result, broadcasting across nside and ipix is supported
automatically.

Fixes #72, closes #75.
  • Loading branch information
@lpsinger
lpsinger committed Oct 31, 2018
1 parent 93c9b25 commit 87bf0c1
Show file tree
Hide file tree
Showing 11 changed files with 640 additions and 773 deletions.
368 changes: 368 additions & 0 deletions astropy_healpix/_core.c
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,368 @@
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

#include <Python.h>
#include <numpy/arrayobject.h>
#include <numpy/ufuncobject.h>
#include "healpix.h"
#include "healpix-utils.h"
#include "interpolation.h"

/* FIXME: Remove this once we drop Python 2 support. */
#include "six.h"

/* FIXME: We need npy_set_floatstatus_invalid(), but unlike most of the Numpy
* C API it is only available on some platforms if you explicitly link against
* Numpy, which is not typically done for building C extensions. This bundled
* header file contains a static definition of _npy_set_floatstatus_invalid().
*/
#include "ieee754.h"


typedef struct {
int64_t (*order_to_xy)(int64_t, int);
int64_t (*xy_to_order)(int64_t, int);
} order_funcs;


static int64_t nside2npix(int nside)
{
return 12 * ((int64_t) nside) * ((int64_t) nside);
}


static int pixel_nside_valid(int64_t pixel, int nside)
{
return pixel >= 0 && pixel < nside2npix(nside);
}


static void healpix_to_lonlat_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
order_funcs *funcs = data;
npy_intp i, n = dimensions[0];

for (i = 0; i < n; i ++)
{
int64_t pixel = *(int64_t *) &args[0][i * steps[0]];
int nside = *(int *) &args[1][i * steps[1]];
double dx = *(double *) &args[2][i * steps[2]];
double dy = *(double *) &args[3][i * steps[3]];
double *lon = (double *) &args[4][i * steps[4]];
double *lat = (double *) &args[5][i * steps[5]];
int64_t xy = -1;

if (pixel_nside_valid(pixel, nside))
xy = funcs->order_to_xy(pixel, nside);

if (xy >= 0)
healpixl_to_radec(xy, nside, dx, dy, lon, lat);
else
{
*lon = *lat = NPY_NAN;
_npy_set_floatstatus_invalid();
}
}
}


static void lonlat_to_healpix_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
order_funcs *funcs = data;
npy_intp i, n = dimensions[0];

for (i = 0; i < n; i ++)
{
double lon = *(double *) &args[0][i * steps[0]];
double lat = *(double *) &args[1][i * steps[1]];
int nside = *(int *) &args[2][i * steps[2]];
int64_t *pixel = (int64_t *) &args[3][i * steps[3]];
double *dx = (double *) &args[4][i * steps[4]];
double *dy = (double *) &args[5][i * steps[5]];
int64_t xy = -1;

xy = radec_to_healpixlf(lon, lat, nside, dx, dy);
if (xy >= 0)
*pixel = funcs->xy_to_order(xy, nside);
else {
*pixel = -1;
*dx = *dy = NPY_NAN;
_npy_set_floatstatus_invalid();
}
}
}


static void nested_to_ring_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
npy_intp i, n = dimensions[0];

for (i = 0; i < n; i ++)
{
int64_t nested = *(int64_t *) &args[0][i * steps[0]];
int nside = *(int *) &args[1][i * steps[1]];
int64_t *ring = (int64_t *) &args[2][i * steps[2]];
int64_t xy = -1;

if (pixel_nside_valid(nested, nside))
xy = healpixl_nested_to_xy(nested, nside);
if (xy >= 0)
*ring = healpixl_xy_to_ring(xy, nside);
else {
*ring = -1;
_npy_set_floatstatus_invalid();
}
}
}


static void ring_to_nested_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
npy_intp i, n = dimensions[0];

for (i = 0; i < n; i ++)
{
int64_t ring = *(int64_t *) &args[0][i * steps[0]];
int nside = *(int *) &args[1][i * steps[1]];
int64_t *nested = (int64_t *) &args[2][i * steps[2]];
int64_t xy = -1;

if (pixel_nside_valid(ring, nside))
xy = healpixl_ring_to_xy(ring, nside);
if (xy >= 0)
*nested = healpixl_xy_to_nested(xy, nside);
else {
*nested = -1;
_npy_set_floatstatus_invalid();
}
}
}


static void bilinear_interpolation_weights_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
npy_intp i, n = dimensions[0];

for (i = 0; i < n; i ++)
{
double lon = *(double *) &args[0][i * steps[0]];
double lat = *(double *) &args[1][i * steps[1]];
int nside = *(int *) &args[2][i * steps[2]];
int64_t indices[4];
double weights[4];
int j;

interpolate_weights(lon, lat, indices, weights, nside);
for (j = 0; j < 4; j ++)
{
*(int64_t *) &args[3 + j][i * steps[3 + j]] = indices[j];
*(double *) &args[7 + j][i * steps[7 + j]] = weights[j];
}
}
}


static void neighbours_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
order_funcs *funcs = data;
npy_intp i, n = dimensions[0];

for (i = 0; i < n; i ++)
{
int64_t pixel = *(int64_t *) &args[0][i * steps[0]];
int nside = *(int *) &args[1][i * steps[1]];
int64_t neighbours[] = {-1, -1, -1, -1, -1, -1, -1, -1};
int j;
int64_t xy = -1;

if (pixel_nside_valid(pixel, nside))
xy = funcs->order_to_xy(pixel, nside);
if (xy >= 0)
healpixl_get_neighbours(xy, neighbours, nside);

for (j = 0; j < 8; j ++)
{
int k = 4 - j;
if (k < 0)
k += 8;
xy = neighbours[k];
if (xy >= 0)
pixel = funcs->xy_to_order(xy, nside);
else {
pixel = -1;
_npy_set_floatstatus_invalid();
}
*(int64_t *) &args[2 + j][i * steps[2 + j]] = pixel;
}
}
}


static PyObject *healpix_cone_search(
PyObject *self, PyObject *args, PyObject *kwargs)
{
PyObject *result = NULL;
static const char *kws[] = {"lon", "lat", "radius", "nside", "order", NULL};
double lon, lat, radius;
int nside;
char *order;
int64_t *indices, n_indices;
int64_t *result_data;

if (!PyArg_ParseTupleAndKeywords(
args, kwargs, "dddis", kws, &lon, &lat, &radius, &nside, &order))
return NULL;

n_indices = healpix_rangesearch_radec_simple(
lon, lat, radius, nside, 0, &indices);
if (!indices)
{
PyErr_SetString(
PyExc_RuntimeError, "healpix_rangesearch_radec_simple failed");
goto done;
}

{
npy_intp dims[] = {n_indices};
result = PyArray_SimpleNew(1, dims, NPY_INT64);
}
if (!result)
goto done;

result_data = PyArray_DATA((PyArrayObject *) result);

if (strcmp(order, "nested") == 0)
{
int i;
for (i = 0; i < n_indices; i ++)
result_data[i] = healpixl_xy_to_nested(indices[i], nside);
} else {
int i;
for (i = 0; i < n_indices; i ++)
result_data[i] = healpixl_xy_to_ring(indices[i], nside);
}

done:
free(indices);
return result;
}


static PyMethodDef methods[] = {
{"healpix_cone_search", healpix_cone_search, METH_VARARGS | METH_KEYWORDS, NULL},
{NULL, NULL, 0, NULL}
};


static PyModuleDef moduledef = {
PyModuleDef_HEAD_INIT,
"_core", NULL, -1, methods
};

static const PyUFuncGenericFunction
healpix_to_lonlat_loops [] = {healpix_to_lonlat_loop},
lonlat_to_healpix_loops [] = {lonlat_to_healpix_loop},
nested_to_ring_loops [] = {nested_to_ring_loop},
ring_to_nested_loops [] = {ring_to_nested_loop},
bilinear_interpolation_weights_loops[] = {bilinear_interpolation_weights_loop},
neighbours_loops [] = {neighbours_loop};

static const order_funcs
nested_order_funcs = {healpixl_nested_to_xy, healpixl_xy_to_nested},
ring_order_funcs = {healpixl_ring_to_xy, healpixl_xy_to_ring};

static void *no_ufunc_data[] = {NULL},
*nested_ufunc_data[] = {&nested_order_funcs},
*ring_ufunc_data[] = {&ring_order_funcs};

static const char
healpix_to_lonlat_types[] = {
NPY_INT64, NPY_INT, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE},
lonlat_to_healpix_types[] = {
NPY_DOUBLE, NPY_DOUBLE, NPY_INT, NPY_INT64, NPY_DOUBLE, NPY_DOUBLE},
healpix_to_healpix_types[] = {
NPY_INT64, NPY_INT, NPY_INT64},
bilinear_interpolation_weights_types[] = {
NPY_DOUBLE, NPY_DOUBLE, NPY_INT,
NPY_INT64, NPY_INT64, NPY_INT64, NPY_INT64,
NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE},
neighbours_types[] = {
NPY_INT64, NPY_INT,
NPY_INT64, NPY_INT64, NPY_INT64, NPY_INT64,
NPY_INT64, NPY_INT64, NPY_INT64, NPY_INT64};


PyMODINIT_FUNC PyInit__core(void)
{
PyObject *module;

import_array();
import_umath();

module = PyModule_Create(&moduledef);

PyModule_AddObject(
module, "healpix_nested_to_lonlat", PyUFunc_FromFuncAndData(
healpix_to_lonlat_loops, nested_ufunc_data,
healpix_to_lonlat_types, 1, 4, 2, PyUFunc_None,
"healpix_nested_to_lonlat", NULL, 0));

PyModule_AddObject(
module, "healpix_ring_to_lonlat", PyUFunc_FromFuncAndData(
healpix_to_lonlat_loops, ring_ufunc_data,
healpix_to_lonlat_types, 1, 4, 2, PyUFunc_None,
"healpix_ring_to_lonlat", NULL, 0));

PyModule_AddObject(
module, "lonlat_to_healpix_nested", PyUFunc_FromFuncAndData(
lonlat_to_healpix_loops, nested_ufunc_data,
lonlat_to_healpix_types, 1, 3, 3, PyUFunc_None,
"lonlat_to_healpix_nested", NULL, 0));

PyModule_AddObject(
module, "lonlat_to_healpix_ring", PyUFunc_FromFuncAndData(
lonlat_to_healpix_loops, ring_ufunc_data,
lonlat_to_healpix_types, 1, 3, 3, PyUFunc_None,
"lonlat_to_healpix_ring", NULL, 0));

PyModule_AddObject(
module, "nested_to_ring", PyUFunc_FromFuncAndData(
nested_to_ring_loops, no_ufunc_data,
healpix_to_healpix_types, 1, 2, 1, PyUFunc_None,
"nested_to_ring", NULL, 0));

PyModule_AddObject(
module, "ring_to_nested", PyUFunc_FromFuncAndData(
ring_to_nested_loops, no_ufunc_data,
healpix_to_healpix_types, 1, 2, 1, PyUFunc_None,
"ring_to_nested", NULL, 0));

PyModule_AddObject(
module, "bilinear_interpolation_weights", PyUFunc_FromFuncAndData(
bilinear_interpolation_weights_loops, no_ufunc_data,
bilinear_interpolation_weights_types, 1, 3, 8, PyUFunc_None,
"bilinear_interpolation_weights", NULL, 0));

PyModule_AddObject(
module, "neighbours_nested", PyUFunc_FromFuncAndData(
neighbours_loops, nested_ufunc_data,
neighbours_types, 1, 2, 8, PyUFunc_None,
"neighbours_nested", NULL, 0));

PyModule_AddObject(
module, "neighbours_ring", PyUFunc_FromFuncAndData(
neighbours_loops, ring_ufunc_data,
neighbours_types, 1, 2, 8, PyUFunc_None,
"neighbours_ring", NULL, 0));

return module;
}


/* FIXME: Remove this once we drop Python 2 support. */
SIX_COMPAT_MODULE(_core)
Loading

0 comments on commit 87bf0c1

Please sign in to comment.