WIP: Reimplement core_cython routines as ufuncs

astropy_healpix/core.py

@@ -42,11 +42,13 @@ def _restore_shape(*args, **kwargs):
 
 def _validate_healpix_index(label, healpix_index, nside):
     npix = nside_to_npix(nside)
+    healpix_index = np.asarray(healpix_index)
     if np.any((healpix_index < 0) | (healpix_index > npix - 1)):
         raise ValueError('{0} must be in the range [0:{1}]'.format(label, npix))
 
 
 def _validate_offset(label, offset):
+    offset = np.asarray(offset)
     if np.any((offset < 0) | (offset > 1)):
         raise ValueError('d{0} must be in the range [0:1]'.format(label))
 
@@ -236,7 +238,7 @@ def npix_to_nside(npix):
     return np.round(square_root).astype(int)
 
 
-def healpix_to_lonlat(healpix_index, nside, dx=None, dy=None, order='ring'):
+def healpix_to_lonlat(healpix_index, nside, dx=0.5, dy=0.5, order='ring'):
     """
     Convert HEALPix indices (optionally with offsets) to longitudes/latitudes.
 
@@ -247,7 +249,7 @@ def healpix_to_lonlat(healpix_index, nside, dx=None, dy=None, order='ring'):
     ----------
     healpix_index : int or `~numpy.ndarray`
         HEALPix indices (as a scalar or array)
-    nside : int
+    nside : int or `~numpy.ndarray`
         Number of pixels along the side of each of the 12 top-level HEALPix tiles
     dx, dy : float or `~numpy.ndarray`, optional
         Offsets inside the HEALPix pixel, which must be in the range [0:1],
@@ -263,42 +265,21 @@ def healpix_to_lonlat(healpix_index, nside, dx=None, dy=None, order='ring'):
         The latitude values
     """
 
-    if (dx is None and dy is not None) or (dx is not None and dy is None):
-        raise ValueError('Either both or neither dx and dy must be specified')
-
-    healpix_index = np.asarray(healpix_index, dtype=np.int64)
-
-    if dx is None and dy is not None:
-        dx = 0.5
-    elif dx is not None and dy is None:
-        dy = 0.5
-
-    if dx is not None:
-        dx = np.asarray(dx, dtype=np.float)
-        dy = np.asarray(dy, dtype=np.float)
-        _validate_offset('x', dx)
-        _validate_offset('y', dy)
-        healpix_index, dx, dy = np.broadcast_arrays(healpix_index, dx, dy)
-        dx = dx.ravel()
-        dy = dy.ravel()
-
-    shape = healpix_index.shape
-    healpix_index = healpix_index.ravel()
-    nside = int(nside)
-
     _validate_healpix_index('healpix_index', healpix_index, nside)
     _validate_nside(nside)
+    _validate_offset('x', dx)
+    _validate_offset('y', dy)
     order = _validate_order(order)
 
-    if dx is None:
-        lon, lat = core_cython.healpix_to_lonlat(healpix_index, nside, order)
-    else:
-        lon, lat = core_cython.healpix_with_offset_to_lonlat(healpix_index, dx, dy, nside, order)
+    func = (core_cython.healpix_to_lonlat_ring if order == 'ring' else
+            core_cython.healpix_to_lonlat_nested)
+
+    lon, lat = func(healpix_index, nside, dx, dy)
 
     lon = Longitude(lon, unit=u.rad, copy=False)
     lat = Latitude(lat, unit=u.rad, copy=False)
 
-    return _restore_shape(lon, lat, shape=shape)
+    return lon, lat
 
 
 def lonlat_to_healpix(lon, lat, nside, return_offsets=False, order='ring'):

astropy_healpix/core_cython.pyx

@@ -19,6 +19,9 @@ npy_double = np.double
 npy_intp = np.intp
 npy_int64 = np.int64
 
