aboutsummaryrefslogtreecommitdiffstats
path: root/fg21sim
diff options
context:
space:
mode:
authorAaron LI <aaronly.me@outlook.com>2016-10-25 18:13:08 +0800
committerAaron LI <aaronly.me@outlook.com>2016-10-25 18:13:08 +0800
commitf0682d6f5c196001f8820d0f2dbdf18087285ac7 (patch)
tree055699fff2a8977bfae65a9a6977555850f22683 /fg21sim
parenta78f2e507048c77191b82facbd1ce77fa6bd38cd (diff)
downloadfg21sim-f0682d6f5c196001f8820d0f2dbdf18087285ac7.tar.bz2
healpix.py: Implement "ang2pix_ring" and "pix2ang_ring" with JIT
To optimize "map_grid_to_healpix()" in `grid.py` which uses "healpy.ang2pix()" using `numba.jit`, implement the latter's functionality explicitly with `numba.jit` support. The implementation simply mimic the corresponding functions "ang2pix_ring_z_phi()" and "pix2ang_ring_z_phi()" in HEALPix's "src/C/subs/chealpix.c". Thanks!
Diffstat (limited to 'fg21sim')
-rw-r--r--fg21sim/utils/healpix.py213
1 files changed, 213 insertions, 0 deletions
diff --git a/fg21sim/utils/healpix.py b/fg21sim/utils/healpix.py
index 620d7eb..fe436d4 100644
--- a/fg21sim/utils/healpix.py
+++ b/fg21sim/utils/healpix.py
@@ -352,3 +352,216 @@ def _make_healpix_header(header, nside,
for comment in append_comment:
header.add_comment(comment)
return header
+
+
+@nb.jit(nb.int64(nb.int64), nopython=True)
+def nside2npix(nside):
+ """Calculate the number of pixels for the given Nside resolution.
+
+ NOTE
+ ----
+ This is the JIT-optimized version that replaces the ``healpy.nside2npix``
+ """
+ return 12 * nside * nside
+
+
+@nb.jit(nb.int64(nb.int64, nb.float64, nb.float64), nopython=True)
+def ang2pix_ring_single(nside, theta, phi):
+ """Calculate the pixel indexes in RING ordering scheme for one single
+ pair of angular coordinate on the sphere.
+
+ Parameters
+ ----------
+ theta : float
+ The polar angle (i.e., latitude), θ ∈ [0, π]. (unit: rad)
+ phi : float
+ The azimuthal angle (i.e., longitude), φ ∈ [0, 2π). (unit: rad)
+
+ Returns
+ -------
+ ipix : int
+ The index of the pixel corresponding to the input coordinate.
+
+ NOTE
+ ----
+ * Only support the *RING* ordering scheme
+ * This is the JIT-optimized version that partially replaces the
+ ``healpy.ang2pix``
+
+ References
+ ----------
+ - HEALPix software: ``src/C/subs/chealpix.c``: ``ang2pix_ring_z_phi()``
+ http://healpix.sourceforge.net/
+ """
+ z = np.cos(theta) # colatitude
+ za = np.absolute(z)
+ tt = (2.0 / np.pi) * np.remainder(phi, 2*np.pi) # range: [0, 4)
+ if za <= 2.0/3.0:
+ # Equatorial region
+ temp1 = nside * (tt + 0.5)
+ temp2 = nside * z * 0.75
+ jp = int(temp1 - temp2) # Index of ascending edge line
+ jm = int(temp1 + temp2) # Index of descending edge line
+ # Ring number counted from z=2/3
+ iring = nside + 1 + jp - jm # range: [1, 2n+1]
+ kshift = 1 - (iring & 1) # kshift=1 if ir even, 0 otherwise
+ ip = int((jp + jm - nside + kshift + 1) / 2)
+ ip = np.remainder(ip, 4*nside)
+ ipix = nside * (nside-1) * 2 + (iring-1) * 4 * nside + ip
+ else:
+ # North & South polar caps
+ tp = tt - int(tt)
+ tmp = nside * np.sqrt(3 * (1-za))
+ jp = int(tp * tmp)
+ jm = int((1.0-tp) * tmp)
+ # Ring number counted from the closest pole
+ iring = jp + jm + 1
+ ip = int(tt * iring)
+ ip = np.remainder(ip, 4*iring)
+ #
+ if z > 0:
+ ipix = 2 * iring * (iring-1) + ip
+ else:
+ ipix = 12 * nside * nside - 2 * iring * (iring+1) + ip
+ #
+ return ipix
+
+
+@nb.jit(nb.types.UniTuple(nb.float64, 2)(nb.int64, nb.int64), nopython=True)
+def pix2ang_ring_single(nside, ipix):
+ """Calculate the angular coordinate on the sphere for one pixel index
+ in the RING ordering scheme.
+
+ Parameters
+ ----------
+ ipix : int
+ The index of the HEALPix pixel in RING ordering.
+
+ Returns
+ -------
+ theta : float
+ The polar angle (i.e., latitude), θ ∈ [0, π]. (unit: rad)
+ phi : float
+ The azimuthal angle (i.e., longitude), φ ∈ [0, 2π). (unit: rad)
+
+ NOTE
+ ----
+ * Only support the *RING* ordering scheme
+ * This is the JIT-optimized version that partially replaces the
+ ``healpy.ang2pix``
+
+ References
+ ----------
+ - HEALPix software: ``src/C/subs/chealpix.c``: ``pix2ang_ring_z_phi()``
+ http://healpix.sourceforge.net/
+ """
+ ncap = nside * (nside-1) * 2
+ npix = nside2npix(nside)
+ fact2 = 4.0 / npix
+ if ipix < ncap:
+ # North polar cap
+ tmp = int(np.sqrt(2*ipix + 1 + 0.5))
+ # Ring number counted from the North pole
+ iring = int((tmp + 1) / 2)
+ iphi = (ipix + 1) - 2 * iring * (iring-1)
+ z = 1.0 - iring * iring * fact2
+ phi = (iphi - 0.5) * np.pi / (2 * iring)
+ elif ipix < (npix - ncap):
+ # Equatorial region
+ fact1 = 2 * nside * fact2
+ ip = ipix - ncap
+ # Ring number counted from the North pole
+ iring = int(ip / (4*nside) + nside)
+ iphi = ip % (4*nside) + 1
+ if (iring + nside) % 2 == 1:
+ fodd = 1.0 # (iring+nside) is odd
+ else:
+ fodd = 0.5
+ z = (2*nside - iring) * fact1
+ phi = (iphi - fodd) * np.pi / (2 * nside)
+ else:
+ # South polar cap
+ ip = npix - ipix
+ tmp = int(np.sqrt(2*ip - 1 + 0.5))
+ # Ring number counted from the South pole
+ iring = int((tmp + 1) / 2)
+ iphi = 4 * iring + 1 - (ip - 2 * iring * (iring-1))
+ z = iring * iring * fact2 - 1.0
+ phi = (iphi - 0.5) * np.pi / (2 * iring)
+ #
+ theta = np.arccos(z)
+ return (theta, phi)
+
+
+@nb.jit([nb.int64[:](nb.int64, nb.float64[:], nb.float64[:]),
+ nb.int64[:, :](nb.int64, nb.float64[:, :], nb.float64[:, :])],
+ nopython=True)
+def ang2pix_ring(nside, theta, phi):
+ """Calculate the pixel indexes in RING ordering scheme for each
+ pair of angular coordinates on the sphere.
+
+ Parameters
+ ----------
+ theta : 1D or 2D `~numpy.ndarray`
+ The polar angles (i.e., latitudes), θ ∈ [0, π]. (unit: rad)
+ phi : 1D or 2D `~numpy.ndarray`
+ The azimuthal angles (i.e., longitudes), φ ∈ [0, 2π). (unit: rad)
+
+ Returns
+ -------
+ ipix : 1D or 1D `~numpy.ndarray`
+ The indices of the pixels corresponding to the input coordinates.
+ The shape is the same as the input array.
+
+ NOTE
+ ----
+ * Only support the *RING* ordering scheme
+ * This is the JIT-optimized version that partially replaces the
+ ``healpy.ang2pix``
+ """
+ shape = theta.shape
+ size = theta.size
+ theta = theta.flatten()
+ phi = phi.flatten()
+ ipix = np.zeros(size, dtype=np.int64)
+ for i in range(size):
+ ipix[i] = ang2pix_ring_single(nside, theta[i], phi[i])
+ return ipix.reshape(shape)
+
+
+@nb.jit([nb.types.UniTuple(nb.float64[:], 2)(nb.int64, nb.int64[:]),
+ nb.types.UniTuple(nb.float64[:, :], 2)(nb.int64, nb.int64[:, :])],
+ nopython=True)
+def pix2ang_ring(nside, ipix):
+ """Calculate the angular coordinates on the sphere for each pixel
+ index in the RING ordering scheme.
+
+ Parameters
+ ----------
+ ipix : 1D or 2D `~numpy.ndarray`
+ The indices of the HEALPix pixels in the RING ordering
+
+ Returns
+ -------
+ theta : 1D or 2D `~numpy.ndarray`
+ The polar angles (i.e., latitudes), θ ∈ [0, π]. (unit: rad)
+ phi : 1D or 2D `~numpy.ndarray`
+ The azimuthal angles (i.e., longitudes), φ ∈ [0, 2π). (unit: rad)
+ The shape is the same as the input array.
+
+ NOTE
+ ----
+ * Only support the *RING* ordering scheme
+ * This is the JIT-optimized version that partially replaces the
+ ``healpy.pix2ang``
+ """
+ shape = ipix.shape
+ size = ipix.size
+ ipix = ipix.flatten()
+ theta = np.zeros(size, dtype=np.float64)
+ phi = np.zeros(size, dtype=np.float64)
+ for i in range(size):
+ theta_, phi_ = pix2ang_ring_single(nside, ipix[i])
+ theta[i] = theta_
+ phi[i] = phi_
+ return (theta.reshape(shape), phi.reshape(shape))