Traktor/myenv/Lib/site-packages/mpmath/libmp/libhyper.py
2024-05-26 05:12:46 +02:00

1151 lines
36 KiB
Python

"""
This module implements computation of hypergeometric and related
functions. In particular, it provides code for generic summation
of hypergeometric series. Optimized versions for various special
cases are also provided.
"""
import operator
import math
from .backend import MPZ_ZERO, MPZ_ONE, BACKEND, xrange, exec_
from .libintmath import gcd
from .libmpf import (\
ComplexResult, round_fast, round_nearest,
negative_rnd, bitcount, to_fixed, from_man_exp, from_int, to_int,
from_rational,
fzero, fone, fnone, ftwo, finf, fninf, fnan,
mpf_sign, mpf_add, mpf_abs, mpf_pos,
mpf_cmp, mpf_lt, mpf_le, mpf_gt, mpf_min_max,
mpf_perturb, mpf_neg, mpf_shift, mpf_sub, mpf_mul, mpf_div,
sqrt_fixed, mpf_sqrt, mpf_rdiv_int, mpf_pow_int,
to_rational,
)
from .libelefun import (\
mpf_pi, mpf_exp, mpf_log, pi_fixed, mpf_cos_sin, mpf_cos, mpf_sin,
mpf_sqrt, agm_fixed,
)
from .libmpc import (\
mpc_one, mpc_sub, mpc_mul_mpf, mpc_mul, mpc_neg, complex_int_pow,
mpc_div, mpc_add_mpf, mpc_sub_mpf,
mpc_log, mpc_add, mpc_pos, mpc_shift,
mpc_is_infnan, mpc_zero, mpc_sqrt, mpc_abs,
mpc_mpf_div, mpc_square, mpc_exp
)
from .libintmath import ifac
from .gammazeta import mpf_gamma_int, mpf_euler, euler_fixed
class NoConvergence(Exception):
pass
#-----------------------------------------------------------------------#
# #
# Generic hypergeometric series #
# #
#-----------------------------------------------------------------------#
"""
TODO:
1. proper mpq parsing
2. imaginary z special-cased (also: rational, integer?)
3. more clever handling of series that don't converge because of stupid
upwards rounding
4. checking for cancellation
"""
def make_hyp_summator(key):
"""
Returns a function that sums a generalized hypergeometric series,
for given parameter types (integer, rational, real, complex).
"""
p, q, param_types, ztype = key
pstring = "".join(param_types)
fname = "hypsum_%i_%i_%s_%s_%s" % (p, q, pstring[:p], pstring[p:], ztype)
#print "generating hypsum", fname
have_complex_param = 'C' in param_types
have_complex_arg = ztype == 'C'
have_complex = have_complex_param or have_complex_arg
source = []
add = source.append
aint = []
arat = []
bint = []
brat = []
areal = []
breal = []
acomplex = []
bcomplex = []
#add("wp = prec + 40")
add("MAX = kwargs.get('maxterms', wp*100)")
add("HIGH = MPZ_ONE<<epsshift")
add("LOW = -HIGH")
# Setup code
add("SRE = PRE = one = (MPZ_ONE << wp)")
if have_complex:
add("SIM = PIM = MPZ_ZERO")
if have_complex_arg:
add("xsign, xm, xe, xbc = z[0]")
add("if xsign: xm = -xm")
add("ysign, ym, ye, ybc = z[1]")
add("if ysign: ym = -ym")
else:
add("xsign, xm, xe, xbc = z")
add("if xsign: xm = -xm")
add("offset = xe + wp")
add("if offset >= 0:")
add(" ZRE = xm << offset")
add("else:")
add(" ZRE = xm >> (-offset)")
if have_complex_arg:
add("offset = ye + wp")
add("if offset >= 0:")
add(" ZIM = ym << offset")
add("else:")
add(" ZIM = ym >> (-offset)")
for i, flag in enumerate(param_types):
W = ["A", "B"][i >= p]
if flag == 'Z':
([aint,bint][i >= p]).append(i)
add("%sINT_%i = coeffs[%i]" % (W, i, i))
elif flag == 'Q':
([arat,brat][i >= p]).append(i)
add("%sP_%i, %sQ_%i = coeffs[%i]._mpq_" % (W, i, W, i, i))
elif flag == 'R':
([areal,breal][i >= p]).append(i)
add("xsign, xm, xe, xbc = coeffs[%i]._mpf_" % i)
add("if xsign: xm = -xm")
add("offset = xe + wp")
add("if offset >= 0:")
add(" %sREAL_%i = xm << offset" % (W, i))
add("else:")
add(" %sREAL_%i = xm >> (-offset)" % (W, i))
elif flag == 'C':
([acomplex,bcomplex][i >= p]).append(i)
add("__re, __im = coeffs[%i]._mpc_" % i)
add("xsign, xm, xe, xbc = __re")
add("if xsign: xm = -xm")
add("ysign, ym, ye, ybc = __im")
add("if ysign: ym = -ym")
add("offset = xe + wp")
add("if offset >= 0:")
add(" %sCRE_%i = xm << offset" % (W, i))
add("else:")
add(" %sCRE_%i = xm >> (-offset)" % (W, i))
add("offset = ye + wp")
add("if offset >= 0:")
add(" %sCIM_%i = ym << offset" % (W, i))
add("else:")
add(" %sCIM_%i = ym >> (-offset)" % (W, i))
else:
raise ValueError
l_areal = len(areal)
l_breal = len(breal)
cancellable_real = min(l_areal, l_breal)
noncancellable_real_num = areal[cancellable_real:]
noncancellable_real_den = breal[cancellable_real:]
# LOOP
add("for n in xrange(1,10**8):")
add(" if n in magnitude_check:")
add(" p_mag = bitcount(abs(PRE))")
if have_complex:
add(" p_mag = max(p_mag, bitcount(abs(PIM)))")
add(" magnitude_check[n] = wp-p_mag")
# Real factors
multiplier = " * ".join(["AINT_#".replace("#", str(i)) for i in aint] + \
["AP_#".replace("#", str(i)) for i in arat] + \
["BQ_#".replace("#", str(i)) for i in brat])
divisor = " * ".join(["BINT_#".replace("#", str(i)) for i in bint] + \
["BP_#".replace("#", str(i)) for i in brat] + \
["AQ_#".replace("#", str(i)) for i in arat] + ["n"])
if multiplier:
add(" mul = " + multiplier)
add(" div = " + divisor)
# Check for singular terms
add(" if not div:")
if multiplier:
add(" if not mul:")
add(" break")
add(" raise ZeroDivisionError")
# Update product
if have_complex:
# TODO: when there are several real parameters and just a few complex
# (maybe just the complex argument), we only need to do about
# half as many ops if we accumulate the real factor in a single real variable
for k in range(cancellable_real): add(" PRE = PRE * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
for i in noncancellable_real_num: add(" PRE = (PRE * AREAL_#) >> wp".replace("#", str(i)))
for i in noncancellable_real_den: add(" PRE = (PRE << wp) // BREAL_#".replace("#", str(i)))
for k in range(cancellable_real): add(" PIM = PIM * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
for i in noncancellable_real_num: add(" PIM = (PIM * AREAL_#) >> wp".replace("#", str(i)))
for i in noncancellable_real_den: add(" PIM = (PIM << wp) // BREAL_#".replace("#", str(i)))
if multiplier:
if have_complex_arg:
add(" PRE, PIM = (mul*(PRE*ZRE-PIM*ZIM))//div, (mul*(PIM*ZRE+PRE*ZIM))//div")
add(" PRE >>= wp")
add(" PIM >>= wp")
else:
add(" PRE = ((mul * PRE * ZRE) >> wp) // div")
add(" PIM = ((mul * PIM * ZRE) >> wp) // div")
else:
if have_complex_arg:
add(" PRE, PIM = (PRE*ZRE-PIM*ZIM)//div, (PIM*ZRE+PRE*ZIM)//div")
add(" PRE >>= wp")
add(" PIM >>= wp")
else:
add(" PRE = ((PRE * ZRE) >> wp) // div")
add(" PIM = ((PIM * ZRE) >> wp) // div")
for i in acomplex:
add(" PRE, PIM = PRE*ACRE_#-PIM*ACIM_#, PIM*ACRE_#+PRE*ACIM_#".replace("#", str(i)))
add(" PRE >>= wp")
add(" PIM >>= wp")
for i in bcomplex:
add(" mag = BCRE_#*BCRE_#+BCIM_#*BCIM_#".replace("#", str(i)))
add(" re = PRE*BCRE_# + PIM*BCIM_#".replace("#", str(i)))
add(" im = PIM*BCRE_# - PRE*BCIM_#".replace("#", str(i)))
add(" PRE = (re << wp) // mag".replace("#", str(i)))
add(" PIM = (im << wp) // mag".replace("#", str(i)))
else:
for k in range(cancellable_real): add(" PRE = PRE * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
for i in noncancellable_real_num: add(" PRE = (PRE * AREAL_#) >> wp".replace("#", str(i)))
for i in noncancellable_real_den: add(" PRE = (PRE << wp) // BREAL_#".replace("#", str(i)))
if multiplier:
add(" PRE = ((PRE * mul * ZRE) >> wp) // div")
else:
add(" PRE = ((PRE * ZRE) >> wp) // div")
# Add product to sum
if have_complex:
add(" SRE += PRE")
add(" SIM += PIM")
add(" if (HIGH > PRE > LOW) and (HIGH > PIM > LOW):")
add(" break")
else:
add(" SRE += PRE")
add(" if HIGH > PRE > LOW:")
add(" break")
#add(" from mpmath import nprint, log, ldexp")
#add(" nprint([n, log(abs(PRE),2), ldexp(PRE,-wp)])")
add(" if n > MAX:")
add(" raise NoConvergence('Hypergeometric series converges too slowly. Try increasing maxterms.')")
# +1 all parameters for next loop
for i in aint: add(" AINT_# += 1".replace("#", str(i)))
for i in bint: add(" BINT_# += 1".replace("#", str(i)))
for i in arat: add(" AP_# += AQ_#".replace("#", str(i)))
for i in brat: add(" BP_# += BQ_#".replace("#", str(i)))
for i in areal: add(" AREAL_# += one".replace("#", str(i)))
for i in breal: add(" BREAL_# += one".replace("#", str(i)))
for i in acomplex: add(" ACRE_# += one".replace("#", str(i)))
for i in bcomplex: add(" BCRE_# += one".replace("#", str(i)))
if have_complex:
add("a = from_man_exp(SRE, -wp, prec, 'n')")
add("b = from_man_exp(SIM, -wp, prec, 'n')")
add("if SRE:")
add(" if SIM:")
add(" magn = max(a[2]+a[3], b[2]+b[3])")
add(" else:")
add(" magn = a[2]+a[3]")
add("elif SIM:")
add(" magn = b[2]+b[3]")
add("else:")
add(" magn = -wp+1")
add("return (a, b), True, magn")
else:
add("a = from_man_exp(SRE, -wp, prec, 'n')")
add("if SRE:")
add(" magn = a[2]+a[3]")
add("else:")
add(" magn = -wp+1")
add("return a, False, magn")
source = "\n".join((" " + line) for line in source)
source = ("def %s(coeffs, z, prec, wp, epsshift, magnitude_check, **kwargs):\n" % fname) + source
namespace = {}
exec_(source, globals(), namespace)
#print source
return source, namespace[fname]
if BACKEND == 'sage':
def make_hyp_summator(key):
"""
Returns a function that sums a generalized hypergeometric series,
for given parameter types (integer, rational, real, complex).
"""
from sage.libs.mpmath.ext_main import hypsum_internal
p, q, param_types, ztype = key
def _hypsum(coeffs, z, prec, wp, epsshift, magnitude_check, **kwargs):
return hypsum_internal(p, q, param_types, ztype, coeffs, z,
prec, wp, epsshift, magnitude_check, kwargs)
return "(none)", _hypsum
#-----------------------------------------------------------------------#
# #
# Error functions #
# #
#-----------------------------------------------------------------------#
# TODO: mpf_erf should call mpf_erfc when appropriate (currently
# only the converse delegation is implemented)
def mpf_erf(x, prec, rnd=round_fast):
sign, man, exp, bc = x
if not man:
if x == fzero: return fzero
if x == finf: return fone
if x== fninf: return fnone
return fnan
size = exp + bc
lg = math.log
# The approximation erf(x) = 1 is accurate to > x^2 * log(e,2) bits
if size > 3 and 2*(size-1) + 0.528766 > lg(prec,2):
if sign:
return mpf_perturb(fnone, 0, prec, rnd)
else:
return mpf_perturb(fone, 1, prec, rnd)
# erf(x) ~ 2*x/sqrt(pi) close to 0
if size < -prec:
# 2*x
x = mpf_shift(x,1)
c = mpf_sqrt(mpf_pi(prec+20), prec+20)
# TODO: interval rounding
return mpf_div(x, c, prec, rnd)
wp = prec + abs(size) + 25
# Taylor series for erf, fixed-point summation
t = abs(to_fixed(x, wp))
t2 = (t*t) >> wp
s, term, k = t, 12345, 1
while term:
t = ((t * t2) >> wp) // k
term = t // (2*k+1)
if k & 1:
s -= term
else:
s += term
k += 1
s = (s << (wp+1)) // sqrt_fixed(pi_fixed(wp), wp)
if sign:
s = -s
return from_man_exp(s, -wp, prec, rnd)
# If possible, we use the asymptotic series for erfc.
# This is an alternating divergent asymptotic series, so
# the error is at most equal to the first omitted term.
# Here we check if the smallest term is small enough
# for a given x and precision
def erfc_check_series(x, prec):
n = to_int(x)
if n**2 * 1.44 > prec:
return True
return False
def mpf_erfc(x, prec, rnd=round_fast):
sign, man, exp, bc = x
if not man:
if x == fzero: return fone
if x == finf: return fzero
if x == fninf: return ftwo
return fnan
wp = prec + 20
mag = bc+exp
# Preserve full accuracy when exponent grows huge
wp += max(0, 2*mag)
regular_erf = sign or mag < 2
if regular_erf or not erfc_check_series(x, wp):
if regular_erf:
return mpf_sub(fone, mpf_erf(x, prec+10, negative_rnd[rnd]), prec, rnd)
# 1-erf(x) ~ exp(-x^2), increase prec to deal with cancellation
n = to_int(x)+1
return mpf_sub(fone, mpf_erf(x, prec + int(n**2*1.44) + 10), prec, rnd)
s = term = MPZ_ONE << wp
term_prev = 0
t = (2 * to_fixed(x, wp) ** 2) >> wp
k = 1
while 1:
term = ((term * (2*k - 1)) << wp) // t
if k > 4 and term > term_prev or not term:
break
if k & 1:
s -= term
else:
s += term
term_prev = term
#print k, to_str(from_man_exp(term, -wp, 50), 10)
k += 1
s = (s << wp) // sqrt_fixed(pi_fixed(wp), wp)
s = from_man_exp(s, -wp, wp)
z = mpf_exp(mpf_neg(mpf_mul(x,x,wp),wp),wp)
y = mpf_div(mpf_mul(z, s, wp), x, prec, rnd)
return y
#-----------------------------------------------------------------------#
# #
# Exponential integrals #
# #
#-----------------------------------------------------------------------#
def ei_taylor(x, prec):
s = t = x
k = 2
while t:
t = ((t*x) >> prec) // k
s += t // k
k += 1
return s
def complex_ei_taylor(zre, zim, prec):
_abs = abs
sre = tre = zre
sim = tim = zim
k = 2
while _abs(tre) + _abs(tim) > 5:
tre, tim = ((tre*zre-tim*zim)//k)>>prec, ((tre*zim+tim*zre)//k)>>prec
sre += tre // k
sim += tim // k
k += 1
return sre, sim
def ei_asymptotic(x, prec):
one = MPZ_ONE << prec
x = t = ((one << prec) // x)
s = one + x
k = 2
while t:
t = (k*t*x) >> prec
s += t
k += 1
return s
def complex_ei_asymptotic(zre, zim, prec):
_abs = abs
one = MPZ_ONE << prec
M = (zim*zim + zre*zre) >> prec
# 1 / z
xre = tre = (zre << prec) // M
xim = tim = ((-zim) << prec) // M
sre = one + xre
sim = xim
k = 2
while _abs(tre) + _abs(tim) > 1000:
#print tre, tim
tre, tim = ((tre*xre-tim*xim)*k)>>prec, ((tre*xim+tim*xre)*k)>>prec
sre += tre
sim += tim
k += 1
if k > prec:
raise NoConvergence
return sre, sim
def mpf_ei(x, prec, rnd=round_fast, e1=False):
if e1:
x = mpf_neg(x)
sign, man, exp, bc = x
if e1 and not sign:
if x == fzero:
return finf
raise ComplexResult("E1(x) for x < 0")
if man:
xabs = 0, man, exp, bc
xmag = exp+bc
wp = prec + 20
can_use_asymp = xmag > wp
if not can_use_asymp:
if exp >= 0:
xabsint = man << exp
else:
xabsint = man >> (-exp)
can_use_asymp = xabsint > int(wp*0.693) + 10
if can_use_asymp:
if xmag > wp:
v = fone
else:
v = from_man_exp(ei_asymptotic(to_fixed(x, wp), wp), -wp)
v = mpf_mul(v, mpf_exp(x, wp), wp)
v = mpf_div(v, x, prec, rnd)
else:
wp += 2*int(to_int(xabs))
u = to_fixed(x, wp)
v = ei_taylor(u, wp) + euler_fixed(wp)
t1 = from_man_exp(v,-wp)
t2 = mpf_log(xabs,wp)
v = mpf_add(t1, t2, prec, rnd)
else:
if x == fzero: v = fninf
elif x == finf: v = finf
elif x == fninf: v = fzero
else: v = fnan
if e1:
v = mpf_neg(v)
return v
def mpc_ei(z, prec, rnd=round_fast, e1=False):
if e1:
z = mpc_neg(z)
a, b = z
asign, aman, aexp, abc = a
bsign, bman, bexp, bbc = b
if b == fzero:
if e1:
x = mpf_neg(mpf_ei(a, prec, rnd))
if not asign:
y = mpf_neg(mpf_pi(prec, rnd))
else:
y = fzero
return x, y
else:
return mpf_ei(a, prec, rnd), fzero
if a != fzero:
if not aman or not bman:
return (fnan, fnan)
wp = prec + 40
amag = aexp+abc
bmag = bexp+bbc
zmag = max(amag, bmag)
can_use_asymp = zmag > wp
if not can_use_asymp:
zabsint = abs(to_int(a)) + abs(to_int(b))
can_use_asymp = zabsint > int(wp*0.693) + 20
try:
if can_use_asymp:
if zmag > wp:
v = fone, fzero
else:
zre = to_fixed(a, wp)
zim = to_fixed(b, wp)
vre, vim = complex_ei_asymptotic(zre, zim, wp)
v = from_man_exp(vre, -wp), from_man_exp(vim, -wp)
v = mpc_mul(v, mpc_exp(z, wp), wp)
v = mpc_div(v, z, wp)
if e1:
v = mpc_neg(v, prec, rnd)
else:
x, y = v
if bsign:
v = mpf_pos(x, prec, rnd), mpf_sub(y, mpf_pi(wp), prec, rnd)
else:
v = mpf_pos(x, prec, rnd), mpf_add(y, mpf_pi(wp), prec, rnd)
return v
except NoConvergence:
pass
#wp += 2*max(0,zmag)
wp += 2*int(to_int(mpc_abs(z, 5)))
zre = to_fixed(a, wp)
zim = to_fixed(b, wp)
vre, vim = complex_ei_taylor(zre, zim, wp)
vre += euler_fixed(wp)
v = from_man_exp(vre,-wp), from_man_exp(vim,-wp)
if e1:
u = mpc_log(mpc_neg(z),wp)
else:
u = mpc_log(z,wp)
v = mpc_add(v, u, prec, rnd)
if e1:
v = mpc_neg(v)
return v
def mpf_e1(x, prec, rnd=round_fast):
return mpf_ei(x, prec, rnd, True)
def mpc_e1(x, prec, rnd=round_fast):
return mpc_ei(x, prec, rnd, True)
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)
if n_orig & 1:
im = mpf_neg(im)
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 mpf_ci_si_taylor(x, wp, which=0):
"""
0 - Ci(x) - (euler+log(x))
1 - Si(x)
"""
x = to_fixed(x, wp)
x2 = -(x*x) >> wp
if which == 0:
s, t, k = 0, (MPZ_ONE<<wp), 2
else:
s, t, k = x, x, 3
while t:
t = (t*x2//(k*(k-1)))>>wp
s += t//k
k += 2
return from_man_exp(s, -wp)
def mpc_ci_si_taylor(re, im, wp, which=0):
# The following code is only designed for small arguments,
# and not too small arguments (for relative accuracy)
if re[1]:
mag = re[2]+re[3]
elif im[1]:
mag = im[2]+im[3]
if im[1]:
mag = max(mag, im[2]+im[3])
if mag > 2 or mag < -wp:
raise NotImplementedError
wp += (2-mag)
zre = to_fixed(re, wp)
zim = to_fixed(im, wp)
z2re = (zim*zim-zre*zre)>>wp
z2im = (-2*zre*zim)>>wp
tre = zre
tim = zim
one = MPZ_ONE<<wp
if which == 0:
sre, sim, tre, tim, k = 0, 0, (MPZ_ONE<<wp), 0, 2
else:
sre, sim, tre, tim, k = zre, zim, zre, zim, 3
while max(abs(tre), abs(tim)) > 2:
f = k*(k-1)
tre, tim = ((tre*z2re-tim*z2im)//f)>>wp, ((tre*z2im+tim*z2re)//f)>>wp
sre += tre//k
sim += tim//k
k += 2
return from_man_exp(sre, -wp), from_man_exp(sim, -wp)
def mpf_ci_si(x, prec, rnd=round_fast, which=2):
"""
Calculation of Ci(x), Si(x) for real x.
which = 0 -- returns (Ci(x), -)
which = 1 -- returns (Si(x), -)
which = 2 -- returns (Ci(x), Si(x))
Note: if x < 0, Ci(x) needs an additional imaginary term, pi*i.
"""
wp = prec + 20
sign, man, exp, bc = x
ci, si = None, None
if not man:
if x == fzero:
return (fninf, fzero)
if x == fnan:
return (x, x)
ci = fzero
if which != 0:
if x == finf:
si = mpf_shift(mpf_pi(prec, rnd), -1)
if x == fninf:
si = mpf_neg(mpf_shift(mpf_pi(prec, negative_rnd[rnd]), -1))
return (ci, si)
# For small x: Ci(x) ~ euler + log(x), Si(x) ~ x
mag = exp+bc
if mag < -wp:
if which != 0:
si = mpf_perturb(x, 1-sign, prec, rnd)
if which != 1:
y = mpf_euler(wp)
xabs = mpf_abs(x)
ci = mpf_add(y, mpf_log(xabs, wp), prec, rnd)
return ci, si
# For huge x: Ci(x) ~ sin(x)/x, Si(x) ~ pi/2
elif mag > wp:
if which != 0:
if sign:
si = mpf_neg(mpf_pi(prec, negative_rnd[rnd]))
else:
si = mpf_pi(prec, rnd)
si = mpf_shift(si, -1)
if which != 1:
ci = mpf_div(mpf_sin(x, wp), x, prec, rnd)
return ci, si
else:
wp += abs(mag)
# Use an asymptotic series? The smallest value of n!/x^n
# occurs for n ~ x, where the magnitude is ~ exp(-x).
asymptotic = mag-1 > math.log(wp, 2)
# Case 1: convergent series near 0
if not asymptotic:
if which != 0:
si = mpf_pos(mpf_ci_si_taylor(x, wp, 1), prec, rnd)
if which != 1:
ci = mpf_ci_si_taylor(x, wp, 0)
ci = mpf_add(ci, mpf_euler(wp), wp)
ci = mpf_add(ci, mpf_log(mpf_abs(x), wp), prec, rnd)
return ci, si
x = mpf_abs(x)
# Case 2: asymptotic series for x >> 1
xf = to_fixed(x, wp)
xr = (MPZ_ONE<<(2*wp)) // xf # 1/x
s1 = (MPZ_ONE << wp)
s2 = xr
t = xr
k = 2
while t:
t = -t
t = (t*xr*k)>>wp
k += 1
s1 += t
t = (t*xr*k)>>wp
k += 1
s2 += t
s1 = from_man_exp(s1, -wp)
s2 = from_man_exp(s2, -wp)
s1 = mpf_div(s1, x, wp)
s2 = mpf_div(s2, x, wp)
cos, sin = mpf_cos_sin(x, wp)
# Ci(x) = sin(x)*s1-cos(x)*s2
# Si(x) = pi/2-cos(x)*s1-sin(x)*s2
if which != 0:
si = mpf_add(mpf_mul(cos, s1), mpf_mul(sin, s2), wp)
si = mpf_sub(mpf_shift(mpf_pi(wp), -1), si, wp)
if sign:
si = mpf_neg(si)
si = mpf_pos(si, prec, rnd)
if which != 1:
ci = mpf_sub(mpf_mul(sin, s1), mpf_mul(cos, s2), prec, rnd)
return ci, si
def mpf_ci(x, prec, rnd=round_fast):
if mpf_sign(x) < 0:
raise ComplexResult
return mpf_ci_si(x, prec, rnd, 0)[0]
def mpf_si(x, prec, rnd=round_fast):
return mpf_ci_si(x, prec, rnd, 1)[1]
def mpc_ci(z, prec, rnd=round_fast):
re, im = z
if im == fzero:
ci = mpf_ci_si(re, prec, rnd, 0)[0]
if mpf_sign(re) < 0:
return (ci, mpf_pi(prec, rnd))
return (ci, fzero)
wp = prec + 20
cre, cim = mpc_ci_si_taylor(re, im, wp, 0)
cre = mpf_add(cre, mpf_euler(wp), wp)
ci = mpc_add((cre, cim), mpc_log(z, wp), prec, rnd)
return ci
def mpc_si(z, prec, rnd=round_fast):
re, im = z
if im == fzero:
return (mpf_ci_si(re, prec, rnd, 1)[1], fzero)
wp = prec + 20
z = mpc_ci_si_taylor(re, im, wp, 1)
return mpc_pos(z, prec, rnd)
#-----------------------------------------------------------------------#
# #
# Bessel functions #
# #
#-----------------------------------------------------------------------#
# A Bessel function of the first kind of integer order, J_n(x), is
# given by the power series
# oo
# ___ k 2 k + n
# \ (-1) / x \
# J_n(x) = ) ----------- | - |
# /___ k! (k + n)! \ 2 /
# k = 0
# Simplifying the quotient between two successive terms gives the
# ratio x^2 / (-4*k*(k+n)). Hence, we only need one full-precision
# multiplication and one division by a small integer per term.
# The complex version is very similar, the only difference being
# that the multiplication is actually 4 multiplies.
# In the general case, we have
# J_v(x) = (x/2)**v / v! * 0F1(v+1, (-1/4)*z**2)
# TODO: for extremely large x, we could use an asymptotic
# trigonometric approximation.
# TODO: recompute at higher precision if the fixed-point mantissa
# is very small
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 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 mpf_agm(a, b, prec, rnd=round_fast):
"""
Computes the arithmetic-geometric mean agm(a,b) for
nonnegative mpf values a, b.
"""
asign, aman, aexp, abc = a
bsign, bman, bexp, bbc = b
if asign or bsign:
raise ComplexResult("agm of a negative number")
# Handle inf, nan or zero in either operand
if not (aman and bman):
if a == fnan or b == fnan:
return fnan
if a == finf:
if b == fzero:
return fnan
return finf
if b == finf:
if a == fzero:
return fnan
return finf
# agm(0,x) = agm(x,0) = 0
return fzero
wp = prec + 20
amag = aexp+abc
bmag = bexp+bbc
mag_delta = amag - bmag
# Reduce to roughly the same magnitude using floating-point AGM
abs_mag_delta = abs(mag_delta)
if abs_mag_delta > 10:
while abs_mag_delta > 10:
a, b = mpf_shift(mpf_add(a,b,wp),-1), \
mpf_sqrt(mpf_mul(a,b,wp),wp)
abs_mag_delta //= 2
asign, aman, aexp, abc = a
bsign, bman, bexp, bbc = b
amag = aexp+abc
bmag = bexp+bbc
mag_delta = amag - bmag
#print to_float(a), to_float(b)
# Use agm(a,b) = agm(x*a,x*b)/x to obtain a, b ~= 1
min_mag = min(amag,bmag)
max_mag = max(amag,bmag)
n = 0
# If too small, we lose precision when going to fixed-point
if min_mag < -8:
n = -min_mag
# If too large, we waste time using fixed-point with large numbers
elif max_mag > 20:
n = -max_mag
if n:
a = mpf_shift(a, n)
b = mpf_shift(b, n)
#print to_float(a), to_float(b)
af = to_fixed(a, wp)
bf = to_fixed(b, wp)
g = agm_fixed(af, bf, wp)
return from_man_exp(g, -wp-n, prec, rnd)
def mpf_agm1(a, prec, rnd=round_fast):
"""
Computes the arithmetic-geometric mean agm(1,a) for a nonnegative
mpf value a.
"""
return mpf_agm(fone, a, prec, rnd)
def mpc_agm(a, b, prec, rnd=round_fast):
"""
Complex AGM.
TODO:
* check that convergence works as intended
* optimize
* select a nonarbitrary branch
"""
if mpc_is_infnan(a) or mpc_is_infnan(b):
return fnan, fnan
if mpc_zero in (a, b):
return fzero, fzero
if mpc_neg(a) == b:
return fzero, fzero
wp = prec+20
eps = mpf_shift(fone, -wp+10)
while 1:
a1 = mpc_shift(mpc_add(a, b, wp), -1)
b1 = mpc_sqrt(mpc_mul(a, b, wp), wp)
a, b = a1, b1
size = mpf_min_max([mpc_abs(a,10), mpc_abs(b,10)])[1]
err = mpc_abs(mpc_sub(a, b, 10), 10)
if size == fzero or mpf_lt(err, mpf_mul(eps, size)):
return a
def mpc_agm1(a, prec, rnd=round_fast):
return mpc_agm(mpc_one, a, prec, rnd)
def mpf_ellipk(x, prec, rnd=round_fast):
if not x[1]:
if x == fzero:
return mpf_shift(mpf_pi(prec, rnd), -1)
if x == fninf:
return fzero
if x == fnan:
return x
if x == fone:
return finf
# TODO: for |x| << 1/2, one could use fall back to
# pi/2 * hyp2f1_rat((1,2),(1,2),(1,1), x)
wp = prec + 15
# Use K(x) = pi/2/agm(1,a) where a = sqrt(1-x)
# The sqrt raises ComplexResult if x > 0
a = mpf_sqrt(mpf_sub(fone, x, wp), wp)
v = mpf_agm1(a, wp)
r = mpf_div(mpf_pi(wp), v, prec, rnd)
return mpf_shift(r, -1)
def mpc_ellipk(z, prec, rnd=round_fast):
re, im = z
if im == fzero:
if re == finf:
return mpc_zero
if mpf_le(re, fone):
return mpf_ellipk(re, prec, rnd), fzero
wp = prec + 15
a = mpc_sqrt(mpc_sub(mpc_one, z, wp), wp)
v = mpc_agm1(a, wp)
r = mpc_mpf_div(mpf_pi(wp), v, prec, rnd)
return mpc_shift(r, -1)
def mpf_ellipe(x, prec, rnd=round_fast):
# http://functions.wolfram.com/EllipticIntegrals/
# EllipticK/20/01/0001/
# E = (1-m)*(K'(m)*2*m + K(m))
sign, man, exp, bc = x
if not man:
if x == fzero:
return mpf_shift(mpf_pi(prec, rnd), -1)
if x == fninf:
return finf
if x == fnan:
return x
if x == finf:
raise ComplexResult
if x == fone:
return fone
wp = prec+20
mag = exp+bc
if mag < -wp:
return mpf_shift(mpf_pi(prec, rnd), -1)
# Compute a finite difference for K'
p = max(mag, 0) - wp
h = mpf_shift(fone, p)
K = mpf_ellipk(x, 2*wp)
Kh = mpf_ellipk(mpf_sub(x, h), 2*wp)
Kdiff = mpf_shift(mpf_sub(K, Kh), -p)
t = mpf_sub(fone, x)
b = mpf_mul(Kdiff, mpf_shift(x,1), wp)
return mpf_mul(t, mpf_add(K, b), prec, rnd)
def mpc_ellipe(z, prec, rnd=round_fast):
re, im = z
if im == fzero:
if re == finf:
return (fzero, finf)
if mpf_le(re, fone):
return mpf_ellipe(re, prec, rnd), fzero
wp = prec + 15
mag = mpc_abs(z, 1)
p = max(mag[2]+mag[3], 0) - wp
h = mpf_shift(fone, p)
K = mpc_ellipk(z, 2*wp)
Kh = mpc_ellipk(mpc_add_mpf(z, h, 2*wp), 2*wp)
Kdiff = mpc_shift(mpc_sub(Kh, K, wp), -p)
t = mpc_sub(mpc_one, z, wp)
b = mpc_mul(Kdiff, mpc_shift(z,1), wp)
return mpc_mul(t, mpc_add(K, b, wp), prec, rnd)