ai-content-maker/.venv/Lib/site-packages/numba/np/random/distributions.py

721 lines
20 KiB
Python

"""
Algorithmic implementations for generating different types
of random distributions.
"""
import numpy as np
from numba.core.extending import register_jitable
from numba.np.random._constants import (wi_double, ki_double,
ziggurat_nor_r, fi_double,
wi_float, ki_float,
ziggurat_nor_inv_r_f,
ziggurat_nor_r_f, fi_float,
we_double, ke_double,
ziggurat_exp_r, fe_double,
we_float, ke_float,
ziggurat_exp_r_f, fe_float,
INT64_MAX, ziggurat_nor_inv_r)
from numba.np.random.generator_core import (next_double, next_float,
next_uint32, next_uint64)
from numba import float32, int64
# All of the following implementations are direct translations from:
# https://github.com/numpy/numpy/blob/7cfef93c77599bd387ecc6a15d186c5a46024dac/numpy/random/src/distributions/distributions.c
@register_jitable
def np_log1p(x):
return np.log1p(x)
@register_jitable
def np_log1pf(x):
return np.log1p(float32(x))
@register_jitable
def random_rayleigh(bitgen, mode):
return mode * np.sqrt(2.0 * random_standard_exponential(bitgen))
@register_jitable
def np_expm1(x):
return np.expm1(x)
@register_jitable
def random_standard_normal(bitgen):
while 1:
r = next_uint64(bitgen)
idx = r & 0xff
r >>= 8
sign = r & 0x1
rabs = (r >> 1) & 0x000fffffffffffff
x = rabs * wi_double[idx]
if (sign & 0x1):
x = -x
if rabs < ki_double[idx]:
return x
if idx == 0:
while 1:
xx = -ziggurat_nor_inv_r * np.log1p(-next_double(bitgen))
yy = -np.log1p(-next_double(bitgen))
if (yy + yy > xx * xx):
if ((rabs >> 8) & 0x1):
return -(ziggurat_nor_r + xx)
else:
return ziggurat_nor_r + xx
else:
if (((fi_double[idx - 1] - fi_double[idx]) *
next_double(bitgen) + fi_double[idx]) <
np.exp(-0.5 * x * x)):
return x
@register_jitable
def random_standard_normal_f(bitgen):
while 1:
r = next_uint32(bitgen)
idx = r & 0xff
sign = (r >> 8) & 0x1
rabs = (r >> 9) & 0x0007fffff
x = float32(float32(rabs) * wi_float[idx])
if (sign & 0x1):
x = -x
if (rabs < ki_float[idx]):
return x
if (idx == 0):
while 1:
xx = float32(-ziggurat_nor_inv_r_f *
np_log1pf(-next_float(bitgen)))
yy = float32(-np_log1pf(-next_float(bitgen)))
if (float32(yy + yy) > float32(xx * xx)):
if ((rabs >> 8) & 0x1):
return -float32(ziggurat_nor_r_f + xx)
else:
return float32(ziggurat_nor_r_f + xx)
else:
if (((fi_float[idx - 1] - fi_float[idx]) * next_float(bitgen) +
fi_float[idx]) < float32(np.exp(-float32(0.5) * x * x))):
return x
@register_jitable
def random_standard_exponential(bitgen):
while 1:
ri = next_uint64(bitgen)
ri >>= 3
idx = ri & 0xFF
ri >>= 8
x = ri * we_double[idx]
if (ri < ke_double[idx]):
return x
else:
if idx == 0:
return ziggurat_exp_r - np_log1p(-next_double(bitgen))
elif ((fe_double[idx - 1] - fe_double[idx]) * next_double(bitgen) +
fe_double[idx] < np.exp(-x)):
return x
@register_jitable
def random_standard_exponential_f(bitgen):
while 1:
ri = next_uint32(bitgen)
ri >>= 1
idx = ri & 0xFF
ri >>= 8
x = float32(float32(ri) * we_float[idx])
if (ri < ke_float[idx]):
return x
else:
if (idx == 0):
return float32(ziggurat_exp_r_f -
float32(np_log1pf(-next_float(bitgen))))
elif ((fe_float[idx - 1] - fe_float[idx]) * next_float(bitgen) +
fe_float[idx] < float32(np.exp(float32(-x)))):
return x
@register_jitable
def random_standard_exponential_inv(bitgen):
return -np_log1p(-next_double(bitgen))
@register_jitable
def random_standard_exponential_inv_f(bitgen):
return -np.log(float32(1.0) - next_float(bitgen))
@register_jitable
def random_standard_gamma(bitgen, shape):
if (shape == 1.0):
return random_standard_exponential(bitgen)
elif (shape == 0.0):
return 0.0
elif (shape < 1.0):
while 1:
U = next_double(bitgen)
V = random_standard_exponential(bitgen)
if (U <= 1.0 - shape):
X = pow(U, 1. / shape)
if (X <= V):
return X
else:
Y = -np.log((1 - U) / shape)
X = pow(1.0 - shape + shape * Y, 1. / shape)
if (X <= (V + Y)):
return X
else:
b = shape - 1. / 3.
c = 1. / np.sqrt(9 * b)
while 1:
while 1:
X = random_standard_normal(bitgen)
V = 1.0 + c * X
if (V > 0.0):
break
V = V * V * V
U = next_double(bitgen)
if (U < 1.0 - 0.0331 * (X * X) * (X * X)):
return (b * V)
if (np.log(U) < 0.5 * X * X + b * (1. - V + np.log(V))):
return (b * V)
@register_jitable
def random_standard_gamma_f(bitgen, shape):
f32_one = float32(1.0)
shape = float32(shape)
if (shape == f32_one):
return random_standard_exponential_f(bitgen)
elif (shape == float32(0.0)):
return float32(0.0)
elif (shape < f32_one):
while 1:
U = next_float(bitgen)
V = random_standard_exponential_f(bitgen)
if (U <= f32_one - shape):
X = float32(pow(U, float32(f32_one / shape)))
if (X <= V):
return X
else:
Y = float32(-np.log(float32((f32_one - U) / shape)))
X = float32(pow(f32_one - shape + float32(shape * Y),
float32(f32_one / shape)))
if (X <= (V + Y)):
return X
else:
b = shape - f32_one / float32(3.0)
c = float32(f32_one / float32(np.sqrt(float32(9.0) * b)))
while 1:
while 1:
X = float32(random_standard_normal_f(bitgen))
V = float32(f32_one + c * X)
if (V > float32(0.0)):
break
V = float32(V * V * V)
U = next_float(bitgen)
if (U < f32_one - float32(0.0331) * (X * X) * (X * X)):
return float32(b * V)
if (np.log(U) < float32(0.5) * X * X + b *
(f32_one - V + np.log(V))):
return float32(b * V)
@register_jitable
def random_normal(bitgen, loc, scale):
scaled_normal = scale * random_standard_normal(bitgen)
return loc + scaled_normal
@register_jitable
def random_normal_f(bitgen, loc, scale):
scaled_normal = float32(scale * random_standard_normal_f(bitgen))
return float32(loc + scaled_normal)
@register_jitable
def random_exponential(bitgen, scale):
return scale * random_standard_exponential(bitgen)
@register_jitable
def random_uniform(bitgen, lower, range):
scaled_uniform = range * next_double(bitgen)
return lower + scaled_uniform
@register_jitable
def random_gamma(bitgen, shape, scale):
return scale * random_standard_gamma(bitgen, shape)
@register_jitable
def random_gamma_f(bitgen, shape, scale):
return float32(scale * random_standard_gamma_f(bitgen, shape))
@register_jitable
def random_beta(bitgen, a, b):
if a <= 1.0 and b <= 1.0:
while 1:
U = next_double(bitgen)
V = next_double(bitgen)
X = pow(U, 1.0 / a)
Y = pow(V, 1.0 / b)
XpY = X + Y
if XpY <= 1.0 and XpY > 0.0:
if (X + Y > 0):
return X / XpY
else:
logX = np.log(U) / a
logY = np.log(V) / b
logM = min(logX, logY)
logX -= logM
logY -= logM
return np.exp(logX - np.log(np.exp(logX) + np.exp(logY)))
else:
Ga = random_standard_gamma(bitgen, a)
Gb = random_standard_gamma(bitgen, b)
return Ga / (Ga + Gb)
@register_jitable
def random_chisquare(bitgen, df):
return 2.0 * random_standard_gamma(bitgen, df / 2.0)
@register_jitable
def random_f(bitgen, dfnum, dfden):
return ((random_chisquare(bitgen, dfnum) * dfden) /
(random_chisquare(bitgen, dfden) * dfnum))
@register_jitable
def random_standard_cauchy(bitgen):
return random_standard_normal(bitgen) / random_standard_normal(bitgen)
@register_jitable
def random_pareto(bitgen, a):
return np_expm1(random_standard_exponential(bitgen) / a)
@register_jitable
def random_weibull(bitgen, a):
if (a == 0.0):
return 0.0
return pow(random_standard_exponential(bitgen), 1. / a)
@register_jitable
def random_power(bitgen, a):
return pow(-np_expm1(-random_standard_exponential(bitgen)), 1. / a)
@register_jitable
def random_laplace(bitgen, loc, scale):
U = next_double(bitgen)
while U <= 0:
U = next_double(bitgen)
if (U >= 0.5):
U = loc - scale * np.log(2.0 - U - U)
elif (U > 0.0):
U = loc + scale * np.log(U + U)
return U
@register_jitable
def random_logistic(bitgen, loc, scale):
U = next_double(bitgen)
while U <= 0.0:
U = next_double(bitgen)
return loc + scale * np.log(U / (1.0 - U))
@register_jitable
def random_lognormal(bitgen, mean, sigma):
return np.exp(random_normal(bitgen, mean, sigma))
@register_jitable
def random_standard_t(bitgen, df):
num = random_standard_normal(bitgen)
denom = random_standard_gamma(bitgen, df / 2)
return np.sqrt(df / 2) * num / np.sqrt(denom)
@register_jitable
def random_wald(bitgen, mean, scale):
mu_2l = mean / (2 * scale)
Y = random_standard_normal(bitgen)
Y = mean * Y * Y
X = mean + mu_2l * (Y - np.sqrt(4 * scale * Y + Y * Y))
U = next_double(bitgen)
if (U <= mean / (mean + X)):
return X
else:
return mean * mean / X
@register_jitable
def random_geometric_search(bitgen, p):
X = 1
sum = prod = p
q = 1.0 - p
U = next_double(bitgen)
while (U > sum):
prod *= q
sum += prod
X = X + 1
return X
@register_jitable
def random_geometric_inversion(bitgen, p):
return np.ceil(-random_standard_exponential(bitgen) / np.log1p(-p))
@register_jitable
def random_geometric(bitgen, p):
if (p >= 0.333333333333333333333333):
return random_geometric_search(bitgen, p)
else:
return random_geometric_inversion(bitgen, p)
@register_jitable
def random_zipf(bitgen, a):
am1 = a - 1.0
b = pow(2.0, am1)
while 1:
U = 1.0 - next_double(bitgen)
V = next_double(bitgen)
X = np.floor(pow(U, -1.0 / am1))
if (X > INT64_MAX or X < 1.0):
continue
T = pow(1.0 + 1.0 / X, am1)
if (V * X * (T - 1.0) / (b - 1.0) <= T / b):
return X
@register_jitable
def random_triangular(bitgen, left, mode,
right):
base = right - left
leftbase = mode - left
ratio = leftbase / base
leftprod = leftbase * base
rightprod = (right - mode) * base
U = next_double(bitgen)
if (U <= ratio):
return left + np.sqrt(U * leftprod)
else:
return right - np.sqrt((1.0 - U) * rightprod)
@register_jitable
def random_loggam(x):
a = [8.333333333333333e-02, -2.777777777777778e-03,
7.936507936507937e-04, -5.952380952380952e-04,
8.417508417508418e-04, -1.917526917526918e-03,
6.410256410256410e-03, -2.955065359477124e-02,
1.796443723688307e-01, -1.39243221690590e+00]
if ((x == 1.0) or (x == 2.0)):
return 0.0
elif (x < 7.0):
n = int(7 - x)
else:
n = 0
x0 = x + n
x2 = (1.0 / x0) * (1.0 / x0)
# /* log(2 * M_PI) */
lg2pi = 1.8378770664093453e+00
gl0 = a[9]
for k in range(0, 9):
gl0 *= x2
gl0 += a[8 - k]
gl = gl0 / x0 + 0.5 * lg2pi + (x0 - 0.5) * np.log(x0) - x0
if (x < 7.0):
for k in range(1, n + 1):
gl = gl - np.log(x0 - 1.0)
x0 = x0 - 1.0
return gl
@register_jitable
def random_poisson_mult(bitgen, lam):
enlam = np.exp(-lam)
X = 0
prod = 1.0
while (1):
U = next_double(bitgen)
prod *= U
if (prod > enlam):
X += 1
else:
return X
@register_jitable
def random_poisson_ptrs(bitgen, lam):
slam = np.sqrt(lam)
loglam = np.log(lam)
b = 0.931 + 2.53 * slam
a = -0.059 + 0.02483 * b
invalpha = 1.1239 + 1.1328 / (b - 3.4)
vr = 0.9277 - 3.6224 / (b - 2)
while (1):
U = next_double(bitgen) - 0.5
V = next_double(bitgen)
us = 0.5 - np.fabs(U)
k = int((2 * a / us + b) * U + lam + 0.43)
if ((us >= 0.07) and (V <= vr)):
return k
if ((k < 0) or ((us < 0.013) and (V > us))):
continue
# /* log(V) == log(0.0) ok here */
# /* if U==0.0 so that us==0.0, log is ok since always returns */
if ((np.log(V) + np.log(invalpha) - np.log(a / (us * us) + b)) <=
(-lam + k * loglam - random_loggam(k + 1))):
return k
@register_jitable
def random_poisson(bitgen, lam):
if (lam >= 10):
return random_poisson_ptrs(bitgen, lam)
elif (lam == 0):
return 0
else:
return random_poisson_mult(bitgen, lam)
@register_jitable
def random_negative_binomial(bitgen, n, p):
Y = random_gamma(bitgen, n, (1 - p) / p)
return random_poisson(bitgen, Y)
@register_jitable
def random_noncentral_chisquare(bitgen, df, nonc):
if np.isnan(nonc):
return np.nan
if nonc == 0:
return random_chisquare(bitgen, df)
if 1 < df:
Chi2 = random_chisquare(bitgen, df - 1)
n = random_standard_normal(bitgen) + np.sqrt(nonc)
return Chi2 + n * n
else:
i = random_poisson(bitgen, nonc / 2.0)
return random_chisquare(bitgen, df + 2 * i)
@register_jitable
def random_noncentral_f(bitgen, dfnum, dfden, nonc):
t = random_noncentral_chisquare(bitgen, dfnum, nonc) * dfden
return t / (random_chisquare(bitgen, dfden) * dfnum)
@register_jitable
def random_logseries(bitgen, p):
r = np_log1p(-p)
while 1:
V = next_double(bitgen)
if (V >= p):
return 1
U = next_double(bitgen)
q = -np.expm1(r * U)
if (V <= q * q):
result = int64(np.floor(1 + np.log(V) / np.log(q)))
if result < 1 or V == 0.0:
continue
else:
return result
if (V >= q):
return 1
else:
return 2
@register_jitable
def random_binomial_btpe(bitgen, n, p):
r = min(p, 1.0 - p)
q = 1.0 - r
fm = n * r + r
m = int(np.floor(fm))
p1 = int(np.floor(2.195 * np.sqrt(n * r * q) - 4.6 * q) + 0.5)
xm = m + 0.5
xl = xm - p1
xr = xm + p1
c = 0.134 + 20.5 / (15.3 + m)
a = (fm - xl) / (fm - xl * r)
laml = a * (1.0 + a / 2.0)
a = (xr - fm) / (xr * q)
lamr = a * (1.0 + a / 2.0)
p2 = p1 * (1.0 + 2.0 * c)
p3 = p2 + c / laml
p4 = p3 + c / lamr
case = 10
y = k = 0
while 1:
if case == 10:
nrq = n * r * q
u = next_double(bitgen) * p4
v = next_double(bitgen)
if (u > p1):
case = 20
continue
y = int(np.floor(xm - p1 * v + u))
case = 60
continue
elif case == 20:
if (u > p2):
case = 30
continue
x = xl + (u - p1) / c
v = v * c + 1.0 - np.fabs(m - x + 0.5) / p1
if (v > 1.0):
case = 10
continue
y = int(np.floor(x))
case = 50
continue
elif case == 30:
if (u > p3):
case = 40
continue
y = int(np.floor(xl + np.log(v) / laml))
if ((y < 0) or (v == 0.0)):
case = 10
continue
v = v * (u - p2) * laml
case = 50
continue
elif case == 40:
y = int(np.floor(xr - np.log(v) / lamr))
if ((y > n) or (v == 0.0)):
case = 10
continue
v = v * (u - p3) * lamr
case = 50
continue
elif case == 50:
k = abs(y - m)
if ((k > 20) and (k < ((nrq) / 2.0 - 1))):
case = 52
continue
s = r / q
a = s * (n + 1)
F = 1.0
if (m < y):
for i in range(m + 1, y + 1):
F = F * (a / i - s)
elif (m > y):
for i in range(y + 1, m + 1):
F = F / (a / i - s)
if (v > F):
case = 10
continue
case = 60
continue
elif case == 52:
rho = (k / (nrq)) * \
((k * (k / 3.0 + 0.625) + 0.16666666666666666) /
nrq + 0.5)
t = -k * k / (2 * nrq)
A = np.log(v)
if (A < (t - rho)):
case = 60
continue
if (A > (t + rho)):
case = 10
continue
x1 = y + 1
f1 = m + 1
z = n + 1 - m
w = n - y + 1
x2 = x1 * x1
f2 = f1 * f1
z2 = z * z
w2 = w * w
if (A > (xm * np.log(f1 / x1) + (n - m + 0.5) * np.log(z / w) +
(y - m) * np.log(w * r / (x1 * q)) +
(13680. - (462. - (132. - (99. - 140. / f2) / f2) / f2)
/ f2) / f1 / 166320. +
(13680. - (462. - (132. - (99. - 140. / z2) / z2) / z2)
/ z2) / z / 166320. +
(13680. - (462. - (132. - (99. - 140. / x2) / x2) / x2)
/ x2) / x1 / 166320. +
(13680. - (462. - (132. - (99. - 140. / w2) / w2) / w2)
/ w2) / w / 66320.)):
case = 10
continue
elif case == 60:
if (p > 0.5):
y = n - y
return y
@register_jitable
def random_binomial_inversion(bitgen, n, p):
q = 1.0 - p
qn = np.exp(n * np.log(q))
_np = n * p
bound = min(n, _np + 10.0 * np.sqrt(_np * q + 1))
X = 0
px = qn
U = next_double(bitgen)
while (U > px):
X = X + 1
if (X > bound):
X = 0
px = qn
U = next_double(bitgen)
else:
U -= px
px = ((n - X + 1) * p * px) / (X * q)
return X
@register_jitable
def random_binomial(bitgen, n, p):
if ((n == 0) or (p == 0.0)):
return 0
if (p <= 0.5):
if (p * n <= 30.0):
return random_binomial_inversion(bitgen, n, p)
else:
return random_binomial_btpe(bitgen, n, p)
else:
q = 1.0 - p
if (q * n <= 30.0):
return n - random_binomial_inversion(bitgen, n, q)
else:
return n - random_binomial_btpe(bitgen, n, q)