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