Comparison between Healpy functions and a Numba-based implementation

numbahpix.py

# -*- encoding: utf-8 -*-

import numpy as np
import math
from numba import njit, vectorize, int32, int64, float32, float64


@njit
def normalize_angle(ang):
    result = ang
    while result >= math.pi:
        result -= 2 * math.pi

    while result < 0:
        result += 2 * math.pi

    return result


@vectorize(
    [
        int32(int32, float32, float32),
        int32(int32, float64, float64),
        int32(int64, float32, float32),
        int32(int64, float64, float64),
    ]
)
def ang2pix(nside, theta, phi):
    # Put the value of class fields in local variables to help the
    # JIT compiler
    nside = nside
    nside_times_4 = 4 * nside
    pixels_per_face = nside * nside
    ncap = 2 * (pixels_per_face - nside)
    npix = 12 * nside * nside

    z = math.cos(theta)
    z_abs = abs(z)
    scaled_phi = normalize_angle(phi)
    tt = phi / (0.5 * math.pi)

    if z_abs <= 2 / 3:
        jp = math.floor(nside * (0.5 + tt - z * 0.75))
        jm = math.floor(nside * (0.5 + tt + z * 0.75))

        ir = nside + 1 + jp - jm
        kshift = 0
        if ir % 2 == 0:
            kshift = 1

        ip = math.floor((jp + jm - nside + kshift + 1) / 2) + 1
        if ip > nside_times_4:
            ip -= nside_times_4

        ipix1 = ncap + nside_times_4 * (ir - 1) + ip
    else:
        tp = tt - math.floor(tt)
        tmp = math.sqrt(3 * (1 - z_abs))

        jp = math.floor(nside * tp * tmp)
        jm = math.floor(nside * (1 - tp) * tmp)

        ir = jp + jm + 1
        ip = math.floor(tt * ir) + 1
        if ip > 4 * ir:
            ip -= 4 * ir

        ipix1 = 2 * ir * (ir - 1) + ip
        if z <= 0.0:
            ipix1 = npix - 2 * ir * (ir + 1) + ip

    return ipix1 - 1


if __name__ == "__main__":
    import healpy

    print("I'm checking the consistency between Numba and Healpy...")
    nside_choices = 2 ** np.arange(10)
    print("   testing NSIDE in {}".format(", ".join([str(x) for x in nside_choices])))
    for i in range(10000):
        nside = np.random.choice(nside_choices)
        theta = np.random.rand() * np.pi
        phi = np.random.rand() * 2 * np.pi
        assert ang2pix(nside, theta, phi) == healpy.ang2pix(
            nside, theta, phi
        ), f"Error for NSIDE={nside}, angle is ({theta}, {phi})"
    print("...done, they agree")

    import timeit

    print("Number\thealpy.ang2pix\tNumba")
    for num in 10 ** np.arange(6):
        num = int(num)
        theta = np.random.rand(num) * np.pi
        phi = np.random.rand(num) * 2 * np.pi

        healpy_time = timeit.timeit(
            lambda: healpy.ang2pix(num, theta, phi), number=1000
        )

        numba_time = timeit.timeit(lambda: ang2pix(num, theta, phi), number=1000)
        print(f"{num}\t{healpy_time:.4e}\t{numba_time:.4e}")