# Benchmark in Jupyter with the jit decorator

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)

from transonic import wait_for_all_extensions
from transonic.util import print_versions, timeit_verbose

print_versions()

Transonic 0.4.12
Pythran 0.11.0
Numba 0.54.1
Cython 0.29.24

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)

# 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()

Pythranizing /home/pierre/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy.py
Pythranizing /home/pierre/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy_loops.py
File /home/pierre/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy_loops_9f81a0650bb71a866e10757958292685_e3e427414afb66bd6d7d9a2b46ba1060.cpython-39-x86_64-linux-gnu.so created by pythran
File /home/pierre/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy_afb23984aa5312381ff8120d62ec54ff_812f8c2346a1f947c4504d48c4d7a4e0.cpython-39-x86_64-linux-gnu.so created by pythran

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 * norm
norm = 0.000206 s
fxfy_numba                       : 0.919 * norm
fxfy_loops_numba                 : 0.798 * norm
fxfy_pythran                     : 0.289 * norm
fxfy_loops_pythran               : 0.801 * norm