+np.import_array()
+np.import_ufunc()
+
 
 def _validate_order(str order):
     # We also support upper-case, to support directly the values
@@ -33,103 +36,59 @@ def _validate_order(str order):
 
 
 @cython.boundscheck(False)
-def healpix_to_lonlat(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
-                      int nside, str order):
-    """
-    Convert HEALPix indices to longitudes/latitudes.
-
-    This returns the longitudes/latitudes of the center of the HEALPix pixels.
-    If you also want to provide relative offsets inside the pixels, see
-    :func:`healpix_with_offset_to_lonlat`.
-
-    Parameters
-    ----------
-    healpix_index : `~numpy.ndarray`
-        1-D array of HEALPix indices
-    nside : int
-        Number of pixels along the side of each of the 12 top-level HEALPix tiles
-    order : { 'nested' | 'ring' }
-        Order of HEALPix pixels
-
-    Returns
-    -------
-    lon, lat : `~numpy.ndarray`
-        1-D arrays of longitude and latitude in radians
-    """
-
-    cdef intp_t n = healpix_index.shape[0]
+cdef void healpix_to_lonlat_ring_loop(char** args, intp_t* dimensions,
+                                      intp_t* steps, void* data) nogil:
+    cdef intp_t n = dimensions[0]
     cdef intp_t i
-    cdef int64_t xy_index
-    cdef double dx, dy;
-    cdef np.ndarray[double_t, ndim=1, mode="c"] lon = np.zeros(n, dtype=npy_double)
-    cdef np.ndarray[double_t, ndim=1, mode="c"] lat = np.zeros(n, dtype=npy_double)
+    cdef int64_t xy_index, healpix_index, nside
+    cdef double dx, dy
+    cdef double* lon
+    cdef double* lat
 
-    dx = 0.5
-    dy = 0.5
+    for i in prange(n, schedule='static'):
+        healpix_index = (<int64_t*> &args[0][i * steps[0]])[0]
+        nside = (<int64_t*> &args[1][i * steps[1]])[0]
+        dx = (<double*> &args[2][i * steps[2]])[0]
+        dy = (<double*> &args[3][i * steps[3]])[0]
+        lon = <double*> &args[4][i * steps[5]]
+        lat = <double*> &args[5][i * steps[5]]
+        xy_index = healpixl_ring_to_xy(healpix_index, nside)
+        healpixl_to_radec(xy_index, nside, dx, dy, lon, lat)
 
-    order = _validate_order(order)
 
-    if order == 'nested':
-        for i in prange(n, nogil=True, schedule='static'):
-            xy_index = healpixl_nested_to_xy(healpix_index[i], nside)
-            healpixl_to_radec(xy_index, nside, dx, dy, &lon[i], &lat[i])
-    elif order == 'ring':
-        for i in prange(n, nogil=True, schedule='static'):
-            xy_index = healpixl_ring_to_xy(healpix_index[i], nside)
-            healpixl_to_radec(xy_index, nside, dx, dy, &lon[i], &lat[i])
-
-    return lon, lat
+healpix_to_lonlat_ring = np.PyUFunc_FromFuncAndData(
+    [healpix_to_lonlat_ring_loop], [NULL],
+    [np.NPY_INT64, np.NPY_INT64,
+     np.NPY_DOUBLE, np.NPY_DOUBLE, np.NPY_DOUBLE, np.NPY_DOUBLE],
+    1, 4, 2, 0, "healpix_to_lonlat_ring", "", 0)
 
 
 @cython.boundscheck(False)
