Hi amazing team!
I am trying to overcome a bottleneck in a process by implementing Numba’s njit. There are two simple functions in this bottleneck, both ready to work with Numba. See the test with and with njit below:
import numpy as np
from numba import njit # nb.__version__ returns 0.55.1
def weib_2params_sampl(shape, scale, num_samples=1):
arr=np.random.weibull(shape, size=num_samples)
return scale*arr
def weib_2params(y, iters=1000, eps=1e-6):
## Checkers
# Check for distribution ndim
if len(y.shape) != 1:
print("Weibull 2 params fit Error. Function requires 'y' to have ndim==1")
return np.nan, np.nan
# Check for nans ~ automatically erase them
if np.isnan(y).any():
y = np.ravel(y)[~np.isnan(y)]
## 2 params-Weibull fit k via MLE
# Initialization
ln_y = np.log(y)
k = 1.
k_t_1 = k
for t in range(iters):
y_k = y ** k
y_k_ln_y = y_k * ln_y
ff = np.sum(y_k_ln_y)
fg = np.sum(y_k)
f = ff / fg - np.mean(ln_y) - (1. / k)
# Calculate second derivative d^2f/dk^2
ff_prime = np.sum(y_k_ln_y * ln_y)
fg_prime = ff
f_prime = (ff_prime/fg - (ff/fg * fg_prime/fg)) + (1. / (k*k))
# Newton-Raphson method k = k - f(k;x)/f'(k;x)
k -= f/f_prime
if np.isnan(f):
print("Weibull 2 params fit warning. Fitting the 2 parameters Weibull is impossible. Returning NaN for shape and scale")
return np.nan, np.nan
if abs(k - k_t_1) < eps:
break
# Update the MLE
k_t_1 = k
lam = np.mean(y ** k) ** (1.0 / k)
return k, lam
def test():
arr = weib_2params_sampl(*np.random.uniform(.1,6, 2), num_samples=15000)
weib_2params(arr)
######## Numba implementation
@jit(nopython=True)
def weib_2params_sampl_nb(shape, scale, num_samples=1):
return scale * np.random.weibull(shape, size=num_samples)
@jit(nopython=True)
def weib_2params_nb(y, iters=1000, eps=1e-6):
## Checkers
# Check for distribution ndim
if len(y.shape) != 1:
print("Weibull 2 params fit Error. Function requires 'y' to have ndim==1")
return np.nan, np.nan
# Check for nans ~ automatically erase them
if np.isnan(y).any():
y = np.ravel(y)[~np.isnan(y)]
## 2 params-Weibull fit k via MLE
# Initialization
ln_y = np.log(y)
k = 1.
k_t_1 = k
for t in range(iters):
y_k = y ** k
y_k_ln_y = y_k * ln_y
ff = np.sum(y_k_ln_y)
fg = np.sum(y_k)
f = ff / fg - np.mean(ln_y) - (1. / k)
# Calculate second derivative d^2f/dk^2
ff_prime = np.sum(y_k_ln_y * ln_y)
fg_prime = ff
f_prime = (ff_prime/fg - (ff/fg * fg_prime/fg)) + (1. / (k*k))
# Newton-Raphson method k = k - f(k;x)/f'(k;x)
k -= f/f_prime
if np.isnan(f):
print("Weibull 2 params fit warning. Fitting the 2 parameters Weibull is impossible. Returning NaN for shape and scale")
return np.nan, np.nan
if abs(k - k_t_1) < eps:
break
# Update the MLE
k_t_1 = k
lam = np.mean(y ** k) ** (1.0 / k)
return k, lam
@jit(nopython=True)
def test_nb():
arr = weib_2params_sampl_nb(*np.random.uniform(.1,6, 2), num_samples=15000)
weib_2params_nb(arr)
The results from running test are (in notebook):
%timeit test
"21.3 ns ± 0.353 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)"
%timeit test_nb # with numba already compiled
"23.3 ns ± 0.491 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)"
Actually, both functions “weib_2params_sampl_nb” and “weib_2params_nb” are a bit slower and the Python/Numpy versions
Is there anything I can do to increase speed?
Thanks a lot guys!
Best!