Benchmark in Jupyter with the jit decorator

[1]:
import numpy as np
from transonic import jit


def fxfy(ft, fn, theta):
    sin_theta = np.sin(theta)
    cos_theta = np.cos(theta)
    fx = cos_theta * ft - sin_theta * fn
    fy = sin_theta * ft + cos_theta * fn
    return fx, fy


def fxfy_loops(ft, fn, theta):
    n0 = theta.size
    fx = np.empty_like(ft)
    fy = np.empty_like(fn)
    for index in range(n0):
        sin_theta = np.sin(theta[index])
        cos_theta = np.cos(theta[index])
        fx[index] = cos_theta * ft[index] - sin_theta * fn[index]
        fy[index] = sin_theta * ft[index] + cos_theta * fn[index]
    return fx, fy


fxfy_pythran = jit(backend="pythran", native=True, xsimd=True)(fxfy)
fxfy_numba = jit(backend="numba")(fxfy)
fxfy_loops_pythran = jit(backend="pythran", native=True, xsimd=True)(fxfy_loops)
fxfy_loops_numba = jit(backend="numba")(fxfy_loops)
[2]:
from transonic import wait_for_all_extensions
from transonic.util import print_versions, timeit_verbose

print_versions()
Transonic 0.4.1
Pythran 0.9.3post1
Numba 0.46.0
Cython 0.29.13
[3]:
theta = np.linspace(0, 2 * np.pi, 10000)
ft = 2.5 * theta
fv = 1.5 * theta
loc = locals()

out = fxfy(ft, fv, theta)
out_loops = fxfy_loops(ft, fv, theta)
assert np.allclose(out, out_loops)
[4]:
# warmup
fxfy_pythran(ft, fv, theta)
fxfy_loops_pythran(ft, fv, theta)
fxfy_numba(ft, fv, theta)
fxfy_loops_numba(ft, fv, theta)

wait_for_all_extensions()
INFO: Schedule pythranization of file /home/users/augier3pi/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy.py
INFO: Schedule pythranization of file /home/users/augier3pi/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy_loops.py
compile extension
compile extension
[5]:
norm = timeit_verbose("fxfy(ft, fv, theta)", globals=loc)

for backend in ("numba", "pythran"):
    timeit_verbose(f"fxfy_{backend}(ft, fv, theta)", globals=loc, norm=norm)
    timeit_verbose(
        f"fxfy_loops_{backend}(ft, fv, theta)", globals=loc, norm=norm
    )
fxfy                             : 1.000 * norm
norm = 6.90e-04 s
fxfy_numba                       : 0.952 * norm
fxfy_loops_numba                 : 0.776 * norm
fxfy_pythran                     : 0.152 * norm
fxfy_loops_pythran               : 0.784 * norm