-def healpix_with_offset_to_lonlat(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
-                                  np.ndarray[double_t, ndim=1, mode="c"] dx,
-                                  np.ndarray[double_t, ndim=1, mode="c"] dy,
-                                  int nside, str order):
-    """
-    Convert HEALPix indices to longitudes/latitudes
-
-    This function takes relative offsets in x and y inside the HEALPix pixels.
-    If you are only interested in the centers of the pixels, see
-    `healpix_to_lonlat`.
-
-    Parameters
-    ----------
-    healpix_index : `~numpy.ndarray`
-        1-D array of HEALPix indices
-    dx, dy : `~numpy.ndarray`
-        1-D arrays of offsets inside the HEALPix pixel, which must be in the
-        range [0:1] (0.5 is the center of the HEALPix pixels)
-    nside : int
-        Number of pixels along the side of each of the 12 top-level HEALPix tiles
-    order : { 'nested' | 'ring' }
-        Order of HEALPix pixels
-
-
-    Returns
-    -------
-    lon, lat : `~numpy.ndarray`
-        1-D arrays of longitude and latitude in radians
-    """
-
-    cdef intp_t n = healpix_index.shape[0]
+cdef void healpix_to_lonlat_nested_loop(
+        char** args, intp_t* dimensions, intp_t* steps, void* data) nogil:
+    cdef intp_t n = dimensions[0]
     cdef intp_t i
-    cdef int64_t xy_index
-    cdef np.ndarray[double_t, ndim=1, mode="c"] lon = np.zeros(n, dtype=npy_double)
-    cdef np.ndarray[double_t, ndim=1, mode="c"] lat = np.zeros(n, dtype=npy_double)
-
-    order = _validate_order(order)
-
-    if order == 'nested':
-        for i in prange(n, nogil=True, schedule='static'):
-            xy_index = healpixl_nested_to_xy(healpix_index[i], nside)
-            healpixl_to_radec(xy_index, nside, dx[i], dy[i], &lon[i], &lat[i])
-    elif order == 'ring':
-        for i in prange(n, nogil=True, schedule='static'):
-            xy_index = healpixl_ring_to_xy(healpix_index[i], nside)
-            healpixl_to_radec(xy_index, nside, dx[i], dy[i], &lon[i], &lat[i])
-
-    return lon, lat
+    cdef int64_t xy_index, healpix_index, nside
+    cdef double dx, dy
+    cdef double* lon
+    cdef double* lat
+
+    for i in prange(n, schedule='static'):
+        healpix_index = (<int64_t*> &args[0][i * steps[0]])[0]
+        nside = (<int64_t*> &args[1][i * steps[1]])[0]
+        dx = (<double*> &args[2][i * steps[2]])[0]
+        dy = (<double*> &args[3][i * steps[3]])[0]
+        lon = <double*> &args[4][i * steps[5]]
+        lat = <double*> &args[5][i * steps[5]]
+        xy_index = healpixl_nested_to_xy(healpix_index, nside)
+        healpixl_to_radec(xy_index, nside, dx, dy, lon, lat)
+
+
+healpix_to_lonlat_nested = np.PyUFunc_FromFuncAndData(
+    [healpix_to_lonlat_nested_loop], [NULL],
+    [np.NPY_INT64, np.NPY_INT64, np.NPY_DOUBLE, np.NPY_DOUBLE,
+     np.NPY_DOUBLE, np.NPY_DOUBLE],
+    1, 4, 2, 0, "healpix_to_lonlat_nested", "", 0)
 
 
 @cython.boundscheck(False)

astropy_healpix/high_level.py

@@ -73,7 +73,7 @@ def npix(self):
         """
         return nside_to_npix(self.nside)
 
-    def healpix_to_lonlat(self, healpix_index, dx=None, dy=None):
+    def healpix_to_lonlat(self, healpix_index, dx=0.5, dy=0.5):
         """
         Convert HEALPix indices (optionally with offsets) to longitudes/latitudes
 
@@ -269,7 +269,7 @@ def neighbours(self, healpix_index):
         """
         return neighbours(healpix_index, self.nside, order=self.order)
 
-    def healpix_to_skycoord(self, healpix_index, dx=None, dy=None):
+    def healpix_to_skycoord(self, healpix_index, dx=0.5, dy=0.5):
         """
         Convert HEALPix indices (optionally with offsets) to celestial coordinates.
 

