diff --git a/astropy_healpix/core.py b/astropy_healpix/core.py index d374ff8..dc80435 100644 --- a/astropy_healpix/core.py +++ b/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'): diff --git a/astropy_healpix/core_cython.pyx b/astropy_healpix/core_cython.pyx index 0a807ea..76dd14b 100644 --- a/astropy_healpix/core_cython.pyx +++ b/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 = ( &args[0][i * steps[0]])[0] + nside = ( &args[1][i * steps[1]])[0] + dx = ( &args[2][i * steps[2]])[0] + dy = ( &args[3][i * steps[3]])[0] + lon = &args[4][i * steps[5]] + lat = &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 = ( &args[0][i * steps[0]])[0] + nside = ( &args[1][i * steps[1]])[0] + dx = ( &args[2][i * steps[2]])[0] + dy = ( &args[3][i * steps[3]])[0] + lon = &args[4][i * steps[5]] + lat = &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) diff --git a/astropy_healpix/high_level.py b/astropy_healpix/high_level.py index 31bfc6b..8ff6a63 100644 --- a/astropy_healpix/high_level.py +++ b/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 diff --git a/astropy_healpix/tests/test_core_cython.py b/astropy_healpix/tests/test_core_cython.py index 9977794..ec9263c 100644 --- a/astropy_healpix/tests/test_core_cython.py +++ b/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)