mpmath

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

8
import math
9

10
from .backend import MPZ_ONE, MPZ_ZERO
11
from .gammazeta import euler_fixed, mpf_euler, mpf_gamma_int
12
from .libelefun import (agm_fixed, mpf_cos_sin, mpf_exp, mpf_log, mpf_pi,
13
                        mpf_sin, mpf_sqrt, pi_fixed)
14
from .libintmath import ifac, sqrt_fixed
15
from .libmpc import (complex_int_pow, mpc_abs, mpc_add, mpc_add_mpf, mpc_div,
16
                     mpc_exp, mpc_is_infnan, mpc_log, mpc_mpf_div, mpc_mul,
17
                     mpc_neg, mpc_one, mpc_pos, mpc_shift, mpc_sqrt, mpc_sub,
18
                     mpc_zero)
19
from .libmpf import (ComplexResult, finf, fnan, fninf, fnone, fone, from_int,
20
                     from_man_exp, from_rational, ftwo, fzero, mpf_abs,
21
                     mpf_add, mpf_div, mpf_le, mpf_lt, mpf_min_max, mpf_mul,
22
                     mpf_neg, mpf_perturb, mpf_pos, mpf_pow_int, mpf_shift,
23
                     mpf_sign, mpf_sqrt, mpf_sub, negative_rnd, round_fast,
24
                     to_fixed, to_int)
25

26

27
class NoConvergence(Exception):
28
    pass
29

30

31
#-----------------------------------------------------------------------#
32
#                                                                       #
33
#                     Generic hypergeometric series                     #
34
#                                                                       #
35
#-----------------------------------------------------------------------#
36

37
"""
38
TODO:
39

40
1. proper mpq parsing
41
2. imaginary z special-cased (also: rational, integer?)
42
3. more clever handling of series that don't converge because of stupid
43
   upwards rounding
44
4. checking for cancellation
45

46
"""
47

48
def make_hyp_summator(key):
49
    """
50
    Returns a function that sums a generalized hypergeometric series,
51
    for given parameter types (integer, rational, real, complex).
52

53
    """
54
    p, q, param_types, ztype = key
55

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

60
    have_complex_param = 'C' in param_types
61
    have_complex_arg = ztype == 'C'
62
    have_complex = have_complex_param or have_complex_arg
63

64
    source = []
65
    add = source.append
66

67
    aint = []
68
    arat = []
69
    bint = []
70
    brat = []
71
    areal = []
72
    breal = []
73
    acomplex = []
74
    bcomplex = []
75

76
    #add("wp = prec + 40")
77
    add("MAX = kwargs.get('maxterms', wp*100)")
78
    add("HIGH = MPZ_ONE<<epsshift")
79
    add("LOW = -HIGH")
80

81
    # Setup code
82
    add("SRE = PRE = one = (MPZ_ONE << wp)")
83
    if have_complex:
84
        add("SIM = PIM = MPZ_ZERO")
85

86
    if have_complex_arg:
87
        add("xsign, xm, xe, xbc = z[0]")
88
        add("if xsign: xm = -xm")
89
        add("ysign, ym, ye, ybc = z[1]")
90
        add("if ysign: ym = -ym")
91
    else:
92
        add("xsign, xm, xe, xbc = z")
93
        add("if xsign: xm = -xm")
94

95
    add("offset = xe + wp")
96
    add("if offset >= 0:")
97
    add("    ZRE = xm << offset")
98
    add("else:")
99
    add("    ZRE = xm >> (-offset)")
100
    if have_complex_arg:
101
        add("offset = ye + wp")
102
        add("if offset >= 0:")
103
        add("    ZIM = ym << offset")
104
        add("else:")
105
        add("    ZIM = ym >> (-offset)")
106

107
    for i, flag in enumerate(param_types):
108
        W = ["A", "B"][i >= p]
109
        if flag == 'Z':
110
            ([aint,bint][i >= p]).append(i)
111
            add("%sINT_%i = coeffs[%i]" % (W, i, i))
112
        elif flag == 'Q':
113
            ([arat,brat][i >= p]).append(i)
114
            add("%sP_%i, %sQ_%i = coeffs[%i].numerator, coeffs[%i].denominator" % (W, i, W, i, i, i))
115
        elif flag == 'R':
116
            ([areal,breal][i >= p]).append(i)
117
            add("xsign, xm, xe, xbc = coeffs[%i]._mpf_" % i)
118
            add("if xsign: xm = -xm")
119
            add("offset = xe + wp")
120
            add("if offset >= 0:")
121
            add("    %sREAL_%i = xm << offset" % (W, i))
122
            add("else:")
123
            add("    %sREAL_%i = xm >> (-offset)" % (W, i))
124
        elif flag == 'C':
125
            ([acomplex,bcomplex][i >= p]).append(i)
126
            add("__re, __im = coeffs[%i]._mpc_" % i)
127
            add("xsign, xm, xe, xbc = __re")
128
            add("if xsign: xm = -xm")
129
            add("ysign, ym, ye, ybc = __im")
130
            add("if ysign: ym = -ym")
131

132
            add("offset = xe + wp")
133
            add("if offset >= 0:")
134
            add("    %sCRE_%i = xm << offset" % (W, i))
135
            add("else:")
136
            add("    %sCRE_%i = xm >> (-offset)" % (W, i))
137
            add("offset = ye + wp")
138
            add("if offset >= 0:")
139
            add("    %sCIM_%i = ym << offset" % (W, i))
140
            add("else:")
141
            add("    %sCIM_%i = ym >> (-offset)" % (W, i))
142
        else:
143
            raise ValueError
144

145
    l_areal = len(areal)
146
    l_breal = len(breal)
147
    cancellable_real = min(l_areal, l_breal)
148
    noncancellable_real_num = areal[cancellable_real:]
149
    noncancellable_real_den = breal[cancellable_real:]
150

151
    # LOOP
152
    add("for n in range(1,10**8):")