astropy_healpix/tests/test_core_cython.py

@@ -37,7 +37,9 @@ def get_test_indices(nside):
 def test_roundtrip_healpix_no_offsets(order, nside_power, capfd):
     nside = 2 ** nside_power
     index = get_test_indices(nside)
-    lon, lat = core_cython.healpix_to_lonlat(index, nside, order)
+    func = (core_cython.healpix_to_lonlat_ring if order == 'ring' else
+            core_cython.healpix_to_lonlat_nested)
+    lon, lat = func(index, nside, 0.5, 0.5)
     index_new = core_cython.lonlat_to_healpix(lon, lat, nside, order)
     assert_equal(index, index_new)
     out, err = capfd.readouterr()
@@ -50,7 +52,9 @@ def test_roundtrip_healpix_with_offsets(order, nside_power, capfd):
     index = get_test_indices(nside)
     dx = np.random.random(index.shape)
     dy = np.random.random(index.shape)
-    lon, lat = core_cython.healpix_with_offset_to_lonlat(index, dx, dy, nside, order)
+    func = (core_cython.healpix_to_lonlat_ring if order == 'ring' else
+            core_cython.healpix_to_lonlat_nested)
+    lon, lat = func(index, nside, dx, dy)
     index_new, dx_new, dy_new = core_cython.lonlat_to_healpix_with_offset(lon, lat, nside, order)
     assert_equal(index, index_new)
     assert_allclose(dx, dx_new, atol=1e-8)
@@ -97,7 +101,9 @@ def test_healpix_cone_search(order, nside_power, capfd):
     dx = np.repeat(dx, n_inside, axis=0).ravel()
     dy = np.repeat(dy, n_inside, axis=0).ravel()
     index_inside = np.repeat(index_inside, 4).ravel()
-    lon, lat = core_cython.healpix_with_offset_to_lonlat(index_inside, dx, dy, nside, order)
+    func = (core_cython.healpix_to_lonlat_ring if order == 'ring' else
+            core_cython.healpix_to_lonlat_nested)
+    lon, lat = func(index_inside, nside, dx, dy)
     lon, lat = lon.reshape((n_inside, 4)), lat.reshape((n_inside, 4))
     sep = angular_separation(lon0 * u.deg, lat0 * u.deg, lon * u.rad, lat * u.rad)
     sep = sep.min(axis=1)
@@ -114,16 +120,16 @@ def test_regression_healpix_to_lonlat_sqrt():
     # and return a negative longitude.
 
     index = np.array([9007199120523263], dtype=np.int64)
-    lon, lat = core_cython.healpix_to_lonlat(index, 2**26, order='ring')
+    lon, lat = core_cython.healpix_to_lonlat_ring(index, 2**26, 0.5, 0.5)
     assert_allclose(lon, 6.283185295476241, rtol=1e-14)
     assert_allclose(lat, 0.729727669554970, rtol=1e-14)
 
     index = np.array([720575940916150240], dtype=np.int64)
-    lon, lat = core_cython.healpix_to_lonlat(index, 2**28, order='ring')
+    lon, lat = core_cython.healpix_to_lonlat_ring(index, 2**28, 0.5, 0.5)
     assert_allclose(lon, 6.283185122851909, rtol=1e-14)
     assert_allclose(lat, -0.729727656226966, rtol=1e-14)
 
     index = np.array([180143985363255292], dtype=np.int64)
-    lon, lat = core_cython.healpix_to_lonlat(index, 2**27, order='ring')
+    lon, lat = core_cython.healpix_to_lonlat_ring(index, 2**27, 0.5, 0.5)
     assert_allclose(lon, 6.283185266217880, rtol=1e-14)
     assert_allclose(lat, -0.729727656226966, rtol=1e-14)