libhyper.py 36 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150
  1. """
  2. This module implements computation of hypergeometric and related
  3. functions. In particular, it provides code for generic summation
  4. of hypergeometric series. Optimized versions for various special
  5. cases are also provided.
  6. """
  7. import operator
  8. import math
  9. from .backend import MPZ_ZERO, MPZ_ONE, BACKEND, xrange, exec_
  10. from .libintmath import gcd
  11. from .libmpf import (\
  12. ComplexResult, round_fast, round_nearest,
  13. negative_rnd, bitcount, to_fixed, from_man_exp, from_int, to_int,
  14. from_rational,
  15. fzero, fone, fnone, ftwo, finf, fninf, fnan,
  16. mpf_sign, mpf_add, mpf_abs, mpf_pos,
  17. mpf_cmp, mpf_lt, mpf_le, mpf_gt, mpf_min_max,
  18. mpf_perturb, mpf_neg, mpf_shift, mpf_sub, mpf_mul, mpf_div,
  19. sqrt_fixed, mpf_sqrt, mpf_rdiv_int, mpf_pow_int,
  20. to_rational,
  21. )
  22. from .libelefun import (\
  23. mpf_pi, mpf_exp, mpf_log, pi_fixed, mpf_cos_sin, mpf_cos, mpf_sin,
  24. mpf_sqrt, agm_fixed,
  25. )
  26. from .libmpc import (\
  27. mpc_one, mpc_sub, mpc_mul_mpf, mpc_mul, mpc_neg, complex_int_pow,
  28. mpc_div, mpc_add_mpf, mpc_sub_mpf,
  29. mpc_log, mpc_add, mpc_pos, mpc_shift,
  30. mpc_is_infnan, mpc_zero, mpc_sqrt, mpc_abs,
  31. mpc_mpf_div, mpc_square, mpc_exp
  32. )
  33. from .libintmath import ifac
  34. from .gammazeta import mpf_gamma_int, mpf_euler, euler_fixed
  35. class NoConvergence(Exception):
  36. pass
  37. #-----------------------------------------------------------------------#
  38. # #
  39. # Generic hypergeometric series #
  40. # #
  41. #-----------------------------------------------------------------------#
  42. """
  43. TODO:
  44. 1. proper mpq parsing
  45. 2. imaginary z special-cased (also: rational, integer?)
  46. 3. more clever handling of series that don't converge because of stupid
  47. upwards rounding
  48. 4. checking for cancellation
  49. """
  50. def make_hyp_summator(key):
  51. """
  52. Returns a function that sums a generalized hypergeometric series,
  53. for given parameter types (integer, rational, real, complex).
  54. """
  55. p, q, param_types, ztype = key
  56. pstring = "".join(param_types)
  57. fname = "hypsum_%i_%i_%s_%s_%s" % (p, q, pstring[:p], pstring[p:], ztype)
  58. #print "generating hypsum", fname
  59. have_complex_param = 'C' in param_types
  60. have_complex_arg = ztype == 'C'
  61. have_complex = have_complex_param or have_complex_arg
  62. source = []
  63. add = source.append
  64. aint = []
  65. arat = []
  66. bint = []
  67. brat = []
  68. areal = []
  69. breal = []
  70. acomplex = []
  71. bcomplex = []
  72. #add("wp = prec + 40")
  73. add("MAX = kwargs.get('maxterms', wp*100)")
  74. add("HIGH = MPZ_ONE<<epsshift")
  75. add("LOW = -HIGH")
  76. # Setup code
  77. add("SRE = PRE = one = (MPZ_ONE << wp)")
  78. if have_complex:
  79. add("SIM = PIM = MPZ_ZERO")
  80. if have_complex_arg:
  81. add("xsign, xm, xe, xbc = z[0]")
  82. add("if xsign: xm = -xm")
  83. add("ysign, ym, ye, ybc = z[1]")
  84. add("if ysign: ym = -ym")
  85. else:
  86. add("xsign, xm, xe, xbc = z")
  87. add("if xsign: xm = -xm")
  88. add("offset = xe + wp")
  89. add("if offset >= 0:")
  90. add(" ZRE = xm << offset")
  91. add("else:")
  92. add(" ZRE = xm >> (-offset)")
  93. if have_complex_arg:
  94. add("offset = ye + wp")
  95. add("if offset >= 0:")
  96. add(" ZIM = ym << offset")
  97. add("else:")
  98. add(" ZIM = ym >> (-offset)")
  99. for i, flag in enumerate(param_types):
  100. W = ["A", "B"][i >= p]
  101. if flag == 'Z':
  102. ([aint,bint][i >= p]).append(i)
  103. add("%sINT_%i = coeffs[%i]" % (W, i, i))
  104. elif flag == 'Q':
  105. ([arat,brat][i >= p]).append(i)
  106. add("%sP_%i, %sQ_%i = coeffs[%i]._mpq_" % (W, i, W, i, i))
  107. elif flag == 'R':
  108. ([areal,breal][i >= p]).append(i)
  109. add("xsign, xm, xe, xbc = coeffs[%i]._mpf_" % i)
  110. add("if xsign: xm = -xm")
  111. add("offset = xe + wp")
  112. add("if offset >= 0:")
  113. add(" %sREAL_%i = xm << offset" % (W, i))
  114. add("else:")
  115. add(" %sREAL_%i = xm >> (-offset)" % (W, i))
  116. elif flag == 'C':
  117. ([acomplex,bcomplex][i >= p]).append(i)
  118. add("__re, __im = coeffs[%i]._mpc_" % i)
  119. add("xsign, xm, xe, xbc = __re")
  120. add("if xsign: xm = -xm")
  121. add("ysign, ym, ye, ybc = __im")
  122. add("if ysign: ym = -ym")
  123. add("offset = xe + wp")
  124. add("if offset >= 0:")
  125. add(" %sCRE_%i = xm << offset" % (W, i))
  126. add("else:")
  127. add(" %sCRE_%i = xm >> (-offset)" % (W, i))
  128. add("offset = ye + wp")
  129. add("if offset >= 0:")
  130. add(" %sCIM_%i = ym << offset" % (W, i))
  131. add("else:")
  132. add(" %sCIM_%i = ym >> (-offset)" % (W, i))
  133. else:
  134. raise ValueError
  135. l_areal = len(areal)
  136. l_breal = len(breal)
  137. cancellable_real = min(l_areal, l_breal)
  138. noncancellable_real_num = areal[cancellable_real:]
  139. noncancellable_real_den = breal[cancellable_real:]
  140. # LOOP
  141. add("for n in xrange(1,10**8):")
  142. add(" if n in magnitude_check:")
  143. add(" p_mag = bitcount(abs(PRE))")
  144. if have_complex:
  145. add(" p_mag = max(p_mag, bitcount(abs(PIM)))")
  146. add(" magnitude_check[n] = wp-p_mag")
  147. # Real factors
  148. multiplier = " * ".join(["AINT_#".replace("#", str(i)) for i in aint] + \
  149. ["AP_#".replace("#", str(i)) for i in arat] + \
  150. ["BQ_#".replace("#", str(i)) for i in brat])
  151. divisor = " * ".join(["BINT_#".replace("#", str(i)) for i in bint] + \
  152. ["BP_#".replace("#", str(i)) for i in brat] + \
  153. ["AQ_#".replace("#", str(i)) for i in arat] + ["n"])
  154. if multiplier:
  155. add(" mul = " + multiplier)
  156. add(" div = " + divisor)
  157. # Check for singular terms
  158. add(" if not div:")
  159. if multiplier:
  160. add(" if not mul:")
  161. add(" break")
  162. add(" raise ZeroDivisionError")
  163. # Update product
  164. if have_complex:
  165. # TODO: when there are several real parameters and just a few complex
  166. # (maybe just the complex argument), we only need to do about
  167. # half as many ops if we accumulate the real factor in a single real variable
  168. for k in range(cancellable_real): add(" PRE = PRE * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
  169. for i in noncancellable_real_num: add(" PRE = (PRE * AREAL_#) >> wp".replace("#", str(i)))
  170. for i in noncancellable_real_den: add(" PRE = (PRE << wp) // BREAL_#".replace("#", str(i)))
  171. for k in range(cancellable_real): add(" PIM = PIM * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
  172. for i in noncancellable_real_num: add(" PIM = (PIM * AREAL_#) >> wp".replace("#", str(i)))
  173. for i in noncancellable_real_den: add(" PIM = (PIM << wp) // BREAL_#".replace("#", str(i)))
  174. if multiplier:
  175. if have_complex_arg:
  176. add(" PRE, PIM = (mul*(PRE*ZRE-PIM*ZIM))//div, (mul*(PIM*ZRE+PRE*ZIM))//div")
  177. add(" PRE >>= wp")
  178. add(" PIM >>= wp")
  179. else:
  180. add(" PRE = ((mul * PRE * ZRE) >> wp) // div")
  181. add(" PIM = ((mul * PIM * ZRE) >> wp) // div")
  182. else:
  183. if have_complex_arg:
  184. add(" PRE, PIM = (PRE*ZRE-PIM*ZIM)//div, (PIM*ZRE+PRE*ZIM)//div")
  185. add(" PRE >>= wp")
  186. add(" PIM >>= wp")
  187. else:
  188. add(" PRE = ((PRE * ZRE) >> wp) // div")
  189. add(" PIM = ((PIM * ZRE) >> wp) // div")
  190. for i in acomplex:
  191. add(" PRE, PIM = PRE*ACRE_#-PIM*ACIM_#, PIM*ACRE_#+PRE*ACIM_#".replace("#", str(i)))
  192. add(" PRE >>= wp")
  193. add(" PIM >>= wp")
  194. for i in bcomplex:
  195. add(" mag = BCRE_#*BCRE_#+BCIM_#*BCIM_#".replace("#", str(i)))
  196. add(" re = PRE*BCRE_# + PIM*BCIM_#".replace("#", str(i)))
  197. add(" im = PIM*BCRE_# - PRE*BCIM_#".replace("#", str(i)))
  198. add(" PRE = (re << wp) // mag".replace("#", str(i)))
  199. add(" PIM = (im << wp) // mag".replace("#", str(i)))
  200. else:
  201. for k in range(cancellable_real): add(" PRE = PRE * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
  202. for i in noncancellable_real_num: add(" PRE = (PRE * AREAL_#) >> wp".replace("#", str(i)))
  203. for i in noncancellable_real_den: add(" PRE = (PRE << wp) // BREAL_#".replace("#", str(i)))
  204. if multiplier:
  205. add(" PRE = ((PRE * mul * ZRE) >> wp) // div")
  206. else:
  207. add(" PRE = ((PRE * ZRE) >> wp) // div")
  208. # Add product to sum
  209. if have_complex:
  210. add(" SRE += PRE")
  211. add(" SIM += PIM")
  212. add(" if (HIGH > PRE > LOW) and (HIGH > PIM > LOW):")
  213. add(" break")
  214. else:
  215. add(" SRE += PRE")
  216. add(" if HIGH > PRE > LOW:")
  217. add(" break")
  218. #add(" from mpmath import nprint, log, ldexp")
  219. #add(" nprint([n, log(abs(PRE),2), ldexp(PRE,-wp)])")
  220. add(" if n > MAX:")
  221. add(" raise NoConvergence('Hypergeometric series converges too slowly. Try increasing maxterms.')")
  222. # +1 all parameters for next loop
  223. for i in aint: add(" AINT_# += 1".replace("#", str(i)))
  224. for i in bint: add(" BINT_# += 1".replace("#", str(i)))
  225. for i in arat: add(" AP_# += AQ_#".replace("#", str(i)))
  226. for i in brat: add(" BP_# += BQ_#".replace("#", str(i)))
  227. for i in areal: add(" AREAL_# += one".replace("#", str(i)))
  228. for i in breal: add(" BREAL_# += one".replace("#", str(i)))
  229. for i in acomplex: add(" ACRE_# += one".replace("#", str(i)))
  230. for i in bcomplex: add(" BCRE_# += one".replace("#", str(i)))
  231. if have_complex:
  232. add("a = from_man_exp(SRE, -wp, prec, 'n')")
  233. add("b = from_man_exp(SIM, -wp, prec, 'n')")
  234. add("if SRE:")
  235. add(" if SIM:")
  236. add(" magn = max(a[2]+a[3], b[2]+b[3])")
  237. add(" else:")
  238. add(" magn = a[2]+a[3]")
  239. add("elif SIM:")
  240. add(" magn = b[2]+b[3]")
  241. add("else:")
  242. add(" magn = -wp+1")
  243. add("return (a, b), True, magn")
  244. else:
  245. add("a = from_man_exp(SRE, -wp, prec, 'n')")
  246. add("if SRE:")
  247. add(" magn = a[2]+a[3]")
  248. add("else:")
  249. add(" magn = -wp+1")
  250. add("return a, False, magn")
  251. source = "\n".join((" " + line) for line in source)
  252. source = ("def %s(coeffs, z, prec, wp, epsshift, magnitude_check, **kwargs):\n" % fname) + source
  253. namespace = {}
  254. exec_(source, globals(), namespace)
  255. #print source
  256. return source, namespace[fname]
  257. if BACKEND == 'sage':
  258. def make_hyp_summator(key):
  259. """
  260. Returns a function that sums a generalized hypergeometric series,
  261. for given parameter types (integer, rational, real, complex).
  262. """
  263. from sage.libs.mpmath.ext_main import hypsum_internal
  264. p, q, param_types, ztype = key
  265. def _hypsum(coeffs, z, prec, wp, epsshift, magnitude_check, **kwargs):
  266. return hypsum_internal(p, q, param_types, ztype, coeffs, z,
  267. prec, wp, epsshift, magnitude_check, kwargs)
  268. return "(none)", _hypsum
  269. #-----------------------------------------------------------------------#
  270. # #
  271. # Error functions #
  272. # #
  273. #-----------------------------------------------------------------------#
  274. # TODO: mpf_erf should call mpf_erfc when appropriate (currently
  275. # only the converse delegation is implemented)
  276. def mpf_erf(x, prec, rnd=round_fast):
  277. sign, man, exp, bc = x
  278. if not man:
  279. if x == fzero: return fzero
  280. if x == finf: return fone
  281. if x== fninf: return fnone
  282. return fnan
  283. size = exp + bc
  284. lg = math.log
  285. # The approximation erf(x) = 1 is accurate to > x^2 * log(e,2) bits
  286. if size > 3 and 2*(size-1) + 0.528766 > lg(prec,2):
  287. if sign:
  288. return mpf_perturb(fnone, 0, prec, rnd)
  289. else:
  290. return mpf_perturb(fone, 1, prec, rnd)
  291. # erf(x) ~ 2*x/sqrt(pi) close to 0
  292. if size < -prec:
  293. # 2*x
  294. x = mpf_shift(x,1)
  295. c = mpf_sqrt(mpf_pi(prec+20), prec+20)
  296. # TODO: interval rounding
  297. return mpf_div(x, c, prec, rnd)
  298. wp = prec + abs(size) + 25
  299. # Taylor series for erf, fixed-point summation
  300. t = abs(to_fixed(x, wp))
  301. t2 = (t*t) >> wp
  302. s, term, k = t, 12345, 1
  303. while term:
  304. t = ((t * t2) >> wp) // k
  305. term = t // (2*k+1)
  306. if k & 1:
  307. s -= term
  308. else:
  309. s += term
  310. k += 1
  311. s = (s << (wp+1)) // sqrt_fixed(pi_fixed(wp), wp)
  312. if sign:
  313. s = -s
  314. return from_man_exp(s, -wp, prec, rnd)
  315. # If possible, we use the asymptotic series for erfc.
  316. # This is an alternating divergent asymptotic series, so
  317. # the error is at most equal to the first omitted term.
  318. # Here we check if the smallest term is small enough
  319. # for a given x and precision
  320. def erfc_check_series(x, prec):
  321. n = to_int(x)
  322. if n**2 * 1.44 > prec:
  323. return True
  324. return False
  325. def mpf_erfc(x, prec, rnd=round_fast):
  326. sign, man, exp, bc = x
  327. if not man:
  328. if x == fzero: return fone
  329. if x == finf: return fzero
  330. if x == fninf: return ftwo
  331. return fnan
  332. wp = prec + 20
  333. mag = bc+exp
  334. # Preserve full accuracy when exponent grows huge
  335. wp += max(0, 2*mag)
  336. regular_erf = sign or mag < 2
  337. if regular_erf or not erfc_check_series(x, wp):
  338. if regular_erf:
  339. return mpf_sub(fone, mpf_erf(x, prec+10, negative_rnd[rnd]), prec, rnd)
  340. # 1-erf(x) ~ exp(-x^2), increase prec to deal with cancellation
  341. n = to_int(x)+1
  342. return mpf_sub(fone, mpf_erf(x, prec + int(n**2*1.44) + 10), prec, rnd)
  343. s = term = MPZ_ONE << wp
  344. term_prev = 0
  345. t = (2 * to_fixed(x, wp) ** 2) >> wp
  346. k = 1
  347. while 1:
  348. term = ((term * (2*k - 1)) << wp) // t
  349. if k > 4 and term > term_prev or not term:
  350. break
  351. if k & 1:
  352. s -= term
  353. else:
  354. s += term
  355. term_prev = term
  356. #print k, to_str(from_man_exp(term, -wp, 50), 10)
  357. k += 1
  358. s = (s << wp) // sqrt_fixed(pi_fixed(wp), wp)
  359. s = from_man_exp(s, -wp, wp)
  360. z = mpf_exp(mpf_neg(mpf_mul(x,x,wp),wp),wp)
  361. y = mpf_div(mpf_mul(z, s, wp), x, prec, rnd)
  362. return y
  363. #-----------------------------------------------------------------------#
  364. # #
  365. # Exponential integrals #
  366. # #
  367. #-----------------------------------------------------------------------#
  368. def ei_taylor(x, prec):
  369. s = t = x
  370. k = 2
  371. while t:
  372. t = ((t*x) >> prec) // k
  373. s += t // k
  374. k += 1
  375. return s
  376. def complex_ei_taylor(zre, zim, prec):
  377. _abs = abs
  378. sre = tre = zre
  379. sim = tim = zim
  380. k = 2
  381. while _abs(tre) + _abs(tim) > 5:
  382. tre, tim = ((tre*zre-tim*zim)//k)>>prec, ((tre*zim+tim*zre)//k)>>prec
  383. sre += tre // k
  384. sim += tim // k
  385. k += 1
  386. return sre, sim
  387. def ei_asymptotic(x, prec):
  388. one = MPZ_ONE << prec
  389. x = t = ((one << prec) // x)
  390. s = one + x
  391. k = 2
  392. while t:
  393. t = (k*t*x) >> prec
  394. s += t
  395. k += 1
  396. return s
  397. def complex_ei_asymptotic(zre, zim, prec):
  398. _abs = abs
  399. one = MPZ_ONE << prec
  400. M = (zim*zim + zre*zre) >> prec
  401. # 1 / z
  402. xre = tre = (zre << prec) // M
  403. xim = tim = ((-zim) << prec) // M
  404. sre = one + xre
  405. sim = xim
  406. k = 2
  407. while _abs(tre) + _abs(tim) > 1000:
  408. #print tre, tim
  409. tre, tim = ((tre*xre-tim*xim)*k)>>prec, ((tre*xim+tim*xre)*k)>>prec
  410. sre += tre
  411. sim += tim
  412. k += 1
  413. if k > prec:
  414. raise NoConvergence
  415. return sre, sim
  416. def mpf_ei(x, prec, rnd=round_fast, e1=False):
  417. if e1:
  418. x = mpf_neg(x)
  419. sign, man, exp, bc = x
  420. if e1 and not sign:
  421. if x == fzero:
  422. return finf
  423. raise ComplexResult("E1(x) for x < 0")
  424. if man:
  425. xabs = 0, man, exp, bc
  426. xmag = exp+bc
  427. wp = prec + 20
  428. can_use_asymp = xmag > wp
  429. if not can_use_asymp:
  430. if exp >= 0:
  431. xabsint = man << exp
  432. else:
  433. xabsint = man >> (-exp)
  434. can_use_asymp = xabsint > int(wp*0.693) + 10
  435. if can_use_asymp:
  436. if xmag > wp:
  437. v = fone
  438. else:
  439. v = from_man_exp(ei_asymptotic(to_fixed(x, wp), wp), -wp)
  440. v = mpf_mul(v, mpf_exp(x, wp), wp)
  441. v = mpf_div(v, x, prec, rnd)
  442. else:
  443. wp += 2*int(to_int(xabs))
  444. u = to_fixed(x, wp)
  445. v = ei_taylor(u, wp) + euler_fixed(wp)
  446. t1 = from_man_exp(v,-wp)
  447. t2 = mpf_log(xabs,wp)
  448. v = mpf_add(t1, t2, prec, rnd)
  449. else:
  450. if x == fzero: v = fninf
  451. elif x == finf: v = finf
  452. elif x == fninf: v = fzero
  453. else: v = fnan
  454. if e1:
  455. v = mpf_neg(v)
  456. return v
  457. def mpc_ei(z, prec, rnd=round_fast, e1=False):
  458. if e1:
  459. z = mpc_neg(z)
  460. a, b = z
  461. asign, aman, aexp, abc = a
  462. bsign, bman, bexp, bbc = b
  463. if b == fzero:
  464. if e1:
  465. x = mpf_neg(mpf_ei(a, prec, rnd))
  466. if not asign:
  467. y = mpf_neg(mpf_pi(prec, rnd))
  468. else:
  469. y = fzero
  470. return x, y
  471. else:
  472. return mpf_ei(a, prec, rnd), fzero
  473. if a != fzero:
  474. if not aman or not bman:
  475. return (fnan, fnan)
  476. wp = prec + 40
  477. amag = aexp+abc
  478. bmag = bexp+bbc
  479. zmag = max(amag, bmag)
  480. can_use_asymp = zmag > wp
  481. if not can_use_asymp:
  482. zabsint = abs(to_int(a)) + abs(to_int(b))
  483. can_use_asymp = zabsint > int(wp*0.693) + 20
  484. try:
  485. if can_use_asymp:
  486. if zmag > wp:
  487. v = fone, fzero
  488. else:
  489. zre = to_fixed(a, wp)
  490. zim = to_fixed(b, wp)
  491. vre, vim = complex_ei_asymptotic(zre, zim, wp)
  492. v = from_man_exp(vre, -wp), from_man_exp(vim, -wp)
  493. v = mpc_mul(v, mpc_exp(z, wp), wp)
  494. v = mpc_div(v, z, wp)
  495. if e1:
  496. v = mpc_neg(v, prec, rnd)
  497. else:
  498. x, y = v
  499. if bsign:
  500. v = mpf_pos(x, prec, rnd), mpf_sub(y, mpf_pi(wp), prec, rnd)
  501. else:
  502. v = mpf_pos(x, prec, rnd), mpf_add(y, mpf_pi(wp), prec, rnd)
  503. return v
  504. except NoConvergence:
  505. pass
  506. #wp += 2*max(0,zmag)
  507. wp += 2*int(to_int(mpc_abs(z, 5)))
  508. zre = to_fixed(a, wp)
  509. zim = to_fixed(b, wp)
  510. vre, vim = complex_ei_taylor(zre, zim, wp)
  511. vre += euler_fixed(wp)
  512. v = from_man_exp(vre,-wp), from_man_exp(vim,-wp)
  513. if e1:
  514. u = mpc_log(mpc_neg(z),wp)
  515. else:
  516. u = mpc_log(z,wp)
  517. v = mpc_add(v, u, prec, rnd)
  518. if e1:
  519. v = mpc_neg(v)
  520. return v
  521. def mpf_e1(x, prec, rnd=round_fast):
  522. return mpf_ei(x, prec, rnd, True)
  523. def mpc_e1(x, prec, rnd=round_fast):
  524. return mpc_ei(x, prec, rnd, True)
  525. def mpf_expint(n, x, prec, rnd=round_fast, gamma=False):
  526. """
  527. E_n(x), n an integer, x real
  528. With gamma=True, computes Gamma(n,x) (upper incomplete gamma function)
  529. Returns (real, None) if real, otherwise (real, imag)
  530. The imaginary part is an optional branch cut term
  531. """
  532. sign, man, exp, bc = x
  533. if not man:
  534. if gamma:
  535. if x == fzero:
  536. # Actually gamma function pole
  537. if n <= 0:
  538. return finf, None
  539. return mpf_gamma_int(n, prec, rnd), None
  540. if x == finf:
  541. return fzero, None
  542. # TODO: could return finite imaginary value at -inf
  543. return fnan, fnan
  544. else:
  545. if x == fzero:
  546. if n > 1:
  547. return from_rational(1, n-1, prec, rnd), None
  548. else:
  549. return finf, None
  550. if x == finf:
  551. return fzero, None
  552. return fnan, fnan
  553. n_orig = n
  554. if gamma:
  555. n = 1-n
  556. wp = prec + 20
  557. xmag = exp + bc
  558. # Beware of near-poles
  559. if xmag < -10:
  560. raise NotImplementedError
  561. nmag = bitcount(abs(n))
  562. have_imag = n > 0 and sign
  563. negx = mpf_neg(x)
  564. # Skip series if direct convergence
  565. if n == 0 or 2*nmag - xmag < -wp:
  566. if gamma:
  567. v = mpf_exp(negx, wp)
  568. re = mpf_mul(v, mpf_pow_int(x, n_orig-1, wp), prec, rnd)
  569. else:
  570. v = mpf_exp(negx, wp)
  571. re = mpf_div(v, x, prec, rnd)
  572. else:
  573. # Finite number of terms, or...
  574. can_use_asymptotic_series = -3*wp < n <= 0
  575. # ...large enough?
  576. if not can_use_asymptotic_series:
  577. xi = abs(to_int(x))
  578. m = min(max(1, xi-n), 2*wp)
  579. siz = -n*nmag + (m+n)*bitcount(abs(m+n)) - m*xmag - (144*m//100)
  580. tol = -wp-10
  581. can_use_asymptotic_series = siz < tol
  582. if can_use_asymptotic_series:
  583. r = ((-MPZ_ONE) << (wp+wp)) // to_fixed(x, wp)
  584. m = n
  585. t = r*m
  586. s = MPZ_ONE << wp
  587. while m and t:
  588. s += t
  589. m += 1
  590. t = (m*r*t) >> wp
  591. v = mpf_exp(negx, wp)
  592. if gamma:
  593. # ~ exp(-x) * x^(n-1) * (1 + ...)
  594. v = mpf_mul(v, mpf_pow_int(x, n_orig-1, wp), wp)
  595. else:
  596. # ~ exp(-x)/x * (1 + ...)
  597. v = mpf_div(v, x, wp)
  598. re = mpf_mul(v, from_man_exp(s, -wp), prec, rnd)
  599. elif n == 1:
  600. re = mpf_neg(mpf_ei(negx, prec, rnd))
  601. elif n > 0 and n < 3*wp:
  602. T1 = mpf_neg(mpf_ei(negx, wp))
  603. if gamma:
  604. if n_orig & 1:
  605. T1 = mpf_neg(T1)
  606. else:
  607. T1 = mpf_mul(T1, mpf_pow_int(negx, n-1, wp), wp)
  608. r = t = to_fixed(x, wp)
  609. facs = [1] * (n-1)
  610. for k in range(1,n-1):
  611. facs[k] = facs[k-1] * k
  612. facs = facs[::-1]
  613. s = facs[0] << wp
  614. for k in range(1, n-1):
  615. if k & 1:
  616. s -= facs[k] * t
  617. else:
  618. s += facs[k] * t
  619. t = (t*r) >> wp
  620. T2 = from_man_exp(s, -wp, wp)
  621. T2 = mpf_mul(T2, mpf_exp(negx, wp))
  622. if gamma:
  623. T2 = mpf_mul(T2, mpf_pow_int(x, n_orig, wp), wp)
  624. R = mpf_add(T1, T2)
  625. re = mpf_div(R, from_int(ifac(n-1)), prec, rnd)
  626. else:
  627. raise NotImplementedError
  628. if have_imag:
  629. M = from_int(-ifac(n-1))
  630. if gamma:
  631. im = mpf_div(mpf_pi(wp), M, prec, rnd)
  632. if n_orig & 1:
  633. im = mpf_neg(im)
  634. else:
  635. im = mpf_div(mpf_mul(mpf_pi(wp), mpf_pow_int(negx, n_orig-1, wp), wp), M, prec, rnd)
  636. return re, im
  637. else:
  638. return re, None
  639. def mpf_ci_si_taylor(x, wp, which=0):
  640. """
  641. 0 - Ci(x) - (euler+log(x))
  642. 1 - Si(x)
  643. """
  644. x = to_fixed(x, wp)
  645. x2 = -(x*x) >> wp
  646. if which == 0:
  647. s, t, k = 0, (MPZ_ONE<<wp), 2
  648. else:
  649. s, t, k = x, x, 3
  650. while t:
  651. t = (t*x2//(k*(k-1)))>>wp
  652. s += t//k
  653. k += 2
  654. return from_man_exp(s, -wp)
  655. def mpc_ci_si_taylor(re, im, wp, which=0):
  656. # The following code is only designed for small arguments,
  657. # and not too small arguments (for relative accuracy)
  658. if re[1]:
  659. mag = re[2]+re[3]
  660. elif im[1]:
  661. mag = im[2]+im[3]
  662. if im[1]:
  663. mag = max(mag, im[2]+im[3])
  664. if mag > 2 or mag < -wp:
  665. raise NotImplementedError
  666. wp += (2-mag)
  667. zre = to_fixed(re, wp)
  668. zim = to_fixed(im, wp)
  669. z2re = (zim*zim-zre*zre)>>wp
  670. z2im = (-2*zre*zim)>>wp
  671. tre = zre
  672. tim = zim
  673. one = MPZ_ONE<<wp
  674. if which == 0:
  675. sre, sim, tre, tim, k = 0, 0, (MPZ_ONE<<wp), 0, 2
  676. else:
  677. sre, sim, tre, tim, k = zre, zim, zre, zim, 3
  678. while max(abs(tre), abs(tim)) > 2:
  679. f = k*(k-1)
  680. tre, tim = ((tre*z2re-tim*z2im)//f)>>wp, ((tre*z2im+tim*z2re)//f)>>wp
  681. sre += tre//k
  682. sim += tim//k
  683. k += 2
  684. return from_man_exp(sre, -wp), from_man_exp(sim, -wp)
  685. def mpf_ci_si(x, prec, rnd=round_fast, which=2):
  686. """
  687. Calculation of Ci(x), Si(x) for real x.
  688. which = 0 -- returns (Ci(x), -)
  689. which = 1 -- returns (Si(x), -)
  690. which = 2 -- returns (Ci(x), Si(x))
  691. Note: if x < 0, Ci(x) needs an additional imaginary term, pi*i.
  692. """
  693. wp = prec + 20
  694. sign, man, exp, bc = x
  695. ci, si = None, None
  696. if not man:
  697. if x == fzero:
  698. return (fninf, fzero)
  699. if x == fnan:
  700. return (x, x)
  701. ci = fzero
  702. if which != 0:
  703. if x == finf:
  704. si = mpf_shift(mpf_pi(prec, rnd), -1)
  705. if x == fninf:
  706. si = mpf_neg(mpf_shift(mpf_pi(prec, negative_rnd[rnd]), -1))
  707. return (ci, si)
  708. # For small x: Ci(x) ~ euler + log(x), Si(x) ~ x
  709. mag = exp+bc
  710. if mag < -wp:
  711. if which != 0:
  712. si = mpf_perturb(x, 1-sign, prec, rnd)
  713. if which != 1:
  714. y = mpf_euler(wp)
  715. xabs = mpf_abs(x)
  716. ci = mpf_add(y, mpf_log(xabs, wp), prec, rnd)
  717. return ci, si
  718. # For huge x: Ci(x) ~ sin(x)/x, Si(x) ~ pi/2
  719. elif mag > wp:
  720. if which != 0:
  721. if sign:
  722. si = mpf_neg(mpf_pi(prec, negative_rnd[rnd]))
  723. else:
  724. si = mpf_pi(prec, rnd)
  725. si = mpf_shift(si, -1)
  726. if which != 1:
  727. ci = mpf_div(mpf_sin(x, wp), x, prec, rnd)
  728. return ci, si
  729. else:
  730. wp += abs(mag)
  731. # Use an asymptotic series? The smallest value of n!/x^n
  732. # occurs for n ~ x, where the magnitude is ~ exp(-x).
  733. asymptotic = mag-1 > math.log(wp, 2)
  734. # Case 1: convergent series near 0
  735. if not asymptotic:
  736. if which != 0:
  737. si = mpf_pos(mpf_ci_si_taylor(x, wp, 1), prec, rnd)
  738. if which != 1:
  739. ci = mpf_ci_si_taylor(x, wp, 0)
  740. ci = mpf_add(ci, mpf_euler(wp), wp)
  741. ci = mpf_add(ci, mpf_log(mpf_abs(x), wp), prec, rnd)
  742. return ci, si
  743. x = mpf_abs(x)
  744. # Case 2: asymptotic series for x >> 1
  745. xf = to_fixed(x, wp)
  746. xr = (MPZ_ONE<<(2*wp)) // xf # 1/x
  747. s1 = (MPZ_ONE << wp)
  748. s2 = xr
  749. t = xr
  750. k = 2
  751. while t:
  752. t = -t
  753. t = (t*xr*k)>>wp
  754. k += 1
  755. s1 += t
  756. t = (t*xr*k)>>wp
  757. k += 1
  758. s2 += t
  759. s1 = from_man_exp(s1, -wp)
  760. s2 = from_man_exp(s2, -wp)
  761. s1 = mpf_div(s1, x, wp)
  762. s2 = mpf_div(s2, x, wp)
  763. cos, sin = mpf_cos_sin(x, wp)
  764. # Ci(x) = sin(x)*s1-cos(x)*s2
  765. # Si(x) = pi/2-cos(x)*s1-sin(x)*s2
  766. if which != 0:
  767. si = mpf_add(mpf_mul(cos, s1), mpf_mul(sin, s2), wp)
  768. si = mpf_sub(mpf_shift(mpf_pi(wp), -1), si, wp)
  769. if sign:
  770. si = mpf_neg(si)
  771. si = mpf_pos(si, prec, rnd)
  772. if which != 1:
  773. ci = mpf_sub(mpf_mul(sin, s1), mpf_mul(cos, s2), prec, rnd)
  774. return ci, si
  775. def mpf_ci(x, prec, rnd=round_fast):
  776. if mpf_sign(x) < 0:
  777. raise ComplexResult
  778. return mpf_ci_si(x, prec, rnd, 0)[0]
  779. def mpf_si(x, prec, rnd=round_fast):
  780. return mpf_ci_si(x, prec, rnd, 1)[1]
  781. def mpc_ci(z, prec, rnd=round_fast):
  782. re, im = z
  783. if im == fzero:
  784. ci = mpf_ci_si(re, prec, rnd, 0)[0]
  785. if mpf_sign(re) < 0:
  786. return (ci, mpf_pi(prec, rnd))
  787. return (ci, fzero)
  788. wp = prec + 20
  789. cre, cim = mpc_ci_si_taylor(re, im, wp, 0)
  790. cre = mpf_add(cre, mpf_euler(wp), wp)
  791. ci = mpc_add((cre, cim), mpc_log(z, wp), prec, rnd)
  792. return ci
  793. def mpc_si(z, prec, rnd=round_fast):
  794. re, im = z
  795. if im == fzero:
  796. return (mpf_ci_si(re, prec, rnd, 1)[1], fzero)
  797. wp = prec + 20
  798. z = mpc_ci_si_taylor(re, im, wp, 1)
  799. return mpc_pos(z, prec, rnd)
  800. #-----------------------------------------------------------------------#
  801. # #
  802. # Bessel functions #
  803. # #
  804. #-----------------------------------------------------------------------#
  805. # A Bessel function of the first kind of integer order, J_n(x), is
  806. # given by the power series
  807. # oo
  808. # ___ k 2 k + n
  809. # \ (-1) / x \
  810. # J_n(x) = ) ----------- | - |
  811. # /___ k! (k + n)! \ 2 /
  812. # k = 0
  813. # Simplifying the quotient between two successive terms gives the
  814. # ratio x^2 / (-4*k*(k+n)). Hence, we only need one full-precision
  815. # multiplication and one division by a small integer per term.
  816. # The complex version is very similar, the only difference being
  817. # that the multiplication is actually 4 multiplies.
  818. # In the general case, we have
  819. # J_v(x) = (x/2)**v / v! * 0F1(v+1, (-1/4)*z**2)
  820. # TODO: for extremely large x, we could use an asymptotic
  821. # trigonometric approximation.
  822. # TODO: recompute at higher precision if the fixed-point mantissa
  823. # is very small
  824. def mpf_besseljn(n, x, prec, rounding=round_fast):
  825. prec += 50
  826. negate = n < 0 and n & 1
  827. mag = x[2]+x[3]
  828. n = abs(n)
  829. wp = prec + 20 + n*bitcount(n)
  830. if mag < 0:
  831. wp -= n * mag
  832. x = to_fixed(x, wp)
  833. x2 = (x**2) >> wp
  834. if not n:
  835. s = t = MPZ_ONE << wp
  836. else:
  837. s = t = (x**n // ifac(n)) >> ((n-1)*wp + n)
  838. k = 1
  839. while t:
  840. t = ((t * x2) // (-4*k*(k+n))) >> wp
  841. s += t
  842. k += 1
  843. if negate:
  844. s = -s
  845. return from_man_exp(s, -wp, prec, rounding)
  846. def mpc_besseljn(n, z, prec, rounding=round_fast):
  847. negate = n < 0 and n & 1
  848. n = abs(n)
  849. origprec = prec
  850. zre, zim = z
  851. mag = max(zre[2]+zre[3], zim[2]+zim[3])
  852. prec += 20 + n*bitcount(n) + abs(mag)
  853. if mag < 0:
  854. prec -= n * mag
  855. zre = to_fixed(zre, prec)
  856. zim = to_fixed(zim, prec)
  857. z2re = (zre**2 - zim**2) >> prec
  858. z2im = (zre*zim) >> (prec-1)
  859. if not n:
  860. sre = tre = MPZ_ONE << prec
  861. sim = tim = MPZ_ZERO
  862. else:
  863. re, im = complex_int_pow(zre, zim, n)
  864. sre = tre = (re // ifac(n)) >> ((n-1)*prec + n)
  865. sim = tim = (im // ifac(n)) >> ((n-1)*prec + n)
  866. k = 1
  867. while abs(tre) + abs(tim) > 3:
  868. p = -4*k*(k+n)
  869. tre, tim = tre*z2re - tim*z2im, tim*z2re + tre*z2im
  870. tre = (tre // p) >> prec
  871. tim = (tim // p) >> prec
  872. sre += tre
  873. sim += tim
  874. k += 1
  875. if negate:
  876. sre = -sre
  877. sim = -sim
  878. re = from_man_exp(sre, -prec, origprec, rounding)
  879. im = from_man_exp(sim, -prec, origprec, rounding)
  880. return (re, im)
  881. def mpf_agm(a, b, prec, rnd=round_fast):
  882. """
  883. Computes the arithmetic-geometric mean agm(a,b) for
  884. nonnegative mpf values a, b.
  885. """
  886. asign, aman, aexp, abc = a
  887. bsign, bman, bexp, bbc = b
  888. if asign or bsign:
  889. raise ComplexResult("agm of a negative number")
  890. # Handle inf, nan or zero in either operand
  891. if not (aman and bman):
  892. if a == fnan or b == fnan:
  893. return fnan
  894. if a == finf:
  895. if b == fzero:
  896. return fnan
  897. return finf
  898. if b == finf:
  899. if a == fzero:
  900. return fnan
  901. return finf
  902. # agm(0,x) = agm(x,0) = 0
  903. return fzero
  904. wp = prec + 20
  905. amag = aexp+abc
  906. bmag = bexp+bbc
  907. mag_delta = amag - bmag
  908. # Reduce to roughly the same magnitude using floating-point AGM
  909. abs_mag_delta = abs(mag_delta)
  910. if abs_mag_delta > 10:
  911. while abs_mag_delta > 10:
  912. a, b = mpf_shift(mpf_add(a,b,wp),-1), \
  913. mpf_sqrt(mpf_mul(a,b,wp),wp)
  914. abs_mag_delta //= 2
  915. asign, aman, aexp, abc = a
  916. bsign, bman, bexp, bbc = b
  917. amag = aexp+abc
  918. bmag = bexp+bbc
  919. mag_delta = amag - bmag
  920. #print to_float(a), to_float(b)
  921. # Use agm(a,b) = agm(x*a,x*b)/x to obtain a, b ~= 1
  922. min_mag = min(amag,bmag)
  923. max_mag = max(amag,bmag)
  924. n = 0
  925. # If too small, we lose precision when going to fixed-point
  926. if min_mag < -8:
  927. n = -min_mag
  928. # If too large, we waste time using fixed-point with large numbers
  929. elif max_mag > 20:
  930. n = -max_mag
  931. if n:
  932. a = mpf_shift(a, n)
  933. b = mpf_shift(b, n)
  934. #print to_float(a), to_float(b)
  935. af = to_fixed(a, wp)
  936. bf = to_fixed(b, wp)
  937. g = agm_fixed(af, bf, wp)
  938. return from_man_exp(g, -wp-n, prec, rnd)
  939. def mpf_agm1(a, prec, rnd=round_fast):
  940. """
  941. Computes the arithmetic-geometric mean agm(1,a) for a nonnegative
  942. mpf value a.
  943. """
  944. return mpf_agm(fone, a, prec, rnd)
  945. def mpc_agm(a, b, prec, rnd=round_fast):
  946. """
  947. Complex AGM.
  948. TODO:
  949. * check that convergence works as intended
  950. * optimize
  951. * select a nonarbitrary branch
  952. """
  953. if mpc_is_infnan(a) or mpc_is_infnan(b):
  954. return fnan, fnan
  955. if mpc_zero in (a, b):
  956. return fzero, fzero
  957. if mpc_neg(a) == b:
  958. return fzero, fzero
  959. wp = prec+20
  960. eps = mpf_shift(fone, -wp+10)
  961. while 1:
  962. a1 = mpc_shift(mpc_add(a, b, wp), -1)
  963. b1 = mpc_sqrt(mpc_mul(a, b, wp), wp)
  964. a, b = a1, b1
  965. size = mpf_min_max([mpc_abs(a,10), mpc_abs(b,10)])[1]
  966. err = mpc_abs(mpc_sub(a, b, 10), 10)
  967. if size == fzero or mpf_lt(err, mpf_mul(eps, size)):
  968. return a
  969. def mpc_agm1(a, prec, rnd=round_fast):
  970. return mpc_agm(mpc_one, a, prec, rnd)
  971. def mpf_ellipk(x, prec, rnd=round_fast):
  972. if not x[1]:
  973. if x == fzero:
  974. return mpf_shift(mpf_pi(prec, rnd), -1)
  975. if x == fninf:
  976. return fzero
  977. if x == fnan:
  978. return x
  979. if x == fone:
  980. return finf
  981. # TODO: for |x| << 1/2, one could use fall back to
  982. # pi/2 * hyp2f1_rat((1,2),(1,2),(1,1), x)
  983. wp = prec + 15
  984. # Use K(x) = pi/2/agm(1,a) where a = sqrt(1-x)
  985. # The sqrt raises ComplexResult if x > 0
  986. a = mpf_sqrt(mpf_sub(fone, x, wp), wp)
  987. v = mpf_agm1(a, wp)
  988. r = mpf_div(mpf_pi(wp), v, prec, rnd)
  989. return mpf_shift(r, -1)
  990. def mpc_ellipk(z, prec, rnd=round_fast):
  991. re, im = z
  992. if im == fzero:
  993. if re == finf:
  994. return mpc_zero
  995. if mpf_le(re, fone):
  996. return mpf_ellipk(re, prec, rnd), fzero
  997. wp = prec + 15
  998. a = mpc_sqrt(mpc_sub(mpc_one, z, wp), wp)
  999. v = mpc_agm1(a, wp)
  1000. r = mpc_mpf_div(mpf_pi(wp), v, prec, rnd)
  1001. return mpc_shift(r, -1)
  1002. def mpf_ellipe(x, prec, rnd=round_fast):
  1003. # http://functions.wolfram.com/EllipticIntegrals/
  1004. # EllipticK/20/01/0001/
  1005. # E = (1-m)*(K'(m)*2*m + K(m))
  1006. sign, man, exp, bc = x
  1007. if not man:
  1008. if x == fzero:
  1009. return mpf_shift(mpf_pi(prec, rnd), -1)
  1010. if x == fninf:
  1011. return finf
  1012. if x == fnan:
  1013. return x
  1014. if x == finf:
  1015. raise ComplexResult
  1016. if x == fone:
  1017. return fone
  1018. wp = prec+20
  1019. mag = exp+bc
  1020. if mag < -wp:
  1021. return mpf_shift(mpf_pi(prec, rnd), -1)
  1022. # Compute a finite difference for K'
  1023. p = max(mag, 0) - wp
  1024. h = mpf_shift(fone, p)
  1025. K = mpf_ellipk(x, 2*wp)
  1026. Kh = mpf_ellipk(mpf_sub(x, h), 2*wp)
  1027. Kdiff = mpf_shift(mpf_sub(K, Kh), -p)
  1028. t = mpf_sub(fone, x)
  1029. b = mpf_mul(Kdiff, mpf_shift(x,1), wp)
  1030. return mpf_mul(t, mpf_add(K, b), prec, rnd)
  1031. def mpc_ellipe(z, prec, rnd=round_fast):
  1032. re, im = z
  1033. if im == fzero:
  1034. if re == finf:
  1035. return (fzero, finf)
  1036. if mpf_le(re, fone):
  1037. return mpf_ellipe(re, prec, rnd), fzero
  1038. wp = prec + 15
  1039. mag = mpc_abs(z, 1)
  1040. p = max(mag[2]+mag[3], 0) - wp
  1041. h = mpf_shift(fone, p)
  1042. K = mpc_ellipk(z, 2*wp)
  1043. Kh = mpc_ellipk(mpc_add_mpf(z, h, 2*wp), 2*wp)
  1044. Kdiff = mpc_shift(mpc_sub(Kh, K, wp), -p)
  1045. t = mpc_sub(mpc_one, z, wp)
  1046. b = mpc_mul(Kdiff, mpc_shift(z,1), wp)
  1047. return mpc_mul(t, mpc_add(K, b, wp), prec, rnd)