""" 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<= 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 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< 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)