Spaces:
Sleeping
Sleeping
""" | |
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) | |