153

154
    add("    if n in magnitude_check:")
155
    add("        p_mag = PRE.bit_length()")
156
    if have_complex:
157
        add("        p_mag = max(p_mag, PIM.bit_length())")
158
    add("        magnitude_check[n] = wp-p_mag")
159

160
    # Real factors
161
    multiplier = " * ".join(["AINT_#".replace("#", str(i)) for i in aint] + \
162
                            ["AP_#".replace("#", str(i)) for i in arat] + \
163
                            ["BQ_#".replace("#", str(i)) for i in brat])
164

165
    divisor    = " * ".join(["BINT_#".replace("#", str(i)) for i in bint] + \
166
                            ["BP_#".replace("#", str(i)) for i in brat] + \
167
                            ["AQ_#".replace("#", str(i)) for i in arat] + ["n"])
168

169
    if multiplier:
170
        add("    mul = " + multiplier)
171
    add("    div = " + divisor)
172

173
    # Check for singular terms
174
    add("    if not div:")
175
    if multiplier:
176
        add("        if not mul:")
177
        add("            break")
178
    add("        raise ZeroDivisionError")
179

180
    # Update product
181
    if have_complex:
182

183
        # TODO: when there are several real parameters and just a few complex
184
        # (maybe just the complex argument), we only need to do about
185
        # half as many ops if we accumulate the real factor in a single real variable
