chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,14 @@
|
||||
from . import functions
|
||||
# Hack to update methods
|
||||
from . import factorials
|
||||
from . import hypergeometric
|
||||
from . import expintegrals
|
||||
from . import bessel
|
||||
from . import orthogonal
|
||||
from . import theta
|
||||
from . import elliptic
|
||||
from . import signals
|
||||
from . import zeta
|
||||
from . import rszeta
|
||||
from . import zetazeros
|
||||
from . import qfunctions
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,437 @@
|
||||
from .functions import defun, defun_wrapped
|
||||
|
||||
@defun_wrapped
|
||||
def _erf_complex(ctx, z):
|
||||
re_z = ctx.re(z)
|
||||
if re_z > 2:
|
||||
nz = ctx.fneg(z, exact=True)
|
||||
v = ctx._erf_complex(nz)
|
||||
return ctx.fneg(v, exact=True)
|
||||
elif re_z < -2:
|
||||
v = ctx._erfc_complex(ctx.fneg(z, exact=True)) - 1
|
||||
else:
|
||||
z2 = ctx.square_exp_arg(z, -1)
|
||||
v = (2/ctx.sqrt(ctx.pi))*z * ctx.hyp1f1((1,2),(3,2), z2)
|
||||
if not re_z:
|
||||
v = ctx._im(v)*ctx.j
|
||||
return v
|
||||
|
||||
@defun_wrapped
|
||||
def _erfc_complex(ctx, z):
|
||||
re_z = ctx.re(z)
|
||||
if re_z > 2:
|
||||
z2 = ctx.square_exp_arg(z)
|
||||
nz2 = ctx.fneg(z2, exact=True)
|
||||
v = ctx.exp(nz2)/ctx.sqrt(ctx.pi) * ctx.hyperu((1,2),(1,2), z2)
|
||||
else:
|
||||
v = 1 - ctx._erf_complex(z)
|
||||
if not re_z:
|
||||
v = 1+ctx._im(v)*ctx.j
|
||||
return v
|
||||
|
||||
@defun
|
||||
def erf(ctx, z):
|
||||
z = ctx.convert(z)
|
||||
if ctx._is_real_type(z):
|
||||
try:
|
||||
return ctx._erf(z)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
if ctx._is_complex_type(z) and not z.imag:
|
||||
try:
|
||||
return ctx.mpc(ctx._erf(z.real))
|
||||
except NotImplementedError:
|
||||
pass
|
||||
return ctx._erf_complex(z)
|
||||
|
||||
@defun
|
||||
def erfc(ctx, z):
|
||||
z = ctx.convert(z)
|
||||
if ctx._is_real_type(z):
|
||||
try:
|
||||
return ctx._erfc(z)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
if ctx._is_complex_type(z) and not z.imag:
|
||||
try:
|
||||
return ctx.mpc(ctx._erfc(z.real))
|
||||
except NotImplementedError:
|
||||
pass
|
||||
return ctx._erfc_complex(z)
|
||||
|
||||
@defun
|
||||
def square_exp_arg(ctx, z, mult=1, reciprocal=False):
|
||||
prec = ctx.prec*4+20
|
||||
if reciprocal:
|
||||
z2 = ctx.fmul(z, z, prec=prec)
|
||||
z2 = ctx.fdiv(ctx.one, z2, prec=prec)
|
||||
else:
|
||||
z2 = ctx.fmul(z, z, prec=prec)
|
||||
if mult != 1:
|
||||
z2 = ctx.fmul(z2, mult, exact=True)
|
||||
return z2
|
||||
|
||||
@defun_wrapped
|
||||
def erfi(ctx, z):
|
||||
if not z:
|
||||
return z
|
||||
z2 = ctx.square_exp_arg(z)
|
||||
v = (2/ctx.sqrt(ctx.pi)*z) * ctx.hyp1f1((1,2), (3,2), z2)
|
||||
if not ctx._re(z):
|
||||
v = ctx._im(v)*ctx.j
|
||||
return v
|
||||
|
||||
@defun_wrapped
|
||||
def erfinv(ctx, x):
|
||||
xre = ctx._re(x)
|
||||
if (xre != x) or (xre < -1) or (xre > 1):
|
||||
return ctx.bad_domain("erfinv(x) is defined only for -1 <= x <= 1")
|
||||
x = xre
|
||||
#if ctx.isnan(x): return x
|
||||
if not x: return x
|
||||
if x == 1: return ctx.inf
|
||||
if x == -1: return ctx.ninf
|
||||
if abs(x) < 0.9:
|
||||
a = 0.53728*x**3 + 0.813198*x
|
||||
else:
|
||||
# An asymptotic formula
|
||||
u = ctx.ln(2/ctx.pi/(abs(x)-1)**2)
|
||||
a = ctx.sign(x) * ctx.sqrt(u - ctx.ln(u))/ctx.sqrt(2)
|
||||
ctx.prec += 10
|
||||
return ctx.findroot(lambda t: ctx.erf(t)-x, a)
|
||||
|
||||
@defun_wrapped
|
||||
def npdf(ctx, x, mu=0, sigma=1):
|
||||
sigma = ctx.convert(sigma)
|
||||
return ctx.exp(-(x-mu)**2/(2*sigma**2)) / (sigma*ctx.sqrt(2*ctx.pi))
|
||||
|
||||
@defun_wrapped
|
||||
def ncdf(ctx, x, mu=0, sigma=1):
|
||||
a = (x-mu)/(sigma*ctx.sqrt(2))
|
||||
if a < 0:
|
||||
return ctx.erfc(-a)/2
|
||||
else:
|
||||
return (1+ctx.erf(a))/2
|
||||
|
||||
@defun_wrapped
|
||||
def betainc(ctx, a, b, x1=0, x2=1, regularized=False):
|
||||
if x1 == x2:
|
||||
v = ctx.zero
|
||||
elif not x1:
|
||||
if x1 == 0 and x2 == 1:
|
||||
v = ctx.beta(a, b)
|
||||
else:
|
||||
v = x2**a * ctx.hyp2f1(a, 1-b, a+1, x2) / a
|
||||
else:
|
||||
m, d = ctx.nint_distance(a)
|
||||
if m <= 0:
|
||||
if d < -ctx.prec:
|
||||
h = +ctx.eps
|
||||
ctx.prec *= 2
|
||||
a += h
|
||||
elif d < -4:
|
||||
ctx.prec -= d
|
||||
s1 = x2**a * ctx.hyp2f1(a,1-b,a+1,x2)
|
||||
s2 = x1**a * ctx.hyp2f1(a,1-b,a+1,x1)
|
||||
v = (s1 - s2) / a
|
||||
if regularized:
|
||||
v /= ctx.beta(a,b)
|
||||
return v
|
||||
|
||||
@defun
|
||||
def gammainc(ctx, z, a=0, b=None, regularized=False):
|
||||
regularized = bool(regularized)
|
||||
z = ctx.convert(z)
|
||||
if a is None:
|
||||
a = ctx.zero
|
||||
lower_modified = False
|
||||
else:
|
||||
a = ctx.convert(a)
|
||||
lower_modified = a != ctx.zero
|
||||
if b is None:
|
||||
b = ctx.inf
|
||||
upper_modified = False
|
||||
else:
|
||||
b = ctx.convert(b)
|
||||
upper_modified = b != ctx.inf
|
||||
# Complete gamma function
|
||||
if not (upper_modified or lower_modified):
|
||||
if regularized:
|
||||
if ctx.re(z) < 0:
|
||||
return ctx.inf
|
||||
elif ctx.re(z) > 0:
|
||||
return ctx.one
|
||||
else:
|
||||
return ctx.nan
|
||||
return ctx.gamma(z)
|
||||
if a == b:
|
||||
return ctx.zero
|
||||
# Standardize
|
||||
if ctx.re(a) > ctx.re(b):
|
||||
return -ctx.gammainc(z, b, a, regularized)
|
||||
# Generalized gamma
|
||||
if upper_modified and lower_modified:
|
||||
return +ctx._gamma3(z, a, b, regularized)
|
||||
# Upper gamma
|
||||
elif lower_modified:
|
||||
return ctx.upper_gamma(z, a, regularized)
|
||||
# Lower gamma
|
||||
elif upper_modified:
|
||||
return ctx.lower_gamma(z, b, regularized)
|
||||
|
||||
@defun
|
||||
def lower_gamma(ctx, z, b, regularized=False):
|
||||
z = ctx.convert(z)
|
||||
b = ctx.convert(b)
|
||||
# Pole
|
||||
if ctx.isnpint(z):
|
||||
return ctx.inf
|
||||
G = [z] * regularized
|
||||
negb = ctx.fneg(b, exact=True)
|
||||
def h(z):
|
||||
T1 = [ctx.exp(negb), b, z], [1, z, -1], [], G, [1], [1+z], b
|
||||
return (T1,)
|
||||
return ctx.hypercomb(h, [z])
|
||||
|
||||
@defun
|
||||
def upper_gamma(ctx, z, a, regularized=False):
|
||||
z = ctx.convert(z)
|
||||
a = ctx.convert(a)
|
||||
# Fast integer case, when available
|
||||
if ctx.isint(z):
|
||||
try:
|
||||
if regularized:
|
||||
# Gamma pole
|
||||
if ctx.isnpint(z):
|
||||
return ctx.zero
|
||||
orig = ctx.prec
|
||||
try:
|
||||
ctx.prec += 10
|
||||
return ctx._gamma_upper_int(z, a) / ctx.gamma(z)
|
||||
finally:
|
||||
ctx.prec = orig
|
||||
else:
|
||||
return ctx._gamma_upper_int(z, a)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
# hypercomb is unable to detect the exact zeros, so handle them here
|
||||
if z == 2 and a == -1:
|
||||
return (z+a)*0
|
||||
if z == 3 and (a == -1-1j or a == -1+1j):
|
||||
return (z+a)*0
|
||||
nega = ctx.fneg(a, exact=True)
|
||||
G = [z] * regularized
|
||||
# Use 2F0 series when possible; fall back to lower gamma representation
|
||||
try:
|
||||
def h(z):
|
||||
r = z-1
|
||||
return [([ctx.exp(nega), a], [1, r], [], G, [1, -r], [], 1/nega)]
|
||||
return ctx.hypercomb(h, [z], force_series=True)
|
||||
except ctx.NoConvergence:
|
||||
def h(z):
|
||||
T1 = [], [1, z-1], [z], G, [], [], 0
|
||||
T2 = [-ctx.exp(nega), a, z], [1, z, -1], [], G, [1], [1+z], a
|
||||
return T1, T2
|
||||
return ctx.hypercomb(h, [z])
|
||||
|
||||
@defun
|
||||
def _gamma3(ctx, z, a, b, regularized=False):
|
||||
pole = ctx.isnpint(z)
|
||||
if regularized and pole:
|
||||
return ctx.zero
|
||||
try:
|
||||
ctx.prec += 15
|
||||
# We don't know in advance whether it's better to write as a difference
|
||||
# of lower or upper gamma functions, so try both
|
||||
T1 = ctx.gammainc(z, a, regularized=regularized)
|
||||
T2 = ctx.gammainc(z, b, regularized=regularized)
|
||||
R = T1 - T2
|
||||
if ctx.mag(R) - max(ctx.mag(T1), ctx.mag(T2)) > -10:
|
||||
return R
|
||||
if not pole:
|
||||
T1 = ctx.gammainc(z, 0, b, regularized=regularized)
|
||||
T2 = ctx.gammainc(z, 0, a, regularized=regularized)
|
||||
R = T1 - T2
|
||||
# May be ok, but should probably at least print a warning
|
||||
# about possible cancellation
|
||||
if 1: #ctx.mag(R) - max(ctx.mag(T1), ctx.mag(T2)) > -10:
|
||||
return R
|
||||
finally:
|
||||
ctx.prec -= 15
|
||||
raise NotImplementedError
|
||||
|
||||
@defun_wrapped
|
||||
def expint(ctx, n, z):
|
||||
if ctx.isint(n) and ctx._is_real_type(z):
|
||||
try:
|
||||
return ctx._expint_int(n, z)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
if ctx.isnan(n) or ctx.isnan(z):
|
||||
return z*n
|
||||
if z == ctx.inf:
|
||||
return 1/z
|
||||
if z == 0:
|
||||
# integral from 1 to infinity of t^n
|
||||
if ctx.re(n) <= 1:
|
||||
# TODO: reasonable sign of infinity
|
||||
return ctx.inf
|
||||
else:
|
||||
return ctx.one/(n-1)
|
||||
if n == 0:
|
||||
return ctx.exp(-z)/z
|
||||
if n == -1:
|
||||
return ctx.exp(-z)*(z+1)/z**2
|
||||
return z**(n-1) * ctx.upper_gamma(1-n, z)
|
||||
|
||||
@defun_wrapped
|
||||
def li(ctx, z, offset=False):
|
||||
if offset:
|
||||
if z == 2:
|
||||
return ctx.zero
|
||||
return ctx.ei(ctx.ln(z)) - ctx.ei(ctx.ln2)
|
||||
if not z:
|
||||
return z
|
||||
if z == 1:
|
||||
return ctx.ninf
|
||||
return ctx.ei(ctx.ln(z))
|
||||
|
||||
@defun
|
||||
def ei(ctx, z):
|
||||
try:
|
||||
return ctx._ei(z)
|
||||
except NotImplementedError:
|
||||
return ctx._ei_generic(z)
|
||||
|
||||
@defun_wrapped
|
||||
def _ei_generic(ctx, z):
|
||||
# Note: the following is currently untested because mp and fp
|
||||
# both use special-case ei code
|
||||
if z == ctx.inf:
|
||||
return z
|
||||
if z == ctx.ninf:
|
||||
return ctx.zero
|
||||
if ctx.mag(z) > 1:
|
||||
try:
|
||||
r = ctx.one/z
|
||||
v = ctx.exp(z)*ctx.hyper([1,1],[],r,
|
||||
maxterms=ctx.prec, force_series=True)/z
|
||||
im = ctx._im(z)
|
||||
if im > 0:
|
||||
v += ctx.pi*ctx.j
|
||||
if im < 0:
|
||||
v -= ctx.pi*ctx.j
|
||||
return v
|
||||
except ctx.NoConvergence:
|
||||
pass
|
||||
v = z*ctx.hyp2f2(1,1,2,2,z) + ctx.euler
|
||||
if ctx._im(z):
|
||||
v += 0.5*(ctx.log(z) - ctx.log(ctx.one/z))
|
||||
else:
|
||||
v += ctx.log(abs(z))
|
||||
return v
|
||||
|
||||
@defun
|
||||
def e1(ctx, z):
|
||||
try:
|
||||
return ctx._e1(z)
|
||||
except NotImplementedError:
|
||||
return ctx.expint(1, z)
|
||||
|
||||
@defun
|
||||
def ci(ctx, z):
|
||||
try:
|
||||
return ctx._ci(z)
|
||||
except NotImplementedError:
|
||||
return ctx._ci_generic(z)
|
||||
|
||||
@defun_wrapped
|
||||
def _ci_generic(ctx, z):
|
||||
if ctx.isinf(z):
|
||||
if z == ctx.inf: return ctx.zero
|
||||
if z == ctx.ninf: return ctx.pi*1j
|
||||
jz = ctx.fmul(ctx.j,z,exact=True)
|
||||
njz = ctx.fneg(jz,exact=True)
|
||||
v = 0.5*(ctx.ei(jz) + ctx.ei(njz))
|
||||
zreal = ctx._re(z)
|
||||
zimag = ctx._im(z)
|
||||
if zreal == 0:
|
||||
if zimag > 0: v += ctx.pi*0.5j
|
||||
if zimag < 0: v -= ctx.pi*0.5j
|
||||
if zreal < 0:
|
||||
if zimag >= 0: v += ctx.pi*1j
|
||||
if zimag < 0: v -= ctx.pi*1j
|
||||
if ctx._is_real_type(z) and zreal > 0:
|
||||
v = ctx._re(v)
|
||||
return v
|
||||
|
||||
@defun
|
||||
def si(ctx, z):
|
||||
try:
|
||||
return ctx._si(z)
|
||||
except NotImplementedError:
|
||||
return ctx._si_generic(z)
|
||||
|
||||
@defun_wrapped
|
||||
def _si_generic(ctx, z):
|
||||
if ctx.isinf(z):
|
||||
if z == ctx.inf: return 0.5*ctx.pi
|
||||
if z == ctx.ninf: return -0.5*ctx.pi
|
||||
# Suffers from cancellation near 0
|
||||
if ctx.mag(z) >= -1:
|
||||
jz = ctx.fmul(ctx.j,z,exact=True)
|
||||
njz = ctx.fneg(jz,exact=True)
|
||||
v = (-0.5j)*(ctx.ei(jz) - ctx.ei(njz))
|
||||
zreal = ctx._re(z)
|
||||
if zreal > 0:
|
||||
v -= 0.5*ctx.pi
|
||||
if zreal < 0:
|
||||
v += 0.5*ctx.pi
|
||||
if ctx._is_real_type(z):
|
||||
v = ctx._re(v)
|
||||
return v
|
||||
else:
|
||||
return z*ctx.hyp1f2((1,2),(3,2),(3,2),-0.25*z*z)
|
||||
|
||||
@defun_wrapped
|
||||
def chi(ctx, z):
|
||||
nz = ctx.fneg(z, exact=True)
|
||||
v = 0.5*(ctx.ei(z) + ctx.ei(nz))
|
||||
zreal = ctx._re(z)
|
||||
zimag = ctx._im(z)
|
||||
if zimag > 0:
|
||||
v += ctx.pi*0.5j
|
||||
elif zimag < 0:
|
||||
v -= ctx.pi*0.5j
|
||||
elif zreal < 0:
|
||||
v += ctx.pi*1j
|
||||
return v
|
||||
|
||||
@defun_wrapped
|
||||
def shi(ctx, z):
|
||||
# Suffers from cancellation near 0
|
||||
if ctx.mag(z) >= -1:
|
||||
nz = ctx.fneg(z, exact=True)
|
||||
v = 0.5*(ctx.ei(z) - ctx.ei(nz))
|
||||
zimag = ctx._im(z)
|
||||
if zimag > 0: v -= 0.5j*ctx.pi
|
||||
if zimag < 0: v += 0.5j*ctx.pi
|
||||
return v
|
||||
else:
|
||||
return z * ctx.hyp1f2((1,2),(3,2),(3,2),0.25*z*z)
|
||||
|
||||
@defun_wrapped
|
||||
def fresnels(ctx, z):
|
||||
if z == ctx.inf:
|
||||
return ctx.mpf(0.5)
|
||||
if z == ctx.ninf:
|
||||
return ctx.mpf(-0.5)
|
||||
return ctx.pi*z**3/6*ctx.hyp1f2((3,4),(3,2),(7,4),-ctx.pi**2*z**4/16)
|
||||
|
||||
@defun_wrapped
|
||||
def fresnelc(ctx, z):
|
||||
if z == ctx.inf:
|
||||
return ctx.mpf(0.5)
|
||||
if z == ctx.ninf:
|
||||
return ctx.mpf(-0.5)
|
||||
return z*ctx.hyp1f2((1,4),(1,2),(5,4),-ctx.pi**2*z**4/16)
|
||||
@@ -0,0 +1,186 @@
|
||||
from .functions import defun, defun_wrapped
|
||||
|
||||
@defun
|
||||
def gammaprod(ctx, a, b, _infsign=False):
|
||||
a = [ctx.convert(x) for x in a]
|
||||
b = [ctx.convert(x) for x in b]
|
||||
poles_num = []
|
||||
poles_den = []
|
||||
regular_num = []
|
||||
regular_den = []
|
||||
for x in a: [regular_num, poles_num][ctx.isnpint(x)].append(x)
|
||||
for x in b: [regular_den, poles_den][ctx.isnpint(x)].append(x)
|
||||
# One more pole in numerator or denominator gives 0 or inf
|
||||
if len(poles_num) < len(poles_den): return ctx.zero
|
||||
if len(poles_num) > len(poles_den):
|
||||
# Get correct sign of infinity for x+h, h -> 0 from above
|
||||
# XXX: hack, this should be done properly
|
||||
if _infsign:
|
||||
a = [x and x*(1+ctx.eps) or x+ctx.eps for x in poles_num]
|
||||
b = [x and x*(1+ctx.eps) or x+ctx.eps for x in poles_den]
|
||||
return ctx.sign(ctx.gammaprod(a+regular_num,b+regular_den)) * ctx.inf
|
||||
else:
|
||||
return ctx.inf
|
||||
# All poles cancel
|
||||
# lim G(i)/G(j) = (-1)**(i+j) * gamma(1-j) / gamma(1-i)
|
||||
p = ctx.one
|
||||
orig = ctx.prec
|
||||
try:
|
||||
ctx.prec = orig + 15
|
||||
while poles_num:
|
||||
i = poles_num.pop()
|
||||
j = poles_den.pop()
|
||||
p *= (-1)**(i+j) * ctx.gamma(1-j) / ctx.gamma(1-i)
|
||||
for x in regular_num: p *= ctx.gamma(x)
|
||||
for x in regular_den: p /= ctx.gamma(x)
|
||||
finally:
|
||||
ctx.prec = orig
|
||||
return +p
|
||||
|
||||
@defun
|
||||
def beta(ctx, x, y):
|
||||
x = ctx.convert(x)
|
||||
y = ctx.convert(y)
|
||||
if ctx.isinf(y):
|
||||
x, y = y, x
|
||||
if ctx.isinf(x):
|
||||
if x == ctx.inf and not ctx._im(y):
|
||||
if y == ctx.ninf:
|
||||
return ctx.nan
|
||||
if y > 0:
|
||||
return ctx.zero
|
||||
if ctx.isint(y):
|
||||
return ctx.nan
|
||||
if y < 0:
|
||||
return ctx.sign(ctx.gamma(y)) * ctx.inf
|
||||
return ctx.nan
|
||||
xy = ctx.fadd(x, y, prec=2*ctx.prec)
|
||||
return ctx.gammaprod([x, y], [xy])
|
||||
|
||||
@defun
|
||||
def binomial(ctx, n, k):
|
||||
n1 = ctx.fadd(n, 1, prec=2*ctx.prec)
|
||||
k1 = ctx.fadd(k, 1, prec=2*ctx.prec)
|
||||
nk1 = ctx.fsub(n1, k, prec=2*ctx.prec)
|
||||
return ctx.gammaprod([n1], [k1, nk1])
|
||||
|
||||
@defun
|
||||
def rf(ctx, x, n):
|
||||
xn = ctx.fadd(x, n, prec=2*ctx.prec)
|
||||
return ctx.gammaprod([xn], [x])
|
||||
|
||||
@defun
|
||||
def ff(ctx, x, n):
|
||||
x1 = ctx.fadd(x, 1, prec=2*ctx.prec)
|
||||
xn1 = ctx.fadd(ctx.fsub(x, n, prec=2*ctx.prec), 1, prec=2*ctx.prec)
|
||||
return ctx.gammaprod([x1], [xn1])
|
||||
|
||||
@defun_wrapped
|
||||
def fac2(ctx, x):
|
||||
if ctx.isinf(x):
|
||||
if x == ctx.inf:
|
||||
return x
|
||||
return ctx.nan
|
||||
return 2**(x/2)*(ctx.pi/2)**((ctx.cospi(x)-1)/4)*ctx.gamma(x/2+1)
|
||||
|
||||
@defun_wrapped
|
||||
def barnesg(ctx, z):
|
||||
if ctx.isinf(z):
|
||||
if z == ctx.inf:
|
||||
return z
|
||||
return ctx.nan
|
||||
if ctx.isnan(z):
|
||||
return z
|
||||
if ctx.isnpint(z):
|
||||
return z*0
|
||||
# Account for size (would not be needed if computing log(G))
|
||||
if abs(z) > 5:
|
||||
ctx.dps += 2*ctx.log(abs(z),2)
|
||||
# Reflection formula
|
||||
if ctx.re(z) < -ctx.dps:
|
||||
w = 1-z
|
||||
pi2 = 2*ctx.pi
|
||||
u = ctx.expjpi(2*w)
|
||||
v = ctx.j*ctx.pi/12 - ctx.j*ctx.pi*w**2/2 + w*ctx.ln(1-u) - \
|
||||
ctx.j*ctx.polylog(2, u)/pi2
|
||||
v = ctx.barnesg(2-z)*ctx.exp(v)/pi2**w
|
||||
if ctx._is_real_type(z):
|
||||
v = ctx._re(v)
|
||||
return v
|
||||
# Estimate terms for asymptotic expansion
|
||||
# TODO: fixme, obviously
|
||||
N = ctx.dps // 2 + 5
|
||||
G = 1
|
||||
while abs(z) < N or ctx.re(z) < 1:
|
||||
G /= ctx.gamma(z)
|
||||
z += 1
|
||||
z -= 1
|
||||
s = ctx.mpf(1)/12
|
||||
s -= ctx.log(ctx.glaisher)
|
||||
s += z*ctx.log(2*ctx.pi)/2
|
||||
s += (z**2/2-ctx.mpf(1)/12)*ctx.log(z)
|
||||
s -= 3*z**2/4
|
||||
z2k = z2 = z**2
|
||||
for k in range(1, N+1):
|
||||
t = ctx.bernoulli(2*k+2) / (4*k*(k+1)*z2k)
|
||||
if abs(t) < ctx.eps:
|
||||
#print k, N # check how many terms were needed
|
||||
break
|
||||
z2k *= z2
|
||||
s += t
|
||||
#if k == N:
|
||||
# print "warning: series for barnesg failed to converge", ctx.dps
|
||||
return G*ctx.exp(s)
|
||||
|
||||
@defun
|
||||
def superfac(ctx, z):
|
||||
return ctx.barnesg(z+2)
|
||||
|
||||
@defun_wrapped
|
||||
def hyperfac(ctx, z):
|
||||
# XXX: estimate needed extra bits accurately
|
||||
if z == ctx.inf:
|
||||
return z
|
||||
if abs(z) > 5:
|
||||
extra = 4*int(ctx.log(abs(z),2))
|
||||
else:
|
||||
extra = 0
|
||||
ctx.prec += extra
|
||||
if z and ctx.isnpint(z):
|
||||
n = int(ctx.re(z))
|
||||
h = ctx.hyperfac(-n-1)
|
||||
if ((n+1)//2) & 1:
|
||||
h = -h
|
||||
if ctx._is_complex_type(z):
|
||||
return h + 0j
|
||||
return h
|
||||
zp1 = z+1
|
||||
# Wrong branch cut
|
||||
#v = ctx.gamma(zp1)**z
|
||||
#ctx.prec -= extra
|
||||
#return v / ctx.barnesg(zp1)
|
||||
v = ctx.exp(z*ctx.loggamma(zp1))
|
||||
ctx.prec -= extra
|
||||
return v / ctx.barnesg(zp1)
|
||||
|
||||
'''
|
||||
@defun
|
||||
def psi0(ctx, z):
|
||||
"""Shortcut for psi(0,z) (the digamma function)"""
|
||||
return ctx.psi(0, z)
|
||||
|
||||
@defun
|
||||
def psi1(ctx, z):
|
||||
"""Shortcut for psi(1,z) (the trigamma function)"""
|
||||
return ctx.psi(1, z)
|
||||
|
||||
@defun
|
||||
def psi2(ctx, z):
|
||||
"""Shortcut for psi(2,z) (the tetragamma function)"""
|
||||
return ctx.psi(2, z)
|
||||
|
||||
@defun
|
||||
def psi3(ctx, z):
|
||||
"""Shortcut for psi(3,z) (the pentagamma function)"""
|
||||
return ctx.psi(3, z)
|
||||
'''
|
||||
@@ -0,0 +1,691 @@
|
||||
class SpecialFunctions:
|
||||
"""
|
||||
This class implements special functions using high-level code.
|
||||
|
||||
Elementary and some other functions (e.g. gamma function, basecase
|
||||
hypergeometric series) are assumed to be predefined by the context as
|
||||
"builtins" or "low-level" functions.
|
||||
"""
|
||||
defined_functions = {}
|
||||
|
||||
def __init__(self):
|
||||
cls = self.__class__
|
||||
for name in cls.defined_functions:
|
||||
f, wrap = cls.defined_functions[name]
|
||||
cls._wrap_specfun(name, f, wrap)
|
||||
|
||||
self._misc_const_cache = {}
|
||||
|
||||
self._aliases.update({
|
||||
'phase' : 'arg',
|
||||
'conjugate' : 'conj',
|
||||
'nthroot' : 'root',
|
||||
'polygamma' : 'psi',
|
||||
'hurwitz' : 'zeta',
|
||||
#'digamma' : 'psi0',
|
||||
#'trigamma' : 'psi1',
|
||||
#'tetragamma' : 'psi2',
|
||||
#'pentagamma' : 'psi3',
|
||||
'fibonacci' : 'fib',
|
||||
'factorial' : 'fac',
|
||||
})
|
||||
|
||||
self.zetazero_memoized = self.memoize(self.zetazero)
|
||||
|
||||
# Default -- do nothing
|
||||
@classmethod
|
||||
def _wrap_specfun(cls, name, f, wrap):
|
||||
setattr(cls, name, f)
|
||||
|
||||
# Optional fast versions of common functions in common cases.
|
||||
# If not overridden, default (generic hypergeometric series)
|
||||
# implementations will be used
|
||||
def _besselj(ctx, n, z): raise NotImplementedError
|
||||
def _erf(ctx, z): raise NotImplementedError
|
||||
def _erfc(ctx, z): raise NotImplementedError
|
||||
def _gamma_upper_int(ctx, z, a): raise NotImplementedError
|
||||
def _expint_int(ctx, n, z): raise NotImplementedError
|
||||
def _zeta(ctx, s): raise NotImplementedError
|
||||
def _zetasum_fast(ctx, s, a, n, derivatives, reflect): raise NotImplementedError
|
||||
def _ei(ctx, z): raise NotImplementedError
|
||||
def _e1(ctx, z): raise NotImplementedError
|
||||
def _ci(ctx, z): raise NotImplementedError
|
||||
def _si(ctx, z): raise NotImplementedError
|
||||
def _altzeta(ctx, s): raise NotImplementedError
|
||||
|
||||
def defun_wrapped(f):
|
||||
SpecialFunctions.defined_functions[f.__name__] = f, True
|
||||
return f
|
||||
|
||||
def defun(f):
|
||||
SpecialFunctions.defined_functions[f.__name__] = f, False
|
||||
return f
|
||||
|
||||
def defun_static(f):
|
||||
setattr(SpecialFunctions, f.__name__, f)
|
||||
return f
|
||||
|
||||
@defun_wrapped
|
||||
def cot(ctx, z): return ctx.one / ctx.tan(z)
|
||||
|
||||
@defun_wrapped
|
||||
def sec(ctx, z): return ctx.one / ctx.cos(z)
|
||||
|
||||
@defun_wrapped
|
||||
def csc(ctx, z): return ctx.one / ctx.sin(z)
|
||||
|
||||
@defun_wrapped
|
||||
def coth(ctx, z): return ctx.one / ctx.tanh(z)
|
||||
|
||||
@defun_wrapped
|
||||
def sech(ctx, z): return ctx.one / ctx.cosh(z)
|
||||
|
||||
@defun_wrapped
|
||||
def csch(ctx, z): return ctx.one / ctx.sinh(z)
|
||||
|
||||
@defun_wrapped
|
||||
def acot(ctx, z):
|
||||
if not z:
|
||||
return ctx.pi * 0.5
|
||||
else:
|
||||
return ctx.atan(ctx.zero if ctx.isinf(z) else ctx.one / z)
|
||||
|
||||
@defun_wrapped
|
||||
def asec(ctx, z): return ctx.acos(ctx.zero if ctx.isinf(z) else ctx.one / z)
|
||||
|
||||
@defun_wrapped
|
||||
def acsc(ctx, z): return ctx.asin(ctx.zero if ctx.isinf(z) else ctx.one / z)
|
||||
|
||||
@defun_wrapped
|
||||
def acoth(ctx, z):
|
||||
if not z:
|
||||
return ctx.pi * 0.5j
|
||||
else:
|
||||
return ctx.atanh(ctx.zero if ctx.isinf(z) else ctx.one / z)
|
||||
|
||||
|
||||
@defun_wrapped
|
||||
def asech(ctx, z): return ctx.acosh(ctx.one / z)
|
||||
|
||||
@defun_wrapped
|
||||
def acsch(ctx, z): return ctx.asinh(ctx.one / z)
|
||||
|
||||
@defun
|
||||
def sign(ctx, x):
|
||||
x = ctx.convert(x)
|
||||
if not x or ctx.isnan(x):
|
||||
return x
|
||||
if ctx._is_real_type(x):
|
||||
if x > 0:
|
||||
return ctx.one
|
||||
else:
|
||||
return -ctx.one
|
||||
return x / abs(x)
|
||||
|
||||
@defun
|
||||
def agm(ctx, a, b=1):
|
||||
if b == 1:
|
||||
return ctx.agm1(a)
|
||||
a = ctx.convert(a)
|
||||
b = ctx.convert(b)
|
||||
return ctx._agm(a, b)
|
||||
|
||||
@defun_wrapped
|
||||
def sinc(ctx, x):
|
||||
if ctx.isinf(x):
|
||||
return 1/x
|
||||
if not x:
|
||||
return x+1
|
||||
return ctx.sin(x)/x
|
||||
|
||||
@defun_wrapped
|
||||
def sincpi(ctx, x):
|
||||
if ctx.isinf(x):
|
||||
return 1/x
|
||||
if not x:
|
||||
return x+1
|
||||
return ctx.sinpi(x)/(ctx.pi*x)
|
||||
|
||||
# TODO: tests; improve implementation
|
||||
@defun_wrapped
|
||||
def expm1(ctx, x):
|
||||
if not x:
|
||||
return ctx.zero
|
||||
# exp(x) - 1 ~ x
|
||||
if ctx.mag(x) < -ctx.prec:
|
||||
return x + 0.5*x**2
|
||||
# TODO: accurately eval the smaller of the real/imag parts
|
||||
return ctx.sum_accurately(lambda: iter([ctx.exp(x),-1]),1)
|
||||
|
||||
@defun_wrapped
|
||||
def log1p(ctx, x):
|
||||
if not x:
|
||||
return ctx.zero
|
||||
LOG1P_EXTRAPREC = 10 # ctx._wrap_specfun()
|
||||
# Note that all cases could by handled by log(1+c) provided the
|
||||
# add is done exactly. Our aim here is to be much faster than that,
|
||||
# especially when |c| is small.
|
||||
c = ctx.convert(x)
|
||||
cmag = ctx.mag(c)
|
||||
a, b = c.real, c.imag
|
||||
wp = ctx.prec
|
||||
if cmag >= -wp:
|
||||
# |c| isn't very small. We call log(1+c) instead, but
|
||||
# are careful about the precision used by the add. The
|
||||
# real part of the result is log(|c+1|). That's
|
||||
# determined by 1 + 2*a + a**2 + b**2, and the add has
|
||||
# to preserve enough info so that no important bits of
|
||||
# that sum are lost. It doesn't matter to this that 2*a,
|
||||
# a**2, etc, are not computed explicitly here: we're
|
||||
# deducing how many bits have to be present in the sum
|
||||
# for log() to "reverse engineer" the value of 2*a +
|
||||
# a**2 + b**2 to `prec` good bits,
|
||||
if cmag < 4:
|
||||
# |c| isn't very small, or large.
|
||||
if ctx.mag(a) > ctx.mag(b):
|
||||
# `a` already contributes the most to c's norm.
|
||||
# After adding 1, it will utterly dominate it.
|
||||
# We only need enough extra precision to avoid
|
||||
# losing any of a's `prec` most significant bits
|
||||
# when addiog, `b**2` is too small to matter.
|
||||
wp *= 2
|
||||
else:
|
||||
# b**2 is the larger of the square terms. The
|
||||
# smallest b can be is about 2**-prec, so the
|
||||
# smallest b**2 can be is about 2**(-2*prec). So
|
||||
# for a bit to matter compared to b**2, it has
|
||||
# to be at least about 2**(-3*prec). Bits of 2*a
|
||||
# (if any) >= 2**(-3*prec) will be preserved if
|
||||
# we use 3*prec bits for the add.
|
||||
wp *= 3
|
||||
# Else (cmag >= 4), |c+1| >= |c| - 1 is so large that
|
||||
# working precision is fine (although that takes some
|
||||
# careful analysis for cmag=4, given that .mag() _may_
|
||||
# return a rexult too large by 2), So leave wp alone.
|
||||
arg = ctx.fadd(1.0, c, prec=wp)
|
||||
result = ctx.log(arg)
|
||||
else:
|
||||
# Else c is "very small", and we use a series expansion,
|
||||
# c - c**2/2. The real part of that is a+(b*b-a*a)/2,
|
||||
# and the imag part b-a*b. Given that cmag < -prec, it
|
||||
# can be shown that "a*b" is numerically insignifcant in
|
||||
# the imag part, and _usually_ the "a*a/2" in the real
|
||||
# part. What remains is cheap to compute. In the real
|
||||
# part, though, if `a` is negative, the remaining
|
||||
# a+b**2/2 can suffer massive cancellation - even total.
|
||||
real = a + b*b*0.5 # usually the real part of the result
|
||||
if (a < 0.0
|
||||
and ctx.mag(real) <= ctx.mag(a) - LOG1P_EXTRAPREC):
|
||||
# The guard bits were lost to cancellation. Rare. At
|
||||
# the contrived
|
||||
# -1.999999873062092e-40+1.999999936531045e-20j
|
||||
# _all_ bits cancel out. Since a ~= -b*b/2 in this
|
||||
# case, and |b| is at largest (worst case) about
|
||||
# 2**-prec, |a| is about 2**(-2*prec), and the true
|
||||
# result may be as small as a**2/2, which is about
|
||||
# 2**(-4*prec), of which we want the leading prec
|
||||
# bits. To get the leading prec bits starting at
|
||||
# 2**(-4**prec) from addends starting at
|
||||
# 2**-(2*prec), we need the subtraction to handle
|
||||
# 3*prec bits (the first 2*prec of which may cancel
|
||||
# to exactly 0).
|
||||
a2 = a*a # only need at worst prec bits
|
||||
b2 = ctx.fmul(b, b, prec=2*wp)
|
||||
diff = ctx.fsub(b2, a2, prec=3*wp)
|
||||
real = a + ctx.ldexp(diff, -1)
|
||||
result = real if ctx._is_real_type(x) else ctx.mpc(real, b)
|
||||
return result
|
||||
|
||||
@defun_wrapped
|
||||
def powm1(ctx, x, y):
|
||||
mag = ctx.mag
|
||||
one = ctx.one
|
||||
w = x**y - one
|
||||
M = mag(w)
|
||||
# Only moderate cancellation
|
||||
if M > -8:
|
||||
return w
|
||||
# Check for the only possible exact cases
|
||||
if not w:
|
||||
if (not y) or (x in (1, -1, 1j, -1j) and ctx.isint(y)):
|
||||
return w
|
||||
x1 = x - one
|
||||
magy = mag(y)
|
||||
lnx = ctx.ln(x)
|
||||
# Small y: x^y - 1 ~ log(x)*y + O(log(x)^2 * y^2)
|
||||
if magy + mag(lnx) < -ctx.prec:
|
||||
return lnx*y + (lnx*y)**2/2
|
||||
# TODO: accurately eval the smaller of the real/imag part
|
||||
return ctx.sum_accurately(lambda: iter([x**y, -1]), 1)
|
||||
|
||||
@defun
|
||||
def _rootof1(ctx, k, n):
|
||||
k = int(k)
|
||||
n = int(n)
|
||||
k %= n
|
||||
if not k:
|
||||
return ctx.one
|
||||
elif 2*k == n:
|
||||
return -ctx.one
|
||||
elif 4*k == n:
|
||||
return ctx.j
|
||||
elif 4*k == 3*n:
|
||||
return -ctx.j
|
||||
return ctx.expjpi(2*ctx.mpf(k)/n)
|
||||
|
||||
@defun
|
||||
def root(ctx, z, n, k=0):
|
||||
n = int(n)
|
||||
z = ctx.convert(z)
|
||||
if k:
|
||||
# Special case: there is an exact real root
|
||||
if (n & 1 and 2*k == n-1) and (not ctx.im(z)) and (ctx.re(z) < 0):
|
||||
return -ctx.root(-z, n)
|
||||
# Multiply by root of unity
|
||||
prec = ctx.prec
|
||||
try:
|
||||
ctx.prec += 10
|
||||
v = ctx.root(z, n, 0) * ctx._rootof1(k, n)
|
||||
finally:
|
||||
ctx.prec = prec
|
||||
return +v
|
||||
return ctx._nthroot(z, n)
|
||||
|
||||
@defun
|
||||
def unitroots(ctx, n, primitive=False):
|
||||
gcd = ctx._gcd
|
||||
prec = ctx.prec
|
||||
try:
|
||||
ctx.prec += 10
|
||||
if primitive:
|
||||
v = [ctx._rootof1(k,n) for k in range(n) if gcd(k,n) == 1]
|
||||
else:
|
||||
# TODO: this can be done *much* faster
|
||||
v = [ctx._rootof1(k,n) for k in range(n)]
|
||||
finally:
|
||||
ctx.prec = prec
|
||||
return [+x for x in v]
|
||||
|
||||
@defun
|
||||
def arg(ctx, x):
|
||||
x = ctx.convert(x)
|
||||
re = ctx._re(x)
|
||||
im = ctx._im(x)
|
||||
return ctx.atan2(im, re)
|
||||
|
||||
@defun
|
||||
def fabs(ctx, x):
|
||||
return abs(ctx.convert(x))
|
||||
|
||||
@defun
|
||||
def re(ctx, x):
|
||||
x = ctx.convert(x)
|
||||
return x.real
|
||||
|
||||
@defun
|
||||
def im(ctx, x):
|
||||
x = ctx.convert(x)
|
||||
return x.imag
|
||||
|
||||
@defun
|
||||
def conj(ctx, x):
|
||||
x = ctx.convert(x)
|
||||
try:
|
||||
return x.conjugate()
|
||||
except AttributeError:
|
||||
return x
|
||||
|
||||
@defun
|
||||
def polar(ctx, z):
|
||||
return (ctx.fabs(z), ctx.arg(z))
|
||||
|
||||
@defun_wrapped
|
||||
def rect(ctx, r, phi):
|
||||
return r * ctx.mpc(*ctx.cos_sin(phi))
|
||||
|
||||
@defun
|
||||
def log(ctx, x, b=None):
|
||||
if b is None:
|
||||
return ctx.ln(x)
|
||||
wp = ctx.prec + 20
|
||||
return ctx.ln(x, prec=wp) / ctx.ln(b, prec=wp)
|
||||
|
||||
@defun
|
||||
def log10(ctx, x):
|
||||
return ctx.log(x, 10)
|
||||
|
||||
@defun
|
||||
def log2(ctx, x):
|
||||
return ctx.log(x, 2)
|
||||
|
||||
@defun
|
||||
def exp2(ctx, x):
|
||||
return ctx.power(2, x)
|
||||
|
||||
@defun
|
||||
def fmod(ctx, x, y):
|
||||
return ctx.convert(x) % ctx.convert(y)
|
||||
|
||||
@defun
|
||||
def degrees(ctx, x):
|
||||
return x / ctx.degree
|
||||
|
||||
@defun
|
||||
def radians(ctx, x):
|
||||
return x * ctx.degree
|
||||
|
||||
def _lambertw_special(ctx, z, k):
|
||||
# W(0,0) = 0; all other branches are singular
|
||||
if not z:
|
||||
if not k:
|
||||
return z
|
||||
return ctx.ninf + z
|
||||
if z == ctx.inf:
|
||||
if k == 0:
|
||||
return z
|
||||
else:
|
||||
return z + 2*k*ctx.pi*ctx.j
|
||||
if z == ctx.ninf:
|
||||
return (-z) + (2*k+1)*ctx.pi*ctx.j
|
||||
# Some kind of nan or complex inf/nan?
|
||||
return ctx.ln(z)
|
||||
|
||||
import math
|
||||
import cmath
|
||||
|
||||
def _lambertw_approx_hybrid(z, k):
|
||||
imag_sign = 0
|
||||
x = float(z.real)
|
||||
y = z.imag
|
||||
if y:
|
||||
imag_sign = (-1) ** (y < 0)
|
||||
y = float(y)
|
||||
# hack to work regardless of whether Python supports -0.0
|
||||
if not y:
|
||||
y = 0.0
|
||||
z = complex(x,y)
|
||||
if k == 0:
|
||||
if -4.0 < y < 4.0 and -1.0 < x < 2.5:
|
||||
if imag_sign:
|
||||
# Taylor series in upper/lower half-plane
|
||||
if y > 1.00: return (0.876+0.645j) + (0.118-0.174j)*(z-(0.75+2.5j))
|
||||
if y > 0.25: return (0.505+0.204j) + (0.375-0.132j)*(z-(0.75+0.5j))
|
||||
if y < -1.00: return (0.876-0.645j) + (0.118+0.174j)*(z-(0.75-2.5j))
|
||||
if y < -0.25: return (0.505-0.204j) + (0.375+0.132j)*(z-(0.75-0.5j))
|
||||
# Taylor series near -1
|
||||
if x < -0.5:
|
||||
if imag_sign >= 0:
|
||||
return (-0.318+1.34j) + (-0.697-0.593j)*(z+1)
|
||||
else:
|
||||
return (-0.318-1.34j) + (-0.697+0.593j)*(z+1)
|
||||
# return real type
|
||||
r = -0.367879441171442
|
||||
if (not imag_sign) and x > r:
|
||||
z = x
|
||||
# Singularity near -1/e
|
||||
if x < -0.2:
|
||||
return -1 + 2.33164398159712*(z-r)**0.5 - 1.81218788563936*(z-r)
|
||||
# Taylor series near 0
|
||||
if x < 0.5: return z
|
||||
# Simple linear approximation
|
||||
return 0.2 + 0.3*z
|
||||
if (not imag_sign) and x > 0.0:
|
||||
L1 = math.log(x); L2 = math.log(L1)
|
||||
else:
|
||||
L1 = cmath.log(z); L2 = cmath.log(L1)
|
||||
elif k == -1:
|
||||
# return real type
|
||||
r = -0.367879441171442
|
||||
if (not imag_sign) and r < x < 0.0:
|
||||
z = x
|
||||
if (imag_sign >= 0) and y < 0.1 and -0.6 < x < -0.2:
|
||||
return -1 - 2.33164398159712*(z-r)**0.5 - 1.81218788563936*(z-r)
|
||||
if (not imag_sign) and -0.2 <= x < 0.0:
|
||||
L1 = math.log(-x)
|
||||
return L1 - math.log(-L1)
|
||||
else:
|
||||
if imag_sign == -1 and (not y) and x < 0.0:
|
||||
L1 = cmath.log(z) - 3.1415926535897932j
|
||||
else:
|
||||
L1 = cmath.log(z) - 6.2831853071795865j
|
||||
L2 = cmath.log(L1)
|
||||
return L1 - L2 + L2/L1 + L2*(L2-2)/(2*L1**2)
|
||||
|
||||
def _lambertw_series(ctx, z, k, tol):
|
||||
"""
|
||||
Return rough approximation for W_k(z) from an asymptotic series,
|
||||
sufficiently accurate for the Halley iteration to converge to
|
||||
the correct value.
|
||||
"""
|
||||
magz = ctx.mag(z)
|
||||
if (-10 < magz < 900) and (-1000 < k < 1000):
|
||||
# Near the branch point at -1/e
|
||||
if magz < 1 and abs(z+0.36787944117144) < 0.05:
|
||||
if k == 0 or (k == -1 and ctx._im(z) >= 0) or \
|
||||
(k == 1 and ctx._im(z) < 0):
|
||||
delta = ctx.sum_accurately(lambda: [z, ctx.exp(-1)])
|
||||
cancellation = -ctx.mag(delta)
|
||||
ctx.prec += cancellation
|
||||
# Use series given in Corless et al.
|
||||
p = ctx.sqrt(2*(ctx.e*z+1))
|
||||
ctx.prec -= cancellation
|
||||
u = {0:ctx.mpf(-1), 1:ctx.mpf(1)}
|
||||
a = {0:ctx.mpf(2), 1:ctx.mpf(-1)}
|
||||
if k != 0:
|
||||
p = -p
|
||||
s = ctx.zero
|
||||
# The series converges, so we could use it directly, but unless
|
||||
# *extremely* close, it is better to just use the first few
|
||||
# terms to get a good approximation for the iteration
|
||||
for l in range(max(2, cancellation)):
|
||||
if l not in u:
|
||||
a[l] = ctx.fsum(u[j]*u[l+1-j] for j in range(2, l))
|
||||
u[l] = (l-1)*(u[l-2]/2+a[l-2]/4)/(l+1)-a[l]/2-u[l-1]/(l+1)
|
||||
term = u[l] * p**l
|
||||
s += term
|
||||
if ctx.mag(term) < -tol:
|
||||
return s, True
|
||||
l += 1
|
||||
ctx.prec += cancellation//2
|
||||
return s, False
|
||||
if k == 0 or k == -1:
|
||||
return _lambertw_approx_hybrid(z, k), False
|
||||
if k == 0:
|
||||
if magz < -1:
|
||||
return z*(1-z), False
|
||||
L1 = ctx.ln(z)
|
||||
L2 = ctx.ln(L1)
|
||||
elif k == -1 and (not ctx._im(z)) and (-0.36787944117144 < ctx._re(z) < 0):
|
||||
L1 = ctx.ln(-z)
|
||||
return L1 - ctx.ln(-L1), False
|
||||
else:
|
||||
# This holds both as z -> 0 and z -> inf.
|
||||
# Relative error is O(1/log(z)).
|
||||
L1 = ctx.ln(z) + 2j*ctx.pi*k
|
||||
L2 = ctx.ln(L1)
|
||||
return L1 - L2 + L2/L1 + L2*(L2-2)/(2*L1**2), False
|
||||
|
||||
@defun
|
||||
def lambertw(ctx, z, k=0):
|
||||
z = ctx.convert(z)
|
||||
k = int(k)
|
||||
if not ctx.isnormal(z):
|
||||
return _lambertw_special(ctx, z, k)
|
||||
prec = ctx.prec
|
||||
ctx.prec += 20 + ctx.mag(k or 1)
|
||||
wp = ctx.prec
|
||||
tol = wp - 5
|
||||
w, done = _lambertw_series(ctx, z, k, tol)
|
||||
if not done:
|
||||
# Use Halley iteration to solve w*exp(w) = z
|
||||
two = ctx.mpf(2)
|
||||
for i in range(100):
|
||||
ew = ctx.exp(w)
|
||||
wew = w*ew
|
||||
wewz = wew-z
|
||||
wn = w - wewz/(wew+ew-(w+two)*wewz/(two*w+two))
|
||||
if ctx.mag(wn-w) <= ctx.mag(wn) - tol:
|
||||
w = wn
|
||||
break
|
||||
else:
|
||||
w = wn
|
||||
if i == 100:
|
||||
ctx.warn("Lambert W iteration failed to converge for z = %s" % z)
|
||||
ctx.prec = prec
|
||||
return +w
|
||||
|
||||
@defun_wrapped
|
||||
def bell(ctx, n, x=1):
|
||||
x = ctx.convert(x)
|
||||
if not n:
|
||||
if ctx.isnan(x):
|
||||
return x
|
||||
return ctx.one
|
||||
if ctx.isinf(x) or ctx.isinf(n) or ctx.isnan(x) or ctx.isnan(n):
|
||||
return x**n
|
||||
if n == 1: return x
|
||||
if n == 2: return x*(x+1)
|
||||
if x == 0: return ctx.sincpi(n)
|
||||
return _polyexp(ctx, n, x, True) / ctx.exp(x)
|
||||
|
||||
def _polyexp(ctx, n, x, extra=False):
|
||||
def _terms():
|
||||
if extra:
|
||||
yield ctx.sincpi(n)
|
||||
t = x
|
||||
k = 1
|
||||
while 1:
|
||||
yield k**n * t
|
||||
k += 1
|
||||
t = t*x/k
|
||||
return ctx.sum_accurately(_terms, check_step=4)
|
||||
|
||||
@defun_wrapped
|
||||
def polyexp(ctx, s, z):
|
||||
if ctx.isinf(z) or ctx.isinf(s) or ctx.isnan(z) or ctx.isnan(s):
|
||||
return z**s
|
||||
if z == 0: return z*s
|
||||
if s == 0: return ctx.expm1(z)
|
||||
if s == 1: return ctx.exp(z)*z
|
||||
if s == 2: return ctx.exp(z)*z*(z+1)
|
||||
return _polyexp(ctx, s, z)
|
||||
|
||||
@defun_wrapped
|
||||
def cyclotomic(ctx, n, z):
|
||||
n = int(n)
|
||||
if n < 0:
|
||||
raise ValueError("n cannot be negative")
|
||||
p = ctx.one
|
||||
if n == 0:
|
||||
return p
|
||||
if n == 1:
|
||||
return z - p
|
||||
if n == 2:
|
||||
return z + p
|
||||
# Use divisor product representation. Unfortunately, this sometimes
|
||||
# includes singularities for roots of unity, which we have to cancel out.
|
||||
# Matching zeros/poles pairwise, we have (1-z^a)/(1-z^b) ~ a/b + O(z-1).
|
||||
a_prod = 1
|
||||
b_prod = 1
|
||||
num_zeros = 0
|
||||
num_poles = 0
|
||||
for d in range(1,n+1):
|
||||
if not n % d:
|
||||
w = ctx.moebius(n//d)
|
||||
# Use powm1 because it is important that we get 0 only
|
||||
# if it really is exactly 0
|
||||
b = -ctx.powm1(z, d)
|
||||
if b:
|
||||
p *= b**w
|
||||
else:
|
||||
if w == 1:
|
||||
a_prod *= d
|
||||
num_zeros += 1
|
||||
elif w == -1:
|
||||
b_prod *= d
|
||||
num_poles += 1
|
||||
#print n, num_zeros, num_poles
|
||||
if num_zeros:
|
||||
if num_zeros > num_poles:
|
||||
p *= 0
|
||||
else:
|
||||
p *= a_prod
|
||||
p /= b_prod
|
||||
return p
|
||||
|
||||
@defun
|
||||
def mangoldt(ctx, n):
|
||||
r"""
|
||||
Evaluates the von Mangoldt function `\Lambda(n) = \log p`
|
||||
if `n = p^k` a power of a prime, and `\Lambda(n) = 0` otherwise.
|
||||
|
||||
**Examples**
|
||||
|
||||
>>> from mpmath import mp, mangoldt, fsum
|
||||
>>> mp.dps = 25
|
||||
>>> mp.pretty = True
|
||||
>>> [mangoldt(n) for n in range(-2,3)]
|
||||
[0.0, 0.0, 0.0, 0.0, 0.6931471805599453094172321]
|
||||
>>> mangoldt(6)
|
||||
0.0
|
||||
>>> mangoldt(7)
|
||||
1.945910149055313305105353
|
||||
>>> mangoldt(8)
|
||||
0.6931471805599453094172321
|
||||
>>> fsum(mangoldt(n) for n in range(101))
|
||||
94.04531122935739224600493
|
||||
>>> fsum(mangoldt(n) for n in range(10001))
|
||||
10013.39669326311478372032
|
||||
|
||||
"""
|
||||
n = int(n)
|
||||
if n < 2:
|
||||
return ctx.zero
|
||||
if n % 2 == 0:
|
||||
# Must be a power of two
|
||||
if n & (n-1) == 0:
|
||||
return +ctx.ln2
|
||||
else:
|
||||
return ctx.zero
|
||||
# TODO: the following could be generalized into a perfect
|
||||
# power testing function
|
||||
# ---
|
||||
# Look for a small factor
|
||||
for p in (3,5,7,11,13,17,19,23,29,31):
|
||||
if not n % p:
|
||||
q, r = n // p, 0
|
||||
while q > 1:
|
||||
q, r = divmod(q, p)
|
||||
if r:
|
||||
return ctx.zero
|
||||
return ctx.ln(p)
|
||||
if ctx.isprime(n):
|
||||
return ctx.ln(n)
|
||||
# Obviously, we could use arbitrary-precision arithmetic for this...
|
||||
if n > 10**30:
|
||||
raise NotImplementedError
|
||||
k = 2
|
||||
while 1:
|
||||
p = int(n**(1./k) + 0.5)
|
||||
if p < 2:
|
||||
return ctx.zero
|
||||
if p ** k == n:
|
||||
if ctx.isprime(p):
|
||||
return ctx.ln(p)
|
||||
k += 1
|
||||
|
||||
@defun
|
||||
def stirling1(ctx, n, k, exact=False):
|
||||
v = ctx._stirling1(int(n), int(k))
|
||||
if exact:
|
||||
return int(v)
|
||||
else:
|
||||
return ctx.mpf(v)
|
||||
|
||||
@defun
|
||||
def stirling2(ctx, n, k, exact=False):
|
||||
v = ctx._stirling2(int(n), int(k))
|
||||
if exact:
|
||||
return int(v)
|
||||
else:
|
||||
return ctx.mpf(v)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,508 @@
|
||||
from ..libmp.backend import MPQ
|
||||
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 = -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, pcfd, mpf, chop, diff, taylor
|
||||
>>> mp.dps = 25
|
||||
>>> mp.pretty = True
|
||||
>>> pcfd(0,0)
|
||||
1.0
|
||||
>>> pcfd(1,0)
|
||||
0.0
|
||||
>>> pcfd(2,0)
|
||||
-1.0
|
||||
>>> pcfd(3,0)
|
||||
0.0
|
||||
>>> pcfd(4,0)
|
||||
3.0
|
||||
>>> pcfd(-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, mpf, pcfu, sqrt, pi, exp, erfc
|
||||
>>> 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-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, pcfu, diff, pcfv, sqrt, pi, chop
|
||||
>>> 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 = MPQ(1,2)
|
||||
r = 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, mpf, pcfw, power, gamma, sqrt, diff
|
||||
>>> 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], [], [], [MPQ(1,2)*a+MPQ(1,4)], [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], [], [], [MPQ(1,2)*a+MPQ(3,4)], \
|
||||
[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 not z and ctx.isint(n) and int(n.real) % 2:
|
||||
return ctx.zero
|
||||
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
|
||||
if kwargs.get('force_series') is None:
|
||||
kwargs['force_series'] = True
|
||||
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
|
||||
if kwargs.get('force_series') is None:
|
||||
kwargs['force_series'] = True
|
||||
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)
|
||||
@@ -0,0 +1,285 @@
|
||||
from .functions import defun, defun_wrapped
|
||||
|
||||
@defun
|
||||
def qp(ctx, a, q=None, n=None, **kwargs):
|
||||
r"""
|
||||
Evaluates the q-Pochhammer symbol (or q-rising factorial)
|
||||
|
||||
.. math ::
|
||||
|
||||
(a; q)_n = \prod_{k=0}^{n-1} (1-a q^k)
|
||||
|
||||
where `n = \infty` is permitted if `|q| < 1`. Called with two arguments,
|
||||
``qp(a,q)`` computes `(a;q)_{\infty}`; with a single argument, ``qp(q)``
|
||||
computes `(q;q)_{\infty}`. The special case
|
||||
|
||||
.. math ::
|
||||
|
||||
\phi(q) = (q; q)_{\infty} = \prod_{k=1}^{\infty} (1-q^k) =
|
||||
\sum_{k=-\infty}^{\infty} (-1)^k q^{(3k^2-k)/2}
|
||||
|
||||
is also known as the Euler function, or (up to a factor `q^{-1/24}`)
|
||||
the Dedekind eta function.
|
||||
|
||||
**Examples**
|
||||
|
||||
If `n` is a positive integer, the function amounts to a finite product::
|
||||
|
||||
>>> from mpmath import (mp, qp, fprod, limit, rf, taylor, findroot,
|
||||
... diffun, mpf, jtheta, pi, root)
|
||||
>>> mp.dps = 25
|
||||
>>> mp.pretty = True
|
||||
>>> qp(2,3,5)
|
||||
-725305.0
|
||||
>>> fprod(1-2*3**k for k in range(5))
|
||||
-725305.0
|
||||
>>> qp(2,3,0)
|
||||
1.0
|
||||
|
||||
Complex arguments are allowed::
|
||||
|
||||
>>> qp(2-1j, 0.75j)
|
||||
(0.4628842231660149089976379 + 4.481821753552703090628793j)
|
||||
|
||||
The regular Pochhammer symbol `(a)_n` is obtained in the
|
||||
following limit as `q \to 1`::
|
||||
|
||||
>>> a, n = 4, 7
|
||||
>>> limit(lambda q: qp(q**a,q,n) / (1-q)**n, 1)
|
||||
604800.0
|
||||
>>> rf(a,n)
|
||||
604800.0
|
||||
|
||||
The Taylor series of the reciprocal Euler function gives
|
||||
the partition function `P(n)`, i.e. the number of ways of writing
|
||||
`n` as a sum of positive integers::
|
||||
|
||||
>>> taylor(lambda q: 1/qp(q), 0, 10)
|
||||
[1.0, 1.0, 2.0, 3.0, 5.0, 7.0, 11.0, 15.0, 22.0, 30.0, 42.0]
|
||||
|
||||
Special values include::
|
||||
|
||||
>>> qp(0)
|
||||
1.0
|
||||
>>> findroot(diffun(qp), -0.4) # location of maximum
|
||||
-0.4112484791779547734440257
|
||||
>>> qp(_)
|
||||
1.228348867038575112586878
|
||||
|
||||
The q-Pochhammer symbol is related to the Jacobi theta functions.
|
||||
For example, the following identity holds::
|
||||
|
||||
>>> q = mpf(0.5) # arbitrary
|
||||
>>> qp(q)
|
||||
0.2887880950866024212788997
|
||||
>>> root(3,-2)*root(q,-24)*jtheta(2,pi/6,root(q,6))
|
||||
0.2887880950866024212788997
|
||||
|
||||
"""
|
||||
a = ctx.convert(a)
|
||||
if n is None:
|
||||
n = ctx.inf
|
||||
else:
|
||||
n = ctx.convert(n)
|
||||
if n < 0:
|
||||
raise ValueError("n cannot be negative")
|
||||
if q is None:
|
||||
q = a
|
||||
else:
|
||||
q = ctx.convert(q)
|
||||
if n == 0:
|
||||
return ctx.one + 0*(a+q)
|
||||
infinite = (n == ctx.inf)
|
||||
same = (a == q)
|
||||
if infinite:
|
||||
if abs(q) >= 1:
|
||||
if same and (q == -1 or q == 1):
|
||||
return ctx.zero * q
|
||||
raise ValueError("q-function only defined for |q| < 1")
|
||||
elif q == 0:
|
||||
return ctx.one - a
|
||||
maxterms = kwargs.get('maxterms', 50*ctx.prec)
|
||||
if infinite and same:
|
||||
# Euler's pentagonal theorem
|
||||
def terms():
|
||||
t = 1
|
||||
yield t
|
||||
k = 1
|
||||
x1 = q
|
||||
x2 = q**2
|
||||
while 1:
|
||||
yield (-1)**k * x1
|
||||
yield (-1)**k * x2
|
||||
x1 *= q**(3*k+1)
|
||||
x2 *= q**(3*k+2)
|
||||
k += 1
|
||||
if k > maxterms:
|
||||
raise ctx.NoConvergence
|
||||
return ctx.sum_accurately(terms)
|
||||
# return ctx.nprod(lambda k: 1-a*q**k, [0,n-1])
|
||||
def factors():
|
||||
k = 0
|
||||
r = ctx.one
|
||||
while 1:
|
||||
yield 1 - a*r
|
||||
r *= q
|
||||
k += 1
|
||||
if k >= n:
|
||||
return
|
||||
if k > maxterms:
|
||||
raise ctx.NoConvergence
|
||||
return ctx.mul_accurately(factors)
|
||||
|
||||
@defun_wrapped
|
||||
def qgamma(ctx, z, q, **kwargs):
|
||||
r"""
|
||||
Evaluates the q-gamma function
|
||||
|
||||
.. math ::
|
||||
|
||||
\Gamma_q(z) = \frac{(q; q)_{\infty}}{(q^z; q)_{\infty}} (1-q)^{1-z}.
|
||||
|
||||
|
||||
**Examples**
|
||||
|
||||
Evaluation for real and complex arguments::
|
||||
|
||||
>>> from mpmath import mp, qgamma, mpf
|
||||
>>> mp.dps = 25
|
||||
>>> mp.pretty = True
|
||||
>>> qgamma(4,0.75)
|
||||
4.046875
|
||||
>>> qgamma(6,6)
|
||||
121226245.0
|
||||
>>> qgamma(3+4j, 0.5j)
|
||||
(0.1663082382255199834630088 + 0.01952474576025952984418217j)
|
||||
|
||||
The q-gamma function satisfies a functional equation similar
|
||||
to that of the ordinary gamma function::
|
||||
|
||||
>>> q = mpf(0.25)
|
||||
>>> z = mpf(2.5)
|
||||
>>> qgamma(z+1,q)
|
||||
1.428277424823760954685912
|
||||
>>> (1-q**z)/(1-q)*qgamma(z,q)
|
||||
1.428277424823760954685912
|
||||
|
||||
"""
|
||||
if abs(q) > 1:
|
||||
return ctx.qgamma(z,1/q)*q**((z-2)*(z-1)*0.5)
|
||||
return ctx.qp(q, q, None, **kwargs) / \
|
||||
ctx.qp(q**z, q, None, **kwargs) * (1-q)**(1-z)
|
||||
|
||||
@defun_wrapped
|
||||
def qfac(ctx, z, q, **kwargs):
|
||||
r"""
|
||||
Evaluates the q-factorial,
|
||||
|
||||
.. math ::
|
||||
|
||||
[n]_q! = (1+q)(1+q+q^2)\cdots(1+q+\cdots+q^{n-1})
|
||||
|
||||
or more generally
|
||||
|
||||
.. math ::
|
||||
|
||||
[z]_q! = \frac{(q;q)_z}{(1-q)^z}.
|
||||
|
||||
**Examples**
|
||||
|
||||
>>> from mpmath import mp, qfac
|
||||
>>> mp.dps = 25
|
||||
>>> mp.pretty = True
|
||||
>>> qfac(0,0)
|
||||
1.0
|
||||
>>> qfac(4,3)
|
||||
2080.0
|
||||
>>> qfac(5,6)
|
||||
121226245.0
|
||||
>>> qfac(1+1j, 2+1j)
|
||||
(0.4370556551322672478613695 + 0.2609739839216039203708921j)
|
||||
|
||||
"""
|
||||
if ctx.isint(z) and ctx._re(z) > 0:
|
||||
n = int(ctx._re(z))
|
||||
return ctx.qp(q, q, n, **kwargs) / (1-q)**n
|
||||
return ctx.qgamma(z+1, q, **kwargs)
|
||||
|
||||
@defun
|
||||
def qhyper(ctx, a_s, b_s, q, z, **kwargs):
|
||||
r"""
|
||||
Evaluates the basic hypergeometric series or hypergeometric q-series
|
||||
|
||||
.. math ::
|
||||
|
||||
\,_r\phi_s \left[\begin{matrix}
|
||||
a_1 & a_2 & \ldots & a_r \\
|
||||
b_1 & b_2 & \ldots & b_s
|
||||
\end{matrix} ; q,z \right] =
|
||||
\sum_{n=0}^\infty
|
||||
\frac{(a_1;q)_n, \ldots, (a_r;q)_n}
|
||||
{(b_1;q)_n, \ldots, (b_s;q)_n}
|
||||
\left((-1)^n q^{n\choose 2}\right)^{1+s-r}
|
||||
\frac{z^n}{(q;q)_n}
|
||||
|
||||
where `(a;q)_n` denotes the q-Pochhammer symbol (see :func:`~mpmath.qp`).
|
||||
|
||||
**Examples**
|
||||
|
||||
Evaluation works for real and complex arguments::
|
||||
|
||||
>>> from mpmath import qhyper, mp, nsum, qp, inf, j
|
||||
>>> mp.dps = 25
|
||||
>>> mp.pretty = True
|
||||
>>> qhyper([0.5], [2.25], 0.25, 4)
|
||||
-0.1975849091263356009534385
|
||||
>>> qhyper([0.5], [2.25], 0.25-0.25j, 4)
|
||||
(2.806330244925716649839237 + 3.568997623337943121769938j)
|
||||
>>> qhyper([1+j], [2,3+0.5j], 0.25, 3+4j)
|
||||
(9.112885171773400017270226 - 1.272756997166375050700388j)
|
||||
|
||||
Comparing with a summation of the defining series, using
|
||||
:func:`~mpmath.nsum`::
|
||||
|
||||
>>> b, q, z = 3, 0.25, 0.5
|
||||
>>> qhyper([], [b], q, z)
|
||||
0.6221136748254495583228324
|
||||
>>> nsum(lambda n: z**n / qp(q,q,n)/qp(b,q,n) * q**(n*(n-1)), [0,inf])
|
||||
0.6221136748254495583228324
|
||||
|
||||
"""
|
||||
#a_s = [ctx._convert_param(a)[0] for a in a_s]
|
||||
#b_s = [ctx._convert_param(b)[0] for b in b_s]
|
||||
#q = ctx._convert_param(q)[0]
|
||||
a_s = [ctx.convert(a) for a in a_s]
|
||||
b_s = [ctx.convert(b) for b in b_s]
|
||||
q = ctx.convert(q)
|
||||
z = ctx.convert(z)
|
||||
r = len(a_s)
|
||||
s = len(b_s)
|
||||
d = 1+s-r
|
||||
maxterms = kwargs.get('maxterms', 50*ctx.prec)
|
||||
def terms():
|
||||
t = ctx.one
|
||||
yield t
|
||||
qk = 1
|
||||
k = 0
|
||||
x = 1
|
||||
while 1:
|
||||
for a in a_s:
|
||||
p = 1 - a*qk
|
||||
t *= p
|
||||
for b in b_s:
|
||||
p = 1 - b*qk
|
||||
if not p:
|
||||
raise ValueError
|
||||
t /= p
|
||||
t *= z
|
||||
x *= (-1)**d * qk ** d
|
||||
qk *= q
|
||||
t /= (1 - qk)
|
||||
k += 1
|
||||
yield t * x
|
||||
if k > maxterms:
|
||||
raise ctx.NoConvergence
|
||||
return ctx.sum_accurately(terms)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,32 @@
|
||||
from .functions import defun_wrapped
|
||||
|
||||
@defun_wrapped
|
||||
def squarew(ctx, t, amplitude=1, period=1):
|
||||
P = period
|
||||
A = amplitude
|
||||
return A*((-1)**ctx.floor(2*t/P))
|
||||
|
||||
@defun_wrapped
|
||||
def trianglew(ctx, t, amplitude=1, period=1):
|
||||
A = amplitude
|
||||
P = period
|
||||
|
||||
return 2*A*(0.5 - ctx.fabs(1 - 2*ctx.frac(t/P + 0.25)))
|
||||
|
||||
@defun_wrapped
|
||||
def sawtoothw(ctx, t, amplitude=1, period=1):
|
||||
A = amplitude
|
||||
P = period
|
||||
return A*ctx.frac(t/P)
|
||||
|
||||
@defun_wrapped
|
||||
def unit_triangle(ctx, t, amplitude=1):
|
||||
A = amplitude
|
||||
if t <= -1 or t >= 1:
|
||||
return ctx.zero
|
||||
return A*(-ctx.fabs(t) + 1)
|
||||
|
||||
@defun_wrapped
|
||||
def sigmoid(ctx, t, amplitude=1):
|
||||
A = amplitude
|
||||
return A / (1 + ctx.exp(-t))
|
||||
@@ -0,0 +1,552 @@
|
||||
from mpmath.libmp.libintmath import jacobi_symbol
|
||||
|
||||
from .functions import defun, defun_wrapped
|
||||
|
||||
|
||||
@defun
|
||||
def _djacobi_theta2(ctx, z, q, nd):
|
||||
# the loops below break when the fixed precision quantities
|
||||
# a and b go to zero;
|
||||
# right shifting small negative numbers by wp one obtains -1, not zero,
|
||||
# so the condition a**2 + b**2 > MIN is used to break the loops.
|
||||
MIN = 2
|
||||
extra1 = 10
|
||||
extra2 = 20
|
||||
if not ctx._im(q) and not ctx._im(z):
|
||||
wp = ctx.prec + extra1
|
||||
x = ctx.to_fixed(ctx._re(q), wp)
|
||||
x2 = (x*x) >> wp
|
||||
a = b = x2
|
||||
c1, s1 = ctx.cos_sin(ctx._re(z), prec=wp)
|
||||
cn = c1 = ctx.to_fixed(c1, wp)
|
||||
sn = s1 = ctx.to_fixed(s1, wp)
|
||||
c2 = (c1*c1 - s1*s1) >> wp
|
||||
s2 = (c1 * s1) >> (wp - 1)
|
||||
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
|
||||
if nd&1:
|
||||
s = s1 + ((a * sn * 3**nd) >> wp)
|
||||
else:
|
||||
s = c1 + ((a * cn * 3**nd) >> wp)
|
||||
n = 2
|
||||
while abs(a) > MIN:
|
||||
b = (b*x2) >> wp
|
||||
a = (a*b) >> wp
|
||||
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
|
||||
if nd&1:
|
||||
s += (a * sn * (2*n+1)**nd) >> wp
|
||||
else:
|
||||
s += (a * cn * (2*n+1)**nd) >> wp
|
||||
n += 1
|
||||
s = -(s << 1)
|
||||
s = ctx.ldexp(s, -wp)
|
||||
# case z real, q complex
|
||||
elif not ctx._im(z):
|
||||
wp = ctx.prec + extra2
|
||||
xre = ctx.to_fixed(ctx._re(q), wp)
|
||||
xim = ctx.to_fixed(ctx._im(q), wp)
|
||||
x2re = (xre*xre - xim*xim) >> wp
|
||||
x2im = (xre*xim) >> (wp - 1)
|
||||
are = bre = x2re
|
||||
aim = bim = x2im
|
||||
c1, s1 = ctx.cos_sin(ctx._re(z), prec=wp)
|
||||
cn = c1 = ctx.to_fixed(c1, wp)
|
||||
sn = s1 = ctx.to_fixed(s1, wp)
|
||||
c2 = (c1*c1 - s1*s1) >> wp
|
||||
s2 = (c1 * s1) >> (wp - 1)
|
||||
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
|
||||
if nd&1:
|
||||
sre = s1 + ((are * sn * 3**nd) >> wp)
|
||||
sim = ((aim * sn * 3**nd) >> wp)
|
||||
else:
|
||||
sre = c1 + ((are * cn * 3**nd) >> wp)
|
||||
sim = ((aim * cn * 3**nd) >> wp)
|
||||
n = 5
|
||||
while are**2 + aim**2 > MIN:
|
||||
bre, bim = (bre * x2re - bim * x2im) >> wp, \
|
||||
(bre * x2im + bim * x2re) >> wp
|
||||
are, aim = (are * bre - aim * bim) >> wp, \
|
||||
(are * bim + aim * bre) >> wp
|
||||
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
|
||||
|
||||
if nd&1:
|
||||
sre += ((are * sn * n**nd) >> wp)
|
||||
sim += ((aim * sn * n**nd) >> wp)
|
||||
else:
|
||||
sre += ((are * cn * n**nd) >> wp)
|
||||
sim += ((aim * cn * n**nd) >> wp)
|
||||
n += 2
|
||||
sre = -(sre << 1)
|
||||
sim = -(sim << 1)
|
||||
sre = ctx.ldexp(sre, -wp)
|
||||
sim = ctx.ldexp(sim, -wp)
|
||||
s = ctx.mpc(sre, sim)
|
||||
# case z complex, q real
|
||||
elif not ctx._im(q):
|
||||
wp = ctx.prec + extra2
|
||||
x = ctx.to_fixed(ctx._re(q), wp)
|
||||
x2 = (x*x) >> wp
|
||||
a = b = x2
|
||||
c1, s1 = ctx.cos_sin(z, prec=wp)
|
||||
cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
|
||||
cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
|
||||
snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
|
||||
snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
|
||||
c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
|
||||
c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
|
||||
s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
|
||||
s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
|
||||
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
|
||||
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
|
||||
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
|
||||
t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
|
||||
cnre = t1
|
||||
cnim = t2
|
||||
snre = t3
|
||||
snim = t4
|
||||
if nd&1:
|
||||
sre = s1re + ((a * snre * 3**nd) >> wp)
|
||||
sim = s1im + ((a * snim * 3**nd) >> wp)
|
||||
else:
|
||||
sre = c1re + ((a * cnre * 3**nd) >> wp)
|
||||
sim = c1im + ((a * cnim * 3**nd) >> wp)
|
||||
n = 5
|
||||
while abs(a) > MIN:
|
||||
b = (b*x2) >> wp
|
||||
a = (a*b) >> wp
|
||||
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
|
||||
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
|
||||
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
|
||||
t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
|
||||
cnre = t1
|
||||
cnim = t2
|
||||
snre = t3
|
||||
snim = t4
|
||||
if nd&1:
|
||||
sre += ((a * snre * n**nd) >> wp)
|
||||
sim += ((a * snim * n**nd) >> wp)
|
||||
else:
|
||||
sre += ((a * cnre * n**nd) >> wp)
|
||||
sim += ((a * cnim * n**nd) >> wp)
|
||||
n += 2
|
||||
sre = -(sre << 1)
|
||||
sim = -(sim << 1)
|
||||
sre = ctx.ldexp(sre, -wp)
|
||||
sim = ctx.ldexp(sim, -wp)
|
||||
s = ctx.mpc(sre, sim)
|
||||
# case z and q complex
|
||||
else:
|
||||
wp = ctx.prec + extra2
|
||||
xre = ctx.to_fixed(ctx._re(q), wp)
|
||||
xim = ctx.to_fixed(ctx._im(q), wp)
|
||||
x2re = (xre*xre - xim*xim) >> wp
|
||||
x2im = (xre*xim) >> (wp - 1)
|
||||
are = bre = x2re
|
||||
aim = bim = x2im
|
||||
c1, s1 = ctx.cos_sin(z, prec=wp)
|
||||
cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
|
||||
cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
|
||||
snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
|
||||
snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
|
||||
c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
|
||||
c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
|
||||
s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
|
||||
s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
|
||||
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
|
||||
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
|
||||
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
|
||||
t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
|
||||
cnre = t1
|
||||
cnim = t2
|
||||
snre = t3
|
||||
snim = t4
|
||||
if nd&1:
|
||||
sre = s1re + (((are * snre - aim * snim) * 3**nd) >> wp)
|
||||
sim = s1im + (((are * snim + aim * snre)* 3**nd) >> wp)
|
||||
else:
|
||||
sre = c1re + (((are * cnre - aim * cnim) * 3**nd) >> wp)
|
||||
sim = c1im + (((are * cnim + aim * cnre)* 3**nd) >> wp)
|
||||
n = 5
|
||||
while are**2 + aim**2 > MIN:
|
||||
bre, bim = (bre * x2re - bim * x2im) >> wp, \
|
||||
(bre * x2im + bim * x2re) >> wp
|
||||
are, aim = (are * bre - aim * bim) >> wp, \
|
||||
(are * bim + aim * bre) >> wp
|
||||
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
|
||||
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
|
||||
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
|
||||
t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
|
||||
cnre = t1
|
||||
cnim = t2
|
||||
snre = t3
|
||||
snim = t4
|
||||
if nd&1:
|
||||
sre += (((are * snre - aim * snim) * n**nd) >> wp)
|
||||
sim += (((aim * snre + are * snim) * n**nd) >> wp)
|
||||
else:
|
||||
sre += (((are * cnre - aim * cnim) * n**nd) >> wp)
|
||||
sim += (((aim * cnre + are * cnim) * n**nd) >> wp)
|
||||
n += 2
|
||||
sre = -(sre << 1)
|
||||
sim = -(sim << 1)
|
||||
sre = ctx.ldexp(sre, -wp)
|
||||
sim = ctx.ldexp(sim, -wp)
|
||||
s = ctx.mpc(sre, sim)
|
||||
s *= ctx.nthroot(q, 4)
|
||||
return (-1)**(1 - (nd&1) + nd//2) * s
|
||||
|
||||
@defun
|
||||
def _djacobi_theta3(ctx, z, q, nd):
|
||||
MIN = 2
|
||||
extra1 = 10
|
||||
extra2 = 20
|
||||
if not ctx._im(q) and not ctx._im(z):
|
||||
s = 0
|
||||
wp = ctx.prec + extra1
|
||||
x = ctx.to_fixed(ctx._re(q), wp)
|
||||
a = (1 << wp)
|
||||
b = x
|
||||
x2 = (x*x) >> wp
|
||||
c1, s1 = ctx.cos_sin(ctx._re(z)*2, prec=wp)
|
||||
c1 = ctx.to_fixed(c1, wp)
|
||||
s1 = ctx.to_fixed(s1, wp)
|
||||
cn = c1
|
||||
sn = s1
|
||||
if nd&1:
|
||||
s += (a * sn) >> wp
|
||||
else:
|
||||
s += (a * cn) >> wp
|
||||
n = 2
|
||||
while True:
|
||||
b = (b*x2) >> wp
|
||||
a = (a*b) >> wp
|
||||
if abs(a) <= MIN:
|
||||
break
|
||||
cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
|
||||
if nd&1:
|
||||
s += (a * sn * n**nd) >> wp
|
||||
else:
|
||||
s += (a * cn * n**nd) >> wp
|
||||
n += 1
|
||||
s = -(s << (nd+1))
|
||||
s = ctx.ldexp(s, -wp)*q
|
||||
# case z real, q complex
|
||||
elif not ctx._im(z):
|
||||
wp = ctx.prec + extra2
|
||||
xre = ctx.to_fixed(ctx._re(q), wp)
|
||||
xim = ctx.to_fixed(ctx._im(q), wp)
|
||||
x2re = (xre*xre - xim*xim) >> wp
|
||||
x2im = (xre*xim) >> (wp - 1)
|
||||
are = (1 << wp)
|
||||
aim = 0
|
||||
bre = xre
|
||||
bim = xim
|
||||
c1, s1 = ctx.cos_sin(ctx._re(z)*2, prec=wp)
|
||||
c1 = ctx.to_fixed(c1, wp)
|
||||
s1 = ctx.to_fixed(s1, wp)
|
||||
cn = c1
|
||||
sn = s1
|
||||
if nd&1:
|
||||
sre = (are * sn) >> wp
|
||||
sim = (aim * sn) >> wp
|
||||
else:
|
||||
sre = (are * cn) >> wp
|
||||
sim = (aim * cn) >> wp
|
||||
n = 2
|
||||
while True:
|
||||
bre, bim = (bre * x2re - bim * x2im) >> wp, \
|
||||
(bre * x2im + bim * x2re) >> wp
|
||||
are, aim = (are * bre - aim * bim) >> wp, \
|
||||
(are * bim + aim * bre) >> wp
|
||||
if are**2 + aim**2 <= MIN:
|
||||
break
|
||||
cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
|
||||
if nd&1:
|
||||
sre += (are * sn * n**nd) >> wp
|
||||
sim += (aim * sn * n**nd) >> wp
|
||||
else:
|
||||
sre += (are * cn * n**nd) >> wp
|
||||
sim += (aim * cn * n**nd) >> wp
|
||||
n += 1
|
||||
sre = -(sre << (nd+1))
|
||||
sim = -(sim << (nd+1))
|
||||
sre = ctx.ldexp(sre, -wp)
|
||||
sim = ctx.ldexp(sim, -wp)
|
||||
s = ctx.mpc(sre, sim)*q
|
||||
# case z complex, q real
|
||||
elif not ctx._im(q):
|
||||
wp = ctx.prec + extra2
|
||||
x = ctx.to_fixed(ctx._re(q), wp)
|
||||
a = (1 << wp)
|
||||
b = x
|
||||
x2 = (x*x) >> wp
|
||||
c1, s1 = ctx.cos_sin(2*z, prec=wp)
|
||||
cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
|
||||
cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
|
||||
snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
|
||||
snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
|
||||
if nd&1:
|
||||
sre = (a * snre) >> wp
|
||||
sim = (a * snim) >> wp
|
||||
else:
|
||||
sre = (a * cnre) >> wp
|
||||
sim = (a * cnim) >> wp
|
||||
n = 2
|
||||
while True:
|
||||
b = (b*x2) >> wp
|
||||
a = (a*b) >> wp
|
||||
if abs(a) <= MIN:
|
||||
break
|
||||
t1 = (cnre*c1re - cnim*c1im - snre*s1re + snim*s1im) >> wp
|
||||
t2 = (cnre*c1im + cnim*c1re - snre*s1im - snim*s1re) >> wp
|
||||
t3 = (snre*c1re - snim*c1im + cnre*s1re - cnim*s1im) >> wp
|
||||
t4 = (snre*c1im + snim*c1re + cnre*s1im + cnim*s1re) >> wp
|
||||
cnre = t1
|
||||
cnim = t2
|
||||
snre = t3
|
||||
snim = t4
|
||||
if nd&1:
|
||||
sre += (a * snre * n**nd) >> wp
|
||||
sim += (a * snim * n**nd) >> wp
|
||||
else:
|
||||
sre += (a * cnre * n**nd) >> wp
|
||||
sim += (a * cnim * n**nd) >> wp
|
||||
n += 1
|
||||
sre = -(sre << (nd+1))
|
||||
sim = -(sim << (nd+1))
|
||||
sre = ctx.ldexp(sre, -wp)
|
||||
sim = ctx.ldexp(sim, -wp)
|
||||
s = ctx.mpc(sre, sim)*q
|
||||
# case z and q complex
|
||||
else:
|
||||
wp = ctx.prec + extra2
|
||||
xre = ctx.to_fixed(ctx._re(q), wp)
|
||||
xim = ctx.to_fixed(ctx._im(q), wp)
|
||||
x2re = (xre*xre - xim*xim) >> wp
|
||||
x2im = (xre*xim) >> (wp - 1)
|
||||
are = (1 << wp)
|
||||
aim = 0
|
||||
bre = xre
|
||||
bim = xim
|
||||
c1, s1 = ctx.cos_sin(2*z, prec=wp)
|
||||
cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
|
||||
cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
|
||||
snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
|
||||
snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
|
||||
if nd&1:
|
||||
sre = (are * snre - aim * snim) >> wp
|
||||
sim = (aim * snre + are * snim) >> wp
|
||||
else:
|
||||
sre = (are * cnre - aim * cnim) >> wp
|
||||
sim = (aim * cnre + are * cnim) >> wp
|
||||
n = 2
|
||||
while True:
|
||||
bre, bim = (bre * x2re - bim * x2im) >> wp, \
|
||||
(bre * x2im + bim * x2re) >> wp
|
||||
are, aim = (are * bre - aim * bim) >> wp, \
|
||||
(are * bim + aim * bre) >> wp
|
||||
if are**2 + aim**2 <= MIN:
|
||||
break
|
||||
t1 = (cnre*c1re - cnim*c1im - snre*s1re + snim*s1im) >> wp
|
||||
t2 = (cnre*c1im + cnim*c1re - snre*s1im - snim*s1re) >> wp
|
||||
t3 = (snre*c1re - snim*c1im + cnre*s1re - cnim*s1im) >> wp
|
||||
t4 = (snre*c1im + snim*c1re + cnre*s1im + cnim*s1re) >> wp
|
||||
cnre = t1
|
||||
cnim = t2
|
||||
snre = t3
|
||||
snim = t4
|
||||
if nd&1:
|
||||
sre += ((are * snre - aim * snim) * n**nd) >> wp
|
||||
sim += ((aim * snre + are * snim) * n**nd) >> wp
|
||||
else:
|
||||
sre += ((are * cnre - aim * cnim) * n**nd) >> wp
|
||||
sim += ((aim * cnre + are * cnim) * n**nd) >> wp
|
||||
n += 1
|
||||
sre = -(sre << (nd+1))
|
||||
sim = -(sim << (nd+1))
|
||||
sre = ctx.ldexp(sre, -wp)
|
||||
sim = ctx.ldexp(sim, -wp)
|
||||
s = ctx.mpc(sre, sim)*q
|
||||
if nd&1:
|
||||
return (-1)**(nd//2) * s
|
||||
else:
|
||||
return (-1)**(1 + nd//2) * s + (ctx.zero if nd else ctx.one)
|
||||
|
||||
@defun
|
||||
def _reduce_psl2z(ctx, z):
|
||||
"""
|
||||
Returns the cumulative transformation matrix, that reduces a complex
|
||||
number z to the fundamental domain of PSL(2, Z), chosen to be
|
||||
|Re(z)| ≤ 0.5 and |z| ≥ 1.
|
||||
"""
|
||||
z = ctx.convert(z)
|
||||
assert z.imag > 0, f"Expected point from upper half-plane, got {ctx.mpc(z)}"
|
||||
|
||||
a = d = 1
|
||||
b = c = 0
|
||||
|
||||
z_orig = z
|
||||
with ctx.extraprec(30):
|
||||
while True:
|
||||
# Translate to center in |Re(z)| ≤ 1/2
|
||||
n = round(z.real)
|
||||
if n:
|
||||
z -= n
|
||||
a -= n*c
|
||||
b -= n*d
|
||||
|
||||
# Maybe apply an inversion
|
||||
if z.real**2 + z.imag**2 < 1:
|
||||
z = -1/z
|
||||
a, c = -c, a
|
||||
b, d = -d, b
|
||||
if abs(z.real) <= 0.5:
|
||||
break
|
||||
else:
|
||||
break
|
||||
|
||||
# Canonicalize matrix
|
||||
if c < 0 or (c == 0 and d < 0):
|
||||
a, b, c, d = -a, -b, -c, -d
|
||||
|
||||
return a, b, c, d
|
||||
|
||||
#
|
||||
# General modular transformations for jtheta()
|
||||
#
|
||||
# References:
|
||||
# * Hans Rademacher (1973), "Topics in Analytic Number Theory",
|
||||
# Springer. Section 81.
|
||||
# * [DLMF]_, §20.7(viii).
|
||||
#
|
||||
|
||||
_T_map = {(0, 0): 1, (0, 1): 2, (1, 0): 4, (1, 1): 3}
|
||||
|
||||
def _jtheta_permutation(n, a, b, c, d):
|
||||
if n == 2:
|
||||
return _T_map[(c%2, d%2)]
|
||||
if n == 3:
|
||||
return _T_map[((a + c)%2, (b + d)%2)]
|
||||
if n == 4:
|
||||
return _T_map[(a%2, b%2)]
|
||||
return 1
|
||||
|
||||
@defun
|
||||
def _jtheta_eps(ctx, n, a, b, c, d):
|
||||
if n != 1:
|
||||
if n == 2:
|
||||
phi = (c - 2)*d - 2 + 2*(1 - c)*((d + 1)%2)
|
||||
elif n == 3:
|
||||
phi = (a + c - 2)*(b + d) - 3 + 2*(1 - a - c)*((b + d + 1)%2)
|
||||
else:
|
||||
phi = (a - 2)*b - 4 + 2*(1 - a)*((b + 1)%2)
|
||||
k = ctx._jtheta_eps(1, -d, b, c, -a)
|
||||
else:
|
||||
if c % 2 == 0:
|
||||
phi = d*(b - c - 1) + 2
|
||||
k = jacobi_symbol(c, d)
|
||||
else:
|
||||
phi = c*(a + d + 1) - 3
|
||||
k = jacobi_symbol(d, c)
|
||||
return ctx.expjpi(ctx.convert(phi)/4)/k
|
||||
|
||||
@defun
|
||||
def _jtheta_needs_modular(ctx, z, q):
|
||||
if not z.imag:
|
||||
return False
|
||||
tau = ctx.taufrom(q=q)
|
||||
assert abs(q) < 1 and tau.imag > 0
|
||||
return abs(tau.real) > 0.5 or tau.real**2 + tau.imag**2 < 1
|
||||
|
||||
@defun
|
||||
def _jtheta_modular(ctx, g, n, z, q, nd):
|
||||
a, b, c, d = g
|
||||
tau = ctx.taufrom(q=q)
|
||||
v = -1/(c*tau + d)
|
||||
alpha = 1j*v*c/ctx.pi
|
||||
|
||||
assert abs(q) < 1 and tau.imag > 0
|
||||
|
||||
new_n = _jtheta_permutation(n, -d, b, c, -a)
|
||||
new_z = z*v
|
||||
new_tau = (a*tau + b)/(c*tau + d)
|
||||
new_q = ctx.qfrom(tau=new_tau)
|
||||
|
||||
assert abs(new_tau.real) <= 0.5 and new_tau.real**2 + new_tau.imag**2 >= 1
|
||||
|
||||
def terms():
|
||||
Him1, Hi = ctx.zero, ctx.one
|
||||
a2 = alpha*2
|
||||
a2z = a2*z
|
||||
for i in range(nd + 1):
|
||||
yield (ctx.binomial(nd, i) * Hi * v**(nd - i)
|
||||
* ctx.jtheta(new_n, new_z, new_q, nd - i))
|
||||
Him1, Hi = Hi, a2z*Hi + a2*i*Him1
|
||||
|
||||
C = ctx._jtheta_eps(n, -d, b, c, -a)*ctx.sqrt(v/1j)
|
||||
X = alpha*z**2
|
||||
return C*ctx.exp(X)*sum(terms())
|
||||
|
||||
@defun
|
||||
def jtheta(ctx, n, z, q, derivative=0):
|
||||
n = int(n)
|
||||
z = ctx.convert(z)
|
||||
q = ctx.convert(q)
|
||||
nd = int(derivative)
|
||||
|
||||
if n not in range(1, 5):
|
||||
raise ValueError("First argument expected to be 1, 2, 3 or 4")
|
||||
if abs(q) >= 1:
|
||||
raise ValueError(f"abs(q) >= 1")
|
||||
|
||||
# We use Fourier series (DLMF, §20.2(i)) to compute functions, when
|
||||
# |q| is not near 1. Else, transform τ to the fundamental
|
||||
# domain (|Re(τ)| ≤ 0.5 and |τ| ≥ 1), applying transformations
|
||||
# of lattice parameter (DLMF, §20.7(viii)).
|
||||
|
||||
if ctx._jtheta_needs_modular(z, q):
|
||||
tau = ctx.taufrom(q=q)
|
||||
g = ctx._reduce_psl2z(tau)
|
||||
|
||||
# Estimate exponential factor
|
||||
c, d = g[2:]
|
||||
extra = 10*(nd + 1) + max(0, ctx.mag(c/(c*tau + d)*z**2))
|
||||
|
||||
return ctx.extraprec(extra, True)(ctx._jtheta_modular)(g, n, z, q, nd)
|
||||
|
||||
# At that point, τ is in the fundamental domain and thus Im(τ) ≥ √3π/2.
|
||||
# Using quasi-periodicity property (see DLMF, §20.2(ii)) brings
|
||||
# z to the domain |Im(z)| ≤ π |Im(τ)|/2.
|
||||
|
||||
if abs(z.imag) > abs(ctx.log(q).real)/2:
|
||||
with ctx.extraprec(10):
|
||||
tau = ctx.taufrom(q=q)
|
||||
tau_pi = tau*ctx.pi
|
||||
k = round(z.imag/tau_pi.imag)
|
||||
assert k != 0
|
||||
beta = -ctx.j*2*k
|
||||
C = q**(k**2)*ctx.exp(beta*z)
|
||||
if n in (1, 4) and k & 1:
|
||||
C = -C
|
||||
new_z = z - k*tau_pi
|
||||
|
||||
def terms():
|
||||
for i in range(nd + 1):
|
||||
yield (ctx.binomial(nd, i) * beta**i
|
||||
* ctx.jtheta(n, new_z, q, nd - i))
|
||||
|
||||
res = C*sum(terms())
|
||||
return +res
|
||||
|
||||
extra = 10 + ctx.prec * nd // 10
|
||||
if z:
|
||||
M = ctx.mag(z)
|
||||
if M > 5 or ((n != 1 if nd else n == 1) and M < -5):
|
||||
extra += 2*abs(M)
|
||||
with ctx.extraprec(extra):
|
||||
if n < 3:
|
||||
z_inner = z - ctx.pi/2 if n == 1 else z
|
||||
res = ctx._djacobi_theta2(z_inner, q, nd)
|
||||
else:
|
||||
q_inner = -q if n == 4 else q
|
||||
res = ctx._djacobi_theta3(z, q_inner, nd)
|
||||
return +res
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user