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