186
        for k in range(cancellable_real): add("    PRE = PRE * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
187
        for i in noncancellable_real_num: add("    PRE = (PRE * AREAL_#) >> wp".replace("#", str(i)))
188
        for i in noncancellable_real_den: add("    PRE = (PRE << wp) // BREAL_#".replace("#", str(i)))
189
        for k in range(cancellable_real): add("    PIM = PIM * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
190
        for i in noncancellable_real_num: add("    PIM = (PIM * AREAL_#) >> wp".replace("#", str(i)))
191
        for i in noncancellable_real_den: add("    PIM = (PIM << wp) // BREAL_#".replace("#", str(i)))
192

193
        if multiplier:
194
            if have_complex_arg:
195
                add("    PRE, PIM = (mul*(PRE*ZRE-PIM*ZIM))//div, (mul*(PIM*ZRE+PRE*ZIM))//div")
196
                add("    PRE >>= wp")
197
                add("    PIM >>= wp")
198
            else:
199
                add("    PRE = ((mul * PRE * ZRE) >> wp) // div")
200
                add("    PIM = ((mul * PIM * ZRE) >> wp) // div")
201
        else:
202
            if have_complex_arg:
203
                add("    PRE, PIM = (PRE*ZRE-PIM*ZIM)//div, (PIM*ZRE+PRE*ZIM)//div")
204
                add("    PRE >>= wp")
205
                add("    PIM >>= wp")
206
            else:
207
                add("    PRE = ((PRE * ZRE) >> wp) // div")
208
                add("    PIM = ((PIM * ZRE) >> wp) // div")
209

210
        for i in acomplex:
211
            add("    PRE, PIM = PRE*ACRE_#-PIM*ACIM_#, PIM*ACRE_#+PRE*ACIM_#".replace("#", str(i)))
212
            add("    PRE >>= wp")
213
            add("    PIM >>= wp")
214

215
        for i in bcomplex:
216
            add("    mag = BCRE_#*BCRE_#+BCIM_#*BCIM_#".replace("#", str(i)))
217
            add("    re = PRE*BCRE_# + PIM*BCIM_#".replace("#", str(i)))
218
            add("    im = PIM*BCRE_# - PRE*BCIM_#".replace("#", str(i)))
219
            add("    PRE = (re << wp) // mag".replace("#", str(i)))
220
            add("    PIM = (im << wp) // mag".replace("#", str(i)))
221

222
    else:
223
        for k in range(cancellable_real): add("    PRE = PRE * AREAL_%i // BREAL_%i" % (areal[k], breal[k]))
224
        for i in noncancellable_real_num: add("    PRE = (PRE * AREAL_#) >> wp".replace("#", str(i)))
225
        for i in noncancellable_real_den: add("    PRE = (PRE << wp) // BREAL_#".replace("#", str(i)))
226
        if multiplier:
227
            add("    PRE = ((PRE * mul * ZRE) >> wp) // div")
228
        else:
229
            add("    PRE = ((PRE * ZRE) >> wp) // div")
230

231
    # Add product to sum
232
    if have_complex:
233
        add("    SRE += PRE")
234
        add("    SIM += PIM")
235
        add("    if (HIGH > PRE > LOW) and (HIGH > PIM > LOW):")
236
        add("        break")
237
    else:
238
        add("    SRE += PRE")
239
        add("    if HIGH > PRE > LOW:")
240
        add("        break")
241

242
    #add("    from mpmath import nprint, log, ldexp")
243
    #add("    nprint([n, log(abs(PRE),2), ldexp(PRE,-wp)])")
244

245
    add("    if n > MAX:")
246
    add("        raise NoConvergence('Hypergeometric series converges too slowly. Try increasing maxterms.')")
247

248
    # +1 all parameters for next loop
249
    for i in aint:     add("    AINT_# += 1".replace("#", str(i)))
250
    for i in bint:     add("    BINT_# += 1".replace("#", str(i)))
251
    for i in arat:     add("    AP_# += AQ_#".replace("#", str(i)))
252
    for i in brat:     add("    BP_# += BQ_#".replace("#", str(i)))
253
    for i in areal:    add("    AREAL_# += one".replace("#", str(i)))
254
    for i in breal:    add("    BREAL_# += one".replace("#", str(i)))
255
    for i in acomplex: add("    ACRE_# += one".replace("#", str(i)))
256
    for i in bcomplex: add("    BCRE_# += one".replace("#", str(i)))
257

258
    if have_complex:
259
        add("a = from_man_exp(SRE, -wp, prec, 'n')")
260
        add("b = from_man_exp(SIM, -wp, prec, 'n')")
261

262
        add("if SRE:")
263
        add("    if SIM:")
264
        add("        magn = max(a[2]+a[3], b[2]+b[3])")
265
        add("    else:")
266
        add("        magn = a[2]+a[3]")
267
        add("elif SIM:")
268
        add("    magn = b[2]+b[3]")
269
        add("else:")
270
        add("    magn = -wp+1")
271

272
        add("return (a, b), True, magn")
273
    else:
274
        add("a = from_man_exp(SRE, -wp, prec, 'n')")
275

276
        add("if SRE:")
277
        add("    magn = a[2]+a[3]")
278
        add("else:")
279
        add("    magn = -wp+1")
280

281
        add("return a, False, magn")
282

283
    source = "\n".join(("    " + line) for line in source)
284
    source = ("def %s(coeffs, z, prec, wp, epsshift, magnitude_check, **kwargs):\n" % fname) + source
285

286
    namespace = {}
287

288
    exec(source, globals(), namespace)
289

290
    #print source
291
    return source, namespace[fname]
292

293

294
#-----------------------------------------------------------------------#
295
#                                                                       #
296
#                              Error functions                          #
297
#                                                                       #
298
#-----------------------------------------------------------------------#
299

300
# TODO: mpf_erf should call mpf_erfc when appropriate (currently
301
#    only the converse delegation is implemented)
302

303
def mpf_erf(x, prec, rnd=round_fast):
304
    sign, man, exp, bc = x
305
    if not man:
306
        if x == fzero: return fzero
307
        if x == finf: return fone
308
        if x== fninf: return fnone
309
        return fnan
310
    size = exp + bc
311
    lg = math.log
312
    # The approximation erf(x) = 1 is accurate to > x^2 * log(e,2) bits
313
    if size > 3 and 2*(size-1) + 0.528766 > lg(prec,2):
314
        if sign:
315
            return mpf_perturb(fnone, 0, prec, rnd)
316
        else:
317
            return mpf_perturb(fone, 1, prec, rnd)
318
    # erf(x) ~ 2*x/sqrt(pi) close to 0
319
    if size < -prec:
320
        # 2*x
321
        x = mpf_shift(x,1)
322
        c = mpf_sqrt(mpf_pi(prec+20), prec+20)
323
        # TODO: interval rounding
324
        return mpf_div(x, c, prec, rnd)
325
    wp = prec + abs(size) + 25
326
    # Taylor series for erf, fixed-point summation
327
    t = abs(to_fixed(x, wp))
328
    t2 = (t*t) >> wp
329
    s, term, k = t, 12345, 1
330
    while term:
331
        t = ((t * t2) >> wp) // k
332
        term = t // (2*k+1)
333
        if k & 1:
334
            s -= term
335
        else:
336
            s += term
337
        k += 1
338
    s = (s << (wp+1)) // sqrt_fixed(pi_fixed(wp), wp)
339
    if sign:
340
        s = -s
341
    return from_man_exp(s, -wp, prec, rnd)
342

343
# If possible, we use the asymptotic series for erfc.
344
# This is an alternating divergent asymptotic series, so
345
# the error is at most equal to the first omitted term.
346
# Here we check if the smallest term is small enough
347
# for a given x and precision
348
def erfc_check_series(x, prec):
349
    n = to_int(x)
350
    if n**2 * 1.44 > prec:
351
        return True
352
    return False
353

354
def mpf_erfc(x, prec, rnd=round_fast):
355
    sign, man, exp, bc = x
356
    if not man:
357
        if x == fzero: return fone
358
        if x == finf: return fzero
359
        if x == fninf: return ftwo
360
        return fnan
361
    wp = prec + 20
362
    mag = bc+exp
363
    # Preserve full accuracy when exponent grows huge
364
    wp += max(0, 2*mag)
365
    regular_erf = sign or mag < 2
366
    if regular_erf or not erfc_check_series(x, wp):
367
        if regular_erf:
368
            return mpf_sub(fone, mpf_erf(x, prec+10, negative_rnd[rnd]), prec, rnd)
369
        # 1-erf(x) ~ exp(-x^2), increase prec to deal with cancellation
370
        n = to_int(x)+1
371
        return mpf_sub(fone, mpf_erf(x, prec + int(n**2*1.44) + 10), prec, rnd)
372
    s = term = MPZ_ONE << wp
373
    term_prev = 0
374
    t = (2 * to_fixed(x, wp) ** 2) >> wp
375
    k = 1
376
    while 1:
377
        term = ((term * (2*k - 1)) << wp) // t
378
        if k > 4 and term > term_prev or not term:
379
            break
380
        if k & 1:
381
            s -= term
382
        else:
383
            s += term
384
        term_prev = term
385
        #print k, to_str(from_man_exp(term, -wp, 50), 10)
386
        k += 1
387
    s = (s << wp) // sqrt_fixed(pi_fixed(wp), wp)
388
    s = from_man_exp(s, -wp, wp)
389
    z = mpf_exp(mpf_neg(mpf_mul(x,x,wp),wp),wp)
390
    y = mpf_div(mpf_mul(z, s, wp), x, prec, rnd)
391
    return y
392

393

394
#-----------------------------------------------------------------------#
395
#                                                                       #
396
#                         Exponential integrals                         #
397
#                                                                       #
398
#-----------------------------------------------------------------------#
399

400
def ei_taylor(x, prec):
401
    s = t = x
402
    k = 2
403
    while t:
404
        t = ((t*x) >> prec) // k
405
        s += t // k
406
        k += 1
407
    return s
408

409
def complex_ei_taylor(zre, zim, prec):
410
    _abs = abs
411
    sre = tre = zre
412
    sim = tim = zim
413
    k = 2
414
    while _abs(tre) + _abs(tim) > 5:
415
        tre, tim = ((tre*zre-tim*zim)//k)>>prec, ((tre*zim+tim*zre)//k)>>prec
416
        sre += tre // k
417
        sim += tim // k
418
        k += 1
419
    return sre, sim
420

421
def ei_asymptotic(x, prec):
422
    one = MPZ_ONE << prec
423
    x = t = ((one << prec) // x)
424
    s = one + x
425
    k = 2
426
    while t:
427
        t = (k*t*x) >> prec
428
        s += t
429
        k += 1
430
    return s
431

432
def complex_ei_asymptotic(zre, zim, prec):
433
    _abs = abs
434
    one = MPZ_ONE << prec
435
    M = (zim*zim + zre*zre) >> prec
436
    # 1 / z
437
    xre = tre = (zre << prec) // M
438
    xim = tim = ((-zim) << prec) // M
439
    sre = one + xre
440
    sim = xim
441
    k = 2
442
    while _abs(tre) + _abs(tim) > 1000:
443
        #print tre, tim
444
        tre, tim = ((tre*xre-tim*xim)*k)>>prec, ((tre*xim+tim*xre)*k)>>prec
445
        sre += tre
446
        sim += tim
447
        k += 1
448
        if k > prec:
449
            raise NoConvergence
450
    return sre, sim
451

452
def mpf_ei(x, prec, rnd=round_fast, e1=False):
453
    if e1:
454
        x = mpf_neg(x)
455
    sign, man, exp, bc = x
456
    if e1 and not sign:
457
        if x == fzero:
458
            return finf
459
        raise ComplexResult("E1(x) for x < 0")
460
    if man:
461
        xabs = 0, man, exp, bc
462
        xmag = exp+bc
463
        wp = prec + 20
464
        can_use_asymp = xmag > wp
465
        if not can_use_asymp:
466
            if exp >= 0:
467
                xabsint = man << exp
468
            else:
469
                xabsint = man >> (-exp)
470
            can_use_asymp = xabsint > int(wp*0.693) + 10
471
        if can_use_asymp:
472
            if xmag > wp:
473
                v = fone
474
            else:
475
                v = from_man_exp(ei_asymptotic(to_fixed(x, wp), wp), -wp)
476
            v = mpf_mul(v, mpf_exp(x, wp), wp)
477
            v = mpf_div(v, x, prec, rnd)
478
        else:
479
            wp += 2*int(to_int(xabs))
480
            u = to_fixed(x, wp)
481
            v = ei_taylor(u, wp) + euler_fixed(wp)
482
            t1 = from_man_exp(v,-wp)
483
            t2 = mpf_log(xabs,wp)
484
            v = mpf_add(t1, t2, prec, rnd)
485
    else:
486
        if x == fzero: v = fninf
487
        elif x == finf: v = finf
488
        elif x == fninf: v = fzero
489
        else: v = fnan
490
    if e1:
491
        v = mpf_neg(v)
492
    return v
493

494
def mpc_ei(z, prec, rnd=round_fast, e1=False):
495
    if e1:
496
        z = mpc_neg(z)
497
    a, b = z
498
    asign, aman, aexp, abc = a
499
    bsign, bman, bexp, bbc = b
500
    if b == fzero:
501
        if e1:
502
            x = mpf_neg(mpf_ei(a, prec, rnd))
503
            if not asign:
504
                y = mpf_neg(mpf_pi(prec, rnd))
505
            else:
506
                y = fzero
507
            return x, y
508
        else:
509
            return mpf_ei(a, prec, rnd), fzero
510
    if a != fzero:
511
        if not aman or not bman:
512
            return (fnan, fnan)
513
    wp = prec + 40
514
    amag = aexp+abc
515
    bmag = bexp+bbc
516
    zmag = max(amag, bmag)
517
    can_use_asymp = zmag > wp
518
    if not can_use_asymp:
519
        zabsint = abs(to_int(a)) + abs(to_int(b))
520
        can_use_asymp = zabsint > int(wp*0.693) + 20
521
    try:
522
        if can_use_asymp:
523
            if zmag > wp:
524
                v = fone, fzero
525
            else:
526
                zre = to_fixed(a, wp)
527
                zim = to_fixed(b, wp)
528
                vre, vim = complex_ei_asymptotic(zre, zim, wp)
529
                v = from_man_exp(vre, -wp), from_man_exp(vim, -wp)
530
            v = mpc_mul(v, mpc_exp(z, wp), wp)
531
            v = mpc_div(v, z, wp)
532
            if e1:
533
                v = mpc_neg(v, prec, rnd)
534
            else:
535
                x, y = v
536
                if bsign:
537
                    v = mpf_pos(x, prec, rnd), mpf_sub(y, mpf_pi(wp), prec, rnd)
538
                else:
539
                    v = mpf_pos(x, prec, rnd), mpf_add(y, mpf_pi(wp), prec, rnd)
540
            return v
541
    except NoConvergence:
542
        pass
543
    #wp += 2*max(0,zmag)
544
    wp += 2*int(to_int(mpc_abs(z, 5)))
545
    zre = to_fixed(a, wp)
546
    zim = to_fixed(b, wp)
547
    vre, vim = complex_ei_taylor(zre, zim, wp)
548
    vre += euler_fixed(wp)
549
    v = from_man_exp(vre,-wp), from_man_exp(vim,-wp)
550
    if e1:
551
        u = mpc_log(mpc_neg(z),wp)
552
    else:
553
        u = mpc_log(z,wp)
554
    v = mpc_add(v, u, prec, rnd)
555
    if e1:
556
        v = mpc_neg(v)
557
    return v
558

559
def mpf_e1(x, prec, rnd=round_fast):
560
    return mpf_ei(x, prec, rnd, True)
561

562
def mpc_e1(x, prec, rnd=round_fast):
563
    return mpc_ei(x, prec, rnd, True)
564

565
def mpf_expint(n, x, prec, rnd=round_fast, gamma=False):
566
    """
567
    E_n(x), n an integer, x real
568

569
    With gamma=True, computes Gamma(n,x)   (upper incomplete gamma function)
570

571
    Returns (real, None) if real, otherwise (real, imag)
572
    The imaginary part is an optional branch cut term
573

574
    """
575
    sign, man, exp, bc = x
576
    if not man:
577
        if gamma:
578
            if x == fzero:
579
                # Actually gamma function pole
580
                if n <= 0:
581
                    return finf, None
582
                return mpf_gamma_int(n, prec, rnd), None
583
            if x == finf:
584
                return fzero, None
585
            # TODO: could return finite imaginary value at -inf
586
            return fnan, fnan
587
        else:
588
            if x == fzero:
589
                if n > 1:
590
                    return from_rational(1, n-1, prec, rnd), None
591
                else:
592
                    return finf, None
593
            if x == finf:
594
                return fzero, None
595
            return fnan, fnan
596
    n_orig = n
597
    if gamma:
598
        n = 1-n
599
    wp = prec + 20
600
    xmag = exp + bc
601
    # Beware of near-poles
602
    if xmag < -10:
603
        raise NotImplementedError
604
    nmag = n.bit_length()
605
    have_imag = n > 0 and sign
606
    negx = mpf_neg(x)
607
    # Skip series if direct convergence
608
    if n == 0 or 2*nmag - xmag < -wp:
609
        if gamma:
610
            v = mpf_exp(negx, wp)
611
            re = mpf_mul(v, mpf_pow_int(x, n_orig-1, wp), prec, rnd)
612
        else:
613
            v = mpf_exp(negx, wp)
614
            re = mpf_div(v, x, prec, rnd)
615
    else:
616
        # Finite number of terms, or...
617
        can_use_asymptotic_series = -3*wp < n <= 0
618
        # ...large enough?
619
        if not can_use_asymptotic_series:
620
            xi = abs(to_int(x))
621
            m = min(max(1, xi-n), 2*wp)
622
            siz = -n*nmag + (m+n)*(m+n).bit_length() - m*xmag - (144*m//100)
623
            tol = -wp-10
624
            can_use_asymptotic_series = siz < tol
625
        if can_use_asymptotic_series:
626
            r = ((-MPZ_ONE) << (wp+wp)) // to_fixed(x, wp)
627
            m = n
628
            t = r*m
629
            s = MPZ_ONE << wp
630
            while m and t:
631
                s += t
632
                m += 1
633
                t = (m*r*t) >> wp
634
            v = mpf_exp(negx, wp)
635
            if gamma:
636
                # ~ exp(-x) * x^(n-1) * (1 + ...)
637
                v = mpf_mul(v, mpf_pow_int(x, n_orig-1, wp), wp)
638
            else:
639
                # ~ exp(-x)/x * (1 + ...)
640
                v = mpf_div(v, x, wp)
641
            re = mpf_mul(v, from_man_exp(s, -wp), prec, rnd)
642
        elif n == 1:
643
            re = mpf_neg(mpf_ei(negx, prec, rnd))
644
        elif n > 0 and n < 3*wp:
645
            T1 = mpf_neg(mpf_ei(negx, wp))
646
            if gamma:
647
                if n_orig & 1:
648
                    T1 = mpf_neg(T1)
649
            else:
650
                T1 = mpf_mul(T1, mpf_pow_int(negx, n-1, wp), wp)
651
            r = t = to_fixed(x, wp)
652
            facs = [1] * (n-1)
653
            for k in range(1,n-1):
654
                facs[k] = facs[k-1] * k
655
            facs = facs[::-1]
656
            s = facs[0] << wp
657
            for k in range(1, n-1):
658
                if k & 1:
659
                    s -= facs[k] * t
660
                else:
661
                    s += facs[k] * t
662
                t = (t*r) >> wp
663
            T2 = from_man_exp(s, -wp, wp)
664
            T2 = mpf_mul(T2, mpf_exp(negx, wp))
665
            if gamma:
666
                T2 = mpf_mul(T2, mpf_pow_int(x, n_orig, wp), wp)
667
            R = mpf_add(T1, T2)
668
            re = mpf_div(R, from_int(ifac(n-1)), prec, rnd)
669
        else:
670
            raise NotImplementedError
671
    if have_imag:
672
        M = from_int(-ifac(n-1))
673
        if gamma:
674
            im = mpf_div(mpf_pi(wp), M, prec, rnd)
675
            if n_orig & 1:
676
                im = mpf_neg(im)
677
        else:
678
            im = mpf_div(mpf_mul(mpf_pi(wp), mpf_pow_int(negx, n_orig-1, wp), wp), M, prec, rnd)
679
        return re, im
680
    else:
681
        return re, None
682

683
def mpf_ci_si_taylor(x, wp, which=0):
684
    """
685
    0 - Ci(x) - (euler+log(x))
686
    1 - Si(x)
687
    """
688
    x = to_fixed(x, wp)
689
    x2 = -(x*x) >> wp
690
    if which == 0:
691
        s, t, k = 0, (MPZ_ONE<<wp), 2
692
    else:
693
        s, t, k = x, x, 3
694
    while t:
695
        t = (t*x2//(k*(k-1)))>>wp
696
        s += t//k
697
        k += 2
698
    return from_man_exp(s, -wp)
699

700
def mpc_ci_si_taylor(re, im, wp, which=0):
701
    # The following code is only designed for small arguments,
702
    # and not too small arguments (for relative accuracy)
703
    if re[1]:
704
        mag = re[2]+re[3]
705
    elif im[1]:
706
        mag = im[2]+im[3]
707
    if im[1]:
708
        mag = max(mag, im[2]+im[3])
709
    if mag > 2 or mag < -wp:
710
        raise NotImplementedError
711
    wp += (2-mag)
712
    zre = to_fixed(re, wp)
713
    zim = to_fixed(im, wp)
714
    z2re = (zim*zim-zre*zre)>>wp
715
    z2im = (-2*zre*zim)>>wp
716
    tre = zre
717
    tim = zim
718
    one = MPZ_ONE<<wp
719
    if which == 0:
720
        sre, sim, tre, tim, k = 0, 0, (MPZ_ONE<<wp), 0, 2
721
    else:
722
        sre, sim, tre, tim, k = zre, zim, zre, zim, 3
723
    while max(abs(tre), abs(tim)) > 2:
724
        f = k*(k-1)
725
        tre, tim = ((tre*z2re-tim*z2im)//f)>>wp, ((tre*z2im+tim*z2re)//f)>>wp
726
        sre += tre//k
727
        sim += tim//k
728
        k += 2
729
    return from_man_exp(sre, -wp), from_man_exp(sim, -wp)
730

731
def mpf_ci_si(x, prec, rnd=round_fast, which=2):
732
    """
733
    Calculation of Ci(x), Si(x) for real x.
734

735
    which = 0 -- returns (Ci(x), -)
736
    which = 1 -- returns (Si(x), -)
737
    which = 2 -- returns (Ci(x), Si(x))
738

739
    Note: if x < 0, Ci(x) needs an additional imaginary term, pi*i.
740
    """
741
    wp = prec + 20
742
    sign, man, exp, bc = x
743
    ci, si = None, None
744
    if not man:
745
        if x == fzero:
746
            return (fninf, fzero)
747
        if x == fnan:
748
            return (x, x)
749
        ci = fzero
750
        if which != 0:
751
            if x == finf:
752
                si = mpf_shift(mpf_pi(prec, rnd), -1)
753
            if x == fninf:
754
                si = mpf_neg(mpf_shift(mpf_pi(prec, negative_rnd[rnd]), -1))
755
        return (ci, si)
756
    # For small x: Ci(x) ~ euler + log(x), Si(x) ~ x
757
    mag = exp+bc
758
    if mag < -wp:
759
        if which != 0:
760
            si = mpf_perturb(x, 1-sign, prec, rnd)
761
        if which != 1:
762
            y = mpf_euler(wp)
763
            xabs = mpf_abs(x)
764
            ci = mpf_add(y, mpf_log(xabs, wp), prec, rnd)
765
        return ci, si
766
    # For huge x: Ci(x) ~ sin(x)/x, Si(x) ~ pi/2
767
    elif mag > wp:
768
        if which != 0:
769
            if sign:
770
                si = mpf_neg(mpf_pi(prec, negative_rnd[rnd]))
771
            else:
772
                si = mpf_pi(prec, rnd)
773
            si = mpf_shift(si, -1)
774
        if which != 1:
775
            ci = mpf_div(mpf_sin(x, wp), x, prec, rnd)
776
        return ci, si
777
    else:
778
        wp += abs(mag)
779
    # Use an asymptotic series? The smallest value of n!/x^n
780
    # occurs for n ~ x, where the magnitude is ~ exp(-x).
781
    asymptotic = mag-1 > math.log(wp, 2)
782
    # Case 1: convergent series near 0
783
    if not asymptotic:
784
        if which != 0:
785
            si = mpf_pos(mpf_ci_si_taylor(x, wp, 1), prec, rnd)
786
        if which != 1:
787
            ci = mpf_ci_si_taylor(x, wp, 0)
788
            ci = mpf_add(ci, mpf_euler(wp), wp)
789
            ci = mpf_add(ci, mpf_log(mpf_abs(x), wp), prec, rnd)
790
        return ci, si
791
    x = mpf_abs(x)
792
    # Case 2: asymptotic series for x >> 1
793
    xf = to_fixed(x, wp)
794
    xr = (MPZ_ONE<<(2*wp)) // xf   # 1/x
795
    s1 = (MPZ_ONE << wp)
796
    s2 = xr
797
    t = xr
798
    k = 2
799
    while t:
800
        t = -t
801
        t = (t*xr*k)>>wp
802
        k += 1
803
        s1 += t
804
        t = (t*xr*k)>>wp
805
        k += 1
806
        s2 += t
807
    s1 = from_man_exp(s1, -wp)
808
    s2 = from_man_exp(s2, -wp)
809
    s1 = mpf_div(s1, x, wp)
810
    s2 = mpf_div(s2, x, wp)
811
    cos, sin = mpf_cos_sin(x, wp)
812
    # Ci(x) = sin(x)*s1-cos(x)*s2
813
    # Si(x) = pi/2-cos(x)*s1-sin(x)*s2
814
    if which != 0:
815
        si = mpf_add(mpf_mul(cos, s1), mpf_mul(sin, s2), wp)
816
        si = mpf_sub(mpf_shift(mpf_pi(wp), -1), si, wp)
817
        if sign:
818
            si = mpf_neg(si)
819
        si = mpf_pos(si, prec, rnd)
820
    if which != 1:
821
        ci = mpf_sub(mpf_mul(sin, s1), mpf_mul(cos, s2), prec, rnd)
822
    return ci, si
823

824
def mpf_ci(x, prec, rnd=round_fast):
825
    if mpf_sign(x) < 0:
826
        raise ComplexResult
827
    return mpf_ci_si(x, prec, rnd, 0)[0]
828

829
def mpf_si(x, prec, rnd=round_fast):
830
    return mpf_ci_si(x, prec, rnd, 1)[1]
831

832
def mpc_ci(z, prec, rnd=round_fast):
833
    re, im = z
834
    if im == fzero:
835
        ci = mpf_ci_si(re, prec, rnd, 0)[0]
836
        if mpf_sign(re) < 0:
837
            return (ci, mpf_pi(prec, rnd))
838
        return (ci, fzero)
839
    wp = prec + 20
840
    cre, cim = mpc_ci_si_taylor(re, im, wp, 0)
841
    cre = mpf_add(cre, mpf_euler(wp), wp)
842
    ci = mpc_add((cre, cim), mpc_log(z, wp), prec, rnd)
843
    return ci
844

845
def mpc_si(z, prec, rnd=round_fast):
846
    re, im = z
847
    if im == fzero:
848
        return (mpf_ci_si(re, prec, rnd, 1)[1], fzero)
849
    wp = prec + 20
850
    z = mpc_ci_si_taylor(re, im, wp, 1)
851
    return mpc_pos(z, prec, rnd)
852

853

854
#-----------------------------------------------------------------------#
855
#                                                                       #
856
#                             Bessel functions                          #
857
#                                                                       #
858
#-----------------------------------------------------------------------#
859

860
# A Bessel function of the first kind of integer order, J_n(x), is
861
# given by the power series
862

863
#             oo
864
#             ___         k         2 k + n
865
#            \        (-1)     / x \
866
#    J_n(x) = )    ----------- | - |
867
#            /___  k! (k + n)! \ 2 /
868
#            k = 0
869

870
# Simplifying the quotient between two successive terms gives the
871
# ratio x^2 / (-4*k*(k+n)). Hence, we only need one full-precision
872
# multiplication and one division by a small integer per term.
873
# The complex version is very similar, the only difference being
874
# that the multiplication is actually 4 multiplies.
875

876
# In the general case, we have
877
# J_v(x) = (x/2)**v / v! * 0F1(v+1, (-1/4)*z**2)
878

879
# TODO: for extremely large x, we could use an asymptotic
880
# trigonometric approximation.
881

882
# TODO: recompute at higher precision if the fixed-point mantissa
883
# is very small
884

885
def mpf_besseljn(n, x, prec, rounding=round_fast):
886
    prec += 50
887
    negate = n < 0 and n & 1
888
    mag = x[2]+x[3]
889
    n = abs(n)
890
    wp = prec + 20 + n*n.bit_length()
891
    if mag < 0:
892
        wp -= n * mag
893
    x = to_fixed(x, wp)
894
    x2 = (x**2) >> wp
895
    if not n:
896
        s = t = MPZ_ONE << wp
897
    else:
898
        s = t = (x**n // ifac(n)) >> ((n-1)*wp + n)
899
    k = 1
900
    while t:
901
        t = ((t * x2) // (-4*k*(k+n))) >> wp
902
        s += t
903
        k += 1
904
    if negate:
905
        s = -s
906
    return from_man_exp(s, -wp, prec, rounding)
907

908
def mpc_besseljn(n, z, prec, rounding=round_fast):
909
    negate = n < 0 and n & 1
910
    n = abs(n)
911
    origprec = prec
912
    zre, zim = z
913
    mag = max(zre[2]+zre[3], zim[2]+zim[3])
914
    prec += 20 + n*n.bit_length() + abs(mag)
915
    if mag < 0:
916
        prec -= n * mag
917
    zre = to_fixed(zre, prec)
918
    zim = to_fixed(zim, prec)
919
    z2re = (zre**2 - zim**2) >> prec
920
    z2im = (zre*zim) >> (prec-1)
921
    if not n:
922
        sre = tre = MPZ_ONE << prec
923
        sim = tim = MPZ_ZERO
924
    else:
925
        re, im = complex_int_pow(zre, zim, n)
926
        sre = tre = (re // ifac(n)) >> ((n-1)*prec + n)
927
        sim = tim = (im // ifac(n)) >> ((n-1)*prec + n)
928
    k = 1
929
    while abs(tre) + abs(tim) > 3:
930
        p = -4*k*(k+n)
931
        tre, tim = tre*z2re - tim*z2im, tim*z2re + tre*z2im
932
        tre = (tre // p) >> prec
933
        tim = (tim // p) >> prec
934
        sre += tre
935
        sim += tim
936
        k += 1
937
    if negate:
938
        sre = -sre
939
        sim = -sim
940
    re = from_man_exp(sre, -prec, origprec, rounding)
941
    im = from_man_exp(sim, -prec, origprec, rounding)
942
    return (re, im)
943

944
def mpf_agm(a, b, prec, rnd=round_fast):
945
    """
946
    Computes the arithmetic-geometric mean agm(a,b) for
947
    nonnegative mpf values a, b.
948
    """
949
    asign, aman, aexp, abc = a
950
    bsign, bman, bexp, bbc = b
951
    if asign or bsign:
952
        raise ComplexResult("agm of a negative number")
953
    # Handle inf, nan or zero in either operand
954
    if not (aman and bman):
955
        if a == fnan or b == fnan:
956
            return fnan
957
        if a == finf:
958
            if b == fzero:
959
                return fnan
960
            return finf
961
        if b == finf:
962
            if a == fzero:
963
                return fnan
964
            return finf
965
        # agm(0,x) = agm(x,0) = 0
966
        return fzero
967
    wp = prec + 20
968
    amag = aexp+abc
969
    bmag = bexp+bbc
970
    mag_delta = amag - bmag
971
    # Reduce to roughly the same magnitude using floating-point AGM
972
    abs_mag_delta = abs(mag_delta)
973
    if abs_mag_delta > 10:
974
        while abs_mag_delta > 10:
975
            a, b = mpf_shift(mpf_add(a,b,wp),-1), \
976
                mpf_sqrt(mpf_mul(a,b,wp),wp)
977
            abs_mag_delta //= 2
978
        asign, aman, aexp, abc = a
979
        bsign, bman, bexp, bbc = b
980
        amag = aexp+abc
981
        bmag = bexp+bbc
982
        mag_delta = amag - bmag
983
    #print to_float(a), to_float(b)
984
    # Use agm(a,b) = agm(x*a,x*b)/x to obtain a, b ~= 1
985
    min_mag = min(amag,bmag)
986
    max_mag = max(amag,bmag)
987
    n = 0
988
    # If too small, we lose precision when going to fixed-point
989
    if min_mag < -8:
990
        n = -min_mag
991
    # If too large, we waste time using fixed-point with large numbers
992
    elif max_mag > 20:
993
        n = -max_mag
994
    if n:
995
        a = mpf_shift(a, n)
996
        b = mpf_shift(b, n)
997
    #print to_float(a), to_float(b)
998
    af = to_fixed(a, wp)
999
    bf = to_fixed(b, wp)
1000
    g = agm_fixed(af, bf, wp)
1001
    return from_man_exp(g, -wp-n, prec, rnd)
1002

1003
def mpf_agm1(a, prec, rnd=round_fast):
1004
    """
1005
    Computes the arithmetic-geometric mean agm(1,a) for a nonnegative
1006
    mpf value a.
1007
    """
1008
    return mpf_agm(fone, a, prec, rnd)
1009

1010
def mpc_agm(a, b, prec, rnd=round_fast):
1011
    """
1012
    Complex AGM.
1013

1014
    TODO:
1015
    * check that convergence works as intended
1016
    * optimize
1017
    * select a nonarbitrary branch
1018
    """
1019
    if mpc_is_infnan(a) or mpc_is_infnan(b):
1020
        return fnan, fnan
1021
    if mpc_zero in (a, b):
1022
        return fzero, fzero
1023
    if mpc_neg(a) == b:
1024
        return fzero, fzero
1025
    wp = prec+20
1026
    eps = mpf_shift(fone, -wp+10)
1027
    while 1:
1028
        a1 = mpc_shift(mpc_add(a, b, wp), -1)
1029
        b1 = mpc_sqrt(mpc_mul(a, b, wp), wp)
1030
        a, b = a1, b1
1031
        size = mpf_min_max([mpc_abs(a,10), mpc_abs(b,10)])[1]
1032
        err = mpc_abs(mpc_sub(a, b, 10), 10)
1033
        if size == fzero or mpf_lt(err, mpf_mul(eps, size)):
1034
            return a
1035

1036
def mpc_agm1(a, prec, rnd=round_fast):
1037
    return mpc_agm(mpc_one, a, prec, rnd)
1038

1039
def mpf_ellipk(x, prec, rnd=round_fast):
1040
    if not x[1]:
1041
        if x == fzero:
1042
            return mpf_shift(mpf_pi(prec, rnd), -1)
1043
        if x == fninf:
1044
            return fzero
1045
        if x == fnan:
1046
            return x
1047
    if x == fone:
1048
        return finf
1049
    # TODO: for |x| << 1/2, one could use fall back to
1050
    # pi/2 * hyp2f1_rat((1,2),(1,2),(1,1), x)
1051
    wp = prec + 15
1052
    # Use K(x) = pi/2/agm(1,a) where a = sqrt(1-x)
1053
    # The sqrt raises ComplexResult if x > 0
1054
    a = mpf_sqrt(mpf_sub(fone, x, wp), wp)
1055
    v = mpf_agm1(a, wp)
1056
    r = mpf_div(mpf_pi(wp), v, prec, rnd)
1057
    return mpf_shift(r, -1)
1058

1059
def mpc_ellipk(z, prec, rnd=round_fast):
1060
    re, im = z
1061
    if im == fzero:
1062
        if re == finf:
1063
            return mpc_zero
1064
        if mpf_le(re, fone):
1065
            return mpf_ellipk(re, prec, rnd), fzero
1066
    wp = prec + 15
1067
    a = mpc_sqrt(mpc_sub(mpc_one, z, wp), wp)
1068
    v = mpc_agm1(a, wp)
1069
    r = mpc_mpf_div(mpf_pi(wp), v, prec, rnd)
1070
    return mpc_shift(r, -1)
1071

1072
def mpf_ellipe(x, prec, rnd=round_fast):
1073
    # http://functions.wolfram.com/EllipticIntegrals/
1074
    # EllipticK/20/01/0001/
1075
    # E = (1-m)*(K'(m)*2*m + K(m))
1076
    sign, man, exp, bc = x
1077
    if not man:
1078
        if x == fzero:
1079
            return mpf_shift(mpf_pi(prec, rnd), -1)
1080
        if x == fninf:
1081
            return finf
1082
        if x == fnan:
1083
            return x
1084
        if x == finf:
1085
            raise ComplexResult
1086
    if x == fone:
1087
        return fone
1088
    wp = prec+20
1089
    mag = exp+bc
1090
    if mag < -wp:
1091
        return mpf_shift(mpf_pi(prec, rnd), -1)
1092
    # Compute a finite difference for K'
1093
    p = max(mag, 0) - wp
1094
    h = mpf_shift(fone, p)
1095
    K = mpf_ellipk(x, 2*wp)
1096
    Kh = mpf_ellipk(mpf_sub(x, h), 2*wp)
1097
    Kdiff = mpf_shift(mpf_sub(K, Kh), -p)
1098
    t = mpf_sub(fone, x)
1099
    b = mpf_mul(Kdiff, mpf_shift(x,1), wp)
1100
    return mpf_mul(t, mpf_add(K, b), prec, rnd)
1101

1102
def mpc_ellipe(z, prec, rnd=round_fast):
1103
    re, im = z
1104
    if im == fzero:
1105
        if re == finf:
1106
            return (fzero, finf)
1107
        if mpf_le(re, fone):
1108
            return mpf_ellipe(re, prec, rnd), fzero
1109
    wp = prec + 15
1110
    mag = mpc_abs(z, 1)
1111
    p = max(mag[2]+mag[3], 0) - wp
1112
    h = mpf_shift(fone, p)
1113
    K = mpc_ellipk(z, 2*wp)
1114
    Kh = mpc_ellipk(mpc_add_mpf(z, h, 2*wp), 2*wp)
1115
    Kdiff = mpc_shift(mpc_sub(Kh, K, wp), -p)
1116
    t = mpc_sub(mpc_one, z, wp)
1117
    b = mpc_mul(Kdiff, mpc_shift(z,1), wp)
1118
    return mpc_mul(t, mpc_add(K, b, wp), prec, rnd)
1119