1from .functions import defun, defun_wrapped
2
3@defun
4def j0(ctx, x):
5    """Computes the Bessel function `J_0(x)`. See :func:`~mpmath.besselj`."""
6    return ctx.besselj(0, x)
7
8@defun
9def j1(ctx, x):
10    """Computes the Bessel function `J_1(x)`.  See :func:`~mpmath.besselj`."""
11    return ctx.besselj(1, x)
12
13@defun
14def besselj(ctx, n, z, derivative=0, **kwargs):
15    if type(n) is int:
16        n_isint = True
17    else:
18        n = ctx.convert(n)
19        n_isint = ctx.isint(n)
20        if n_isint:
21            n = int(ctx._re(n))
22    if n_isint and n < 0:
23        return (-1)**n * ctx.besselj(-n, z, derivative, **kwargs)
24    z = ctx.convert(z)
25    M = ctx.mag(z)
26    if derivative:
27        d = ctx.convert(derivative)
28        # TODO: the integer special-casing shouldn't be necessary.
29        # However, the hypergeometric series gets inaccurate for large d
30        # because of inaccurate pole cancellation at a pole far from
31        # zero (needs to be fixed in hypercomb or hypsum)
32        if ctx.isint(d) and d >= 0:
33            d = int(d)
34            orig = ctx.prec
35            try:
36                ctx.prec += 15
37                v = ctx.fsum((-1)**k * ctx.binomial(d,k) * ctx.besselj(2*k+n-d,z)
38                    for k in range(d+1))
39            finally:
40                ctx.prec = orig
41            v *= ctx.mpf(2)**(-d)
42        else:
43            def h(n,d):
44                r = ctx.fmul(ctx.fmul(z, z, prec=ctx.prec+M), -0.25, exact=True)
45                B = [0.5*(n-d+1), 0.5*(n-d+2)]
46                T = [([2,ctx.pi,z],[d-2*n,0.5,n-d],[],B,[(n+1)*0.5,(n+2)*0.5],B+[n+1],r)]
47                return T
48            v = ctx.hypercomb(h, [n,d], **kwargs)
49    else:
50        # Fast case: J_n(x), n int, appropriate magnitude for fixed-point calculation
51        if (not derivative) and n_isint and abs(M) < 10 and abs(n) < 20:
52            try:
53                return ctx._besselj(n, z)
54            except NotImplementedError:
55                pass
56        if not z:
57            if not n:
58                v = ctx.one + n+z
59            elif ctx.re(n) > 0:
60                v = n*z
61            else:
62                v = ctx.inf + z + n
63        else:
64            #v = 0
65            orig = ctx.prec
66            try:
67                # XXX: workaround for accuracy in low level hypergeometric series
68                # when alternating, large arguments
69                ctx.prec += min(3*abs(M), ctx.prec)
70                w = ctx.fmul(z, 0.5, exact=True)
71                def h(n):
72                    r = ctx.fneg(ctx.fmul(w, w, prec=max(0,ctx.prec+M)), exact=True)
73                    return [([w], [n], [], [n+1], [], [n+1], r)]
74                v = ctx.hypercomb(h, [n], **kwargs)
75            finally:
76                ctx.prec = orig
77        v = +v
78    return v
79
80@defun
81def besseli(ctx, n, z, derivative=0, **kwargs):
82    n = ctx.convert(n)
83    z = ctx.convert(z)
84    if not z:
85        if derivative:
86            raise ValueError
87        if not n:
88            # I(0,0) = 1
89            return 1+n+z
90        if ctx.isint(n):
91            return 0*(n+z)
92        r = ctx.re(n)
93        if r == 0:
94            return ctx.nan*(n+z)
95        elif r > 0:
96            return 0*(n+z)
97        else:
98            return ctx.inf+(n+z)
99    M = ctx.mag(z)
100    if derivative:
101        d = ctx.convert(derivative)
102        def h(n,d):
103            r = ctx.fmul(ctx.fmul(z, z, prec=ctx.prec+M), 0.25, exact=True)
104            B = [0.5*(n-d+1), 0.5*(n-d+2), n+1]
105            T = [([2,ctx.pi,z],[d-2*n,0.5,n-d],[n+1],B,[(n+1)*0.5,(n+2)*0.5],B,r)]
106            return T
107        v = ctx.hypercomb(h, [n,d], **kwargs)
108    else:
109        def h(n):
110            w = ctx.fmul(z, 0.5, exact=True)
111            r = ctx.fmul(w, w, prec=max(0,ctx.prec+M))
112            return [([w], [n], [], [n+1], [], [n+1], r)]
113        v = ctx.hypercomb(h, [n], **kwargs)
114    return v
115
116@defun_wrapped
117def bessely(ctx, n, z, derivative=0, **kwargs):
118    if not z:
119        if derivative:
120            # Not implemented
121            raise ValueError
122        if not n:
123            # ~ log(z/2)
124            return -ctx.inf + (n+z)
125        if ctx.im(n):
126            return ctx.nan * (n+z)
127        r = ctx.re(n)
128        q = n+0.5
129        if ctx.isint(q):
130            if n > 0:
131                return -ctx.inf + (n+z)
132            else:
133                return 0 * (n+z)
134        if r < 0 and int(ctx.floor(q)) % 2:
135            return ctx.inf + (n+z)
136        else:
137            return ctx.ninf + (n+z)
138    # XXX: use hypercomb
139    ctx.prec += 10
140    m, d = ctx.nint_distance(n)
141    if d < -ctx.prec:
142        h = +ctx.eps
143        ctx.prec *= 2
144        n += h
145    elif d < 0:
146        ctx.prec -= d
147    # TODO: avoid cancellation for imaginary arguments
148    cos, sin = ctx.cospi_sinpi(n)
149    return (ctx.besselj(n,z,derivative,**kwargs)*cos - \
150        ctx.besselj(-n,z,derivative,**kwargs))/sin
151
152@defun_wrapped
153def besselk(ctx, n, z, **kwargs):
154    if not z:
155        return ctx.inf
156    M = ctx.mag(z)
157    if M < 1:
158        # Represent as limit definition
159        def h(n):
160            r = (z/2)**2
161            T1 = [z, 2], [-n, n-1], [n], [], [], [1-n], r
162            T2 = [z, 2], [n, -n-1], [-n], [], [], [1+n], r
163            return T1, T2
164    # We could use the limit definition always, but it leads
165    # to very bad cancellation (of exponentially large terms)
166    # for large real z
167    # Instead represent in terms of 2F0
168    else:
169        ctx.prec += M
170        def h(n):
171            return [([ctx.pi/2, z, ctx.exp(-z)], [0.5,-0.5,1], [], [], \
172                [n+0.5, 0.5-n], [], -1/(2*z))]
173    return ctx.hypercomb(h, [n], **kwargs)
174
175@defun_wrapped
176def hankel1(ctx,n,x,**kwargs):
177    return ctx.besselj(n,x,**kwargs) + ctx.j*ctx.bessely(n,x,**kwargs)
178
179@defun_wrapped
180def hankel2(ctx,n,x,**kwargs):
181    return ctx.besselj(n,x,**kwargs) - ctx.j*ctx.bessely(n,x,**kwargs)
182
183@defun_wrapped
184def whitm(ctx,k,m,z,**kwargs):
185    if z == 0:
186        # M(k,m,z) = 0^(1/2+m)
187        if ctx.re(m) > -0.5:
188            return z
189        elif ctx.re(m) < -0.5:
190            return ctx.inf + z
191        else:
192            return ctx.nan * z
193    x = ctx.fmul(-0.5, z, exact=True)
194    y = 0.5+m
195    return ctx.exp(x) * z**y * ctx.hyp1f1(y-k, 1+2*m, z, **kwargs)
196
197@defun_wrapped
198def whitw(ctx,k,m,z,**kwargs):
199    if z == 0:
200        g = abs(ctx.re(m))
201        if g < 0.5:
202            return z
203        elif g > 0.5:
204            return ctx.inf + z
205        else:
206            return ctx.nan * z
207    x = ctx.fmul(-0.5, z, exact=True)
208    y = 0.5+m
209    return ctx.exp(x) * z**y * ctx.hyperu(y-k, 1+2*m, z, **kwargs)
210
211@defun
212def hyperu(ctx, a, b, z, **kwargs):
213    a, atype = ctx._convert_param(a)
214    b, btype = ctx._convert_param(b)
215    z = ctx.convert(z)
216    if not z:
217        if ctx.re(b) <= 1:
218            return ctx.gammaprod([1-b],[a-b+1])
219        else:
220            return ctx.inf + z
221    bb = 1+a-b
222    bb, bbtype = ctx._convert_param(bb)
223    try:
224        orig = ctx.prec
225        try:
226            ctx.prec += 10
227            v = ctx.hypsum(2, 0, (atype, bbtype), [a, bb], -1/z, maxterms=ctx.prec)
228            return v / z**a
229        finally:
230            ctx.prec = orig
231    except ctx.NoConvergence:
232        pass
233    def h(a,b):
234        w = ctx.sinpi(b)
235        T1 = ([ctx.pi,w],[1,-1],[],[a-b+1,b],[a],[b],z)
236        T2 = ([-ctx.pi,w,z],[1,-1,1-b],[],[a,2-b],[a-b+1],[2-b],z)
237        return T1, T2
238    return ctx.hypercomb(h, [a,b], **kwargs)
239
240@defun
241def struveh(ctx,n,z, **kwargs):
242    n = ctx.convert(n)
243    z = ctx.convert(z)
244    # http://functions.wolfram.com/Bessel-TypeFunctions/StruveH/26/01/02/
245    def h(n):
246        return [([z/2, 0.5*ctx.sqrt(ctx.pi)], [n+1, -1], [], [n+1.5], [1], [1.5, n+1.5], -(z/2)**2)]
247    return ctx.hypercomb(h, [n], **kwargs)
248
249@defun
250def struvel(ctx,n,z, **kwargs):
251    n = ctx.convert(n)
252    z = ctx.convert(z)
253    # http://functions.wolfram.com/Bessel-TypeFunctions/StruveL/26/01/02/
254    def h(n):
255        return [([z/2, 0.5*ctx.sqrt(ctx.pi)], [n+1, -1], [], [n+1.5], [1], [1.5, n+1.5], (z/2)**2)]
256    return ctx.hypercomb(h, [n], **kwargs)
257
258def _anger(ctx,which,v,z,**kwargs):
259    v = ctx._convert_param(v)[0]
260    z = ctx.convert(z)
261    def h(v):
262        b = ctx.mpq_1_2
263        u = v*b
264        m = b*3
265        a1,a2,b1,b2 = m-u, m+u, 1-u, 1+u
266        c, s = ctx.cospi_sinpi(u)
267        if which == 0:
268            A, B = [b*z, s], [c]
269        if which == 1:
270            A, B = [b*z, -c], [s]
271        w = ctx.square_exp_arg(z, mult=-0.25)
272        T1 = A, [1, 1], [], [a1,a2], [1], [a1,a2], w
273        T2 = B, [1], [], [b1,b2], [1], [b1,b2], w
274        return T1, T2
275    return ctx.hypercomb(h, [v], **kwargs)
276
277@defun
278def angerj(ctx, v, z, **kwargs):
279    return _anger(ctx, 0, v, z, **kwargs)
280
281@defun
282def webere(ctx, v, z, **kwargs):
283    return _anger(ctx, 1, v, z, **kwargs)
284
285@defun
286def lommels1(ctx, u, v, z, **kwargs):
287    u = ctx._convert_param(u)[0]
288    v = ctx._convert_param(v)[0]
289    z = ctx.convert(z)
290    def h(u,v):
291        b = ctx.mpq_1_2
292        w = ctx.square_exp_arg(z, mult=-0.25)
293        return ([u-v+1, u+v+1, z], [-1, -1, u+1], [], [], [1], \
294            [b*(u-v+3),b*(u+v+3)], w),
295    return ctx.hypercomb(h, [u,v], **kwargs)
296
297@defun
298def lommels2(ctx, u, v, z, **kwargs):
299    u = ctx._convert_param(u)[0]
300    v = ctx._convert_param(v)[0]
301    z = ctx.convert(z)
302    # Asymptotic expansion (GR p. 947) -- need to be careful
303    # not to use for small arguments
304    # def h(u,v):
305    #    b = ctx.mpq_1_2
306    #    w = -(z/2)**(-2)
307    #    return ([z], [u-1], [], [], [b*(1-u+v)], [b*(1-u-v)], w),
308    def h(u,v):
309        b = ctx.mpq_1_2
310        w = ctx.square_exp_arg(z, mult=-0.25)
311        T1 = [u-v+1, u+v+1, z], [-1, -1, u+1], [], [], [1], [b*(u-v+3),b*(u+v+3)], w
312        T2 = [2, z], [u+v-1, -v], [v, b*(u+v+1)], [b*(v-u+1)], [], [1-v], w
313        T3 = [2, z], [u-v-1, v], [-v, b*(u-v+1)], [b*(1-u-v)], [], [1+v], w
314        #c1 = ctx.cospi((u-v)*b)
315        #c2 = ctx.cospi((u+v)*b)
316        #s = ctx.sinpi(v)
317        #r1 = (u-v+1)*b
318        #r2 = (u+v+1)*b
319        #T2 = [c1, s, z, 2], [1, -1, -v, v], [], [-v+1], [], [-v+1], w
320        #T3 = [-c2, s, z, 2], [1, -1, v, -v], [], [v+1], [], [v+1], w
321        #T2 = [c1, s, z, 2], [1, -1, -v, v+u-1], [r1, r2], [-v+1], [], [-v+1], w
322        #T3 = [-c2, s, z, 2], [1, -1, v, -v+u-1], [r1, r2], [v+1], [], [v+1], w
323        return T1, T2, T3
324    return ctx.hypercomb(h, [u,v], **kwargs)
325
326@defun
327def ber(ctx, n, z, **kwargs):
328    n = ctx.convert(n)
329    z = ctx.convert(z)
330    # http://functions.wolfram.com/Bessel-TypeFunctions/KelvinBer2/26/01/02/0001/
331    def h(n):
332        r = -(z/4)**4
333        cos, sin = ctx.cospi_sinpi(-0.75*n)
334        T1 = [cos, z/2], [1, n], [], [n+1], [], [0.5, 0.5*(n+1), 0.5*n+1], r
335        T2 = [sin, z/2], [1, n+2], [], [n+2], [], [1.5, 0.5*(n+3), 0.5*n+1], r
336        return T1, T2
337    return ctx.hypercomb(h, [n], **kwargs)
338
339@defun
340def bei(ctx, n, z, **kwargs):
341    n = ctx.convert(n)
342    z = ctx.convert(z)
343    # http://functions.wolfram.com/Bessel-TypeFunctions/KelvinBei2/26/01/02/0001/
344    def h(n):
345        r = -(z/4)**4
346        cos, sin = ctx.cospi_sinpi(0.75*n)
347        T1 = [cos, z/2], [1, n+2], [], [n+2], [], [1.5, 0.5*(n+3), 0.5*n+1], r
348        T2 = [sin, z/2], [1, n], [], [n+1], [], [0.5, 0.5*(n+1), 0.5*n+1], r
349        return T1, T2
350    return ctx.hypercomb(h, [n], **kwargs)
351
352@defun
353def ker(ctx, n, z, **kwargs):
354    n = ctx.convert(n)
355    z = ctx.convert(z)
356    # http://functions.wolfram.com/Bessel-TypeFunctions/KelvinKer2/26/01/02/0001/
357    def h(n):
358        r = -(z/4)**4
359        cos1, sin1 = ctx.cospi_sinpi(0.25*n)
360        cos2, sin2 = ctx.cospi_sinpi(0.75*n)
361        T1 = [2, z, 4*cos1], [-n-3, n, 1], [-n], [], [], [0.5, 0.5*(1+n), 0.5*(n+2)], r
362        T2 = [2, z, -sin1], [-n-3, 2+n, 1], [-n-1], [], [], [1.5, 0.5*(3+n), 0.5*(n+2)], r
363        T3 = [2, z, 4*cos2], [n-3, -n, 1], [n], [], [], [0.5, 0.5*(1-n), 1-0.5*n], r
364        T4 = [2, z, -sin2], [n-3, 2-n, 1], [n-1], [], [], [1.5, 0.5*(3-n), 1-0.5*n], r
365        return T1, T2, T3, T4
366    return ctx.hypercomb(h, [n], **kwargs)
367
368@defun
369def kei(ctx, n, z, **kwargs):
370    n = ctx.convert(n)
371    z = ctx.convert(z)
372    # http://functions.wolfram.com/Bessel-TypeFunctions/KelvinKei2/26/01/02/0001/
373    def h(n):
374        r = -(z/4)**4
375        cos1, sin1 = ctx.cospi_sinpi(0.75*n)
376        cos2, sin2 = ctx.cospi_sinpi(0.25*n)
377        T1 = [-cos1, 2, z], [1, n-3, 2-n], [n-1], [], [], [1.5, 0.5*(3-n), 1-0.5*n], r
378        T2 = [-sin1, 2, z], [1, n-1, -n], [n], [], [], [0.5, 0.5*(1-n), 1-0.5*n], r
379        T3 = [-sin2, 2, z], [1, -n-1, n], [-n], [], [], [0.5, 0.5*(n+1), 0.5*(n+2)], r
380        T4 = [-cos2, 2, z], [1, -n-3, n+2], [-n-1], [], [], [1.5, 0.5*(n+3), 0.5*(n+2)], r
381        return T1, T2, T3, T4
382    return ctx.hypercomb(h, [n], **kwargs)
383
384# TODO: do this more generically?
385def c_memo(f):
386    name = f.__name__
387    def f_wrapped(ctx):
388        cache = ctx._misc_const_cache
389        prec = ctx.prec
390        p,v = cache.get(name, (-1,0))
391        if p >= prec:
392            return +v
393        else:
394            cache[name] = (prec, f(ctx))
395            return cache[name][1]
396    return f_wrapped
397
398@c_memo
399def _airyai_C1(ctx):
400    return 1 / (ctx.cbrt(9) * ctx.gamma(ctx.mpf(2)/3))
401
402@c_memo
403def _airyai_C2(ctx):
404    return -1 / (ctx.cbrt(3) * ctx.gamma(ctx.mpf(1)/3))
405
406@c_memo
407def _airybi_C1(ctx):
408    return 1 / (ctx.nthroot(3,6) * ctx.gamma(ctx.mpf(2)/3))
409
410@c_memo
411def _airybi_C2(ctx):
412    return ctx.nthroot(3,6) / ctx.gamma(ctx.mpf(1)/3)
413
414def _airybi_n2_inf(ctx):
415    prec = ctx.prec
416    try:
417        v = ctx.power(3,'2/3')*ctx.gamma('2/3')/(2*ctx.pi)
418    finally:
419        ctx.prec = prec
420    return +v
421
422# Derivatives at z = 0
423# TODO: could be expressed more elegantly using triple factorials
424def _airyderiv_0(ctx, z, n, ntype, which):
425    if ntype == 'Z':
426        if n < 0:
427            return z
428        r = ctx.mpq_1_3
429        prec = ctx.prec
430        try:
431            ctx.prec += 10
432            v = ctx.gamma((n+1)*r) * ctx.power(3,n*r) / ctx.pi
433            if which == 0:
434                v *= ctx.sinpi(2*(n+1)*r)
435                v /= ctx.power(3,'2/3')
436            else:
437                v *= abs(ctx.sinpi(2*(n+1)*r))
438                v /= ctx.power(3,'1/6')
439        finally:
440            ctx.prec = prec
441        return +v + z
442    else:
443        # singular (does the limit exist?)
444        raise NotImplementedError
445
446@defun
447def airyai(ctx, z, derivative=0, **kwargs):
448    z = ctx.convert(z)
449    if derivative:
450        n, ntype = ctx._convert_param(derivative)
451    else:
452        n = 0
453    # Values at infinities
454    if not ctx.isnormal(z) and z:
455        if n and ntype == 'Z':
456            if n == -1:
457                if z == ctx.inf:
458                    return ctx.mpf(1)/3 + 1/z
459                if z == ctx.ninf:
460                    return ctx.mpf(-2)/3 + 1/z
461            if n < -1:
462                if z == ctx.inf:
463                    return z
464                if z == ctx.ninf:
465                    return (-1)**n * (-z)
466        if (not n) and z == ctx.inf or z == ctx.ninf:
467            return 1/z
468        # TODO: limits
469        raise ValueError("essential singularity of Ai(z)")
470    # Account for exponential scaling
471    if z:
472        extraprec = max(0, int(1.5*ctx.mag(z)))
473    else:
474        extraprec = 0
475    if n:
476        if n == 1:
477            def h():
478                # http://functions.wolfram.com/03.07.06.0005.01
479                if ctx._re(z) > 4:
480                    ctx.prec += extraprec
481                    w = z**1.5; r = -0.75/w; u = -2*w/3
482                    ctx.prec -= extraprec
483                    C = -ctx.exp(u)/(2*ctx.sqrt(ctx.pi))*ctx.nthroot(z,4)
484                    return ([C],[1],[],[],[(-1,6),(7,6)],[],r),
485                # http://functions.wolfram.com/03.07.26.0001.01
486                else:
487                    ctx.prec += extraprec
488                    w = z**3 / 9
489                    ctx.prec -= extraprec
490                    C1 = _airyai_C1(ctx) * 0.5
491                    C2 = _airyai_C2(ctx)
492                    T1 = [C1,z],[1,2],[],[],[],[ctx.mpq_5_3],w
493                    T2 = [C2],[1],[],[],[],[ctx.mpq_1_3],w
494                    return T1, T2
495            return ctx.hypercomb(h, [], **kwargs)
496        else:
497            if z == 0:
498                return _airyderiv_0(ctx, z, n, ntype, 0)
499            # http://functions.wolfram.com/03.05.20.0004.01
500            def h(n):
501                ctx.prec += extraprec
502                w = z**3/9
503                ctx.prec -= extraprec
504                q13,q23,q43 = ctx.mpq_1_3, ctx.mpq_2_3, ctx.mpq_4_3
505                a1=q13; a2=1; b1=(1-n)*q13; b2=(2-n)*q13; b3=1-n*q13
506                T1 = [3, z], [n-q23, -n], [a1], [b1,b2,b3], \
507                    [a1,a2], [b1,b2,b3], w
508                a1=q23; b1=(2-n)*q13; b2=1-n*q13; b3=(4-n)*q13
509                T2 = [3, z, -z], [n-q43, -n, 1], [a1], [b1,b2,b3], \
510                    [a1,a2], [b1,b2,b3], w
511                return T1, T2
512            v = ctx.hypercomb(h, [n], **kwargs)
513            if ctx._is_real_type(z) and ctx.isint(n):
514                v = ctx._re(v)
515            return v
516    else:
517        def h():
518            if ctx._re(z) > 4:
519                # We could use 1F1, but it results in huge cancellation;
520                # the following expansion is better.
521                # TODO: asymptotic series for derivatives
522                ctx.prec += extraprec
523                w = z**1.5; r = -0.75/w; u = -2*w/3
524                ctx.prec -= extraprec
525                C = ctx.exp(u)/(2*ctx.sqrt(ctx.pi)*ctx.nthroot(z,4))
526                return ([C],[1],[],[],[(1,6),(5,6)],[],r),
527            else:
528                ctx.prec += extraprec
529                w = z**3 / 9
530                ctx.prec -= extraprec
531                C1 = _airyai_C1(ctx)
532                C2 = _airyai_C2(ctx)
533                T1 = [C1],[1],[],[],[],[ctx.mpq_2_3],w
534                T2 = [z*C2],[1],[],[],[],[ctx.mpq_4_3],w
535                return T1, T2
536        return ctx.hypercomb(h, [], **kwargs)
537
538@defun
539def airybi(ctx, z, derivative=0, **kwargs):
540    z = ctx.convert(z)
541    if derivative:
542        n, ntype = ctx._convert_param(derivative)
543    else:
544        n = 0
545    # Values at infinities
546    if not ctx.isnormal(z) and z:
547        if n and ntype == 'Z':
548            if z == ctx.inf:
549                return z
550            if z == ctx.ninf:
551                if n == -1:
552                    return 1/z
553                if n == -2:
554                    return _airybi_n2_inf(ctx)
555                if n < -2:
556                    return (-1)**n * (-z)
557        if not n:
558            if z == ctx.inf:
559                return z
560            if z == ctx.ninf:
561                return 1/z
562        # TODO: limits
563        raise ValueError("essential singularity of Bi(z)")
564    if z:
565        extraprec = max(0, int(1.5*ctx.mag(z)))
566    else:
567        extraprec = 0
568    if n:
569        if n == 1:
570            # http://functions.wolfram.com/03.08.26.0001.01
571            def h():
572                ctx.prec += extraprec
573                w = z**3 / 9
574                ctx.prec -= extraprec
575                C1 = _airybi_C1(ctx)*0.5
576                C2 = _airybi_C2(ctx)
577                T1 = [C1,z],[1,2],[],[],[],[ctx.mpq_5_3],w
578                T2 = [C2],[1],[],[],[],[ctx.mpq_1_3],w
579                return T1, T2
580            return ctx.hypercomb(h, [], **kwargs)
581        else:
582            if z == 0:
583                return _airyderiv_0(ctx, z, n, ntype, 1)
584            def h(n):
585                ctx.prec += extraprec
586                w = z**3/9
587                ctx.prec -= extraprec
588                q13,q23,q43 = ctx.mpq_1_3, ctx.mpq_2_3, ctx.mpq_4_3
589                q16 = ctx.mpq_1_6
590                q56 = ctx.mpq_5_6
591                a1=q13; a2=1; b1=(1-n)*q13; b2=(2-n)*q13; b3=1-n*q13
592                T1 = [3, z], [n-q16, -n], [a1], [b1,b2,b3], \
593                    [a1,a2], [b1,b2,b3], w
594                a1=q23; b1=(2-n)*q13; b2=1-n*q13; b3=(4-n)*q13
595                T2 = [3, z], [n-q56, 1-n], [a1], [b1,b2,b3], \
596                    [a1,a2], [b1,b2,b3], w
597                return T1, T2
598            v = ctx.hypercomb(h, [n], **kwargs)
599            if ctx._is_real_type(z) and ctx.isint(n):
600                v = ctx._re(v)
601            return v
602    else:
603        def h():
604            ctx.prec += extraprec
605            w = z**3 / 9
606            ctx.prec -= extraprec
607            C1 = _airybi_C1(ctx)
608            C2 = _airybi_C2(ctx)
609            T1 = [C1],[1],[],[],[],[ctx.mpq_2_3],w
610            T2 = [z*C2],[1],[],[],[],[ctx.mpq_4_3],w
611            return T1, T2
612        return ctx.hypercomb(h, [], **kwargs)
613
614def _airy_zero(ctx, which, k, derivative, complex=False):
615    # Asymptotic formulas are given in DLMF section 9.9
616    def U(t): return t**(2/3.)*(1-7/(t**2*48))
617    def T(t): return t**(2/3.)*(1+5/(t**2*48))
618    k = int(k)
619    if k < 1:
620        raise ValueError("k cannot be less than 1")
621    if not derivative in (0,1):
622        raise ValueError("Derivative should lie between 0 and 1")
623    if which == 0:
624        if derivative:
625            return ctx.findroot(lambda z: ctx.airyai(z,1),
626                -U(3*ctx.pi*(4*k-3)/8))
627        return ctx.findroot(ctx.airyai, -T(3*ctx.pi*(4*k-1)/8))
628    if which == 1 and complex == False:
629        if derivative:
630            return ctx.findroot(lambda z: ctx.airybi(z,1),
631                -U(3*ctx.pi*(4*k-1)/8))
632        return ctx.findroot(ctx.airybi, -T(3*ctx.pi*(4*k-3)/8))
633    if which == 1 and complex == True:
634        if derivative:
635            t = 3*ctx.pi*(4*k-3)/8 + 0.75j*ctx.ln2
636            s = ctx.expjpi(ctx.mpf(1)/3) * T(t)
637            return ctx.findroot(lambda z: ctx.airybi(z,1), s)
638        t = 3*ctx.pi*(4*k-1)/8 + 0.75j*ctx.ln2
639        s = ctx.expjpi(ctx.mpf(1)/3) * U(t)
640        return ctx.findroot(ctx.airybi, s)
641
642@defun
643def airyaizero(ctx, k, derivative=0):
644    return _airy_zero(ctx, 0, k, derivative, False)
645
646@defun
647def airybizero(ctx, k, derivative=0, complex=False):
648    return _airy_zero(ctx, 1, k, derivative, complex)
649
650def _scorer(ctx, z, which, kwargs):
651    z = ctx.convert(z)
652    if ctx.isinf(z):
653        if z == ctx.inf:
654            if which == 0: return 1/z
655            if which == 1: return z
656        if z == ctx.ninf:
657            return 1/z
658        raise ValueError("essential singularity")
659    if z:
660        extraprec = max(0, int(1.5*ctx.mag(z)))
661    else:
662        extraprec = 0
663    if kwargs.get('derivative'):
664        raise NotImplementedError
665    # Direct asymptotic expansions, to avoid
666    # exponentially large cancellation
667    try:
668        if ctx.mag(z) > 3:
669            if which == 0 and abs(ctx.arg(z)) < ctx.pi/3 * 0.999:
670                def h():
671                    return (([ctx.pi,z],[-1,-1],[],[],[(1,3),(2,3),1],[],9/z**3),)
672                return ctx.hypercomb(h, [], maxterms=ctx.prec, force_series=True)
673            if which == 1 and abs(ctx.arg(-z)) < 2*ctx.pi/3 * 0.999:
674                def h():
675                    return (([-ctx.pi,z],[-1,-1],[],[],[(1,3),(2,3),1],[],9/z**3),)
676                return ctx.hypercomb(h, [], maxterms=ctx.prec, force_series=True)
677    except ctx.NoConvergence:
678        pass
679    def h():
680        A = ctx.airybi(z, **kwargs)/3
681        B = -2*ctx.pi
682        if which == 1:
683            A *= 2
684            B *= -1
685        ctx.prec += extraprec
686        w = z**3/9
687        ctx.prec -= extraprec
688        T1 = [A], [1], [], [], [], [], 0
689        T2 = [B,z], [-1,2], [], [], [1], [ctx.mpq_4_3,ctx.mpq_5_3], w
690        return T1, T2
691    return ctx.hypercomb(h, [], **kwargs)
692
693@defun
694def scorergi(ctx, z, **kwargs):
695    return _scorer(ctx, z, 0, kwargs)
696
697@defun
698def scorerhi(ctx, z, **kwargs):
699    return _scorer(ctx, z, 1, kwargs)
700
701@defun_wrapped
702def coulombc(ctx, l, eta, _cache={}):
703    if (l, eta) in _cache and _cache[l,eta][0] >= ctx.prec:
704        return +_cache[l,eta][1]
705    G3 = ctx.loggamma(2*l+2)
706    G1 = ctx.loggamma(1+l+ctx.j*eta)
707    G2 = ctx.loggamma(1+l-ctx.j*eta)
708    v = 2**l * ctx.exp((-ctx.pi*eta+G1+G2)/2 - G3)
709    if not (ctx.im(l) or ctx.im(eta)):
710        v = ctx.re(v)
711    _cache[l,eta] = (ctx.prec, v)
712    return v
713
714@defun_wrapped
715def coulombf(ctx, l, eta, z, w=1, chop=True, **kwargs):
716    # Regular Coulomb wave function
717    # Note: w can be either 1 or -1; the other may be better in some cases
718    # TODO: check that chop=True chops when and only when it should
719    #ctx.prec += 10
720    def h(l, eta):
721        try:
722            jw = ctx.j*w
723            jwz = ctx.fmul(jw, z, exact=True)
724            jwz2 = ctx.fmul(jwz, -2, exact=True)
725            C = ctx.coulombc(l, eta)
726            T1 = [C, z, ctx.exp(jwz)], [1, l+1, 1], [], [], [1+l+jw*eta], \
727                [2*l+2], jwz2
728        except ValueError:
729            T1 = [0], [-1], [], [], [], [], 0
730        return (T1,)
731    v = ctx.hypercomb(h, [l,eta], **kwargs)
732    if chop and (not ctx.im(l)) and (not ctx.im(eta)) and (not ctx.im(z)) and \
733        (ctx.re(z) >= 0):
734        v = ctx.re(v)
735    return v
736
737@defun_wrapped
738def _coulomb_chi(ctx, l, eta, _cache={}):
739    if (l, eta) in _cache and _cache[l,eta][0] >= ctx.prec:
740        return _cache[l,eta][1]
741    def terms():
742        l2 = -l-1
743        jeta = ctx.j*eta
744        return [ctx.loggamma(1+l+jeta) * (-0.5j),
745            ctx.loggamma(1+l-jeta) * (0.5j),
746            ctx.loggamma(1+l2+jeta) * (0.5j),
747            ctx.loggamma(1+l2-jeta) * (-0.5j),
748            -(l+0.5)*ctx.pi]
749    v = ctx.sum_accurately(terms, 1)
750    _cache[l,eta] = (ctx.prec, v)
751    return v
752
753@defun_wrapped
754def coulombg(ctx, l, eta, z, w=1, chop=True, **kwargs):
755    # Irregular Coulomb wave function
756    # Note: w can be either 1 or -1; the other may be better in some cases
757    # TODO: check that chop=True chops when and only when it should
758    if not ctx._im(l):
759        l = ctx._re(l)  # XXX: for isint
760    def h(l, eta):
761        # Force perturbation for integers and half-integers
762        if ctx.isint(l*2):
763            T1 = [0], [-1], [], [], [], [], 0
764            return (T1,)
765        l2 = -l-1
766        try:
767            chi = ctx._coulomb_chi(l, eta)
768            jw = ctx.j*w
769            s = ctx.sin(chi); c = ctx.cos(chi)
770            C1 = ctx.coulombc(l,eta)
771            C2 = ctx.coulombc(l2,eta)
772            u = ctx.exp(jw*z)
773            x = -2*jw*z
774            T1 = [s, C1, z, u, c], [-1, 1, l+1, 1, 1], [], [], \
775                [1+l+jw*eta], [2*l+2], x
776            T2 = [-s, C2, z, u],   [-1, 1, l2+1, 1],    [], [], \
777                [1+l2+jw*eta], [2*l2+2], x
778            return T1, T2
779        except ValueError:
780            T1 = [0], [-1], [], [], [], [], 0
781            return (T1,)
782    v = ctx.hypercomb(h, [l,eta], **kwargs)
783    if chop and (not ctx._im(l)) and (not ctx._im(eta)) and (not ctx._im(z)) and \
784        (ctx._re(z) >= 0):
785        v = ctx._re(v)
786    return v
787
788def mcmahon(ctx,kind,prime,v,m):
789    """
790    Computes an estimate for the location of the Bessel function zero
791    j_{v,m}, y_{v,m}, j'_{v,m} or y'_{v,m} using McMahon's asymptotic
792    expansion (Abramowitz & Stegun 9.5.12-13, DLMF 20.21(vi)).
793
794    Returns (r,err) where r is the estimated location of the root
795    and err is a positive number estimating the error of the
796    asymptotic expansion.
797    """
798    u = 4*v**2
799    if kind == 1 and not prime: b = (4*m+2*v-1)*ctx.pi/4
800    if kind == 2 and not prime: b = (4*m+2*v-3)*ctx.pi/4
801    if kind == 1 and prime: b = (4*m+2*v-3)*ctx.pi/4
802    if kind == 2 and prime: b = (4*m+2*v-1)*ctx.pi/4
803    if not prime:
804        s1 = b
805        s2 = -(u-1)/(8*b)
806        s3 = -4*(u-1)*(7*u-31)/(3*(8*b)**3)
807        s4 = -32*(u-1)*(83*u**2-982*u+3779)/(15*(8*b)**5)
808        s5 = -64*(u-1)*(6949*u**3-153855*u**2+1585743*u-6277237)/(105*(8*b)**7)
809    if prime:
810        s1 = b
811        s2 = -(u+3)/(8*b)
812        s3 = -4*(7*u**2+82*u-9)/(3*(8*b)**3)
813        s4 = -32*(83*u**3+2075*u**2-3039*u+3537)/(15*(8*b)**5)
814        s5 = -64*(6949*u**4+296492*u**3-1248002*u**2+7414380*u-5853627)/(105*(8*b)**7)
815    terms = [s1,s2,s3,s4,s5]
816    s = s1
817    err = 0.0
818    for i in range(1,len(terms)):
819        if abs(terms[i]) < abs(terms[i-1]):
820            s += terms[i]
821        else:
822            err = abs(terms[i])
823    if i == len(terms)-1:
824        err = abs(terms[-1])
825    return s, err
826
827def generalized_bisection(ctx,f,a,b,n):
828    """
829    Given f known to have exactly n simple roots within [a,b],
830    return a list of n intervals isolating the roots
831    and having opposite signs at the endpoints.
832
833    TODO: this can be optimized, e.g. by reusing evaluation points.
834    """
835    if n < 1:
836        raise ValueError("n cannot be less than 1")
837    N = n+1
838    points = []
839    signs = []
840    while 1:
841        points = ctx.linspace(a,b,N)
842        signs = [ctx.sign(f(x)) for x in points]
843        ok_intervals = [(points[i],points[i+1]) for i in range(N-1) \
844            if signs[i]*signs[i+1] == -1]
845        if len(ok_intervals) == n:
846            return ok_intervals
847        N = N*2
848
849def find_in_interval(ctx, f, ab):
850    return ctx.findroot(f, ab, solver='illinois', verify=False)
851
852def bessel_zero(ctx, kind, prime, v, m, isoltol=0.01, _interval_cache={}):
853    prec = ctx.prec
854    workprec = max(prec, ctx.mag(v), ctx.mag(m))+10
855    try:
856        ctx.prec = workprec
857        v = ctx.mpf(v)
858        m = int(m)
859        prime = int(prime)
860        if v < 0:
861            raise ValueError("v cannot be negative")
862        if m < 1:
863            raise ValueError("m cannot be less than 1")
864        if not prime in (0,1):
865            raise ValueError("prime should lie between 0 and 1")
866        if kind == 1:
867            if prime: f = lambda x: ctx.besselj(v,x,derivative=1)
868            else:     f = lambda x: ctx.besselj(v,x)
869        if kind == 2:
870            if prime: f = lambda x: ctx.bessely(v,x,derivative=1)
871            else:     f = lambda x: ctx.bessely(v,x)
872        # The first root of J' is very close to 0 for small
873        # orders, and this needs to be special-cased
874        if kind == 1 and prime and m == 1:
875            if v == 0:
876                return ctx.zero
877            if v <= 1:
878                # TODO: use v <= j'_{v,1} < y_{v,1}?
879                r = 2*ctx.sqrt(v*(1+v)/(v+2))
880                return find_in_interval(ctx, f, (r/10, 2*r))
881        if (kind,prime,v,m) in _interval_cache:
882            return find_in_interval(ctx, f, _interval_cache[kind,prime,v,m])
883        r, err = mcmahon(ctx, kind, prime, v, m)
884        if err < isoltol:
885            return find_in_interval(ctx, f, (r-isoltol, r+isoltol))
886        # An x such that 0 < x < r_{v,1}
887        if kind == 1 and not prime: low = 2.4
888        if kind == 1 and prime: low = 1.8
889        if kind == 2 and not prime: low = 0.8
890        if kind == 2 and prime: low = 2.0
891        n = m+1
892        while 1:
893            r1, err = mcmahon(ctx, kind, prime, v, n)
894            if err < isoltol:
895                r2, err2 = mcmahon(ctx, kind, prime, v, n+1)
896                intervals = generalized_bisection(ctx, f, low, 0.5*(r1+r2), n)
897                for k, ab in enumerate(intervals):
898                    _interval_cache[kind,prime,v,k+1] = ab
899                return find_in_interval(ctx, f, intervals[m-1])
900            else:
901                n = n*2
902    finally:
903        ctx.prec = prec
904
905@defun
906def besseljzero(ctx, v, m, derivative=0):
907    r"""
908    For a real order `\nu \ge 0` and a positive integer `m`, returns
909    `j_{\nu,m}`, the `m`-th positive zero of the Bessel function of the
910    first kind `J_{\nu}(z)` (see :func:`~mpmath.besselj`). Alternatively,
911    with *derivative=1*, gives the first nonnegative simple zero
912    `j'_{\nu,m}` of `J'_{\nu}(z)`.
913
914    The indexing convention is that used by Abramowitz & Stegun
915    and the DLMF. Note the special case `j'_{0,1} = 0`, while all other
916    zeros are positive. In effect, only simple zeros are counted
917    (all zeros of Bessel functions are simple except possibly `z = 0`)
918    and `j_{\nu,m}` becomes a monotonic function of both `\nu`
919    and `m`.
920
921    The zeros are interlaced according to the inequalities
922
923    .. math ::
924
925        j'_{\nu,k} < j_{\nu,k} < j'_{\nu,k+1}
926
927        j_{\nu,1} < j_{\nu+1,2} < j_{\nu,2} < j_{\nu+1,2} < j_{\nu,3} < \cdots
928
929    **Examples**
930
931    Initial zeros of the Bessel functions `J_0(z), J_1(z), J_2(z)`::
932
933        >>> from mpmath import *
934        >>> mp.dps = 25; mp.pretty = True
935        >>> besseljzero(0,1); besseljzero(0,2); besseljzero(0,3)
936        2.404825557695772768621632
937        5.520078110286310649596604
938        8.653727912911012216954199
939        >>> besseljzero(1,1); besseljzero(1,2); besseljzero(1,3)
940        3.831705970207512315614436
941        7.01558666981561875353705
942        10.17346813506272207718571
943        >>> besseljzero(2,1); besseljzero(2,2); besseljzero(2,3)
944        5.135622301840682556301402
945        8.417244140399864857783614
946        11.61984117214905942709415
947
948    Initial zeros of `J'_0(z), J'_1(z), J'_2(z)`::
949
950        0.0
951        3.831705970207512315614436
952        7.01558666981561875353705
953        >>> besseljzero(1,1,1); besseljzero(1,2,1); besseljzero(1,3,1)
954        1.84118378134065930264363
955        5.331442773525032636884016
956        8.536316366346285834358961
957        >>> besseljzero(2,1,1); besseljzero(2,2,1); besseljzero(2,3,1)
958        3.054236928227140322755932
959        6.706133194158459146634394
960        9.969467823087595793179143
961
962    Zeros with large index::
963
964        >>> besseljzero(0,100); besseljzero(0,1000); besseljzero(0,10000)
965        313.3742660775278447196902
966        3140.807295225078628895545
967        31415.14114171350798533666
968        >>> besseljzero(5,100); besseljzero(5,1000); besseljzero(5,10000)
969        321.1893195676003157339222
970        3148.657306813047523500494
971        31422.9947255486291798943
972        >>> besseljzero(0,100,1); besseljzero(0,1000,1); besseljzero(0,10000,1)
973        311.8018681873704508125112
974        3139.236339643802482833973
975        31413.57032947022399485808
976
977    Zeros of functions with large order::
978
979        >>> besseljzero(50,1)
980        57.11689916011917411936228
981        >>> besseljzero(50,2)
982        62.80769876483536093435393
983        >>> besseljzero(50,100)
984        388.6936600656058834640981
985        >>> besseljzero(50,1,1)
986        52.99764038731665010944037
987        >>> besseljzero(50,2,1)
988        60.02631933279942589882363
989        >>> besseljzero(50,100,1)
990        387.1083151608726181086283
991
992    Zeros of functions with fractional order::
993
994        >>> besseljzero(0.5,1); besseljzero(1.5,1); besseljzero(2.25,4)
995        3.141592653589793238462643
996        4.493409457909064175307881
997        15.15657692957458622921634
998
999    Both `J_{\nu}(z)` and `J'_{\nu}(z)` can be expressed as infinite
1000    products over their zeros::
1001
1002        >>> v,z = 2, mpf(1)
1003        >>> (z/2)**v/gamma(v+1) * \
1004        ...     nprod(lambda k: 1-(z/besseljzero(v,k))**2, [1,inf])
1005        ...
1006        0.1149034849319004804696469
1007        >>> besselj(v,z)
1008        0.1149034849319004804696469
1009        >>> (z/2)**(v-1)/2/gamma(v) * \
1010        ...     nprod(lambda k: 1-(z/besseljzero(v,k,1))**2, [1,inf])
1011        ...
1012        0.2102436158811325550203884
1013        >>> besselj(v,z,1)
1014        0.2102436158811325550203884
1015
1016    """
1017    return +bessel_zero(ctx, 1, derivative, v, m)
1018
1019@defun
1020def besselyzero(ctx, v, m, derivative=0):
1021    r"""
1022    For a real order `\nu \ge 0` and a positive integer `m`, returns
1023    `y_{\nu,m}`, the `m`-th positive zero of the Bessel function of the
1024    second kind `Y_{\nu}(z)` (see :func:`~mpmath.bessely`). Alternatively,
1025    with *derivative=1*, gives the first positive zero `y'_{\nu,m}` of
1026    `Y'_{\nu}(z)`.
1027
1028    The zeros are interlaced according to the inequalities
1029
1030    .. math ::
1031
1032        y_{\nu,k} < y'_{\nu,k} < y_{\nu,k+1}
1033
1034        y_{\nu,1} < y_{\nu+1,2} < y_{\nu,2} < y_{\nu+1,2} < y_{\nu,3} < \cdots
1035
1036    **Examples**
1037
1038    Initial zeros of the Bessel functions `Y_0(z), Y_1(z), Y_2(z)`::
1039
1040        >>> from mpmath import *
1041        >>> mp.dps = 25; mp.pretty = True
1042        >>> besselyzero(0,1); besselyzero(0,2); besselyzero(0,3)
1043        0.8935769662791675215848871
1044        3.957678419314857868375677
1045        7.086051060301772697623625
1046        >>> besselyzero(1,1); besselyzero(1,2); besselyzero(1,3)
1047        2.197141326031017035149034
1048        5.429681040794135132772005
1049        8.596005868331168926429606
1050        >>> besselyzero(2,1); besselyzero(2,2); besselyzero(2,3)
1051        3.384241767149593472701426
1052        6.793807513268267538291167
1053        10.02347797936003797850539
1054
1055    Initial zeros of `Y'_0(z), Y'_1(z), Y'_2(z)`::
1056
1057        >>> besselyzero(0,1,1); besselyzero(0,2,1); besselyzero(0,3,1)
1058        2.197141326031017035149034
1059        5.429681040794135132772005
1060        8.596005868331168926429606
1061        >>> besselyzero(1,1,1); besselyzero(1,2,1); besselyzero(1,3,1)
1062        3.683022856585177699898967
1063        6.941499953654175655751944
1064        10.12340465543661307978775
1065        >>> besselyzero(2,1,1); besselyzero(2,2,1); besselyzero(2,3,1)
1066        5.002582931446063945200176
1067        8.350724701413079526349714
1068        11.57419546521764654624265
1069
1070    Zeros with large index::
1071
1072        >>> besselyzero(0,100); besselyzero(0,1000); besselyzero(0,10000)
1073        311.8034717601871549333419
1074        3139.236498918198006794026
1075        31413.57034538691205229188
1076        >>> besselyzero(5,100); besselyzero(5,1000); besselyzero(5,10000)
1077        319.6183338562782156235062
1078        3147.086508524556404473186
1079        31421.42392920214673402828
1080        >>> besselyzero(0,100,1); besselyzero(0,1000,1); besselyzero(0,10000,1)
1081        313.3726705426359345050449
1082        3140.807136030340213610065
1083        31415.14112579761578220175
1084
1085    Zeros of functions with large order::
1086
1087        >>> besselyzero(50,1)
1088        53.50285882040036394680237
1089        >>> besselyzero(50,2)
1090        60.11244442774058114686022
1091        >>> besselyzero(50,100)
1092        387.1096509824943957706835
1093        >>> besselyzero(50,1,1)
1094        56.96290427516751320063605
1095        >>> besselyzero(50,2,1)
1096        62.74888166945933944036623
1097        >>> besselyzero(50,100,1)
1098        388.6923300548309258355475
1099
1100    Zeros of functions with fractional order::
1101
1102        >>> besselyzero(0.5,1); besselyzero(1.5,1); besselyzero(2.25,4)
1103        1.570796326794896619231322
1104        2.798386045783887136720249
1105        13.56721208770735123376018
1106
1107    """
1108    return +bessel_zero(ctx, 2, derivative, v, m)
1109