You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

494 lines
16 KiB

5 months ago
from .functions import defun, defun_wrapped
def _hermite_param(ctx, n, z, parabolic_cylinder):
"""
Combined calculation of the Hermite polynomial H_n(z) (and its
generalization to complex n) and the parabolic cylinder
function D.
"""
n, ntyp = ctx._convert_param(n)
z = ctx.convert(z)
q = -ctx.mpq_1_2
# For re(z) > 0, 2F0 -- http://functions.wolfram.com/
# HypergeometricFunctions/HermiteHGeneral/06/02/0009/
# Otherwise, there is a reflection formula
# 2F0 + http://functions.wolfram.com/HypergeometricFunctions/
# HermiteHGeneral/16/01/01/0006/
#
# TODO:
# An alternative would be to use
# http://functions.wolfram.com/HypergeometricFunctions/
# HermiteHGeneral/06/02/0006/
#
# Also, the 1F1 expansion
# http://functions.wolfram.com/HypergeometricFunctions/
# HermiteHGeneral/26/01/02/0001/
# should probably be used for tiny z
if not z:
T1 = [2, ctx.pi], [n, 0.5], [], [q*(n-1)], [], [], 0
if parabolic_cylinder:
T1[1][0] += q*n
return T1,
can_use_2f0 = ctx.isnpint(-n) or ctx.re(z) > 0 or \
(ctx.re(z) == 0 and ctx.im(z) > 0)
expprec = ctx.prec*4 + 20
if parabolic_cylinder:
u = ctx.fmul(ctx.fmul(z,z,prec=expprec), -0.25, exact=True)
w = ctx.fmul(z, ctx.sqrt(0.5,prec=expprec), prec=expprec)
else:
w = z
w2 = ctx.fmul(w, w, prec=expprec)
rw2 = ctx.fdiv(1, w2, prec=expprec)
nrw2 = ctx.fneg(rw2, exact=True)
nw = ctx.fneg(w, exact=True)
if can_use_2f0:
T1 = [2, w], [n, n], [], [], [q*n, q*(n-1)], [], nrw2
terms = [T1]
else:
T1 = [2, nw], [n, n], [], [], [q*n, q*(n-1)], [], nrw2
T2 = [2, ctx.pi, nw], [n+2, 0.5, 1], [], [q*n], [q*(n-1)], [1-q], w2
terms = [T1,T2]
# Multiply by prefactor for D_n
if parabolic_cylinder:
expu = ctx.exp(u)
for i in range(len(terms)):
terms[i][1][0] += q*n
terms[i][0].append(expu)
terms[i][1].append(1)
return tuple(terms)
@defun
def hermite(ctx, n, z, **kwargs):
return ctx.hypercomb(lambda: _hermite_param(ctx, n, z, 0), [], **kwargs)
@defun
def pcfd(ctx, n, z, **kwargs):
r"""
Gives the parabolic cylinder function in Whittaker's notation
`D_n(z) = U(-n-1/2, z)` (see :func:`~mpmath.pcfu`).
It solves the differential equation
.. math ::
y'' + \left(n + \frac{1}{2} - \frac{1}{4} z^2\right) y = 0.
and can be represented in terms of Hermite polynomials
(see :func:`~mpmath.hermite`) as
.. math ::
D_n(z) = 2^{-n/2} e^{-z^2/4} H_n\left(\frac{z}{\sqrt{2}}\right).
**Plots**
.. literalinclude :: /plots/pcfd.py
.. image :: /plots/pcfd.png
**Examples**
>>> from mpmath import *
>>> mp.dps = 25; mp.pretty = True
>>> pcfd(0,0); pcfd(1,0); pcfd(2,0); pcfd(3,0)
1.0
0.0
-1.0
0.0
>>> pcfd(4,0); pcfd(-3,0)
3.0
0.6266570686577501256039413
>>> pcfd('1/2', 2+3j)
(-5.363331161232920734849056 - 3.858877821790010714163487j)
>>> pcfd(2, -10)
1.374906442631438038871515e-9
Verifying the differential equation::
>>> n = mpf(2.5)
>>> y = lambda z: pcfd(n,z)
>>> z = 1.75
>>> chop(diff(y,z,2) + (n+0.5-0.25*z**2)*y(z))
0.0
Rational Taylor series expansion when `n` is an integer::
>>> taylor(lambda z: pcfd(5,z), 0, 7)
[0.0, 15.0, 0.0, -13.75, 0.0, 3.96875, 0.0, -0.6015625]
"""
return ctx.hypercomb(lambda: _hermite_param(ctx, n, z, 1), [], **kwargs)
@defun
def pcfu(ctx, a, z, **kwargs):
r"""
Gives the parabolic cylinder function `U(a,z)`, which may be
defined for `\Re(z) > 0` in terms of the confluent
U-function (see :func:`~mpmath.hyperu`) by
.. math ::
U(a,z) = 2^{-\frac{1}{4}-\frac{a}{2}} e^{-\frac{1}{4} z^2}
U\left(\frac{a}{2}+\frac{1}{4},
\frac{1}{2}, \frac{1}{2}z^2\right)
or, for arbitrary `z`,
.. math ::
e^{-\frac{1}{4}z^2} U(a,z) =
U(a,0) \,_1F_1\left(-\tfrac{a}{2}+\tfrac{1}{4};
\tfrac{1}{2}; -\tfrac{1}{2}z^2\right) +
U'(a,0) z \,_1F_1\left(-\tfrac{a}{2}+\tfrac{3}{4};
\tfrac{3}{2}; -\tfrac{1}{2}z^2\right).
**Examples**
Connection to other functions::
>>> from mpmath import *
>>> mp.dps = 25; mp.pretty = True
>>> z = mpf(3)
>>> pcfu(0.5,z)
0.03210358129311151450551963
>>> sqrt(pi/2)*exp(z**2/4)*erfc(z/sqrt(2))
0.03210358129311151450551963
>>> pcfu(0.5,-z)
23.75012332835297233711255
>>> sqrt(pi/2)*exp(z**2/4)*erfc(-z/sqrt(2))
23.75012332835297233711255
>>> pcfu(0.5,-z)
23.75012332835297233711255
>>> sqrt(pi/2)*exp(z**2/4)*erfc(-z/sqrt(2))
23.75012332835297233711255
"""
n, _ = ctx._convert_param(a)
return ctx.pcfd(-n-ctx.mpq_1_2, z)
@defun
def pcfv(ctx, a, z, **kwargs):
r"""
Gives the parabolic cylinder function `V(a,z)`, which can be
represented in terms of :func:`~mpmath.pcfu` as
.. math ::
V(a,z) = \frac{\Gamma(a+\tfrac{1}{2}) (U(a,-z)-\sin(\pi a) U(a,z)}{\pi}.
**Examples**
Wronskian relation between `U` and `V`::
>>> from mpmath import *
>>> mp.dps = 25; mp.pretty = True
>>> a, z = 2, 3
>>> pcfu(a,z)*diff(pcfv,(a,z),(0,1))-diff(pcfu,(a,z),(0,1))*pcfv(a,z)
0.7978845608028653558798921
>>> sqrt(2/pi)
0.7978845608028653558798921
>>> a, z = 2.5, 3
>>> pcfu(a,z)*diff(pcfv,(a,z),(0,1))-diff(pcfu,(a,z),(0,1))*pcfv(a,z)
0.7978845608028653558798921
>>> a, z = 0.25, -1
>>> pcfu(a,z)*diff(pcfv,(a,z),(0,1))-diff(pcfu,(a,z),(0,1))*pcfv(a,z)
0.7978845608028653558798921
>>> a, z = 2+1j, 2+3j
>>> chop(pcfu(a,z)*diff(pcfv,(a,z),(0,1))-diff(pcfu,(a,z),(0,1))*pcfv(a,z))
0.7978845608028653558798921
"""
n, ntype = ctx._convert_param(a)
z = ctx.convert(z)
q = ctx.mpq_1_2
r = ctx.mpq_1_4
if ntype == 'Q' and ctx.isint(n*2):
# Faster for half-integers
def h():
jz = ctx.fmul(z, -1j, exact=True)
T1terms = _hermite_param(ctx, -n-q, z, 1)
T2terms = _hermite_param(ctx, n-q, jz, 1)
for T in T1terms:
T[0].append(1j)
T[1].append(1)
T[3].append(q-n)
u = ctx.expjpi((q*n-r)) * ctx.sqrt(2/ctx.pi)
for T in T2terms:
T[0].append(u)
T[1].append(1)
return T1terms + T2terms
v = ctx.hypercomb(h, [], **kwargs)
if ctx._is_real_type(n) and ctx._is_real_type(z):
v = ctx._re(v)
return v
else:
def h(n):
w = ctx.square_exp_arg(z, -0.25)
u = ctx.square_exp_arg(z, 0.5)
e = ctx.exp(w)
l = [ctx.pi, q, ctx.exp(w)]
Y1 = l, [-q, n*q+r, 1], [r-q*n], [], [q*n+r], [q], u
Y2 = l + [z], [-q, n*q-r, 1, 1], [1-r-q*n], [], [q*n+1-r], [1+q], u
c, s = ctx.cospi_sinpi(r+q*n)
Y1[0].append(s)
Y2[0].append(c)
for Y in (Y1, Y2):
Y[1].append(1)
Y[3].append(q-n)
return Y1, Y2
return ctx.hypercomb(h, [n], **kwargs)
@defun
def pcfw(ctx, a, z, **kwargs):
r"""
Gives the parabolic cylinder function `W(a,z)` defined in (DLMF 12.14).
**Examples**
Value at the origin::
>>> from mpmath import *
>>> mp.dps = 25; mp.pretty = True
>>> a = mpf(0.25)
>>> pcfw(a,0)
0.9722833245718180765617104
>>> power(2,-0.75)*sqrt(abs(gamma(0.25+0.5j*a)/gamma(0.75+0.5j*a)))
0.9722833245718180765617104
>>> diff(pcfw,(a,0),(0,1))
-0.5142533944210078966003624
>>> -power(2,-0.25)*sqrt(abs(gamma(0.75+0.5j*a)/gamma(0.25+0.5j*a)))
-0.5142533944210078966003624
"""
n, _ = ctx._convert_param(a)
z = ctx.convert(z)
def terms():
phi2 = ctx.arg(ctx.gamma(0.5 + ctx.j*n))
phi2 = (ctx.loggamma(0.5+ctx.j*n) - ctx.loggamma(0.5-ctx.j*n))/2j
rho = ctx.pi/8 + 0.5*phi2
# XXX: cancellation computing k
k = ctx.sqrt(1 + ctx.exp(2*ctx.pi*n)) - ctx.exp(ctx.pi*n)
C = ctx.sqrt(k/2) * ctx.exp(0.25*ctx.pi*n)
yield C * ctx.expj(rho) * ctx.pcfu(ctx.j*n, z*ctx.expjpi(-0.25))
yield C * ctx.expj(-rho) * ctx.pcfu(-ctx.j*n, z*ctx.expjpi(0.25))
v = ctx.sum_accurately(terms)
if ctx._is_real_type(n) and ctx._is_real_type(z):
v = ctx._re(v)
return v
"""
Even/odd PCFs. Useful?
@defun
def pcfy1(ctx, a, z, **kwargs):
a, _ = ctx._convert_param(n)
z = ctx.convert(z)
def h():
w = ctx.square_exp_arg(z)
w1 = ctx.fmul(w, -0.25, exact=True)
w2 = ctx.fmul(w, 0.5, exact=True)
e = ctx.exp(w1)
return [e], [1], [], [], [ctx.mpq_1_2*a+ctx.mpq_1_4], [ctx.mpq_1_2], w2
return ctx.hypercomb(h, [], **kwargs)
@defun
def pcfy2(ctx, a, z, **kwargs):
a, _ = ctx._convert_param(n)
z = ctx.convert(z)
def h():
w = ctx.square_exp_arg(z)
w1 = ctx.fmul(w, -0.25, exact=True)
w2 = ctx.fmul(w, 0.5, exact=True)
e = ctx.exp(w1)
return [e, z], [1, 1], [], [], [ctx.mpq_1_2*a+ctx.mpq_3_4], \
[ctx.mpq_3_2], w2
return ctx.hypercomb(h, [], **kwargs)
"""
@defun_wrapped
def gegenbauer(ctx, n, a, z, **kwargs):
# Special cases: a+0.5, a*2 poles
if ctx.isnpint(a):
return 0*(z+n)
if ctx.isnpint(a+0.5):
# TODO: something else is required here
# E.g.: gegenbauer(-2, -0.5, 3) == -12
if ctx.isnpint(n+1):
raise NotImplementedError("Gegenbauer function with two limits")
def h(a):
a2 = 2*a
T = [], [], [n+a2], [n+1, a2], [-n, n+a2], [a+0.5], 0.5*(1-z)
return [T]
return ctx.hypercomb(h, [a], **kwargs)
def h(n):
a2 = 2*a
T = [], [], [n+a2], [n+1, a2], [-n, n+a2], [a+0.5], 0.5*(1-z)
return [T]
return ctx.hypercomb(h, [n], **kwargs)
@defun_wrapped
def jacobi(ctx, n, a, b, x, **kwargs):
if not ctx.isnpint(a):
def h(n):
return (([], [], [a+n+1], [n+1, a+1], [-n, a+b+n+1], [a+1], (1-x)*0.5),)
return ctx.hypercomb(h, [n], **kwargs)
if not ctx.isint(b):
def h(n, a):
return (([], [], [-b], [n+1, -b-n], [-n, a+b+n+1], [b+1], (x+1)*0.5),)
return ctx.hypercomb(h, [n, a], **kwargs)
# XXX: determine appropriate limit
return ctx.binomial(n+a,n) * ctx.hyp2f1(-n,1+n+a+b,a+1,(1-x)/2, **kwargs)
@defun_wrapped
def laguerre(ctx, n, a, z, **kwargs):
# XXX: limits, poles
#if ctx.isnpint(n):
# return 0*(a+z)
def h(a):
return (([], [], [a+n+1], [a+1, n+1], [-n], [a+1], z),)
return ctx.hypercomb(h, [a], **kwargs)
@defun_wrapped
def legendre(ctx, n, x, **kwargs):
if ctx.isint(n):
n = int(n)
# Accuracy near zeros
if (n + (n < 0)) & 1:
if not x:
return x
mag = ctx.mag(x)
if mag < -2*ctx.prec-10:
return x
if mag < -5:
ctx.prec += -mag
return ctx.hyp2f1(-n,n+1,1,(1-x)/2, **kwargs)
@defun
def legenp(ctx, n, m, z, type=2, **kwargs):
# Legendre function, 1st kind
n = ctx.convert(n)
m = ctx.convert(m)
# Faster
if not m:
return ctx.legendre(n, z, **kwargs)
# TODO: correct evaluation at singularities
if type == 2:
def h(n,m):
g = m*0.5
T = [1+z, 1-z], [g, -g], [], [1-m], [-n, n+1], [1-m], 0.5*(1-z)
return (T,)
return ctx.hypercomb(h, [n,m], **kwargs)
if type == 3:
def h(n,m):
g = m*0.5
T = [z+1, z-1], [g, -g], [], [1-m], [-n, n+1], [1-m], 0.5*(1-z)
return (T,)
return ctx.hypercomb(h, [n,m], **kwargs)
raise ValueError("requires type=2 or type=3")
@defun
def legenq(ctx, n, m, z, type=2, **kwargs):
# Legendre function, 2nd kind
n = ctx.convert(n)
m = ctx.convert(m)
z = ctx.convert(z)
if z in (1, -1):
#if ctx.isint(m):
# return ctx.nan
#return ctx.inf # unsigned
return ctx.nan
if type == 2:
def h(n, m):
cos, sin = ctx.cospi_sinpi(m)
s = 2 * sin / ctx.pi
c = cos
a = 1+z
b = 1-z
u = m/2
w = (1-z)/2
T1 = [s, c, a, b], [-1, 1, u, -u], [], [1-m], \
[-n, n+1], [1-m], w
T2 = [-s, a, b], [-1, -u, u], [n+m+1], [n-m+1, m+1], \
[-n, n+1], [m+1], w
return T1, T2
return ctx.hypercomb(h, [n, m], **kwargs)
if type == 3:
# The following is faster when there only is a single series
# Note: not valid for -1 < z < 0 (?)
if abs(z) > 1:
def h(n, m):
T1 = [ctx.expjpi(m), 2, ctx.pi, z, z-1, z+1], \
[1, -n-1, 0.5, -n-m-1, 0.5*m, 0.5*m], \
[n+m+1], [n+1.5], \
[0.5*(2+n+m), 0.5*(1+n+m)], [n+1.5], z**(-2)
return [T1]
return ctx.hypercomb(h, [n, m], **kwargs)
else:
# not valid for 1 < z < inf ?
def h(n, m):
s = 2 * ctx.sinpi(m) / ctx.pi
c = ctx.expjpi(m)
a = 1+z
b = z-1
u = m/2
w = (1-z)/2
T1 = [s, c, a, b], [-1, 1, u, -u], [], [1-m], \
[-n, n+1], [1-m], w
T2 = [-s, c, a, b], [-1, 1, -u, u], [n+m+1], [n-m+1, m+1], \
[-n, n+1], [m+1], w
return T1, T2
return ctx.hypercomb(h, [n, m], **kwargs)
raise ValueError("requires type=2 or type=3")
@defun_wrapped
def chebyt(ctx, n, x, **kwargs):
if (not x) and ctx.isint(n) and int(ctx._re(n)) % 2 == 1:
return x * 0
return ctx.hyp2f1(-n,n,(1,2),(1-x)/2, **kwargs)
@defun_wrapped
def chebyu(ctx, n, x, **kwargs):
if (not x) and ctx.isint(n) and int(ctx._re(n)) % 2 == 1:
return x * 0
return (n+1) * ctx.hyp2f1(-n, n+2, (3,2), (1-x)/2, **kwargs)
@defun
def spherharm(ctx, l, m, theta, phi, **kwargs):
l = ctx.convert(l)
m = ctx.convert(m)
theta = ctx.convert(theta)
phi = ctx.convert(phi)
l_isint = ctx.isint(l)
l_natural = l_isint and l >= 0
m_isint = ctx.isint(m)
if l_isint and l < 0 and m_isint:
return ctx.spherharm(-(l+1), m, theta, phi, **kwargs)
if theta == 0 and m_isint and m < 0:
return ctx.zero * 1j
if l_natural and m_isint:
if abs(m) > l:
return ctx.zero * 1j
# http://functions.wolfram.com/Polynomials/
# SphericalHarmonicY/26/01/02/0004/
def h(l,m):
absm = abs(m)
C = [-1, ctx.expj(m*phi),
(2*l+1)*ctx.fac(l+absm)/ctx.pi/ctx.fac(l-absm),
ctx.sin(theta)**2,
ctx.fac(absm), 2]
P = [0.5*m*(ctx.sign(m)+1), 1, 0.5, 0.5*absm, -1, -absm-1]
return ((C, P, [], [], [absm-l, l+absm+1], [absm+1],
ctx.sin(0.5*theta)**2),)
else:
# http://functions.wolfram.com/HypergeometricFunctions/
# SphericalHarmonicYGeneral/26/01/02/0001/
def h(l,m):
if ctx.isnpint(l-m+1) or ctx.isnpint(l+m+1) or ctx.isnpint(1-m):
return (([0], [-1], [], [], [], [], 0),)
cos, sin = ctx.cos_sin(0.5*theta)
C = [0.5*ctx.expj(m*phi), (2*l+1)/ctx.pi,
ctx.gamma(l-m+1), ctx.gamma(l+m+1),
cos**2, sin**2]
P = [1, 0.5, 0.5, -0.5, 0.5*m, -0.5*m]
return ((C, P, [], [1-m], [-l,l+1], [1-m], sin**2),)
return ctx.hypercomb(h, [l,m], **kwargs)