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""" 7 8import operator 9import math 10 11from .backend import MPZ_ZERO, MPZ_ONE, BACKEND, xrange, exec_ 12 13from .libintmath import gcd 14 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) 26 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) 31 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) 39 40from .libintmath import ifac 41from .gammazeta import mpf_gamma_int, mpf_euler, euler_fixed 42 43class NoConvergence(Exception): 44 pass 45 46 47#-----------------------------------------------------------------------# 48# # 49# Generic hypergeometric series # 50# # 51#-----------------------------------------------------------------------# 52 53""" 54TODO: 55 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 61 62""" 63 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). 68 69 """ 70 p, q, param_types, ztype = key 71 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 75 76 have_complex_param = 'C' in param_types 77 have_complex_arg = ztype == 'C' 78 have_complex = have_complex_param or have_complex_arg 79 80 source = [] 81 add = source.append 82 83 aint = [] 84 arat = [] 85 bint = [] 86 brat = [] 87 areal = [] 88 breal = [] 89 acomplex = [] 90 bcomplex = [] 91 92 #add("wp = prec + 40") 93 add("MAX = kwargs.get('maxterms', wp*100)") 94 add("HIGH = MPZ_ONE<<epsshift") 95 add("LOW = -HIGH") 96 97 # Setup code 98 add("SRE = PRE = one = (MPZ_ONE << wp)") 99 if have_complex: 100 add("SIM = PIM = MPZ_ZERO") 101 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") 110 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)") 122 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") 147 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 160 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:] 166 167 # LOOP 168 add("for n in xrange(1,10**8):") 169 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") 175 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]) 180 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"]) 184 185 if multiplier: 186 add(" mul = " + multiplier) 187 add(" div = " + divisor) 188 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") 195 196 # Update product 197 if have_complex: 198 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))) 208 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") 225 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") 230 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))) 237 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") 246 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") 257 258 #add(" from mpmath import nprint, log, ldexp") 259 #add(" nprint([n, log(abs(PRE),2), ldexp(PRE,-wp)])") 260 261 add(" if n > MAX:") 262 add(" raise NoConvergence('Hypergeometric series converges too slowly. Try increasing maxterms.')") 263 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))) 273 274 if have_complex: 275 add("a = from_man_exp(SRE, -wp, prec, 'n')") 276 add("b = from_man_exp(SIM, -wp, prec, 'n')") 277 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") 287 288 add("return (a, b), True, magn") 289 else: 290 add("a = from_man_exp(SRE, -wp, prec, 'n')") 291 292 add("if SRE:") 293 add(" magn = a[2]+a[3]") 294 add("else:") 295 add(" magn = -wp+1") 296 297 add("return a, False, magn") 298 299 source = "\n".join((" " + line) for line in source) 300 source = ("def %s(coeffs, z, prec, wp, epsshift, magnitude_check, **kwargs):\n" % fname) + source 301 302 namespace = {} 303 304 exec_(source, globals(), namespace) 305 306 #print source 307 return source, namespace[fname] 308 309 310if BACKEND == 'sage': 311 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) 322 323 return "(none)", _hypsum 324 325 326#-----------------------------------------------------------------------# 327# # 328# Error functions # 329# # 330#-----------------------------------------------------------------------# 331 332# TODO: mpf_erf should call mpf_erfc when appropriate (currently 333# only the converse delegation is implemented) 334 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) 374 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 385 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 424 425 426#-----------------------------------------------------------------------# 427# # 428# Exponential integrals # 429# # 430#-----------------------------------------------------------------------# 431 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 440 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 452 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 463 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 483 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 525 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 590 591def mpf_e1(x, prec, rnd=round_fast): 592 return mpf_ei(x, prec, rnd, True) 593 594def mpc_e1(x, prec, rnd=round_fast): 595 return mpc_ei(x, prec, rnd, True) 596 597def mpf_expint(n, x, prec, rnd=round_fast, gamma=False): 598 """ 599 E_n(x), n an integer, x real 600 601 With gamma=True, computes Gamma(n,x) (upper incomplete gamma function) 602 603 Returns (real, None) if real, otherwise (real, imag) 604 The imaginary part is an optional branch cut term 605 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 714 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) 731 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) 762 763def mpf_ci_si(x, prec, rnd=round_fast, which=2): 764 """ 765 Calculation of Ci(x), Si(x) for real x. 766 767 which = 0 -- returns (Ci(x), -) 768 which = 1 -- returns (Si(x), -) 769 which = 2 -- returns (Ci(x), Si(x)) 770 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 855 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] 860 861def mpf_si(x, prec, rnd=round_fast): 862 return mpf_ci_si(x, prec, rnd, 1)[1] 863 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 876 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) 884 885 886#-----------------------------------------------------------------------# 887# # 888# Bessel functions # 889# # 890#-----------------------------------------------------------------------# 891 892# A Bessel function of the first kind of integer order, J_n(x), is 893# given by the power series 894 895# oo 896# ___ k 2 k + n 897# \ (-1) / x \ 898# J_n(x) = ) ----------- | - | 899# /___ k! (k + n)! \ 2 / 900# k = 0 901 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. 907 908# In the general case, we have 909# J_v(x) = (x/2)**v / v! * 0F1(v+1, (-1/4)*z**2) 910 911# TODO: for extremely large x, we could use an asymptotic 912# trigonometric approximation. 913 914# TODO: recompute at higher precision if the fixed-point mantissa 915# is very small 916 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) 939 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) 975 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) 1034 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) 1041 1042def mpc_agm(a, b, prec, rnd=round_fast): 1043 """ 1044 Complex AGM. 1045 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 1067 1068def mpc_agm1(a, prec, rnd=round_fast): 1069 return mpc_agm(mpc_one, a, prec, rnd) 1070 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) 1090 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) 1103 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) 1133 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) 1151