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_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_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 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)