Rewrite core in C using Numpy ufuncs

astropy_healpix/_core.c

@@ -0,0 +1,371 @@
+#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"
+
+
+/* Data structure for storing function pointers for routines that are specific
+ * to the HEALPix ordering scheme. When we create the ufuncs using
+ * PyUFunc_FromFuncAndData, we will set them up to pass a pointer to this
+ * data structure through as the void *data parameter for the loop functions
+ * defined below. */
+typedef struct {
+    int64_t (*order_to_xy)(int64_t, int);
+    int64_t (*xy_to_order)(int64_t, int);
+} order_funcs;
+
+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 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;
+    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;
+    npy_intp dims[1];
+
+    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");
+        return NULL;
+    }
+
+    dims[0] = n_indices;
+    result = PyArray_SimpleNew(1, dims, NPY_INT64);
+    if (result)
+    {
+        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);
+        }
+    }
+
+    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 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)