Skip to content

Commit

Permalink
WIP: 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 astropy#72, closes astropy#75.
  • Loading branch information
@lpsinger
lpsinger committed Oct 30, 2018
1 parent 93c9b25 commit dada4d0
Show file tree
Hide file tree
Showing 7 changed files with 362 additions and 763 deletions.
301 changes: 301 additions & 0 deletions astropy_healpix/_core.c
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,301 @@
#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"


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


static void healpix_to_lonlat_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
order_funcs *funcs = data;

for (npy_intp i = 0; i < dimensions[0]; 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 = funcs->order_to_xy(pixel, nside);
if (xy >= 0)
healpixl_to_radec(xy, nside, dx, dy, lon, lat);
else
*lon = *lat = NPY_NAN;
}
}


static void lonlat_to_healpix_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
order_funcs *funcs = data;

for (npy_intp i = 0; i < dimensions[0]; 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 = 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;
}
}
}


static void nested_to_ring_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
for (npy_intp i = 0; i < dimensions[0]; 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 = healpixl_nested_to_xy(nested, nside);
if (xy >= 0)
*ring = healpixl_xy_to_ring(xy, nside);
else
*ring = -1;
}
}


static void ring_to_nested_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
for (npy_intp i = 0; i < dimensions[0]; 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 = healpixl_ring_to_xy(ring, nside);
if (xy >= 0)
*nested = healpixl_xy_to_nested(xy, nside);
else
*nested = -1;
}
}


static void bilinear_interpolation_weights_loop(
char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
for (npy_intp i = 0; i < dimensions[0]; 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];

interpolate_weights(lon, lat, indices, weights, nside);
for (int 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;

for (npy_intp i = 0; i < dimensions[0]; 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};

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

for (int j = 0; j < 8; j ++)
{
xy = neighbours[(4 - j) % 8];
if (xy >= 0)
pixel = funcs->xy_to_order(xy, nside);
else
pixel = -1;
*(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;

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

int64_t *indices;
int64_t 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;

int64_t *result_data = PyArray_DATA((PyArrayObject *) result);

if (strcmp(order, "nested") == 0)
{
for (int i = 0; i < n_indices; i ++)
result_data[i] = healpixl_xy_to_nested(indices[i], nside);
} else {
for (int 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_DOUBLE, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE,
NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE};


PyMODINIT_FUNC PyInit__core(void)
{
import_array();
import_umath();

PyObject *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;
}
Loading

0 comments on commit dada4d0

Please sign in to comment.