""" ----------------------------------------------------------------------- This module implements gamma- and zeta-related functions: * Bernoulli numbers * Factorials * The gamma function * Polygamma functions * Harmonic numbers * The Riemann zeta function * Constants related to these functions ----------------------------------------------------------------------- """ import math import sys from .backend import xrange from .backend import MPZ, MPZ_ZERO, MPZ_ONE, MPZ_THREE, gmpy from .libintmath import list_primes, ifac, ifac2, moebius from .libmpf import (\ round_floor, round_ceiling, round_down, round_up, round_nearest, round_fast, lshift, sqrt_fixed, isqrt_fast, fzero, fone, fnone, fhalf, ftwo, finf, fninf, fnan, from_int, to_int, to_fixed, from_man_exp, from_rational, mpf_pos, mpf_neg, mpf_abs, mpf_add, mpf_sub, mpf_mul, mpf_mul_int, mpf_div, mpf_sqrt, mpf_pow_int, mpf_rdiv_int, mpf_perturb, mpf_le, mpf_lt, mpf_gt, mpf_shift, negative_rnd, reciprocal_rnd, bitcount, to_float, mpf_floor, mpf_sign, ComplexResult ) from .libelefun import (\ constant_memo, def_mpf_constant, mpf_pi, pi_fixed, ln2_fixed, log_int_fixed, mpf_ln2, mpf_exp, mpf_log, mpf_pow, mpf_cosh, mpf_cos_sin, mpf_cosh_sinh, mpf_cos_sin_pi, mpf_cos_pi, mpf_sin_pi, ln_sqrt2pi_fixed, mpf_ln_sqrt2pi, sqrtpi_fixed, mpf_sqrtpi, cos_sin_fixed, exp_fixed ) from .libmpc import (\ mpc_zero, mpc_one, mpc_half, mpc_two, mpc_abs, mpc_shift, mpc_pos, mpc_neg, mpc_add, mpc_sub, mpc_mul, mpc_div, mpc_add_mpf, mpc_mul_mpf, mpc_div_mpf, mpc_mpf_div, mpc_mul_int, mpc_pow_int, mpc_log, mpc_exp, mpc_pow, mpc_cos_pi, mpc_sin_pi, mpc_reciprocal, mpc_square, mpc_sub_mpf ) # Catalan's constant is computed using Lupas's rapidly convergent series # (listed on http://mathworld.wolfram.com/CatalansConstant.html) # oo # ___ n-1 8n 2 3 2 # 1 \ (-1) 2 (40n - 24n + 3) [(2n)!] (n!) # K = --- ) ----------------------------------------- # 64 /___ 3 2 # n (2n-1) [(4n)!] # n = 1 @constant_memo def catalan_fixed(prec): prec = prec + 20 a = one = MPZ_ONE << prec s, t, n = 0, 1, 1 while t: a *= 32 * n**3 * (2*n-1) a //= (3-16*n+16*n**2)**2 t = a * (-1)**(n-1) * (40*n**2-24*n+3) // (n**3 * (2*n-1)) s += t n += 1 return s >> (20 + 6) # Khinchin's constant is relatively difficult to compute. Here # we use the rational zeta series # oo 2*n-1 # ___ ___ # \ ` zeta(2*n)-1 \ ` (-1)^(k+1) # log(K)*log(2) = ) ------------ ) ---------- # /___. n /___. k # n = 1 k = 1 # which adds half a digit per term. The essential trick for achieving # reasonable efficiency is to recycle both the values of the zeta # function (essentially Bernoulli numbers) and the partial terms of # the inner sum. # An alternative might be to use K = 2*exp[1/log(2) X] where # / 1 1 [ pi*x*(1-x^2) ] # X = | ------ log [ ------------ ]. # / 0 x(1+x) [ sin(pi*x) ] # and integrate numerically. In practice, this seems to be slightly # slower than the zeta series at high precision. @constant_memo def khinchin_fixed(prec): wp = int(prec + prec**0.5 + 15) s = MPZ_ZERO fac = from_int(4) t = ONE = MPZ_ONE << wp pi = mpf_pi(wp) pipow = twopi2 = mpf_shift(mpf_mul(pi, pi, wp), 2) n = 1 while 1: zeta2n = mpf_abs(mpf_bernoulli(2*n, wp)) zeta2n = mpf_mul(zeta2n, pipow, wp) zeta2n = mpf_div(zeta2n, fac, wp) zeta2n = to_fixed(zeta2n, wp) term = (((zeta2n - ONE) * t) // n) >> wp if term < 100: break #if not n % 10: # print n, math.log(int(abs(term))) s += term t += ONE//(2*n+1) - ONE//(2*n) n += 1 fac = mpf_mul_int(fac, (2*n)*(2*n-1), wp) pipow = mpf_mul(pipow, twopi2, wp) s = (s << wp) // ln2_fixed(wp) K = mpf_exp(from_man_exp(s, -wp), wp) K = to_fixed(K, prec) return K # Glaisher's constant is defined as A = exp(1/2 - zeta'(-1)). # One way to compute it would be to perform direct numerical # differentiation, but computing arbitrary Riemann zeta function # values at high precision is expensive. We instead use the formula # A = exp((6 (-zeta'(2))/pi^2 + log 2 pi + gamma)/12) # and compute zeta'(2) from the series representation # oo # ___ # \ log k # -zeta'(2) = ) ----- # /___ 2 # k # k = 2 # This series converges exceptionally slowly, but can be accelerated # using Euler-Maclaurin formula. The important insight is that the # E-M integral can be done in closed form and that the high order # are given by # n / \ # d | log x | a + b log x # --- | ----- | = ----------- # n | 2 | 2 + n # dx \ x / x # where a and b are integers given by a simple recurrence. Note # that just one logarithm is needed. However, lots of integer # logarithms are required for the initial summation. # This algorithm could possibly be turned into a faster algorithm # for general evaluation of zeta(s) or zeta'(s); this should be # looked into. @constant_memo def glaisher_fixed(prec): wp = prec + 30 # Number of direct terms to sum before applying the Euler-Maclaurin # formula to the tail. TODO: choose more intelligently N = int(0.33*prec + 5) ONE = MPZ_ONE << wp # Euler-Maclaurin, step 1: sum log(k)/k**2 for k from 2 to N-1 s = MPZ_ZERO for k in range(2, N): #print k, N s += log_int_fixed(k, wp) // k**2 logN = log_int_fixed(N, wp) #logN = to_fixed(mpf_log(from_int(N), wp+20), wp) # E-M step 2: integral of log(x)/x**2 from N to inf s += (ONE + logN) // N # E-M step 3: endpoint correction term f(N)/2 s += logN // (N**2 * 2) # E-M step 4: the series of derivatives pN = N**3 a = 1 b = -2 j = 3 fac = from_int(2) k = 1 while 1: # D(2*k-1) * B(2*k) / fac(2*k) [D(n) = nth derivative] D = ((a << wp) + b*logN) // pN D = from_man_exp(D, -wp) B = mpf_bernoulli(2*k, wp) term = mpf_mul(B, D, wp) term = mpf_div(term, fac, wp) term = to_fixed(term, wp) if abs(term) < 100: break #if not k % 10: # print k, math.log(int(abs(term)), 10) s -= term # Advance derivative twice a, b, pN, j = b-a*j, -j*b, pN*N, j+1 a, b, pN, j = b-a*j, -j*b, pN*N, j+1 k += 1 fac = mpf_mul_int(fac, (2*k)*(2*k-1), wp) # A = exp((6*s/pi**2 + log(2*pi) + euler)/12) pi = pi_fixed(wp) s *= 6 s = (s << wp) // (pi**2 >> wp) s += euler_fixed(wp) s += to_fixed(mpf_log(from_man_exp(2*pi, -wp), wp), wp) s //= 12 A = mpf_exp(from_man_exp(s, -wp), wp) return to_fixed(A, prec) # Apery's constant can be computed using the very rapidly convergent # series # oo # ___ 2 10 # \ n 205 n + 250 n + 77 (n!) # zeta(3) = ) (-1) ------------------- ---------- # /___ 64 5 # n = 0 ((2n+1)!) @constant_memo def apery_fixed(prec): prec += 20 d = MPZ_ONE << prec term = MPZ(77) << prec n = 1 s = MPZ_ZERO while term: s += term d *= (n**10) d //= (((2*n+1)**5) * (2*n)**5) term = (-1)**n * (205*(n**2) + 250*n + 77) * d n += 1 return s >> (20 + 6) """ Euler's constant (gamma) is computed using the Brent-McMillan formula, gamma ~= I(n)/J(n) - log(n), where I(n) = sum_{k=0,1,2,...} (n**k / k!)**2 * H(k) J(n) = sum_{k=0,1,2,...} (n**k / k!)**2 H(k) = 1 + 1/2 + 1/3 + ... + 1/k The error is bounded by O(exp(-4n)). Choosing n to be a power of two, 2**p, the logarithm becomes particularly easy to calculate.[1] We use the formulation of Algorithm 3.9 in [2] to make the summation more efficient. Reference: [1] Xavier Gourdon & Pascal Sebah, The Euler constant: gamma http://numbers.computation.free.fr/Constants/Gamma/gamma.pdf [2] [BorweinBailey]_ """ @constant_memo def euler_fixed(prec): extra = 30 prec += extra # choose p such that exp(-4*(2**p)) < 2**-n p = int(math.log((prec/4) * math.log(2), 2)) + 1 n = 2**p A = U = -p*ln2_fixed(prec) B = V = MPZ_ONE << prec k = 1 while 1: B = B*n**2//k**2 A = (A*n**2//k + B)//k U += A V += B if max(abs(A), abs(B)) < 100: break k += 1 return (U<<(prec-extra))//V # Use zeta accelerated formulas for the Mertens and twin # prime constants; see # http://mathworld.wolfram.com/MertensConstant.html # http://mathworld.wolfram.com/TwinPrimesConstant.html @constant_memo def mertens_fixed(prec): wp = prec + 20 m = 2 s = mpf_euler(wp) while 1: t = mpf_zeta_int(m, wp) if t == fone: break t = mpf_log(t, wp) t = mpf_mul_int(t, moebius(m), wp) t = mpf_div(t, from_int(m), wp) s = mpf_add(s, t) m += 1 return to_fixed(s, prec) @constant_memo def twinprime_fixed(prec): def I(n): return sum(moebius(d)<<(n//d) for d in xrange(1,n+1) if not n%d)//n wp = 2*prec + 30 res = fone primes = [from_rational(1,p,wp) for p in [2,3,5,7]] ppowers = [mpf_mul(p,p,wp) for p in primes] n = 2 while 1: a = mpf_zeta_int(n, wp) for i in range(4): a = mpf_mul(a, mpf_sub(fone, ppowers[i]), wp) ppowers[i] = mpf_mul(ppowers[i], primes[i], wp) a = mpf_pow_int(a, -I(n), wp) if mpf_pos(a, prec+10, 'n') == fone: break #from libmpf import to_str #print n, to_str(mpf_sub(fone, a), 6) res = mpf_mul(res, a, wp) n += 1 res = mpf_mul(res, from_int(3*15*35), wp) res = mpf_div(res, from_int(4*16*36), wp) return to_fixed(res, prec) mpf_euler = def_mpf_constant(euler_fixed) mpf_apery = def_mpf_constant(apery_fixed) mpf_khinchin = def_mpf_constant(khinchin_fixed) mpf_glaisher = def_mpf_constant(glaisher_fixed) mpf_catalan = def_mpf_constant(catalan_fixed) mpf_mertens = def_mpf_constant(mertens_fixed) mpf_twinprime = def_mpf_constant(twinprime_fixed) #-----------------------------------------------------------------------# # # # Bernoulli numbers # # # #-----------------------------------------------------------------------# MAX_BERNOULLI_CACHE = 3000 r""" Small Bernoulli numbers and factorials are used in numerous summations, so it is critical for speed that sequential computation is fast and that values are cached up to a fairly high threshold. On the other hand, we also want to support fast computation of isolated large numbers. Currently, no such acceleration is provided for integer factorials (though it is for large floating-point factorials, which are computed via gamma if the precision is low enough). For sequential computation of Bernoulli numbers, we use Ramanujan's formula / n + 3 \ B = (A(n) - S(n)) / | | n \ n / where A(n) = (n+3)/3 when n = 0 or 2 (mod 6), A(n) = -(n+3)/6 when n = 4 (mod 6), and [n/6] ___ \ / n + 3 \ S(n) = ) | | * B /___ \ n - 6*k / n-6*k k = 1 For isolated large Bernoulli numbers, we use the Riemann zeta function to calculate a numerical value for B_n. The von Staudt-Clausen theorem can then be used to optionally find the exact value of the numerator and denominator. """ bernoulli_cache = {} f3 = from_int(3) f6 = from_int(6) def bernoulli_size(n): """Accurately estimate the size of B_n (even n > 2 only)""" lgn = math.log(n,2) return int(2.326 + 0.5*lgn + n*(lgn - 4.094)) BERNOULLI_PREC_CUTOFF = bernoulli_size(MAX_BERNOULLI_CACHE) def mpf_bernoulli(n, prec, rnd=None): """Computation of Bernoulli numbers (numerically)""" if n < 2: if n < 0: raise ValueError("Bernoulli numbers only defined for n >= 0") if n == 0: return fone if n == 1: return mpf_neg(fhalf) # For odd n > 1, the Bernoulli numbers are zero if n & 1: return fzero # If precision is extremely high, we can save time by computing # the Bernoulli number at a lower precision that is sufficient to # obtain the exact fraction, round to the exact fraction, and # convert the fraction back to an mpf value at the original precision if prec > BERNOULLI_PREC_CUTOFF and prec > bernoulli_size(n)*1.1 + 1000: p, q = bernfrac(n) return from_rational(p, q, prec, rnd or round_floor) if n > MAX_BERNOULLI_CACHE: return mpf_bernoulli_huge(n, prec, rnd) wp = prec + 30 # Reuse nearby precisions wp += 32 - (prec & 31) cached = bernoulli_cache.get(wp) if cached: numbers, state = cached if n in numbers: if not rnd: return numbers[n] return mpf_pos(numbers[n], prec, rnd) m, bin, bin1 = state if n - m > 10: return mpf_bernoulli_huge(n, prec, rnd) else: if n > 10: return mpf_bernoulli_huge(n, prec, rnd) numbers = {0:fone} m, bin, bin1 = state = [2, MPZ(10), MPZ_ONE] bernoulli_cache[wp] = (numbers, state) while m <= n: #print m case = m % 6 # Accurately estimate size of B_m so we can use # fixed point math without using too much precision szbm = bernoulli_size(m) s = 0 sexp = max(0, szbm) - wp if m < 6: a = MPZ_ZERO else: a = bin1 for j in xrange(1, m//6+1): usign, uman, uexp, ubc = u = numbers[m-6*j] if usign: uman = -uman s += lshift(a*uman, uexp-sexp) # Update inner binomial coefficient j6 = 6*j a *= ((m-5-j6)*(m-4-j6)*(m-3-j6)*(m-2-j6)*(m-1-j6)*(m-j6)) a //= ((4+j6)*(5+j6)*(6+j6)*(7+j6)*(8+j6)*(9+j6)) if case == 0: b = mpf_rdiv_int(m+3, f3, wp) if case == 2: b = mpf_rdiv_int(m+3, f3, wp) if case == 4: b = mpf_rdiv_int(-m-3, f6, wp) s = from_man_exp(s, sexp, wp) b = mpf_div(mpf_sub(b, s, wp), from_int(bin), wp) numbers[m] = b m += 2 # Update outer binomial coefficient bin = bin * ((m+2)*(m+3)) // (m*(m-1)) if m > 6: bin1 = bin1 * ((2+m)*(3+m)) // ((m-7)*(m-6)) state[:] = [m, bin, bin1] return numbers[n] def mpf_bernoulli_huge(n, prec, rnd=None): wp = prec + 10 piprec = wp + int(math.log(n,2)) v = mpf_gamma_int(n+1, wp) v = mpf_mul(v, mpf_zeta_int(n, wp), wp) v = mpf_mul(v, mpf_pow_int(mpf_pi(piprec), -n, wp)) v = mpf_shift(v, 1-n) if not n & 3: v = mpf_neg(v) return mpf_pos(v, prec, rnd or round_fast) def bernfrac(n): r""" Returns a tuple of integers `(p, q)` such that `p/q = B_n` exactly, where `B_n` denotes the `n`-th Bernoulli number. The fraction is always reduced to lowest terms. Note that for `n > 1` and `n` odd, `B_n = 0`, and `(0, 1)` is returned. **Examples** The first few Bernoulli numbers are exactly:: >>> from mpmath import * >>> for n in range(15): ... p, q = bernfrac(n) ... print("%s %s/%s" % (n, p, q)) ... 0 1/1 1 -1/2 2 1/6 3 0/1 4 -1/30 5 0/1 6 1/42 7 0/1 8 -1/30 9 0/1 10 5/66 11 0/1 12 -691/2730 13 0/1 14 7/6 This function works for arbitrarily large `n`:: >>> p, q = bernfrac(10**4) >>> print(q) 2338224387510 >>> print(len(str(p))) 27692 >>> mp.dps = 15 >>> print(mpf(p) / q) -9.04942396360948e+27677 >>> print(bernoulli(10**4)) -9.04942396360948e+27677 .. note :: :func:`~mpmath.bernoulli` computes a floating-point approximation directly, without computing the exact fraction first. This is much faster for large `n`. **Algorithm** :func:`~mpmath.bernfrac` works by computing the value of `B_n` numerically and then using the von Staudt-Clausen theorem [1] to reconstruct the exact fraction. For large `n`, this is significantly faster than computing `B_1, B_2, \ldots, B_2` recursively with exact arithmetic. The implementation has been tested for `n = 10^m` up to `m = 6`. In practice, :func:`~mpmath.bernfrac` appears to be about three times slower than the specialized program calcbn.exe [2] **References** 1. MathWorld, von Staudt-Clausen Theorem: http://mathworld.wolfram.com/vonStaudt-ClausenTheorem.html 2. The Bernoulli Number Page: http://www.bernoulli.org/ """ n = int(n) if n < 3: return [(1, 1), (-1, 2), (1, 6)][n] if n & 1: return (0, 1) q = 1 for k in list_primes(n+1): if not (n % (k-1)): q *= k prec = bernoulli_size(n) + int(math.log(q,2)) + 20 b = mpf_bernoulli(n, prec) p = mpf_mul(b, from_int(q)) pint = to_int(p, round_nearest) return (pint, q) #-----------------------------------------------------------------------# # # # Polygamma functions # # # #-----------------------------------------------------------------------# r""" For all polygamma (psi) functions, we use the Euler-Maclaurin summation formula. It looks slightly different in the m = 0 and m > 0 cases. For m = 0, we have oo ___ B (0) 1 \ 2 k -2 k psi (z) ~ log z + --- - ) ------ z 2 z /___ (2 k)! k = 1 Experiment shows that the minimum term of the asymptotic series reaches 2^(-p) when Re(z) > 0.11*p. So we simply use the recurrence for psi (equivalent, in fact, to summing to the first few terms directly before applying E-M) to obtain z large enough. Since, very crudely, log z ~= 1 for Re(z) > 1, we can use fixed-point arithmetic (if z is extremely large, log(z) itself is a sufficient approximation, so we can stop there already). For Re(z) << 0, we could use recurrence, but this is of course inefficient for large negative z, so there we use the reflection formula instead. For m > 0, we have N - 1 ___ ~~~(m) [ \ 1 ] 1 1 psi (z) ~ [ ) -------- ] + ---------- + -------- + [ /___ m+1 ] m+1 m k = 1 (z+k) ] 2 (z+N) m (z+N) oo ___ B \ 2 k (m+1) (m+2) ... (m+2k-1) + ) ------ ------------------------ /___ (2 k)! m + 2 k k = 1 (z+N) where ~~~ denotes the function rescaled by 1/((-1)^(m+1) m!). Here again N is chosen to make z+N large enough for the minimum term in the last series to become smaller than eps. TODO: the current estimation of N for m > 0 is *very suboptimal*. TODO: implement the reflection formula for m > 0, Re(z) << 0. It is generally a combination of multiple cotangents. Need to figure out a reasonably simple way to generate these formulas on the fly. TODO: maybe use exact algorithms to compute psi for integral and certain rational arguments, as this can be much more efficient. (On the other hand, the availability of these special values provides a convenient way to test the general algorithm.) """ # Harmonic numbers are just shifted digamma functions # We should calculate these exactly when x is an integer # and when doing so is faster. def mpf_harmonic(x, prec, rnd): if x in (fzero, fnan, finf): return x a = mpf_psi0(mpf_add(fone, x, prec+5), prec) return mpf_add(a, mpf_euler(prec+5, rnd), prec, rnd) def mpc_harmonic(z, prec, rnd): if z[1] == fzero: return (mpf_harmonic(z[0], prec, rnd), fzero) a = mpc_psi0(mpc_add_mpf(z, fone, prec+5), prec) return mpc_add_mpf(a, mpf_euler(prec+5, rnd), prec, rnd) def mpf_psi0(x, prec, rnd=round_fast): """ Computation of the digamma function (psi function of order 0) of a real argument. """ sign, man, exp, bc = x wp = prec + 10 if not man: if x == finf: return x if x == fninf or x == fnan: return fnan if x == fzero or (exp >= 0 and sign): raise ValueError("polygamma pole") # Near 0 -- fixed-point arithmetic becomes bad if exp+bc < -5: v = mpf_psi0(mpf_add(x, fone, prec, rnd), prec, rnd) return mpf_sub(v, mpf_div(fone, x, wp, rnd), prec, rnd) # Reflection formula if sign and exp+bc > 3: c, s = mpf_cos_sin_pi(x, wp) q = mpf_mul(mpf_div(c, s, wp), mpf_pi(wp), wp) p = mpf_psi0(mpf_sub(fone, x, wp), wp) return mpf_sub(p, q, prec, rnd) # The logarithmic term is accurate enough if (not sign) and bc + exp > wp: return mpf_log(mpf_sub(x, fone, wp), prec, rnd) # Initial recurrence to obtain a large enough x m = to_int(x) n = int(0.11*wp) + 2 s = MPZ_ZERO x = to_fixed(x, wp) one = MPZ_ONE << wp if m < n: for k in xrange(m, n): s -= (one << wp) // x x += one x -= one # Logarithmic term s += to_fixed(mpf_log(from_man_exp(x, -wp, wp), wp), wp) # Endpoint term in Euler-Maclaurin expansion s += (one << wp) // (2*x) # Euler-Maclaurin remainder sum x2 = (x*x) >> wp t = one prev = 0 k = 1 while 1: t = (t*x2) >> wp bsign, bman, bexp, bbc = mpf_bernoulli(2*k, wp) offset = (bexp + 2*wp) if offset >= 0: term = (bman << offset) // (t*(2*k)) else: term = (bman >> (-offset)) // (t*(2*k)) if k & 1: s -= term else: s += term if k > 2 and term >= prev: break prev = term k += 1 return from_man_exp(s, -wp, wp, rnd) 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 szprev = fzero 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) or mpf_le(szprev, szterm)): break prev = term szprev = szterm k += 1 return s # Currently unoptimized def mpf_psi(m, x, prec, rnd=round_fast): """ Computation of the polygamma function of arbitrary integer order m >= 0, for a real argument x. """ if m == 0: return mpf_psi0(x, prec, rnd=round_fast) return mpc_psi(m, (x, fzero), prec, rnd)[0] 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 im[1]: if im in (finf, fninf, fnan): return (fnan, fnan) if not man: if re == finf and im == fzero: return (fzero, fzero) if re == fnan: return (fnan, 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 #-----------------------------------------------------------------------# # # # Riemann zeta function # # # #-----------------------------------------------------------------------# r""" We use zeta(s) = eta(s) / (1 - 2**(1-s)) and Borwein's approximation n-1 ___ k -1 \ (-1) (d_k - d_n) eta(s) ~= ---- ) ------------------ d_n /___ s k = 0 (k + 1) where k ___ i \ (n + i - 1)! 4 d_k = n ) ---------------. /___ (n - i)! (2i)! i = 0 If s = a + b*I, the absolute error for eta(s) is bounded by 3 (1 + 2|b|) ------------ * exp(|b| pi/2) n (3+sqrt(8)) Disregarding the linear term, we have approximately, log(err) ~= log(exp(1.58*|b|)) - log(5.8**n) log(err) ~= 1.58*|b| - log(5.8)*n log(err) ~= 1.58*|b| - 1.76*n log2(err) ~= 2.28*|b| - 2.54*n So for p bits, we should choose n > (p + 2.28*|b|) / 2.54. References: ----------- Peter Borwein, "An Efficient Algorithm for the Riemann Zeta Function" http://www.cecm.sfu.ca/personal/pborwein/PAPERS/P117.ps http://en.wikipedia.org/wiki/Dirichlet_eta_function """ borwein_cache = {} def borwein_coefficients(n): if n in borwein_cache: return borwein_cache[n] ds = [MPZ_ZERO] * (n+1) d = MPZ_ONE s = ds[0] = MPZ_ONE for i in range(1, n+1): d = d * 4 * (n+i-1) * (n-i+1) d //= ((2*i) * ((2*i)-1)) s += d ds[i] = s borwein_cache[n] = ds return ds ZETA_INT_CACHE_MAX_PREC = 1000 zeta_int_cache = {} def mpf_zeta_int(s, prec, rnd=round_fast): """ Optimized computation of zeta(s) for an integer s. """ wp = prec + 20 s = int(s) if s in zeta_int_cache and zeta_int_cache[s][0] >= wp: return mpf_pos(zeta_int_cache[s][1], prec, rnd) if s < 2: if s == 1: raise ValueError("zeta(1) pole") if not s: return mpf_neg(fhalf) return mpf_div(mpf_bernoulli(-s+1, wp), from_int(s-1), prec, rnd) # 2^-s term vanishes? if s >= wp: return mpf_perturb(fone, 0, prec, rnd) # 5^-s term vanishes? elif s >= wp*0.431: t = one = 1 << wp t += 1 << (wp - s) t += one // (MPZ_THREE ** s) t += 1 << max(0, wp - s*2) return from_man_exp(t, -wp, prec, rnd) else: # Fast enough to sum directly? # Even better, we use the Euler product (idea stolen from pari) m = (float(wp)/(s-1) + 1) if m < 30: needed_terms = int(2.0**m + 1) if needed_terms < int(wp/2.54 + 5) / 10: t = fone for k in list_primes(needed_terms): #print k, needed_terms powprec = int(wp - s*math.log(k,2)) if powprec < 2: break a = mpf_sub(fone, mpf_pow_int(from_int(k), -s, powprec), wp) t = mpf_mul(t, a, wp) return mpf_div(fone, t, wp) # Use Borwein's algorithm n = int(wp/2.54 + 5) d = borwein_coefficients(n) t = MPZ_ZERO s = MPZ(s) for k in xrange(n): t += (((-1)**k * (d[k] - d[n])) << wp) // (k+1)**s t = (t << wp) // (-d[n]) t = (t << wp) // ((1 << wp) - (1 << (wp+1-s))) if (s in zeta_int_cache and zeta_int_cache[s][0] < wp) or (s not in zeta_int_cache): zeta_int_cache[s] = (wp, from_man_exp(t, -wp-wp)) return from_man_exp(t, -wp-wp, prec, rnd) def mpf_zeta(s, prec, rnd=round_fast, alt=0): sign, man, exp, bc = s if not man: if s == fzero: if alt: return fhalf else: return mpf_neg(fhalf) if s == finf: return fone return fnan wp = prec + 20 # First term vanishes? if (not sign) and (exp + bc > (math.log(wp,2) + 2)): return mpf_perturb(fone, alt, prec, rnd) # Optimize for integer arguments elif exp >= 0: if alt: if s == fone: return mpf_ln2(prec, rnd) z = mpf_zeta_int(to_int(s), wp, negative_rnd[rnd]) q = mpf_sub(fone, mpf_pow(ftwo, mpf_sub(fone, s, wp), wp), wp) return mpf_mul(z, q, prec, rnd) else: return mpf_zeta_int(to_int(s), prec, rnd) # Negative: use the reflection formula # Borwein only proves the accuracy bound for x >= 1/2. However, based on # tests, the accuracy without reflection is quite good even some distance # to the left of 1/2. XXX: verify this. if sign: # XXX: could use the separate refl. formula for Dirichlet eta if alt: q = mpf_sub(fone, mpf_pow(ftwo, mpf_sub(fone, s, wp), wp), wp) return mpf_mul(mpf_zeta(s, wp), q, prec, rnd) # XXX: -1 should be done exactly y = mpf_sub(fone, s, 10*wp) a = mpf_gamma(y, wp) b = mpf_zeta(y, wp) c = mpf_sin_pi(mpf_shift(s, -1), wp) wp2 = wp + max(0,exp+bc) pi = mpf_pi(wp+wp2) d = mpf_div(mpf_pow(mpf_shift(pi, 1), s, wp2), pi, wp2) return mpf_mul(a,mpf_mul(b,mpf_mul(c,d,wp),wp),prec,rnd) # Near pole r = mpf_sub(fone, s, wp) asign, aman, aexp, abc = mpf_abs(r) pole_dist = -2*(aexp+abc) if pole_dist > wp: if alt: return mpf_ln2(prec, rnd) else: q = mpf_neg(mpf_div(fone, r, wp)) return mpf_add(q, mpf_euler(wp), prec, rnd) else: wp += max(0, pole_dist) t = MPZ_ZERO #wp += 16 - (prec & 15) # Use Borwein's algorithm n = int(wp/2.54 + 5) d = borwein_coefficients(n) t = MPZ_ZERO sf = to_fixed(s, wp) ln2 = ln2_fixed(wp) for k in xrange(n): u = (-sf*log_int_fixed(k+1, wp, ln2)) >> wp #esign, eman, eexp, ebc = mpf_exp(u, wp) #offset = eexp + wp #if offset >= 0: # w = ((d[k] - d[n]) * eman) << offset #else: # w = ((d[k] - d[n]) * eman) >> (-offset) eman = exp_fixed(u, wp, ln2) w = (d[k] - d[n]) * eman if k & 1: t -= w else: t += w t = t // (-d[n]) t = from_man_exp(t, -wp, wp) if alt: return mpf_pos(t, prec, rnd) else: q = mpf_sub(fone, mpf_pow(ftwo, mpf_sub(fone, s, wp), wp), wp) return mpf_div(t, q, prec, rnd) def mpc_zeta(s, prec, rnd=round_fast, alt=0, force=False): re, im = s if im == fzero: return mpf_zeta(re, prec, rnd, alt), fzero # slow for large s if (not force) and mpf_gt(mpc_abs(s, 10), from_int(prec)): raise NotImplementedError wp = prec + 20 # Near pole r = mpc_sub(mpc_one, s, wp) asign, aman, aexp, abc = mpc_abs(r, 10) pole_dist = -2*(aexp+abc) if pole_dist > wp: if alt: q = mpf_ln2(wp) y = mpf_mul(q, mpf_euler(wp), wp) g = mpf_shift(mpf_mul(q, q, wp), -1) g = mpf_sub(y, g) z = mpc_mul_mpf(r, mpf_neg(g), wp) z = mpc_add_mpf(z, q, wp) return mpc_pos(z, prec, rnd) else: q = mpc_neg(mpc_div(mpc_one, r, wp)) q = mpc_add_mpf(q, mpf_euler(wp), wp) return mpc_pos(q, prec, rnd) else: wp += max(0, pole_dist) # Reflection formula. To be rigorous, we should reflect to the left of # re = 1/2 (see comments for mpf_zeta), but this leads to unnecessary # slowdown for interesting values of s if mpf_lt(re, fzero): # XXX: could use the separate refl. formula for Dirichlet eta if alt: q = mpc_sub(mpc_one, mpc_pow(mpc_two, mpc_sub(mpc_one, s, wp), wp), wp) return mpc_mul(mpc_zeta(s, wp), q, prec, rnd) # XXX: -1 should be done exactly y = mpc_sub(mpc_one, s, 10*wp) a = mpc_gamma(y, wp) b = mpc_zeta(y, wp) c = mpc_sin_pi(mpc_shift(s, -1), wp) rsign, rman, rexp, rbc = re isign, iman, iexp, ibc = im mag = max(rexp+rbc, iexp+ibc) wp2 = wp + max(0, mag) pi = mpf_pi(wp+wp2) pi2 = (mpf_shift(pi, 1), fzero) d = mpc_div_mpf(mpc_pow(pi2, s, wp2), pi, wp2) return mpc_mul(a,mpc_mul(b,mpc_mul(c,d,wp),wp),prec,rnd) n = int(wp/2.54 + 5) n += int(0.9*abs(to_int(im))) d = borwein_coefficients(n) ref = to_fixed(re, wp) imf = to_fixed(im, wp) tre = MPZ_ZERO tim = MPZ_ZERO one = MPZ_ONE << wp one_2wp = MPZ_ONE << (2*wp) critical_line = re == fhalf ln2 = ln2_fixed(wp) pi2 = pi_fixed(wp-1) wp2 = wp+wp for k in xrange(n): log = log_int_fixed(k+1, wp, ln2) # A square root is much cheaper than an exp if critical_line: w = one_2wp // isqrt_fast((k+1) << wp2) else: w = exp_fixed((-ref*log) >> wp, wp) if k & 1: w *= (d[n] - d[k]) else: w *= (d[k] - d[n]) wre, wim = cos_sin_fixed((-imf*log)>>wp, wp, pi2) tre += (w * wre) >> wp tim += (w * wim) >> wp tre //= (-d[n]) tim //= (-d[n]) tre = from_man_exp(tre, -wp, wp) tim = from_man_exp(tim, -wp, wp) if alt: return mpc_pos((tre, tim), prec, rnd) else: q = mpc_sub(mpc_one, mpc_pow(mpc_two, r, wp), wp) return mpc_div((tre, tim), q, prec, rnd) def mpf_altzeta(s, prec, rnd=round_fast): return mpf_zeta(s, prec, rnd, 1) def mpc_altzeta(s, prec, rnd=round_fast): return mpc_zeta(s, prec, rnd, 1) # Not optimized currently mpf_zetasum = None def pow_fixed(x, n, wp): if n == 1: return x y = MPZ_ONE << wp while n: if n & 1: y = (y*x) >> wp n -= 1 x = (x*x) >> wp n //= 2 return y # TODO: optimize / cleanup interface / unify with list_primes sieve_cache = [] primes_cache = [] mult_cache = [] def primesieve(n): global sieve_cache, primes_cache, mult_cache if n < len(sieve_cache): sieve = sieve_cache#[:n+1] primes = primes_cache[:primes_cache.index(max(sieve))+1] mult = mult_cache#[:n+1] return sieve, primes, mult sieve = [0] * (n+1) mult = [0] * (n+1) primes = list_primes(n) for p in primes: #sieve[p::p] = p for k in xrange(p,n+1,p): sieve[k] = p for i, p in enumerate(sieve): if i >= 2: m = 1 n = i // p while not n % p: n //= p m += 1 mult[i] = m sieve_cache = sieve primes_cache = primes mult_cache = mult return sieve, primes, mult def zetasum_sieved(critical_line, sre, sim, a, n, wp): if a < 1: raise ValueError("a cannot be less than 1") sieve, primes, mult = primesieve(a+n) basic_powers = {} one = MPZ_ONE << wp one_2wp = MPZ_ONE << (2*wp) wp2 = wp+wp ln2 = ln2_fixed(wp) pi2 = pi_fixed(wp-1) for p in primes: if p*2 > a+n: break log = log_int_fixed(p, wp, ln2) cos, sin = cos_sin_fixed((-sim*log)>>wp, wp, pi2) if critical_line: u = one_2wp // isqrt_fast(p<>wp, wp) pre = (u*cos) >> wp pim = (u*sin) >> wp basic_powers[p] = [(pre, pim)] tre, tim = pre, pim for m in range(1,int(math.log(a+n,p)+0.01)+1): tre, tim = ((pre*tre-pim*tim)>>wp), ((pim*tre+pre*tim)>>wp) basic_powers[p].append((tre,tim)) xre = MPZ_ZERO xim = MPZ_ZERO if a == 1: xre += one aa = max(a,2) for k in xrange(aa, a+n+1): p = sieve[k] if p in basic_powers: m = mult[k] tre, tim = basic_powers[p][m-1] while 1: k //= p**m if k == 1: break p = sieve[k] m = mult[k] pre, pim = basic_powers[p][m-1] tre, tim = ((pre*tre-pim*tim)>>wp), ((pim*tre+pre*tim)>>wp) else: log = log_int_fixed(k, wp, ln2) cos, sin = cos_sin_fixed((-sim*log)>>wp, wp, pi2) if critical_line: u = one_2wp // isqrt_fast(k<>wp, wp) tre = (u*cos) >> wp tim = (u*sin) >> wp xre += tre xim += tim return xre, xim # Set to something large to disable ZETASUM_SIEVE_CUTOFF = 10 def mpc_zetasum(s, a, n, derivatives, reflect, prec): """ Fast version of mp._zetasum, assuming s = complex, a = integer. """ wp = prec + 10 derivatives = list(derivatives) have_derivatives = derivatives != [0] have_one_derivative = len(derivatives) == 1 # parse s sre, sim = s critical_line = (sre == fhalf) sre = to_fixed(sre, wp) sim = to_fixed(sim, wp) if a > 0 and n > ZETASUM_SIEVE_CUTOFF and not have_derivatives \ and not reflect and (n < 4e7 or sys.maxsize > 2**32): re, im = zetasum_sieved(critical_line, sre, sim, a, n, wp) xs = [(from_man_exp(re, -wp, prec, 'n'), from_man_exp(im, -wp, prec, 'n'))] return xs, [] maxd = max(derivatives) if not have_one_derivative: derivatives = range(maxd+1) # x_d = 0, y_d = 0 xre = [MPZ_ZERO for d in derivatives] xim = [MPZ_ZERO for d in derivatives] if reflect: yre = [MPZ_ZERO for d in derivatives] yim = [MPZ_ZERO for d in derivatives] else: yre = yim = [] one = MPZ_ONE << wp one_2wp = MPZ_ONE << (2*wp) ln2 = ln2_fixed(wp) pi2 = pi_fixed(wp-1) wp2 = wp+wp for w in xrange(a, a+n+1): log = log_int_fixed(w, wp, ln2) cos, sin = cos_sin_fixed((-sim*log)>>wp, wp, pi2) if critical_line: u = one_2wp // isqrt_fast(w<>wp, wp) xterm_re = (u * cos) >> wp xterm_im = (u * sin) >> wp if reflect: reciprocal = (one_2wp // (u*w)) yterm_re = (reciprocal * cos) >> wp yterm_im = (reciprocal * sin) >> wp if have_derivatives: if have_one_derivative: log = pow_fixed(log, maxd, wp) xre[0] += (xterm_re * log) >> wp xim[0] += (xterm_im * log) >> wp if reflect: yre[0] += (yterm_re * log) >> wp yim[0] += (yterm_im * log) >> wp else: t = MPZ_ONE << wp for d in derivatives: xre[d] += (xterm_re * t) >> wp xim[d] += (xterm_im * t) >> wp if reflect: yre[d] += (yterm_re * t) >> wp yim[d] += (yterm_im * t) >> wp t = (t * log) >> wp else: xre[0] += xterm_re xim[0] += xterm_im if reflect: yre[0] += yterm_re yim[0] += yterm_im if have_derivatives: if have_one_derivative: if maxd % 2: xre[0] = -xre[0] xim[0] = -xim[0] if reflect: yre[0] = -yre[0] yim[0] = -yim[0] else: xre = [(-1)**d * xre[d] for d in derivatives] xim = [(-1)**d * xim[d] for d in derivatives] if reflect: yre = [(-1)**d * yre[d] for d in derivatives] yim = [(-1)**d * yim[d] for d in derivatives] xs = [(from_man_exp(xa, -wp, prec, 'n'), from_man_exp(xb, -wp, prec, 'n')) for (xa, xb) in zip(xre, xim)] ys = [(from_man_exp(ya, -wp, prec, 'n'), from_man_exp(yb, -wp, prec, 'n')) for (ya, yb) in zip(yre, yim)] return xs, ys #-----------------------------------------------------------------------# # # # The gamma function (NEW IMPLEMENTATION) # # # #-----------------------------------------------------------------------# # Higher means faster, but more precomputation time MAX_GAMMA_TAYLOR_PREC = 5000 # Need to derive higher bounds for Taylor series to go higher assert MAX_GAMMA_TAYLOR_PREC < 15000 # Use Stirling's series if abs(x) > beta*prec # Important: must be large enough for convergence! GAMMA_STIRLING_BETA = 0.2 SMALL_FACTORIAL_CACHE_SIZE = 150 gamma_taylor_cache = {} gamma_stirling_cache = {} small_factorial_cache = [from_int(ifac(n)) for \ n in range(SMALL_FACTORIAL_CACHE_SIZE+1)] def zeta_array(N, prec): """ zeta(n) = A * pi**n / n! + B where A is a rational number (A = Bernoulli number for n even) and B is an infinite sum over powers of exp(2*pi). (B = 0 for n even). TODO: this is currently only used for gamma, but could be very useful elsewhere. """ extra = 30 wp = prec+extra zeta_values = [MPZ_ZERO] * (N+2) pi = pi_fixed(wp) # STEP 1: one = MPZ_ONE << wp zeta_values[0] = -one//2 f_2pi = mpf_shift(mpf_pi(wp),1) exp_2pi_k = exp_2pi = mpf_exp(f_2pi, wp) # Compute exponential series # Store values of 1/(exp(2*pi*k)-1), # exp(2*pi*k)/(exp(2*pi*k)-1)**2, 1/(exp(2*pi*k)-1)**2 # pi*k*exp(2*pi*k)/(exp(2*pi*k)-1)**2 exps3 = [] k = 1 while 1: tp = wp - 9*k if tp < 1: break # 1/(exp(2*pi*k-1) q1 = mpf_div(fone, mpf_sub(exp_2pi_k, fone, tp), tp) # pi*k*exp(2*pi*k)/(exp(2*pi*k)-1)**2 q2 = mpf_mul(exp_2pi_k, mpf_mul(q1,q1,tp), tp) q1 = to_fixed(q1, wp) q2 = to_fixed(q2, wp) q2 = (k * q2 * pi) >> wp exps3.append((q1, q2)) # Multiply for next round exp_2pi_k = mpf_mul(exp_2pi_k, exp_2pi, wp) k += 1 # Exponential sum for n in xrange(3, N+1, 2): s = MPZ_ZERO k = 1 for e1, e2 in exps3: if n%4 == 3: t = e1 // k**n else: U = (n-1)//4 t = (e1 + e2//U) // k**n if not t: break s += t k += 1 zeta_values[n] = -2*s # Even zeta values B = [mpf_abs(mpf_bernoulli(k,wp)) for k in xrange(N+2)] pi_pow = fpi = mpf_pow_int(mpf_shift(mpf_pi(wp), 1), 2, wp) pi_pow = mpf_div(pi_pow, from_int(4), wp) for n in xrange(2,N+2,2): z = mpf_mul(B[n], pi_pow, wp) zeta_values[n] = to_fixed(z, wp) pi_pow = mpf_mul(pi_pow, fpi, wp) pi_pow = mpf_div(pi_pow, from_int((n+1)*(n+2)), wp) # Zeta sum reciprocal_pi = (one << wp) // pi for n in xrange(3, N+1, 4): U = (n-3)//4 s = zeta_values[4*U+4]*(4*U+7)//4 for k in xrange(1, U+1): s -= (zeta_values[4*k] * zeta_values[4*U+4-4*k]) >> wp zeta_values[n] += (2*s*reciprocal_pi) >> wp for n in xrange(5, N+1, 4): U = (n-1)//4 s = zeta_values[4*U+2]*(2*U+1) for k in xrange(1, 2*U+1): s += ((-1)**k*2*k* zeta_values[2*k] * zeta_values[4*U+2-2*k])>>wp zeta_values[n] += ((s*reciprocal_pi)>>wp)//(2*U) return [x>>extra for x in zeta_values] def gamma_taylor_coefficients(inprec): """ Gives the Taylor coefficients of 1/gamma(1+x) as a list of fixed-point numbers. Enough coefficients are returned to ensure that the series converges to the given precision when x is in [0.5, 1.5]. """ # Reuse nearby cache values (small case) if inprec < 400: prec = inprec + (10-(inprec%10)) elif inprec < 1000: prec = inprec + (30-(inprec%30)) else: prec = inprec if prec in gamma_taylor_cache: return gamma_taylor_cache[prec], prec # Experimentally determined bounds if prec < 1000: N = int(prec**0.76 + 2) else: # Valid to at least 15000 bits N = int(prec**0.787 + 2) # Reuse higher precision values for cprec in gamma_taylor_cache: if cprec > prec: coeffs = [x>>(cprec-prec) for x in gamma_taylor_cache[cprec][-N:]] if inprec < 1000: gamma_taylor_cache[prec] = coeffs return coeffs, prec # Cache at a higher precision (large case) if prec > 1000: prec = int(prec * 1.2) wp = prec + 20 A = [0] * N A[0] = MPZ_ZERO A[1] = MPZ_ONE << wp A[2] = euler_fixed(wp) # SLOW, reference implementation #zeta_values = [0,0]+[to_fixed(mpf_zeta_int(k,wp),wp) for k in xrange(2,N)] zeta_values = zeta_array(N, wp) for k in xrange(3, N): a = (-A[2]*A[k-1])>>wp for j in xrange(2,k): a += ((-1)**j * zeta_values[j] * A[k-j]) >> wp a //= (1-k) A[k] = a A = [a>>20 for a in A] A = A[::-1] A = A[:-1] gamma_taylor_cache[prec] = A #return A, prec return gamma_taylor_coefficients(inprec) def gamma_fixed_taylor(xmpf, x, wp, prec, rnd, type): # Determine nearest multiple of N/2 #n = int(x >> (wp-1)) #steps = (n-1)>>1 nearest_int = ((x >> (wp-1)) + MPZ_ONE) >> 1 one = MPZ_ONE << wp coeffs, cwp = gamma_taylor_coefficients(wp) if nearest_int > 0: r = one for i in xrange(nearest_int-1): x -= one r = (r*x) >> wp x -= one p = MPZ_ZERO for c in coeffs: p = c + ((x*p)>>wp) p >>= (cwp-wp) if type == 0: return from_man_exp((r<> wp x += one p = MPZ_ZERO for c in coeffs: p = c + ((x*p)>>wp) p >>= (cwp-wp) if wp - bitcount(abs(x)) > 10: # pass very close to 0, so do floating-point multiply g = mpf_add(xmpf, from_int(-nearest_int)) # exact r = from_man_exp(p*r,-wp-wp) r = mpf_mul(r, g, wp) if type == 0: return mpf_div(fone, r, prec, rnd) if type == 2: return mpf_pos(r, prec, rnd) if type == 3: return mpf_log(mpf_abs(mpf_div(fone, r, wp)), prec, rnd) else: r = from_man_exp(x*p*r,-3*wp) if type == 0: return mpf_div(fone, r, prec, rnd) if type == 2: return mpf_pos(r, prec, rnd) if type == 3: return mpf_neg(mpf_log(mpf_abs(r), prec, rnd)) def stirling_coefficient(n): if n in gamma_stirling_cache: return gamma_stirling_cache[n] p, q = bernfrac(n) q *= MPZ(n*(n-1)) gamma_stirling_cache[n] = p, q, bitcount(abs(p)), bitcount(q) return gamma_stirling_cache[n] def real_stirling_series(x, prec): """ Sums the rational part of Stirling's expansion, log(sqrt(2*pi)) - z + 1/(12*z) - 1/(360*z^3) + ... """ t = (MPZ_ONE<<(prec+prec)) // x # t = 1/x u = (t*t)>>prec # u = 1/x**2 s = ln_sqrt2pi_fixed(prec) - x # Add initial terms of Stirling's series s += t//12; t = (t*u)>>prec s -= t//360; t = (t*u)>>prec s += t//1260; t = (t*u)>>prec s -= t//1680; t = (t*u)>>prec if not t: return s s += t//1188; t = (t*u)>>prec s -= 691*t//360360; t = (t*u)>>prec s += t//156; t = (t*u)>>prec if not t: return s s -= 3617*t//122400; t = (t*u)>>prec s += 43867*t//244188; t = (t*u)>>prec s -= 174611*t//125400; t = (t*u)>>prec if not t: return s k = 22 # From here on, the coefficients are growing, so we # have to keep t at a roughly constant size usize = bitcount(abs(u)) tsize = bitcount(abs(t)) texp = 0 while 1: p, q, pb, qb = stirling_coefficient(k) term_mag = tsize + pb + texp shift = -texp m = pb - term_mag if m > 0 and shift < m: p >>= m shift -= m m = tsize - term_mag if m > 0 and shift < m: w = t >> m shift -= m else: w = t term = (t*p//q) >> shift if not term: break s += term t = (t*u) >> usize texp -= (prec - usize) k += 2 return s def complex_stirling_series(x, y, prec): # t = 1/z _m = (x*x + y*y) >> prec tre = (x << prec) // _m tim = (-y << prec) // _m # u = 1/z**2 ure = (tre*tre - tim*tim) >> prec uim = tim*tre >> (prec-1) # s = log(sqrt(2*pi)) - z sre = ln_sqrt2pi_fixed(prec) - x sim = -y # Add initial terms of Stirling's series sre += tre//12; sim += tim//12; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) sre -= tre//360; sim -= tim//360; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) sre += tre//1260; sim += tim//1260; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) sre -= tre//1680; sim -= tim//1680; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) if abs(tre) + abs(tim) < 5: return sre, sim sre += tre//1188; sim += tim//1188; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) sre -= 691*tre//360360; sim -= 691*tim//360360; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) sre += tre//156; sim += tim//156; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) if abs(tre) + abs(tim) < 5: return sre, sim sre -= 3617*tre//122400; sim -= 3617*tim//122400; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) sre += 43867*tre//244188; sim += 43867*tim//244188; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) sre -= 174611*tre//125400; sim -= 174611*tim//125400; tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec) if abs(tre) + abs(tim) < 5: return sre, sim k = 22 # From here on, the coefficients are growing, so we # have to keep t at a roughly constant size usize = bitcount(max(abs(ure), abs(uim))) tsize = bitcount(max(abs(tre), abs(tim))) texp = 0 while 1: p, q, pb, qb = stirling_coefficient(k) term_mag = tsize + pb + texp shift = -texp m = pb - term_mag if m > 0 and shift < m: p >>= m shift -= m m = tsize - term_mag if m > 0 and shift < m: wre = tre >> m wim = tim >> m shift -= m else: wre = tre wim = tim termre = (tre*p//q) >> shift termim = (tim*p//q) >> shift if abs(termre) + abs(termim) < 5: break sre += termre sim += termim tre, tim = ((tre*ure - tim*uim)>>usize), \ ((tre*uim + tim*ure)>>usize) texp -= (prec - usize) k += 2 return sre, sim def mpf_gamma(x, prec, rnd='d', type=0): """ This function implements multipurpose evaluation of the gamma function, G(x), as well as the following versions of the same: type = 0 -- G(x) [standard gamma function] type = 1 -- G(x+1) = x*G(x+1) = x! [factorial] type = 2 -- 1/G(x) [reciprocal gamma function] type = 3 -- log(|G(x)|) [log-gamma function, real part] """ # Specal values sign, man, exp, bc = x if not man: if x == fzero: if type == 1: return fone if type == 2: return fzero raise ValueError("gamma function pole") if x == finf: if type == 2: return fzero return finf return fnan # First of all, for log gamma, numbers can be well beyond the fixed-point # range, so we must take care of huge numbers before e.g. trying # to convert x to the nearest integer if type == 3: wp = prec+20 if exp+bc > wp and not sign: return mpf_sub(mpf_mul(x, mpf_log(x, wp), wp), x, prec, rnd) # We strongly want to special-case small integers is_integer = exp >= 0 if is_integer: # Poles if sign: if type == 2: return fzero raise ValueError("gamma function pole") # n = x n = man << exp if n < SMALL_FACTORIAL_CACHE_SIZE: if type == 0: return mpf_pos(small_factorial_cache[n-1], prec, rnd) if type == 1: return mpf_pos(small_factorial_cache[n], prec, rnd) if type == 2: return mpf_div(fone, small_factorial_cache[n-1], prec, rnd) if type == 3: return mpf_log(small_factorial_cache[n-1], prec, rnd) else: # floor(abs(x)) n = int(man >> (-exp)) # Estimate size and precision # Estimate log(gamma(|x|),2) as x*log(x,2) mag = exp + bc gamma_size = n*mag if type == 3: wp = prec + 20 else: wp = prec + bitcount(gamma_size) + 20 # Very close to 0, pole if mag < -wp: if type == 0: return mpf_sub(mpf_div(fone,x, wp),mpf_shift(fone,-wp),prec,rnd) if type == 1: return mpf_sub(fone, x, prec, rnd) if type == 2: return mpf_add(x, mpf_shift(fone,mag-wp), prec, rnd) if type == 3: return mpf_neg(mpf_log(mpf_abs(x), prec, rnd)) # From now on, we assume having a gamma function if type == 1: return mpf_gamma(mpf_add(x, fone), prec, rnd, 0) # Special case integers (those not small enough to be caught above, # but still small enough for an exact factorial to be faster # than an approximate algorithm), and half-integers if exp >= -1: if is_integer: if gamma_size < 10*wp: if type == 0: return from_int(ifac(n-1), prec, rnd) if type == 2: return from_rational(MPZ_ONE, ifac(n-1), prec, rnd) if type == 3: return mpf_log(from_int(ifac(n-1)), prec, rnd) # half-integer if n < 100 or gamma_size < 10*wp: if sign: w = sqrtpi_fixed(wp) if n % 2: f = ifac2(2*n+1) else: f = -ifac2(2*n+1) if type == 0: return mpf_shift(from_rational(w, f, prec, rnd), -wp+n+1) if type == 2: return mpf_shift(from_rational(f, w, prec, rnd), wp-n-1) if type == 3: return mpf_log(mpf_shift(from_rational(w, abs(f), prec, rnd), -wp+n+1), prec, rnd) elif n == 0: if type == 0: return mpf_sqrtpi(prec, rnd) if type == 2: return mpf_div(fone, mpf_sqrtpi(wp), prec, rnd) if type == 3: return mpf_log(mpf_sqrtpi(wp), prec, rnd) else: w = sqrtpi_fixed(wp) w = from_man_exp(w * ifac2(2*n-1), -wp-n) if type == 0: return mpf_pos(w, prec, rnd) if type == 2: return mpf_div(fone, w, prec, rnd) if type == 3: return mpf_log(mpf_abs(w), prec, rnd) # Convert to fixed point offset = exp + wp if offset >= 0: absxman = man << offset else: absxman = man >> (-offset) # For log gamma, provide accurate evaluation for x = 1+eps and 2+eps if type == 3 and not sign: one = MPZ_ONE << wp one_dist = abs(absxman-one) two_dist = abs(absxman-2*one) cancellation = (wp - bitcount(min(one_dist, two_dist))) if cancellation > 10: xsub1 = mpf_sub(fone, x) xsub2 = mpf_sub(ftwo, x) xsub1mag = xsub1[2]+xsub1[3] xsub2mag = xsub2[2]+xsub2[3] if xsub1mag < -wp: return mpf_mul(mpf_euler(wp), mpf_sub(fone, x), prec, rnd) if xsub2mag < -wp: return mpf_mul(mpf_sub(fone, mpf_euler(wp)), mpf_sub(x, ftwo), prec, rnd) # Proceed but increase precision wp += max(-xsub1mag, -xsub2mag) offset = exp + wp if offset >= 0: absxman = man << offset else: absxman = man >> (-offset) # Use Taylor series if appropriate n_for_stirling = int(GAMMA_STIRLING_BETA*wp) if n < max(100, n_for_stirling) and wp < MAX_GAMMA_TAYLOR_PREC: if sign: absxman = -absxman return gamma_fixed_taylor(x, absxman, wp, prec, rnd, type) # Use Stirling's series # First ensure that |x| is large enough for rapid convergence xorig = x # Argument reduction r = 0 if n < n_for_stirling: r = one = MPZ_ONE << wp d = n_for_stirling - n for k in xrange(d): r = (r * absxman) >> wp absxman += one x = xabs = from_man_exp(absxman, -wp) if sign: x = mpf_neg(x) else: xabs = mpf_abs(x) # Asymptotic series y = real_stirling_series(absxman, wp) u = to_fixed(mpf_log(xabs, wp), wp) u = ((absxman - (MPZ_ONE<<(wp-1))) * u) >> wp y += u w = from_man_exp(y, -wp) # Compute final value if sign: # Reflection formula A = mpf_mul(mpf_sin_pi(xorig, wp), xorig, wp) B = mpf_neg(mpf_pi(wp)) if type == 0 or type == 2: A = mpf_mul(A, mpf_exp(w, wp)) if r: B = mpf_mul(B, from_man_exp(r, -wp), wp) if type == 0: return mpf_div(B, A, prec, rnd) if type == 2: return mpf_div(A, B, prec, rnd) if type == 3: if r: B = mpf_mul(B, from_man_exp(r, -wp), wp) A = mpf_add(mpf_log(mpf_abs(A), wp), w, wp) return mpf_sub(mpf_log(mpf_abs(B), wp), A, prec, rnd) else: if type == 0: if r: return mpf_div(mpf_exp(w, wp), from_man_exp(r, -wp), prec, rnd) return mpf_exp(w, prec, rnd) if type == 2: if r: return mpf_div(from_man_exp(r, -wp), mpf_exp(w, wp), prec, rnd) return mpf_exp(mpf_neg(w), prec, rnd) if type == 3: if r: return mpf_sub(w, mpf_log(from_man_exp(r,-wp), wp), prec, rnd) return mpf_pos(w, prec, rnd) def mpc_gamma(z, prec, rnd='d', type=0): a, b = z asign, aman, aexp, abc = a bsign, bman, bexp, bbc = b if b == fzero: # Imaginary part on negative half-axis for log-gamma function if type == 3 and asign: re = mpf_gamma(a, prec, rnd, 3) n = (-aman) >> (-aexp) im = mpf_mul_int(mpf_pi(prec+10), n, prec, rnd) return re, im return mpf_gamma(a, prec, rnd, type), fzero # Some kind of complex inf/nan if (not aman and aexp) or (not bman and bexp): return (fnan, fnan) # Initial working precision wp = prec + 20 amag = aexp+abc bmag = bexp+bbc if aman: mag = max(amag, bmag) else: mag = bmag # Close to 0 if mag < -8: if mag < -wp: # 1/gamma(z) = z + euler*z^2 + O(z^3) v = mpc_add(z, mpc_mul_mpf(mpc_mul(z,z,wp),mpf_euler(wp),wp), wp) if type == 0: return mpc_reciprocal(v, prec, rnd) if type == 1: return mpc_div(z, v, prec, rnd) if type == 2: return mpc_pos(v, prec, rnd) if type == 3: return mpc_log(mpc_reciprocal(v, prec), prec, rnd) elif type != 1: wp += (-mag) # Handle huge log-gamma values; must do this before converting to # a fixed-point value. TODO: determine a precise cutoff of validity # depending on amag and bmag if type == 3 and mag > wp and ((not asign) or (bmag >= amag)): return mpc_sub(mpc_mul(z, mpc_log(z, wp), wp), z, prec, rnd) # From now on, we assume having a gamma function if type == 1: return mpc_gamma((mpf_add(a, fone), b), prec, rnd, 0) an = abs(to_int(a)) bn = abs(to_int(b)) absn = max(an, bn) gamma_size = absn*mag if type == 3: pass else: wp += bitcount(gamma_size) # Reflect to the right half-plane. Note that Stirling's expansion # is valid in the left half-plane too, as long as we're not too close # to the real axis, but in order to use this argument reduction # in the negative direction must be implemented. #need_reflection = asign and ((bmag < 0) or (amag-bmag > 4)) need_reflection = asign zorig = z if need_reflection: z = mpc_neg(z) asign, aman, aexp, abc = a = z[0] bsign, bman, bexp, bbc = b = z[1] # Imaginary part very small compared to real one? yfinal = 0 balance_prec = 0 if bmag < -10: # Check z ~= 1 and z ~= 2 for loggamma if type == 3: zsub1 = mpc_sub_mpf(z, fone) if zsub1[0] == fzero: cancel1 = -bmag else: cancel1 = -max(zsub1[0][2]+zsub1[0][3], bmag) if cancel1 > wp: pi = mpf_pi(wp) x = mpc_mul_mpf(zsub1, pi, wp) x = mpc_mul(x, x, wp) x = mpc_div_mpf(x, from_int(12), wp) y = mpc_mul_mpf(zsub1, mpf_neg(mpf_euler(wp)), wp) yfinal = mpc_add(x, y, wp) if not need_reflection: return mpc_pos(yfinal, prec, rnd) elif cancel1 > 0: wp += cancel1 zsub2 = mpc_sub_mpf(z, ftwo) if zsub2[0] == fzero: cancel2 = -bmag else: cancel2 = -max(zsub2[0][2]+zsub2[0][3], bmag) if cancel2 > wp: pi = mpf_pi(wp) t = mpf_sub(mpf_mul(pi, pi), from_int(6)) x = mpc_mul_mpf(mpc_mul(zsub2, zsub2, wp), t, wp) x = mpc_div_mpf(x, from_int(12), wp) y = mpc_mul_mpf(zsub2, mpf_sub(fone, mpf_euler(wp)), wp) yfinal = mpc_add(x, y, wp) if not need_reflection: return mpc_pos(yfinal, prec, rnd) elif cancel2 > 0: wp += cancel2 if bmag < -wp: # Compute directly from the real gamma function. pp = 2*(wp+10) aabs = mpf_abs(a) eps = mpf_shift(fone, amag-wp) x1 = mpf_gamma(aabs, pp, type=type) x2 = mpf_gamma(mpf_add(aabs, eps), pp, type=type) xprime = mpf_div(mpf_sub(x2, x1, pp), eps, pp) y = mpf_mul(b, xprime, prec, rnd) yfinal = (x1, y) # Note: we still need to use the reflection formula for # near-poles, and the correct branch of the log-gamma function if not need_reflection: return mpc_pos(yfinal, prec, rnd) else: balance_prec += (-bmag) wp += balance_prec n_for_stirling = int(GAMMA_STIRLING_BETA*wp) need_reduction = absn < n_for_stirling afix = to_fixed(a, wp) bfix = to_fixed(b, wp) r = 0 if not yfinal: zprered = z # Argument reduction if absn < n_for_stirling: absn = complex(an, bn) d = int((1 + n_for_stirling**2 - bn**2)**0.5 - an) rre = one = MPZ_ONE << wp rim = MPZ_ZERO for k in xrange(d): rre, rim = ((afix*rre-bfix*rim)>>wp), ((afix*rim + bfix*rre)>>wp) afix += one r = from_man_exp(rre, -wp), from_man_exp(rim, -wp) a = from_man_exp(afix, -wp) z = a, b yre, yim = complex_stirling_series(afix, bfix, wp) # (z-1/2)*log(z) + S lre, lim = mpc_log(z, wp) lre = to_fixed(lre, wp) lim = to_fixed(lim, wp) yre = ((lre*afix - lim*bfix)>>wp) - (lre>>1) + yre yim = ((lre*bfix + lim*afix)>>wp) - (lim>>1) + yim y = from_man_exp(yre, -wp), from_man_exp(yim, -wp) if r and type == 3: # If re(z) > 0 and abs(z) <= 4, the branches of loggamma(z) # and log(gamma(z)) coincide. Otherwise, use the zeroth order # Stirling expansion to compute the correct imaginary part. y = mpc_sub(y, mpc_log(r, wp), wp) zfa = to_float(zprered[0]) zfb = to_float(zprered[1]) zfabs = math.hypot(zfa,zfb) #if not (zfa > 0.0 and zfabs <= 4): yfb = to_float(y[1]) u = math.atan2(zfb, zfa) if zfabs <= 0.5: gi = 0.577216*zfb - u else: gi = -zfb - 0.5*u + zfa*u + zfb*math.log(zfabs) n = int(math.floor((gi-yfb)/(2*math.pi)+0.5)) y = (y[0], mpf_add(y[1], mpf_mul_int(mpf_pi(wp), 2*n, wp), wp)) if need_reflection: if type == 0 or type == 2: A = mpc_mul(mpc_sin_pi(zorig, wp), zorig, wp) B = (mpf_neg(mpf_pi(wp)), fzero) if yfinal: if type == 2: A = mpc_div(A, yfinal, wp) else: A = mpc_mul(A, yfinal, wp) else: A = mpc_mul(A, mpc_exp(y, wp), wp) if r: B = mpc_mul(B, r, wp) if type == 0: return mpc_div(B, A, prec, rnd) if type == 2: return mpc_div(A, B, prec, rnd) # Reflection formula for the log-gamma function with correct branch # http://functions.wolfram.com/GammaBetaErf/LogGamma/16/01/01/0006/ # LogGamma[z] == -LogGamma[-z] - Log[-z] + # Sign[Im[z]] Floor[Re[z]] Pi I + Log[Pi] - # Log[Sin[Pi (z - Floor[Re[z]])]] - # Pi I (1 - Abs[Sign[Im[z]]]) Abs[Floor[Re[z]]] if type == 3: if yfinal: s1 = mpc_neg(yfinal) else: s1 = mpc_neg(y) # s -= log(-z) s1 = mpc_sub(s1, mpc_log(mpc_neg(zorig), wp), wp) # floor(re(z)) rezfloor = mpf_floor(zorig[0]) imzsign = mpf_sign(zorig[1]) pi = mpf_pi(wp) t = mpf_mul(pi, rezfloor) t = mpf_mul_int(t, imzsign, wp) s1 = (s1[0], mpf_add(s1[1], t, wp)) s1 = mpc_add_mpf(s1, mpf_log(pi, wp), wp) t = mpc_sin_pi(mpc_sub_mpf(zorig, rezfloor), wp) t = mpc_log(t, wp) s1 = mpc_sub(s1, t, wp) # Note: may actually be unused, because we fall back # to the mpf_ function for real arguments if not imzsign: t = mpf_mul(pi, mpf_floor(rezfloor), wp) s1 = (s1[0], mpf_sub(s1[1], t, wp)) return mpc_pos(s1, prec, rnd) else: if type == 0: if r: return mpc_div(mpc_exp(y, wp), r, prec, rnd) return mpc_exp(y, prec, rnd) if type == 2: if r: return mpc_div(r, mpc_exp(y, wp), prec, rnd) return mpc_exp(mpc_neg(y), prec, rnd) if type == 3: return mpc_pos(y, prec, rnd) def mpf_factorial(x, prec, rnd='d'): return mpf_gamma(x, prec, rnd, 1) def mpc_factorial(x, prec, rnd='d'): return mpc_gamma(x, prec, rnd, 1) def mpf_rgamma(x, prec, rnd='d'): return mpf_gamma(x, prec, rnd, 2) def mpc_rgamma(x, prec, rnd='d'): return mpc_gamma(x, prec, rnd, 2) def mpf_loggamma(x, prec, rnd='d'): sign, man, exp, bc = x if sign: raise ComplexResult return mpf_gamma(x, prec, rnd, 3) def mpc_loggamma(z, prec, rnd='d'): a, b = z asign, aman, aexp, abc = a bsign, bman, bexp, bbc = b if b == fzero and asign: re = mpf_gamma(a, prec, rnd, 3) n = (-aman) >> (-aexp) im = mpf_mul_int(mpf_pi(prec+10), n, prec, rnd) return re, im return mpc_gamma(z, prec, rnd, 3) def mpf_gamma_int(n, prec, rnd=round_fast): if n < SMALL_FACTORIAL_CACHE_SIZE: return mpf_pos(small_factorial_cache[n-1], prec, rnd) return mpf_gamma(from_int(n), prec, rnd)