def int_pow_fixed(y, n, prec):
"""n-th power of a fixed point number with precision prec
Returns the power in the form man, exp,
man * 2**exp ~= y**n
"""
if n == 2:
return (y*y), 0
bc = bitcount(y)
exp = 0
workprec = 2 * (prec + 4*bitcount(n) + 4)
_, pm, pe, pbc = fone
while 1:
if n & 1:
pm = pm*y
pe = pe+exp
pbc += bc - 2
pbc = pbc + bctable[int(pm >> pbc)]
if pbc > workprec:
pm = pm >> (pbc-workprec)
pe += pbc - workprec
pbc = workprec
n -= 1
if not n:
break
y = y*y
exp = exp+exp
bc = bc + bc - 2
bc = bc + bctable[int(y >> bc)]
if bc > workprec:
y = y >> (bc-workprec)
exp += bc - workprec
bc = workprec
n = n // 2
return pm, pe
def int_pow_fixed(y, n, prec):
"""n-th power of a fixed point number with precision prec
Returns the power in the form man, exp,
man * 2**exp ~= y**n
"""
if n == 2:
return (y*y), 0
bc = bitcount(y)
exp = 0
workprec = 2 * (prec + 4*bitcount(n) + 4)
_, pm, pe, pbc = fone
while 1:
if n & 1:
pm = pm*y
pe = pe+exp
pbc += bc - 2
pbc = pbc + bctable[int(pm >> pbc)]
if pbc > workprec:
pm = pm >> (pbc-workprec)
pe += pbc - workprec
pbc = workprec
n -= 1
if not n:
break
y = y*y
exp = exp+exp
bc = bc + bc - 2
bc = bc + bctable[int(y >> bc)]
if bc > workprec:
y = y >> (bc-workprec)
exp += bc - workprec
bc = workprec
n = n // 2
return pm, pe
def mpc_besseljn(n, z, prec):
negate = n < 0 and n & 1
n = abs(n)
origprec = prec
prec += 20 + bitcount(abs(n))
zre, zim = z
zre = to_fixed(zre, prec)
zim = to_fixed(zim, prec)
z2re = (zre**2 - zim**2) >> prec
z2im = (zre*zim) >> (prec-1)
if not n:
sre = tre = MP_ONE << prec
sim = tim = MP_ZERO
else:
re, im = complex_int_pow(zre, zim, n)
sre = tre = (re // int_fac(n)) >> ((n-1)*prec + n)
sim = tim = (im // int_fac(n)) >> ((n-1)*prec + n)
k = 1
while abs(tre) + abs(tim) > 3:
p = -4*k*(k+n)
tre, tim = tre*z2re - tim*z2im, tim*z2re + tre*z2im
tre = (tre // p) >> prec
tim = (tim // p) >> prec
sre += tre
sim += tim
k += 1
if negate:
sre = -sre
sim = -sim
re = from_man_exp(sre, -prec, origprec, round_nearest)
im = from_man_exp(sim, -prec, origprec, round_nearest)
return (re, im)
def mpc_besseljn(n, z, prec, rounding=round_fast):
negate = n < 0 and n & 1
n = abs(n)
origprec = prec
prec += 20 + bitcount(abs(n))
zre, zim = z
zre = to_fixed(zre, prec)
zim = to_fixed(zim, prec)
z2re = (zre**2 - zim**2) >> prec
z2im = (zre * zim) >> (prec - 1)
if not n:
sre = tre = MP_ONE << prec
sim = tim = MP_ZERO
else:
re, im = complex_int_pow(zre, zim, n)
sre = tre = (re // int_fac(n)) >> ((n - 1) * prec + n)
sim = tim = (im // int_fac(n)) >> ((n - 1) * prec + n)
k = 1
while abs(tre) + abs(tim) > 3:
p = -4 * k * (k + n)
tre, tim = tre * z2re - tim * z2im, tim * z2re + tre * z2im
tre = (tre // p) >> prec
tim = (tim // p) >> prec
sre += tre
sim += tim
k += 1
if negate:
sre = -sre
sim = -sim
re = from_man_exp(sre, -prec, origprec, rounding)
im = from_man_exp(sim, -prec, origprec, rounding)
return (re, im)
def mpc_besseljn(n, z, prec, rounding=round_fast):
negate = n < 0 and n & 1
n = abs(n)
origprec = prec
zre, zim = z
mag = max(zre[2]+zre[3], zim[2]+zim[3])
prec += 20 + n*bitcount(n) + abs(mag)
if mag < 0:
prec -= n * mag
zre = to_fixed(zre, prec)
zim = to_fixed(zim, prec)
z2re = (zre**2 - zim**2) >> prec
z2im = (zre*zim) >> (prec-1)
if not n:
sre = tre = MPZ_ONE << prec
sim = tim = MPZ_ZERO
else:
re, im = complex_int_pow(zre, zim, n)
sre = tre = (re // ifac(n)) >> ((n-1)*prec + n)
sim = tim = (im // ifac(n)) >> ((n-1)*prec + n)
k = 1
while abs(tre) + abs(tim) > 3:
p = -4*k*(k+n)
tre, tim = tre*z2re - tim*z2im, tim*z2re + tre*z2im
tre = (tre // p) >> prec
tim = (tim // p) >> prec
sre += tre
sim += tim
k += 1
if negate:
sre = -sre
sim = -sim
re = from_man_exp(sre, -prec, origprec, rounding)
im = from_man_exp(sim, -prec, origprec, rounding)
return (re, im)
def log_taylor_get_cached(n, prec):
# Taylor series with caching wins up to huge precisions
# To avoid unnecessary precomputation at low precision, we
# do it in steps
# Round to next power of 2
prec2 = (1<<(bitcount(prec-1))) + 20
dprec = prec2 - prec
if (n, prec2) in log_taylor_cache:
a, atan_a = log_taylor_cache[n, prec2]
else:
a = n << (prec2 - LOG_TAYLOR_SHIFT)
atan_a = log_newton(a, prec2)
log_taylor_cache[n, prec2] = (a, atan_a)
return (a >> dprec), (atan_a >> dprec)
def atan_taylor_get_cached(n, prec):
# Taylor series with caching wins up to huge precisions
# To avoid unnecessary precomputation at low precision, we
# do it in steps
# Round to next power of 2
prec2 = (1<<(bitcount(prec-1))) + 20
dprec = prec2 - prec
if (n, prec2) in atan_taylor_cache:
a, atan_a = atan_taylor_cache[n, prec2]
else:
a = n << (prec2 - ATAN_TAYLOR_SHIFT)
atan_a = atan_newton(a, prec2)
atan_taylor_cache[n, prec2] = (a, atan_a)
return (a >> dprec), (atan_a >> dprec)
def mpf_besseljn(n, x, prec):
negate = n < 0 and n & 1
n = abs(n)
origprec = prec
prec += 20 + bitcount(abs(n))
x = to_fixed(x, prec)
x2 = (x**2) >> prec
if not n:
s = t = MP_ONE << prec
else:
s = t = (x**n // int_fac(n)) >> ((n-1)*prec + n)
k = 1
while t:
t = ((t * x2) // (-4*k*(k+n))) >> prec
s += t
k += 1
if negate:
s = -s
return from_man_exp(s, -prec, origprec, round_nearest)
def mpf_besseljn(n, x, prec, rounding=round_fast):
negate = n < 0 and n & 1
n = abs(n)
origprec = prec
prec += 20 + bitcount(abs(n))
x = to_fixed(x, prec)
x2 = (x**2) >> prec
if not n:
s = t = MP_ONE << prec
else:
s = t = (x**n // int_fac(n)) >> ((n - 1) * prec + n)
k = 1
while t:
t = ((t * x2) // (-4 * k * (k + n))) >> prec
s += t
k += 1
if negate:
s = -s
return from_man_exp(s, -prec, origprec, rounding)
def log_int_fixed(n, prec, ln2=None):
"""
Fast computation of log(n), caching the value for small n,
intended for zeta sums.
"""
if n in log_int_cache:
value, vprec = log_int_cache[n]
if vprec >= prec:
return value >> (vprec - prec)
wp = prec + 10
if wp <= LOG_TAYLOR_SHIFT:
if ln2 is None:
ln2 = ln2_fixed(wp)
r = bitcount(n)
x = n << (wp - r)
v = log_taylor_cached(x, wp) + r * ln2
else:
v = to_fixed(mpf_log(from_int(n), wp + 5), wp)
if n < MAX_LOG_INT_CACHE:
log_int_cache[n] = (v, wp)
return v >> (wp - prec)
def log_int_fixed(n, prec, ln2=None):
"""
Fast computation of log(n), caching the value for small n,
intended for zeta sums.
"""
if n in log_int_cache:
value, vprec = log_int_cache[n]
if vprec >= prec:
return value >> (vprec - prec)
wp = prec + 10
if wp <= LOG_TAYLOR_SHIFT:
if ln2 is None:
ln2 = ln2_fixed(wp)
r = bitcount(n)
x = n << (wp - r)
v = log_taylor_cached(x, wp) + r * ln2
else:
v = to_fixed(mpf_log(from_int(n), wp + 5), wp)
if n < MAX_LOG_INT_CACHE:
log_int_cache[n] = (v, wp)
return v >> (wp - prec)
def mpf_besseljn(n, x, prec, rounding=round_fast):
prec += 50
negate = n < 0 and n & 1
mag = x[2]+x[3]
n = abs(n)
wp = prec + 20 + n*bitcount(n)
if mag < 0:
wp -= n * mag
x = to_fixed(x, wp)
x2 = (x**2) >> wp
if not n:
s = t = MPZ_ONE << wp
else:
s = t = (x**n // ifac(n)) >> ((n-1)*wp + n)
k = 1
while t:
t = ((t * x2) // (-4*k*(k+n))) >> wp
s += t
k += 1
if negate:
s = -s
return from_man_exp(s, -wp, prec, rounding)
def mpf_besseljn(n, x, prec, rounding=round_fast):
prec += 50
negate = n < 0 and n & 1
mag = x[2] + x[3]
n = abs(n)
wp = prec + 20 + n * bitcount(n)
if mag < 0:
wp -= n * mag
x = to_fixed(x, wp)
x2 = (x**2) >> wp
if not n:
s = t = MP_ONE << wp
else:
s = t = (x**n // int_fac(n)) >> ((n - 1) * wp + n)
k = 1
while t:
t = ((t * x2) // (-4 * k * (k + n))) >> wp
s += t
k += 1
if negate:
s = -s
return from_man_exp(s, -wp, prec, rounding)
def mpf_fibonacci(x, prec, rnd=round_fast):
sign, man, exp, bc = x
if not man:
if x == fninf:
return fnan
return x
# F(2^n) ~= 2^(2^n)
size = abs(exp+bc)
if exp >= 0:
# Exact
if size < 10 or size <= bitcount(prec):
return from_int(ifib(to_int(x)), prec, rnd)
# Use the modified Binet formula
wp = prec + size + 20
a = mpf_phi(wp)
b = mpf_add(mpf_shift(a, 1), fnone, wp)
u = mpf_pow(a, x, wp)
v = mpf_cos_pi(x, wp)
v = mpf_div(v, u, wp)
u = mpf_sub(u, v, wp)
u = mpf_div(u, b, prec, rnd)
return u
def mpf_exp(x, prec, rnd=round_fast):
sign, man, exp, bc = x
if not man:
if not exp:
return fone
if x == fninf:
return fzero
return x
mag = bc+exp
# Fast handling e**n. TODO: the best cutoff depends on both the
# size of n and the precision.
if prec > 600 and exp >= 0:
e = mpf_e(prec+10+int(1.45*mag))
return mpf_pow_int(e, (-1)**sign *(man<<exp), prec, rnd)
if mag < -prec-10:
return mpf_perturb(fone, sign, prec, rnd)
# extra precision needs to be similar in magnitude to log_2(|x|)
# for the modulo reduction, plus r for the error from squaring r times
wp = prec + max(0, mag)
if wp < 300:
r = int(2*wp**0.4)
if mag < 0:
r = max(1, r + mag)
wp += r + 20
t = to_fixed(x, wp)
# abs(x) > 1?
if mag > 1:
lg2 = ln2_fixed(wp)
n, t = divmod(t, lg2)
else:
n = 0
man = exp_series(t, wp, r)
else:
use_newton = False
# put a bound on exp to avoid infinite recursion in exp_newton
# TODO find a good bound
if wp > LIM_EXP_SERIES2 and exp < 1000:
if mag > 0:
use_newton = True
elif mag <= 0 and -mag <= ns_exp[-1]:
i = bisect(ns_exp, -mag-1)
if i < len(ns_exp):
wp0 = precs_exp[i]
if wp > wp0:
use_newton = True
if not use_newton:
r = int(0.7 * wp**0.5)
if mag < 0:
r = max(1, r + mag)
wp += r + 20
t = to_fixed(x, wp)
if mag > 1:
lg2 = ln2_fixed(wp)
n, t = divmod(t, lg2)
else:
n = 0
man = exp_series2(t, wp, r)
else:
# if x is very small or very large use
# exp(x + m) = exp(x) * e**m
if mag > LIM_MAG:
wp += mag*10 + 100
n = int(mag * math.log(2)) + 1
x = mpf_sub(x, from_int(n, wp), wp)
elif mag <= 0:
wp += -mag*10 + 100
if mag < 0:
n = int(-mag * math.log(2)) + 1
x = mpf_add(x, from_int(n, wp), wp)
res = exp_newton(x, wp)
sign, man, exp, bc = res
if mag < 0:
t = mpf_pow_int(mpf_e(wp), n, wp)
res = mpf_div(res, t, wp)
sign, man, exp, bc = res
if mag > LIM_MAG:
t = mpf_pow_int(mpf_e(wp), n, wp)
res = mpf_mul(res, t, wp)
sign, man, exp, bc = res
return normalize(sign, man, exp, bc, prec, rnd)
bc = bitcount(man)
return normalize(0, man, int(-wp+n), bc, prec, rnd)
def f(prec, rnd=round_fast):
wp = prec + 20
v = fixed(wp)
if rnd in (round_up, round_ceiling):
v += 1
return normalize(0, v, -wp, bitcount(v), prec, rnd)
def mpf_log(x, prec, rnd=round_fast):
"""
Compute the natural logarithm of the mpf value x. If x is negative,
ComplexResult is raised.
"""
sign, man, exp, bc = x
#------------------------------------------------------------------
# Handle special values
if not man:
if x == fzero: return fninf
if x == finf: return finf
if x == fnan: return fnan
if sign:
raise ComplexResult("logarithm of a negative number")
wp = prec + 20
#------------------------------------------------------------------
# Handle log(2^n) = log(n)*2.
# Here we catch the only possible exact value, log(1) = 0
if man == 1:
if not exp:
return fzero
return from_man_exp(exp*ln2_fixed(wp), -wp, prec, rnd)
mag = exp+bc
abs_mag = abs(mag)
#------------------------------------------------------------------
# Handle x = 1+eps, where log(x) ~ x. We need to check for
# cancellation when moving to fixed-point math and compensate
# by increasing the precision. Note that abs_mag in (0, 1) <=>
# 0.5 < x < 2 and x != 1
if abs_mag <= 1:
# Calculate t = x-1 to measure distance from 1 in bits
tsign = 1-abs_mag
if tsign:
tman = (MP_ONE<<bc) - man
else:
tman = man - (MP_ONE<<(bc-1))
tbc = bitcount(tman)
cancellation = bc - tbc
if cancellation > wp:
t = normalize(tsign, tman, abs_mag-bc, tbc, tbc, 'n')
return mpf_perturb(t, tsign, prec, rnd)
else:
wp += cancellation
# TODO: if close enough to 1, we could use Taylor series
# even in the AGM precision range, since the Taylor series
# converges rapidly
#------------------------------------------------------------------
# Another special case:
# n*log(2) is a good enough approximation
if abs_mag > 10000:
if bitcount(abs_mag) > wp:
return from_man_exp(exp*ln2_fixed(wp), -wp, prec, rnd)
#------------------------------------------------------------------
# General case.
# Perform argument reduction using log(x) = log(x*2^n) - n*log(2):
# If we are in the Taylor precision range, choose magnitude 0 or 1.
# If we are in the AGM precision range, choose magnitude -m for
# some large m; benchmarking on one machine showed m = prec/20 to be
# optimal between 1000 and 100,000 digits.
if wp <= LOG_TAYLOR_PREC:
m = log_taylor_cached(lshift(man, wp-bc), wp)
if mag:
m += mag*ln2_fixed(wp)
else:
optimal_mag = -wp//LOG_AGM_MAG_PREC_RATIO
n = optimal_mag - mag
x = mpf_shift(x, n)
wp += (-optimal_mag)
m = -log_agm(to_fixed(x, wp), wp)
m -= n*ln2_fixed(wp)
return from_man_exp(m, -wp, prec, rnd)
def mpci_gamma(z, prec, type=0):
(a1, a2), (b1, b2) = z
# Real case
if b1 == b2 == fzero and (type != 3 or mpf_gt(a1, fzero)):
return mpi_gamma(z, prec, type), mpi_zero
# Estimate precision
wp = prec + 20
if type != 3:
amag = a2[2] + a2[3]
bmag = b2[2] + b2[3]
if a2 != fzero:
mag = max(amag, bmag)
else:
mag = bmag
an = abs(to_int(a2))
bn = abs(to_int(b2))
absn = max(an, bn)
gamma_size = max(0, absn * mag)
wp += bitcount(gamma_size)
# Assume type != 1
if type == 1:
(a1, a2) = mpi_add((a1, a2), mpi_one, wp)
z = (a1, a2), (b1, b2)
type = 0
# Avoid non-monotonic region near the negative real axis
if mpf_lt(a1, gamma_min_b):
if mpi_overlap((b1, b2), (gamma_mono_imag_a, gamma_mono_imag_b)):
# TODO: reflection formula
#if mpf_lt(a2, mpf_shift(fone,-1)):
# znew = mpci_sub((mpi_one,mpi_zero),z,wp)
# ...
# Recurrence:
# gamma(z) = gamma(z+1)/z
znew = mpi_add((a1, a2), mpi_one, wp), (b1, b2)
if type == 0:
return mpci_div(mpci_gamma(znew, prec + 2, 0), z, prec)
if type == 2:
return mpci_mul(mpci_gamma(znew, prec + 2, 2), z, prec)
if type == 3:
return mpci_sub(mpci_gamma(znew, prec + 2, 3),
mpci_log(z, prec + 2), prec)
# Use monotonicity (except for a small region close to the
# origin and near poles)
# upper half-plane
if mpf_ge(b1, fzero):
minre = mpc_loggamma((a1, b2), wp, round_floor)
maxre = mpc_loggamma((a2, b1), wp, round_ceiling)
minim = mpc_loggamma((a1, b1), wp, round_floor)
maxim = mpc_loggamma((a2, b2), wp, round_ceiling)
# lower half-plane
elif mpf_le(b2, fzero):
minre = mpc_loggamma((a1, b1), wp, round_floor)
maxre = mpc_loggamma((a2, b2), wp, round_ceiling)
minim = mpc_loggamma((a2, b1), wp, round_floor)
maxim = mpc_loggamma((a1, b2), wp, round_ceiling)
# crosses real axis
else:
maxre = mpc_loggamma((a2, fzero), wp, round_ceiling)
# stretches more into the lower half-plane
if mpf_gt(mpf_neg(b1), b2):
minre = mpc_loggamma((a1, b1), wp, round_ceiling)
else:
minre = mpc_loggamma((a1, b2), wp, round_ceiling)
minim = mpc_loggamma((a2, b1), wp, round_floor)
maxim = mpc_loggamma((a2, b2), wp, round_floor)
w = (minre[0], maxre[0]), (minim[1], maxim[1])
if type == 3:
return mpi_pos(w[0], prec), mpi_pos(w[1], prec)
if type == 2:
w = mpci_neg(w)
return mpci_exp(w, prec)
def calc_cos_sin(which, y, swaps, prec, cos_rnd, sin_rnd):
"""
Simultaneous computation of cos and sin (internal function).
"""
y, wp = y
swap_cos_sin, cos_sign, sin_sign = swaps
if swap_cos_sin:
which_compute = -which
else:
which_compute = which
# XXX: assumes no swaps
if not y:
return fone, fzero
# Tiny nonzero argument
if wp > prec*2 + 30:
y = from_man_exp(y, -wp)
if swap_cos_sin:
cos_rnd, sin_rnd = sin_rnd, cos_rnd
cos_sign, sin_sign = sin_sign, cos_sign
if cos_sign: cos = mpf_perturb(fnone, 0, prec, cos_rnd)
else: cos = mpf_perturb(fone, 1, prec, cos_rnd)
if sin_sign: sin = mpf_perturb(mpf_neg(y), 0, prec, sin_rnd)
else: sin = mpf_perturb(y, 1, prec, sin_rnd)
if swap_cos_sin:
cos, sin = sin, cos
return cos, sin
# Use standard Taylor series
if prec < 600:
if which_compute == 0:
sin = sin_taylor(y, wp)
# only need to evaluate one of the series
cos = isqrt_fast((MP_ONE<<(2*wp)) - sin*sin)
elif which_compute == 1:
sin = 0
cos = cos_taylor(y, wp)
elif which_compute == -1:
sin = sin_taylor(y, wp)
cos = 0
# Use exp(i*x) with Brent's trick
else:
r = int(0.137 * prec**0.579)
ep = r+20
cos, sin = expi_series(y<<ep, wp+ep, r)
cos >>= ep
sin >>= ep
if swap_cos_sin:
cos, sin = sin, cos
if cos_rnd is not round_nearest:
# Round and set correct signs
# XXX: this logic needs a second look
ONE = MP_ONE << wp
if cos_sign:
cos += (-1)**(cos_rnd in (round_ceiling, round_down))
cos = min(ONE, cos)
else:
cos += (-1)**(cos_rnd in (round_ceiling, round_up))
cos = min(ONE, cos)
if sin_sign:
sin += (-1)**(sin_rnd in (round_ceiling, round_down))
sin = min(ONE, sin)
else:
sin += (-1)**(sin_rnd in (round_ceiling, round_up))
sin = min(ONE, sin)
if which != -1:
cos = normalize(cos_sign, cos, -wp, bitcount(cos), prec, cos_rnd)
if which != 1:
sin = normalize(sin_sign, sin, -wp, bitcount(sin), prec, sin_rnd)
return cos, sin
def exponential_series(x, prec, type=0):
"""
Taylor series for cosh/sinh or cos/sin.
type = 0 -- returns exp(x) (slightly faster than cosh+sinh)
type = 1 -- returns (cosh(x), sinh(x))
type = 2 -- returns (cos(x), sin(x))
"""
if x < 0:
x = -x
sign = 1
else:
sign = 0
r = int(0.5 * prec ** 0.5)
xmag = bitcount(x) - prec
r = max(0, xmag + r)
extra = 10 + 2 * max(r, -xmag)
wp = prec + extra
x <<= extra - r
one = MPZ_ONE << wp
alt = type == 2
if prec < EXP_SERIES_U_CUTOFF:
x2 = a = (x * x) >> wp
x4 = (x2 * x2) >> wp
s0 = s1 = MPZ_ZERO
k = 2
while a:
a //= (k - 1) * k
s0 += a
k += 2
a //= (k - 1) * k
s1 += a
k += 2
a = (a * x4) >> wp
s1 = (x2 * s1) >> wp
if alt:
c = s1 - s0 + one
else:
c = s1 + s0 + one
else:
u = int(0.3 * prec ** 0.35)
x2 = a = (x * x) >> wp
xpowers = [one, x2]
for i in xrange(1, u):
xpowers.append((xpowers[-1] * x2) >> wp)
sums = [MPZ_ZERO] * u
k = 2
while a:
for i in xrange(u):
a //= (k - 1) * k
if alt and k & 2:
sums[i] -= a
else:
sums[i] += a
k += 2
a = (a * xpowers[-1]) >> wp
for i in xrange(1, u):
sums[i] = (sums[i] * xpowers[i]) >> wp
c = sum(sums) + one
if type == 0:
s = isqrt_fast(c * c - (one << wp))
if sign:
v = c - s
else:
v = c + s
for i in xrange(r):
v = (v * v) >> wp
return v >> extra
else:
# Repeatedly apply the double-angle formula
# cosh(2*x) = 2*cosh(x)^2 - 1
# cos(2*x) = 2*cos(x)^2 - 1
pshift = wp - 1
for i in xrange(r):
c = ((c * c) >> pshift) - one
# With the abs, this is the same for sinh and sin
s = isqrt_fast(abs((one << wp) - c * c))
if sign:
s = -s
return (c >> extra), (s >> extra)
def mpf_expint(n, x, prec, rnd=round_fast, gamma=False):
"""
E_n(x), n an integer, x real
With gamma=True, computes Gamma(n,x) (upper incomplete gamma function)
Returns (real, None) if real, otherwise (real, imag)
The imaginary part is an optional branch cut term
"""
sign, man, exp, bc = x
if not man:
if gamma:
if x == fzero:
# Actually gamma function pole
if n <= 0:
return finf
return mpf_gamma_int(n, prec, rnd)
if x == finf:
return fzero, None
# TODO: could return finite imaginary value at -inf
return fnan, fnan
else:
if x == fzero:
if n > 1:
return from_rational(1, n - 1, prec, rnd), None
else:
return finf, None
if x == finf:
return fzero, None
return fnan, fnan
n_orig = n
if gamma:
n = 1 - n
wp = prec + 20
xmag = exp + bc
# Beware of near-poles
if xmag < -10:
raise NotImplementedError
nmag = bitcount(abs(n))
have_imag = n > 0 and sign
negx = mpf_neg(x)
# Skip series if direct convergence
if n == 0 or 2 * nmag - xmag < -wp:
if gamma:
v = mpf_exp(negx, wp)
re = mpf_mul(v, mpf_pow_int(x, n_orig - 1, wp), prec, rnd)
else:
v = mpf_exp(negx, wp)
re = mpf_div(v, x, prec, rnd)
else:
# Finite number of terms, or...
can_use_asymptotic_series = -3 * wp < n <= 0
# ...large enough?
if not can_use_asymptotic_series:
xi = abs(to_int(x))
m = min(max(1, xi - n), 2 * wp)
siz = -n * nmag + (m + n) * bitcount(abs(m + n)) - m * xmag - (
144 * m // 100)
tol = -wp - 10
can_use_asymptotic_series = siz < tol
if can_use_asymptotic_series:
r = ((-MP_ONE) << (wp + wp)) // to_fixed(x, wp)
m = n
t = r * m
s = MP_ONE << wp
while m and t:
s += t
m += 1
t = (m * r * t) >> wp
v = mpf_exp(negx, wp)
if gamma:
# ~ exp(-x) * x^(n-1) * (1 + ...)
v = mpf_mul(v, mpf_pow_int(x, n_orig - 1, wp), wp)
else:
# ~ exp(-x)/x * (1 + ...)
v = mpf_div(v, x, wp)
re = mpf_mul(v, from_man_exp(s, -wp), prec, rnd)
elif n == 1:
re = mpf_neg(mpf_ei(negx, prec, rnd))
elif n > 0 and n < 3 * wp:
T1 = mpf_neg(mpf_ei(negx, wp))
if gamma:
if n_orig & 1:
T1 = mpf_neg(T1)
else:
T1 = mpf_mul(T1, mpf_pow_int(negx, n - 1, wp), wp)
r = t = to_fixed(x, wp)
facs = [1] * (n - 1)
for k in range(1, n - 1):
facs[k] = facs[k - 1] * k
facs = facs[::-1]
s = facs[0] << wp
for k in range(1, n - 1):
if k & 1:
s -= facs[k] * t
else:
s += facs[k] * t
t = (t * r) >> wp
T2 = from_man_exp(s, -wp, wp)
T2 = mpf_mul(T2, mpf_exp(negx, wp))
if gamma:
T2 = mpf_mul(T2, mpf_pow_int(x, n_orig, wp), wp)
R = mpf_add(T1, T2)
re = mpf_div(R, from_int(int_fac(n - 1)), prec, rnd)
else:
raise NotImplementedError
if have_imag:
M = from_int(-int_fac(n - 1))
if gamma:
im = mpf_div(mpf_pi(wp), M, prec, rnd)
else:
im = mpf_div(
mpf_mul(mpf_pi(wp), mpf_pow_int(negx, n_orig - 1, wp), wp), M,
prec, rnd)
return re, im
else:
return re, None
def mpci_gamma(z, prec, type=0):
(a1,a2), (b1,b2) = z
# Real case
if b1 == b2 == fzero and (type != 3 or mpf_gt(a1,fzero)):
return mpi_gamma(z, prec, type), mpi_zero
# Estimate precision
wp = prec+20
if type != 3:
amag = a2[2]+a2[3]
bmag = b2[2]+b2[3]
if a2 != fzero:
mag = max(amag, bmag)
else:
mag = bmag
an = abs(to_int(a2))
bn = abs(to_int(b2))
absn = max(an, bn)
gamma_size = max(0,absn*mag)
wp += bitcount(gamma_size)
# Assume type != 1
if type == 1:
(a1,a2) = mpi_add((a1,a2), mpi_one, wp); z = (a1,a2), (b1,b2)
type = 0
# Avoid non-monotonic region near the negative real axis
if mpf_lt(a1, gamma_min_b):
if mpi_overlap((b1,b2), (gamma_mono_imag_a, gamma_mono_imag_b)):
# TODO: reflection formula
#if mpf_lt(a2, mpf_shift(fone,-1)):
# znew = mpci_sub((mpi_one,mpi_zero),z,wp)
# ...
# Recurrence:
# gamma(z) = gamma(z+1)/z
znew = mpi_add((a1,a2), mpi_one, wp), (b1,b2)
if type == 0: return mpci_div(mpci_gamma(znew, prec+2, 0), z, prec)
if type == 2: return mpci_mul(mpci_gamma(znew, prec+2, 2), z, prec)
if type == 3: return mpci_sub(mpci_gamma(znew, prec+2, 3), mpci_log(z,prec+2), prec)
# Use monotonicity (except for a small region close to the
# origin and near poles)
# upper half-plane
if mpf_ge(b1, fzero):
minre = mpc_loggamma((a1,b2), wp, round_floor)
maxre = mpc_loggamma((a2,b1), wp, round_ceiling)
minim = mpc_loggamma((a1,b1), wp, round_floor)
maxim = mpc_loggamma((a2,b2), wp, round_ceiling)
# lower half-plane
elif mpf_le(b2, fzero):
minre = mpc_loggamma((a1,b1), wp, round_floor)
maxre = mpc_loggamma((a2,b2), wp, round_ceiling)
minim = mpc_loggamma((a2,b1), wp, round_floor)
maxim = mpc_loggamma((a1,b2), wp, round_ceiling)
# crosses real axis
else:
maxre = mpc_loggamma((a2,fzero), wp, round_ceiling)
# stretches more into the lower half-plane
if mpf_gt(mpf_neg(b1), b2):
minre = mpc_loggamma((a1,b1), wp, round_ceiling)
else:
minre = mpc_loggamma((a1,b2), wp, round_ceiling)
minim = mpc_loggamma((a2,b1), wp, round_floor)
maxim = mpc_loggamma((a2,b2), wp, round_floor)
w = (minre[0], maxre[0]), (minim[1], maxim[1])
if type == 3:
return mpi_pos(w[0], prec), mpi_pos(w[1], prec)
if type == 2:
w = mpci_neg(w)
return mpci_exp(w, prec)
def mpf_cos_sin(x, prec, rnd=round_fast, which=0, pi=False):
"""
which:
0 -- return cos(x), sin(x)
1 -- return cos(x)
2 -- return sin(x)
3 -- return tan(x)
if pi=True, compute for pi*x
"""
sign, man, exp, bc = x
if not man:
if exp:
c, s = fnan, fnan
else:
c, s = fone, fzero
if which == 0: return c, s
if which == 1: return c
if which == 2: return s
if which == 3: return s
mag = bc + exp
wp = prec + 10
# Extremely small?
if mag < 0:
if mag < -wp:
if pi:
x = mpf_mul(x, mpf_pi(wp))
c = mpf_perturb(fone, 1, prec, rnd)
s = mpf_perturb(x, 1 - sign, prec, rnd)
if which == 0: return c, s
if which == 1: return c
if which == 2: return s
if which == 3: return mpf_perturb(x, sign, prec, rnd)
if pi:
if exp >= -1:
if exp == -1:
c = fzero
s = (fone, fnone)[bool(man & 2) ^ sign]
elif exp == 0:
c, s = (fnone, fzero)
else:
c, s = (fone, fzero)
if which == 0: return c, s
if which == 1: return c
if which == 2: return s
if which == 3: return mpf_div(s, c, prec, rnd)
# Subtract nearest half-integer (= mod by pi/2)
n = ((man >> (-exp - 2)) + 1) >> 1
man = man - (n << (-exp - 1))
mag2 = bitcount(man) + exp
wp = prec + 10 - mag2
offset = exp + wp
if offset >= 0:
t = man << offset
else:
t = man >> (-offset)
t = (t * pi_fixed(wp)) >> wp
else:
t, n, wp = mod_pi2(man, exp, mag, wp)
c, s = cos_sin_basecase(t, wp)
m = n & 3
if m == 1: c, s = -s, c
elif m == 2: c, s = -c, -s
elif m == 3: c, s = s, -c
if sign:
s = -s
if which == 0:
c = from_man_exp(c, -wp, prec, rnd)
s = from_man_exp(s, -wp, prec, rnd)
return c, s
if which == 1:
return from_man_exp(c, -wp, prec, rnd)
if which == 2:
return from_man_exp(s, -wp, prec, rnd)
if which == 3:
return from_rational(s, c, prec, rnd)
def mpf_cos_sin(x, prec, rnd=round_fast, which=0, pi=False):
"""
which:
0 -- return cos(x), sin(x)
1 -- return cos(x)
2 -- return sin(x)
3 -- return tan(x)
if pi=True, compute for pi*x
"""
sign, man, exp, bc = x
if not man:
if exp:
c, s = fnan, fnan
else:
c, s = fone, fzero
if which == 0:
return c, s
if which == 1:
return c
if which == 2:
return s
if which == 3:
return s
mag = bc + exp
wp = prec + 10
# Extremely small?
if mag < 0:
if mag < -wp:
if pi:
x = mpf_mul(x, mpf_pi(wp))
c = mpf_perturb(fone, 1, prec, rnd)
s = mpf_perturb(x, 1 - sign, prec, rnd)
if which == 0:
return c, s
if which == 1:
return c
if which == 2:
return s
if which == 3:
return mpf_perturb(x, sign, prec, rnd)
if pi:
if exp >= -1:
if exp == -1:
c = fzero
s = (fone, fnone)[bool(man & 2) ^ sign]
elif exp == 0:
c, s = (fnone, fzero)
else:
c, s = (fone, fzero)
if which == 0:
return c, s
if which == 1:
return c
if which == 2:
return s
if which == 3:
return mpf_div(s, c, prec, rnd)
# Subtract nearest half-integer (= mod by pi/2)
n = ((man >> (-exp - 2)) + 1) >> 1
man = man - (n << (-exp - 1))
mag2 = bitcount(man) + exp
wp = prec + 10 - mag2
offset = exp + wp
if offset >= 0:
t = man << offset
else:
t = man >> (-offset)
t = (t * pi_fixed(wp)) >> wp
else:
t, n, wp = mod_pi2(man, exp, mag, wp)
c, s = cos_sin_basecase(t, wp)
m = n & 3
if m == 1:
c, s = -s, c
elif m == 2:
c, s = -c, -s
elif m == 3:
c, s = s, -c
if sign:
s = -s
if which == 0:
c = from_man_exp(c, -wp, prec, rnd)
s = from_man_exp(s, -wp, prec, rnd)
return c, s
if which == 1:
return from_man_exp(c, -wp, prec, rnd)
if which == 2:
return from_man_exp(s, -wp, prec, rnd)
if which == 3:
return from_rational(s, c, prec, rnd)
def mpf_expint(n, x, prec, rnd=round_fast, gamma=False):
"""
E_n(x), n an integer, x real
With gamma=True, computes Gamma(n,x) (upper incomplete gamma function)
Returns (real, None) if real, otherwise (real, imag)
The imaginary part is an optional branch cut term
"""
sign, man, exp, bc = x
if not man:
if gamma:
if x == fzero:
# Actually gamma function pole
if n <= 0:
return finf, None
return mpf_gamma_int(n, prec, rnd), None
if x == finf:
return fzero, None
# TODO: could return finite imaginary value at -inf
return fnan, fnan
else:
if x == fzero:
if n > 1:
return from_rational(1, n-1, prec, rnd), None
else:
return finf, None
if x == finf:
return fzero, None
return fnan, fnan
n_orig = n
if gamma:
n = 1-n
wp = prec + 20
xmag = exp + bc
# Beware of near-poles
if xmag < -10:
raise NotImplementedError
nmag = bitcount(abs(n))
have_imag = n > 0 and sign
negx = mpf_neg(x)
# Skip series if direct convergence
if n == 0 or 2*nmag - xmag < -wp:
if gamma:
v = mpf_exp(negx, wp)
re = mpf_mul(v, mpf_pow_int(x, n_orig-1, wp), prec, rnd)
else:
v = mpf_exp(negx, wp)
re = mpf_div(v, x, prec, rnd)
else:
# Finite number of terms, or...
can_use_asymptotic_series = -3*wp < n <= 0
# ...large enough?
if not can_use_asymptotic_series:
xi = abs(to_int(x))
m = min(max(1, xi-n), 2*wp)
siz = -n*nmag + (m+n)*bitcount(abs(m+n)) - m*xmag - (144*m//100)
tol = -wp-10
can_use_asymptotic_series = siz < tol
if can_use_asymptotic_series:
r = ((-MPZ_ONE) << (wp+wp)) // to_fixed(x, wp)
m = n
t = r*m
s = MPZ_ONE << wp
while m and t:
s += t
m += 1
t = (m*r*t) >> wp
v = mpf_exp(negx, wp)
if gamma:
# ~ exp(-x) * x^(n-1) * (1 + ...)
v = mpf_mul(v, mpf_pow_int(x, n_orig-1, wp), wp)
else:
# ~ exp(-x)/x * (1 + ...)
v = mpf_div(v, x, wp)
re = mpf_mul(v, from_man_exp(s, -wp), prec, rnd)
elif n == 1:
re = mpf_neg(mpf_ei(negx, prec, rnd))
elif n > 0 and n < 3*wp:
T1 = mpf_neg(mpf_ei(negx, wp))
if gamma:
if n_orig & 1:
T1 = mpf_neg(T1)
else:
T1 = mpf_mul(T1, mpf_pow_int(negx, n-1, wp), wp)
r = t = to_fixed(x, wp)
facs = [1] * (n-1)
for k in range(1,n-1):
facs[k] = facs[k-1] * k
facs = facs[::-1]
s = facs[0] << wp
for k in range(1, n-1):
if k & 1:
s -= facs[k] * t
else:
s += facs[k] * t
t = (t*r) >> wp
T2 = from_man_exp(s, -wp, wp)
T2 = mpf_mul(T2, mpf_exp(negx, wp))
if gamma:
T2 = mpf_mul(T2, mpf_pow_int(x, n_orig, wp), wp)
R = mpf_add(T1, T2)
re = mpf_div(R, from_int(ifac(n-1)), prec, rnd)
else:
raise NotImplementedError
if have_imag:
M = from_int(-ifac(n-1))
if gamma:
im = mpf_div(mpf_pi(wp), M, prec, rnd)
else:
im = mpf_div(mpf_mul(mpf_pi(wp), mpf_pow_int(negx, n_orig-1, wp), wp), M, prec, rnd)
return re, im
else:
return re, None
def f(prec, rnd=round_fast):
wp = prec + 20
v = fixed(wp)
if rnd in (round_up, round_ceiling):
v += 1
return normalize(0, v, -wp, bitcount(v), prec, rnd)
return cache[n]
# Use Taylor series with caching up to this prec
LOG_TAYLOR_PREC = 2500
# Cache log values in steps of size 2^-N
LOG_TAYLOR_SHIFT = 9
# prec/size ratio of x for fastest convergence in AGM formula
LOG_AGM_MAG_PREC_RATIO = 20
log_taylor_cache = {}
# ~= next power of two + 20
cache_prec_steps = [22,22]
for k in xrange(1, bitcount(LOG_TAYLOR_PREC)+1):
cache_prec_steps += [min(2**k,LOG_TAYLOR_PREC)+20] * 2**(k-1)
def agm_fixed(a, b, prec):
"""
Fixed-point computation of agm(a,b), assuming
a, b both close to unit magnitude.
"""
i = 0
while 1:
anew = (a+b)>>1
if i > 4 and abs(a-anew) < 8:
return a
b = isqrt_fast(a*b)
a = anew
i += 1
def calc_cos_sin(which, y, swaps, prec, cos_rnd, sin_rnd):
"""
Simultaneous computation of cos and sin (internal function).
"""
y, wp = y
swap_cos_sin, cos_sign, sin_sign = swaps
if swap_cos_sin:
which_compute = -which
else:
which_compute = which
# XXX: assumes no swaps
if not y:
return fone, fzero
# Tiny nonzero argument
if wp > prec*2 + 30:
y = from_man_exp(y, -wp)
if swap_cos_sin:
cos_rnd, sin_rnd = sin_rnd, cos_rnd
cos_sign, sin_sign = sin_sign, cos_sign
if cos_sign: cos = mpf_perturb(fnone, 0, prec, cos_rnd)
else: cos = mpf_perturb(fone, 1, prec, cos_rnd)
if sin_sign: sin = mpf_perturb(mpf_neg(y), 0, prec, sin_rnd)
else: sin = mpf_perturb(y, 1, prec, sin_rnd)
if swap_cos_sin:
cos, sin = sin, cos
return cos, sin
# Use standard Taylor series
if prec < 600:
if which_compute == 0:
sin = sin_taylor(y, wp)
# only need to evaluate one of the series
cos = sqrt_fixed((1<<wp) - ((sin*sin)>>wp), wp)
elif which_compute == 1:
sin = 0
cos = cos_taylor(y, wp)
elif which_compute == -1:
sin = sin_taylor(y, wp)
cos = 0
# Use exp(i*x) with Brent's trick
else:
r = int(0.137 * prec**0.579)
ep = r+20
cos, sin = expi_series(y<<ep, wp+ep, r)
cos >>= ep
sin >>= ep
if swap_cos_sin:
cos, sin = sin, cos
if cos_rnd is not round_nearest:
# Round and set correct signs
# XXX: this logic needs a second look
ONE = MP_ONE << wp
if cos_sign:
cos += (-1)**(cos_rnd in (round_ceiling, round_down))
cos = min(ONE, cos)
else:
cos += (-1)**(cos_rnd in (round_ceiling, round_up))
cos = min(ONE, cos)
if sin_sign:
sin += (-1)**(sin_rnd in (round_ceiling, round_down))
sin = min(ONE, sin)
else:
sin += (-1)**(sin_rnd in (round_ceiling, round_up))
sin = min(ONE, sin)
if which != -1:
cos = normalize(cos_sign, cos, -wp, bitcount(cos), prec, cos_rnd)
if which != 1:
sin = normalize(sin_sign, sin, -wp, bitcount(sin), prec, sin_rnd)
return cos, sin
def exponential_series(x, prec, type=0):
"""
Taylor series for cosh/sinh or cos/sin.
type = 0 -- returns exp(x) (slightly faster than cosh+sinh)
type = 1 -- returns (cosh(x), sinh(x))
type = 2 -- returns (cos(x), sin(x))
"""
if x < 0:
x = -x
sign = 1
else:
sign = 0
r = int(0.5 * prec**0.5)
xmag = bitcount(x) - prec
r = max(0, xmag + r)
extra = 10 + 2 * max(r, -xmag)
wp = prec + extra
x <<= (extra - r)
one = MPZ_ONE << wp
alt = (type == 2)
if prec < EXP_SERIES_U_CUTOFF:
x2 = a = (x * x) >> wp
x4 = (x2 * x2) >> wp
s0 = s1 = MPZ_ZERO
k = 2
while a:
a //= (k - 1) * k
s0 += a
k += 2
a //= (k - 1) * k
s1 += a
k += 2
a = (a * x4) >> wp
s1 = (x2 * s1) >> wp
if alt:
c = s1 - s0 + one
else:
c = s1 + s0 + one
else:
u = int(0.3 * prec**0.35)
x2 = a = (x * x) >> wp
xpowers = [one, x2]
for i in xrange(1, u):
xpowers.append((xpowers[-1] * x2) >> wp)
sums = [MPZ_ZERO] * u
k = 2
while a:
for i in xrange(u):
a //= (k - 1) * k
if alt and k & 2: sums[i] -= a
else: sums[i] += a
k += 2
a = (a * xpowers[-1]) >> wp
for i in xrange(1, u):
sums[i] = (sums[i] * xpowers[i]) >> wp
c = sum(sums) + one
if type == 0:
s = isqrt_fast(c * c - (one << wp))
if sign:
v = c - s
else:
v = c + s
for i in xrange(r):
v = (v * v) >> wp
return v >> extra
else:
# Repeatedly apply the double-angle formula
# cosh(2*x) = 2*cosh(x)^2 - 1
# cos(2*x) = 2*cos(x)^2 - 1
pshift = wp - 1
for i in xrange(r):
c = ((c * c) >> pshift) - one
# With the abs, this is the same for sinh and sin
s = isqrt_fast(abs((one << wp) - c * c))
if sign:
s = -s
return (c >> extra), (s >> extra)
