def mpc_psi(m, z, prec, rnd=round_fast):
"""
Computation of the polygamma function of arbitrary integer order
m >= 0, for a complex argument z.
"""
if m == 0:
return mpc_psi0(z, prec, rnd)
re, im = z
wp = prec + 20
sign, man, exp, bc = re
if not man:
if re == finf and im == fzero:
return (fzero, fzero)
if re == fnan:
return fnan
# Recurrence
w = to_int(re)
n = int(0.4 * wp + 4 * m)
s = mpc_zero
if w < n:
for k in xrange(w, n):
t = mpc_pow_int(z, -m - 1, wp)
s = mpc_add(s, t, wp)
z = mpc_add_mpf(z, fone, wp)
zm = mpc_pow_int(z, -m, wp)
z2 = mpc_pow_int(z, -2, wp)
# 1/m*(z+N)^m
integral_term = mpc_div_mpf(zm, from_int(m), wp)
s = mpc_add(s, integral_term, wp)
# 1/2*(z+N)^(-(m+1))
s = mpc_add(s, mpc_mul_mpf(mpc_div(zm, z, wp), fhalf, wp), wp)
a = m + 1
b = 2
k = 1
# Important: we want to sum up to the *relative* error,
# not the absolute error, because psi^(m)(z) might be tiny
magn = mpc_abs(s, 10)
magn = magn[2] + magn[3]
eps = mpf_shift(fone, magn - wp + 2)
while 1:
zm = mpc_mul(zm, z2, wp)
bern = mpf_bernoulli(2 * k, wp)
scal = mpf_mul_int(bern, a, wp)
scal = mpf_div(scal, from_int(b), wp)
term = mpc_mul_mpf(zm, scal, wp)
s = mpc_add(s, term, wp)
szterm = mpc_abs(term, 10)
if k > 2 and mpf_le(szterm, eps):
break
#print k, to_str(szterm, 10), to_str(eps, 10)
a *= (m + 2 * k) * (m + 2 * k + 1)
b *= (2 * k + 1) * (2 * k + 2)
k += 1
# Scale and sign factor
v = mpc_mul_mpf(s, mpf_gamma(from_int(m + 1), wp), prec, rnd)
if not (m & 1):
v = mpf_neg(v[0]), mpf_neg(v[1])
return v
def mpc_psi(m, z, prec, rnd=round_fast):
"""
Computation of the polygamma function of arbitrary integer order
m >= 0, for a complex argument z.
"""
if m == 0:
return mpc_psi0(z, prec, rnd)
re, im = z
wp = prec + 20
sign, man, exp, bc = re
if not man:
if re == finf and im == fzero:
return (fzero, fzero)
if re == fnan:
return fnan
# Recurrence
w = to_int(re)
n = int(0.4*wp + 4*m)
s = mpc_zero
if w < n:
for k in xrange(w, n):
t = mpc_pow_int(z, -m-1, wp)
s = mpc_add(s, t, wp)
z = mpc_add_mpf(z, fone, wp)
zm = mpc_pow_int(z, -m, wp)
z2 = mpc_pow_int(z, -2, wp)
# 1/m*(z+N)^m
integral_term = mpc_div_mpf(zm, from_int(m), wp)
s = mpc_add(s, integral_term, wp)
# 1/2*(z+N)^(-(m+1))
s = mpc_add(s, mpc_mul_mpf(mpc_div(zm, z, wp), fhalf, wp), wp)
a = m + 1
b = 2
k = 1
# Important: we want to sum up to the *relative* error,
# not the absolute error, because psi^(m)(z) might be tiny
magn = mpc_abs(s, 10)
magn = magn[2]+magn[3]
eps = mpf_shift(fone, magn-wp+2)
while 1:
zm = mpc_mul(zm, z2, wp)
bern = mpf_bernoulli(2*k, wp)
scal = mpf_mul_int(bern, a, wp)
scal = mpf_div(scal, from_int(b), wp)
term = mpc_mul_mpf(zm, scal, wp)
s = mpc_add(s, term, wp)
szterm = mpc_abs(term, 10)
if k > 2 and mpf_le(szterm, eps):
break
#print k, to_str(szterm, 10), to_str(eps, 10)
a *= (m+2*k)*(m+2*k+1)
b *= (2*k+1)*(2*k+2)
k += 1
# Scale and sign factor
v = mpc_mul_mpf(s, mpf_gamma(from_int(m+1), wp), prec, rnd)
if not (m & 1):
v = mpf_neg(v[0]), mpf_neg(v[1])
return v
def mpc_ellipk(z, prec, rnd=round_fast):
re, im = z
if im == fzero:
if re == finf:
return mpc_zero
if mpf_le(re, fone):
return mpf_ellipk(re, prec, rnd), fzero
wp = prec + 15
a = mpc_sqrt(mpc_sub(mpc_one, z, wp), wp)
v = mpc_agm1(a, wp)
r = mpc_mpf_div(mpf_pi(wp), v, prec, rnd)
return mpc_shift(r, -1)
def mpc_ellipk(z, prec, rnd=round_fast):
re, im = z
if im == fzero:
if re == finf:
return mpc_zero
if mpf_le(re, fone):
return mpf_ellipk(re, prec, rnd), fzero
wp = prec + 15
a = mpc_sqrt(mpc_sub(mpc_one, z, wp), wp)
v = mpc_agm1(a, wp)
r = mpc_mpf_div(mpf_pi(wp), v, prec, rnd)
return mpc_shift(r, -1)
def mpc_psi0(z, prec, rnd=round_fast):
"""
Computation of the digamma function (psi function of order 0)
of a complex argument.
"""
re, im = z
# Fall back to the real case
if im == fzero:
return (mpf_psi0(re, prec, rnd), fzero)
wp = prec + 20
sign, man, exp, bc = re
# Reflection formula
if sign and exp+bc > 3:
c = mpc_cos_pi(z, wp)
s = mpc_sin_pi(z, wp)
q = mpc_mul_mpf(mpc_div(c, s, wp), mpf_pi(wp), wp)
p = mpc_psi0(mpc_sub(mpc_one, z, wp), wp)
return mpc_sub(p, q, prec, rnd)
# Just the logarithmic term
if (not sign) and bc + exp > wp:
return mpc_log(mpc_sub(z, mpc_one, wp), prec, rnd)
# Initial recurrence to obtain a large enough z
w = to_int(re)
n = int(0.11*wp) + 2
s = mpc_zero
if w < n:
for k in xrange(w, n):
s = mpc_sub(s, mpc_reciprocal(z, wp), wp)
z = mpc_add_mpf(z, fone, wp)
z = mpc_sub(z, mpc_one, wp)
# Logarithmic and endpoint term
s = mpc_add(s, mpc_log(z, wp), wp)
s = mpc_add(s, mpc_div(mpc_half, z, wp), wp)
# Euler-Maclaurin remainder sum
z2 = mpc_square(z, wp)
t = mpc_one
prev = mpc_zero
k = 1
eps = mpf_shift(fone, -wp+2)
while 1:
t = mpc_mul(t, z2, wp)
bern = mpf_bernoulli(2*k, wp)
term = mpc_mpf_div(bern, mpc_mul_int(t, 2*k, wp), wp)
s = mpc_sub(s, term, wp)
szterm = mpc_abs(term, 10)
if k > 2 and mpf_le(szterm, eps):
break
prev = term
k += 1
return s
def mpc_psi0(z, prec, rnd=round_fast):
"""
Computation of the digamma function (psi function of order 0)
of a complex argument.
"""
re, im = z
# Fall back to the real case
if im == fzero:
return (mpf_psi0(re, prec, rnd), fzero)
wp = prec + 20
sign, man, exp, bc = re
# Reflection formula
if sign and exp + bc > 3:
c = mpc_cos_pi(z, wp)
s = mpc_sin_pi(z, wp)
q = mpc_mul(mpc_div(c, s, wp), (mpf_pi(wp), fzero), wp)
p = mpc_psi0(mpc_sub(mpc_one, z, wp), wp)
return mpc_sub(p, q, prec, rnd)
# Just the logarithmic term
if (not sign) and bc + exp > wp:
return mpc_log(mpc_sub(z, mpc_one, wp), prec, rnd)
# Initial recurrence to obtain a large enough z
w = to_int(re)
n = int(0.11 * wp) + 2
s = mpc_zero
if w < n:
for k in xrange(w, n):
s = mpc_sub(s, mpc_div(mpc_one, z, wp), wp)
z = mpc_add_mpf(z, fone, wp)
z = mpc_sub(z, mpc_one, wp)
# Logarithmic and endpoint term
s = mpc_add(s, mpc_log(z, wp), wp)
s = mpc_add(s, mpc_div(mpc_half, z, wp), wp)
# Euler-Maclaurin remainder sum
z2 = mpc_mul(z, z, wp)
t = mpc_one
prev = mpc_zero
k = 1
eps = mpf_shift(fone, -wp + 2)
while 1:
t = mpc_mul(t, z2, wp)
bern = mpf_bernoulli(2 * k, wp)
term = mpc_div((bern, fzero), mpc_mul_int(t, 2 * k, wp), wp)
s = mpc_sub(s, term, wp)
szterm = mpc_abs(term, 10)
if k > 2 and mpf_le(szterm, eps):
break
prev = term
k += 1
return s
def mpi_square(s, prec=0):
sa, sb = s
if mpf_ge(sa, fzero):
a = mpf_mul(sa, sa, prec, round_floor)
b = mpf_mul(sb, sb, prec, round_ceiling)
elif mpf_le(sb, fzero):
a = mpf_mul(sb, sb, prec, round_floor)
b = mpf_mul(sa, sa, prec, round_ceiling)
else:
sa = mpf_neg(sa)
sa, sb = mpf_min_max([sa, sb])
a = fzero
b = mpf_mul(sb, sb, prec, round_ceiling)
return a, b
def mpi_square(s, prec=0):
sa, sb = s
if mpf_ge(sa, fzero):
a = mpf_mul(sa, sa, prec, round_floor)
b = mpf_mul(sb, sb, prec, round_ceiling)
elif mpf_le(sb, fzero):
a = mpf_mul(sb, sb, prec, round_floor)
b = mpf_mul(sa, sa, prec, round_ceiling)
else:
sa = mpf_neg(sa)
sa, sb = mpf_min_max([sa, sb])
a = fzero
b = mpf_mul(sb, sb, prec, round_ceiling)
return a, b
def mpi_atan2(y, x, prec):
ya, yb = y
xa, xb = x
# Constrained to the real line
if ya == yb == fzero:
if mpf_ge(xa, fzero):
return mpi_zero
return mpi_pi(prec)
# Right half-plane
if mpf_ge(xa, fzero):
if mpf_ge(ya, fzero):
a = mpf_atan2(ya, xb, prec, round_floor)
else:
a = mpf_atan2(ya, xa, prec, round_floor)
if mpf_ge(yb, fzero):
b = mpf_atan2(yb, xa, prec, round_ceiling)
else:
b = mpf_atan2(yb, xb, prec, round_ceiling)
# Upper half-plane
elif mpf_ge(ya, fzero):
b = mpf_atan2(ya, xa, prec, round_ceiling)
if mpf_le(xb, fzero):
a = mpf_atan2(yb, xb, prec, round_floor)
else:
a = mpf_atan2(ya, xb, prec, round_floor)
# Lower half-plane
elif mpf_le(yb, fzero):
a = mpf_atan2(yb, xa, prec, round_floor)
if mpf_le(xb, fzero):
b = mpf_atan2(ya, xb, prec, round_ceiling)
else:
b = mpf_atan2(yb, xb, prec, round_ceiling)
# Covering the origin
else:
b = mpf_pi(prec, round_ceiling)
a = mpf_neg(b)
return a, b
def mpi_atan2(y, x, prec):
ya, yb = y
xa, xb = x
# Constrained to the real line
if ya == yb == fzero:
if mpf_ge(xa, fzero):
return mpi_zero
return mpi_pi(prec)
# Right half-plane
if mpf_ge(xa, fzero):
if mpf_ge(ya, fzero):
a = mpf_atan2(ya, xb, prec, round_floor)
else:
a = mpf_atan2(ya, xa, prec, round_floor)
if mpf_ge(yb, fzero):
b = mpf_atan2(yb, xa, prec, round_ceiling)
else:
b = mpf_atan2(yb, xb, prec, round_ceiling)
# Upper half-plane
elif mpf_ge(ya, fzero):
b = mpf_atan2(ya, xa, prec, round_ceiling)
if mpf_le(xb, fzero):
a = mpf_atan2(yb, xb, prec, round_floor)
else:
a = mpf_atan2(ya, xb, prec, round_floor)
# Lower half-plane
elif mpf_le(yb, fzero):
a = mpf_atan2(yb, xa, prec, round_floor)
if mpf_le(xb, fzero):
b = mpf_atan2(ya, xb, prec, round_ceiling)
else:
b = mpf_atan2(yb, xb, prec, round_ceiling)
# Covering the origin
else:
b = mpf_pi(prec, round_ceiling)
a = mpf_neg(b)
return a, b
def mpc_ellipe(z, prec, rnd=round_fast):
re, im = z
if im == fzero:
if re == finf:
return (fzero, finf)
if mpf_le(re, fone):
return mpf_ellipe(re, prec, rnd), fzero
wp = prec + 15
mag = mpc_abs(z, 1)
p = max(mag[2]+mag[3], 0) - wp
h = mpf_shift(fone, p)
K = mpc_ellipk(z, 2*wp)
Kh = mpc_ellipk(mpc_add_mpf(z, h, 2*wp), 2*wp)
Kdiff = mpc_shift(mpc_sub(Kh, K, wp), -p)
t = mpc_sub(mpc_one, z, wp)
b = mpc_mul(Kdiff, mpc_shift(z,1), wp)
return mpc_mul(t, mpc_add(K, b, wp), prec, rnd)
def mpc_ellipe(z, prec, rnd=round_fast):
re, im = z
if im == fzero:
if re == finf:
return (fzero, finf)
if mpf_le(re, fone):
return mpf_ellipe(re, prec, rnd), fzero
wp = prec + 15
mag = mpc_abs(z, 1)
p = max(mag[2] + mag[3], 0) - wp
h = mpf_shift(fone, p)
K = mpc_ellipk(z, 2 * wp)
Kh = mpc_ellipk(mpc_add_mpf(z, h, 2 * wp), 2 * wp)
Kdiff = mpc_shift(mpc_sub(Kh, K, wp), -p)
t = mpc_sub(mpc_one, z, wp)
b = mpc_mul(Kdiff, mpc_shift(z, 1), wp)
return mpc_mul(t, mpc_add(K, b, wp), prec, rnd)
def MIN(x, y):
if mpf_le(x, y):
return x
return y
def mpci_gamma(z, prec, type=0):
(a1, a2), (b1, b2) = z
# Real case
if b1 == b2 == fzero and (type != 3 or mpf_gt(a1, fzero)):
return mpi_gamma(z, prec, type), mpi_zero
# Estimate precision
wp = prec + 20
if type != 3:
amag = a2[2] + a2[3]
bmag = b2[2] + b2[3]
if a2 != fzero:
mag = max(amag, bmag)
else:
mag = bmag
an = abs(to_int(a2))
bn = abs(to_int(b2))
absn = max(an, bn)
gamma_size = max(0, absn * mag)
wp += bitcount(gamma_size)
# Assume type != 1
if type == 1:
(a1, a2) = mpi_add((a1, a2), mpi_one, wp)
z = (a1, a2), (b1, b2)
type = 0
# Avoid non-monotonic region near the negative real axis
if mpf_lt(a1, gamma_min_b):
if mpi_overlap((b1, b2), (gamma_mono_imag_a, gamma_mono_imag_b)):
# TODO: reflection formula
#if mpf_lt(a2, mpf_shift(fone,-1)):
# znew = mpci_sub((mpi_one,mpi_zero),z,wp)
# ...
# Recurrence:
# gamma(z) = gamma(z+1)/z
znew = mpi_add((a1, a2), mpi_one, wp), (b1, b2)
if type == 0:
return mpci_div(mpci_gamma(znew, prec + 2, 0), z, prec)
if type == 2:
return mpci_mul(mpci_gamma(znew, prec + 2, 2), z, prec)
if type == 3:
return mpci_sub(mpci_gamma(znew, prec + 2, 3),
mpci_log(z, prec + 2), prec)
# Use monotonicity (except for a small region close to the
# origin and near poles)
# upper half-plane
if mpf_ge(b1, fzero):
minre = mpc_loggamma((a1, b2), wp, round_floor)
maxre = mpc_loggamma((a2, b1), wp, round_ceiling)
minim = mpc_loggamma((a1, b1), wp, round_floor)
maxim = mpc_loggamma((a2, b2), wp, round_ceiling)
# lower half-plane
elif mpf_le(b2, fzero):
minre = mpc_loggamma((a1, b1), wp, round_floor)
maxre = mpc_loggamma((a2, b2), wp, round_ceiling)
minim = mpc_loggamma((a2, b1), wp, round_floor)
maxim = mpc_loggamma((a1, b2), wp, round_ceiling)
# crosses real axis
else:
maxre = mpc_loggamma((a2, fzero), wp, round_ceiling)
# stretches more into the lower half-plane
if mpf_gt(mpf_neg(b1), b2):
minre = mpc_loggamma((a1, b1), wp, round_ceiling)
else:
minre = mpc_loggamma((a1, b2), wp, round_ceiling)
minim = mpc_loggamma((a2, b1), wp, round_floor)
maxim = mpc_loggamma((a2, b2), wp, round_floor)
w = (minre[0], maxre[0]), (minim[1], maxim[1])
if type == 3:
return mpi_pos(w[0], prec), mpi_pos(w[1], prec)
if type == 2:
w = mpci_neg(w)
return mpci_exp(w, prec)
def mpi_le(s, t):
sa, sb = s
ta, tb = t
if mpf_le(sb, ta): return True
if mpf_gt(sa, tb): return False
return None
def mpi_le(s, t):
sa, sb = s
ta, tb = t
if mpf_le(sb, ta): return True
if mpf_gt(sa, tb): return False
return None
def MIN(x, y):
if mpf_le(x, y):
return x
return y
def mpci_gamma(z, prec, type=0):
(a1,a2), (b1,b2) = z
# Real case
if b1 == b2 == fzero and (type != 3 or mpf_gt(a1,fzero)):
return mpi_gamma(z, prec, type), mpi_zero
# Estimate precision
wp = prec+20
if type != 3:
amag = a2[2]+a2[3]
bmag = b2[2]+b2[3]
if a2 != fzero:
mag = max(amag, bmag)
else:
mag = bmag
an = abs(to_int(a2))
bn = abs(to_int(b2))
absn = max(an, bn)
gamma_size = max(0,absn*mag)
wp += bitcount(gamma_size)
# Assume type != 1
if type == 1:
(a1,a2) = mpi_add((a1,a2), mpi_one, wp); z = (a1,a2), (b1,b2)
type = 0
# Avoid non-monotonic region near the negative real axis
if mpf_lt(a1, gamma_min_b):
if mpi_overlap((b1,b2), (gamma_mono_imag_a, gamma_mono_imag_b)):
# TODO: reflection formula
#if mpf_lt(a2, mpf_shift(fone,-1)):
# znew = mpci_sub((mpi_one,mpi_zero),z,wp)
# ...
# Recurrence:
# gamma(z) = gamma(z+1)/z
znew = mpi_add((a1,a2), mpi_one, wp), (b1,b2)
if type == 0: return mpci_div(mpci_gamma(znew, prec+2, 0), z, prec)
if type == 2: return mpci_mul(mpci_gamma(znew, prec+2, 2), z, prec)
if type == 3: return mpci_sub(mpci_gamma(znew, prec+2, 3), mpci_log(z,prec+2), prec)
# Use monotonicity (except for a small region close to the
# origin and near poles)
# upper half-plane
if mpf_ge(b1, fzero):
minre = mpc_loggamma((a1,b2), wp, round_floor)
maxre = mpc_loggamma((a2,b1), wp, round_ceiling)
minim = mpc_loggamma((a1,b1), wp, round_floor)
maxim = mpc_loggamma((a2,b2), wp, round_ceiling)
# lower half-plane
elif mpf_le(b2, fzero):
minre = mpc_loggamma((a1,b1), wp, round_floor)
maxre = mpc_loggamma((a2,b2), wp, round_ceiling)
minim = mpc_loggamma((a2,b1), wp, round_floor)
maxim = mpc_loggamma((a1,b2), wp, round_ceiling)
# crosses real axis
else:
maxre = mpc_loggamma((a2,fzero), wp, round_ceiling)
# stretches more into the lower half-plane
if mpf_gt(mpf_neg(b1), b2):
minre = mpc_loggamma((a1,b1), wp, round_ceiling)
else:
minre = mpc_loggamma((a1,b2), wp, round_ceiling)
minim = mpc_loggamma((a2,b1), wp, round_floor)
maxim = mpc_loggamma((a2,b2), wp, round_floor)
w = (minre[0], maxre[0]), (minim[1], maxim[1])
if type == 3:
return mpi_pos(w[0], prec), mpi_pos(w[1], prec)
if type == 2:
w = mpci_neg(w)
return mpci_exp(w, prec)
