Coverage for /usr/lib/python3/dist-packages/mpmath/libmp/gammazeta.py: 7%
1367 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""
2-----------------------------------------------------------------------
3This module implements gamma- and zeta-related functions:
5* Bernoulli numbers
6* Factorials
7* The gamma function
8* Polygamma functions
9* Harmonic numbers
10* The Riemann zeta function
11* Constants related to these functions
13-----------------------------------------------------------------------
14"""
16import math
17import sys
19from .backend import xrange
20from .backend import MPZ, MPZ_ZERO, MPZ_ONE, MPZ_THREE, gmpy
22from .libintmath import list_primes, ifac, ifac2, moebius
24from .libmpf import (\
25 round_floor, round_ceiling, round_down, round_up,
26 round_nearest, round_fast,
27 lshift, sqrt_fixed, isqrt_fast,
28 fzero, fone, fnone, fhalf, ftwo, finf, fninf, fnan,
29 from_int, to_int, to_fixed, from_man_exp, from_rational,
30 mpf_pos, mpf_neg, mpf_abs, mpf_add, mpf_sub,
31 mpf_mul, mpf_mul_int, mpf_div, mpf_sqrt, mpf_pow_int,
32 mpf_rdiv_int,
33 mpf_perturb, mpf_le, mpf_lt, mpf_gt, mpf_shift,
34 negative_rnd, reciprocal_rnd,
35 bitcount, to_float, mpf_floor, mpf_sign, ComplexResult
36)
38from .libelefun import (\
39 constant_memo,
40 def_mpf_constant,
41 mpf_pi, pi_fixed, ln2_fixed, log_int_fixed, mpf_ln2,
42 mpf_exp, mpf_log, mpf_pow, mpf_cosh,
43 mpf_cos_sin, mpf_cosh_sinh, mpf_cos_sin_pi, mpf_cos_pi, mpf_sin_pi,
44 ln_sqrt2pi_fixed, mpf_ln_sqrt2pi, sqrtpi_fixed, mpf_sqrtpi,
45 cos_sin_fixed, exp_fixed
46)
48from .libmpc import (\
49 mpc_zero, mpc_one, mpc_half, mpc_two,
50 mpc_abs, mpc_shift, mpc_pos, mpc_neg,
51 mpc_add, mpc_sub, mpc_mul, mpc_div,
52 mpc_add_mpf, mpc_mul_mpf, mpc_div_mpf, mpc_mpf_div,
53 mpc_mul_int, mpc_pow_int,
54 mpc_log, mpc_exp, mpc_pow,
55 mpc_cos_pi, mpc_sin_pi,
56 mpc_reciprocal, mpc_square,
57 mpc_sub_mpf
58)
62# Catalan's constant is computed using Lupas's rapidly convergent series
63# (listed on http://mathworld.wolfram.com/CatalansConstant.html)
64# oo
65# ___ n-1 8n 2 3 2
66# 1 \ (-1) 2 (40n - 24n + 3) [(2n)!] (n!)
67# K = --- ) -----------------------------------------
68# 64 /___ 3 2
69# n (2n-1) [(4n)!]
70# n = 1
72@constant_memo
73def catalan_fixed(prec):
74 prec = prec + 20
75 a = one = MPZ_ONE << prec
76 s, t, n = 0, 1, 1
77 while t:
78 a *= 32 * n**3 * (2*n-1)
79 a //= (3-16*n+16*n**2)**2
80 t = a * (-1)**(n-1) * (40*n**2-24*n+3) // (n**3 * (2*n-1))
81 s += t
82 n += 1
83 return s >> (20 + 6)
85# Khinchin's constant is relatively difficult to compute. Here
86# we use the rational zeta series
88# oo 2*n-1
89# ___ ___
90# \ ` zeta(2*n)-1 \ ` (-1)^(k+1)
91# log(K)*log(2) = ) ------------ ) ----------
92# /___. n /___. k
93# n = 1 k = 1
95# which adds half a digit per term. The essential trick for achieving
96# reasonable efficiency is to recycle both the values of the zeta
97# function (essentially Bernoulli numbers) and the partial terms of
98# the inner sum.
100# An alternative might be to use K = 2*exp[1/log(2) X] where
102# / 1 1 [ pi*x*(1-x^2) ]
103# X = | ------ log [ ------------ ].
104# / 0 x(1+x) [ sin(pi*x) ]
106# and integrate numerically. In practice, this seems to be slightly
107# slower than the zeta series at high precision.
109@constant_memo
110def khinchin_fixed(prec):
111 wp = int(prec + prec**0.5 + 15)
112 s = MPZ_ZERO
113 fac = from_int(4)
114 t = ONE = MPZ_ONE << wp
115 pi = mpf_pi(wp)
116 pipow = twopi2 = mpf_shift(mpf_mul(pi, pi, wp), 2)
117 n = 1
118 while 1:
119 zeta2n = mpf_abs(mpf_bernoulli(2*n, wp))
120 zeta2n = mpf_mul(zeta2n, pipow, wp)
121 zeta2n = mpf_div(zeta2n, fac, wp)
122 zeta2n = to_fixed(zeta2n, wp)
123 term = (((zeta2n - ONE) * t) // n) >> wp
124 if term < 100:
125 break
126 #if not n % 10:
127 # print n, math.log(int(abs(term)))
128 s += term
129 t += ONE//(2*n+1) - ONE//(2*n)
130 n += 1
131 fac = mpf_mul_int(fac, (2*n)*(2*n-1), wp)
132 pipow = mpf_mul(pipow, twopi2, wp)
133 s = (s << wp) // ln2_fixed(wp)
134 K = mpf_exp(from_man_exp(s, -wp), wp)
135 K = to_fixed(K, prec)
136 return K
139# Glaisher's constant is defined as A = exp(1/2 - zeta'(-1)).
140# One way to compute it would be to perform direct numerical
141# differentiation, but computing arbitrary Riemann zeta function
142# values at high precision is expensive. We instead use the formula
144# A = exp((6 (-zeta'(2))/pi^2 + log 2 pi + gamma)/12)
146# and compute zeta'(2) from the series representation
148# oo
149# ___
150# \ log k
151# -zeta'(2) = ) -----
152# /___ 2
153# k
154# k = 2
156# This series converges exceptionally slowly, but can be accelerated
157# using Euler-Maclaurin formula. The important insight is that the
158# E-M integral can be done in closed form and that the high order
159# are given by
161# n / \
162# d | log x | a + b log x
163# --- | ----- | = -----------
164# n | 2 | 2 + n
165# dx \ x / x
167# where a and b are integers given by a simple recurrence. Note
168# that just one logarithm is needed. However, lots of integer
169# logarithms are required for the initial summation.
171# This algorithm could possibly be turned into a faster algorithm
172# for general evaluation of zeta(s) or zeta'(s); this should be
173# looked into.
175@constant_memo
176def glaisher_fixed(prec):
177 wp = prec + 30
178 # Number of direct terms to sum before applying the Euler-Maclaurin
179 # formula to the tail. TODO: choose more intelligently
180 N = int(0.33*prec + 5)
181 ONE = MPZ_ONE << wp
182 # Euler-Maclaurin, step 1: sum log(k)/k**2 for k from 2 to N-1
183 s = MPZ_ZERO
184 for k in range(2, N):
185 #print k, N
186 s += log_int_fixed(k, wp) // k**2
187 logN = log_int_fixed(N, wp)
188 #logN = to_fixed(mpf_log(from_int(N), wp+20), wp)
189 # E-M step 2: integral of log(x)/x**2 from N to inf
190 s += (ONE + logN) // N
191 # E-M step 3: endpoint correction term f(N)/2
192 s += logN // (N**2 * 2)
193 # E-M step 4: the series of derivatives
194 pN = N**3
195 a = 1
196 b = -2
197 j = 3
198 fac = from_int(2)
199 k = 1
200 while 1:
201 # D(2*k-1) * B(2*k) / fac(2*k) [D(n) = nth derivative]
202 D = ((a << wp) + b*logN) // pN
203 D = from_man_exp(D, -wp)
204 B = mpf_bernoulli(2*k, wp)
205 term = mpf_mul(B, D, wp)
206 term = mpf_div(term, fac, wp)
207 term = to_fixed(term, wp)
208 if abs(term) < 100:
209 break
210 #if not k % 10:
211 # print k, math.log(int(abs(term)), 10)
212 s -= term
213 # Advance derivative twice
214 a, b, pN, j = b-a*j, -j*b, pN*N, j+1
215 a, b, pN, j = b-a*j, -j*b, pN*N, j+1
216 k += 1
217 fac = mpf_mul_int(fac, (2*k)*(2*k-1), wp)
218 # A = exp((6*s/pi**2 + log(2*pi) + euler)/12)
219 pi = pi_fixed(wp)
220 s *= 6
221 s = (s << wp) // (pi**2 >> wp)
222 s += euler_fixed(wp)
223 s += to_fixed(mpf_log(from_man_exp(2*pi, -wp), wp), wp)
224 s //= 12
225 A = mpf_exp(from_man_exp(s, -wp), wp)
226 return to_fixed(A, prec)
228# Apery's constant can be computed using the very rapidly convergent
229# series
230# oo
231# ___ 2 10
232# \ n 205 n + 250 n + 77 (n!)
233# zeta(3) = ) (-1) ------------------- ----------
234# /___ 64 5
235# n = 0 ((2n+1)!)
237@constant_memo
238def apery_fixed(prec):
239 prec += 20
240 d = MPZ_ONE << prec
241 term = MPZ(77) << prec
242 n = 1
243 s = MPZ_ZERO
244 while term:
245 s += term
246 d *= (n**10)
247 d //= (((2*n+1)**5) * (2*n)**5)
248 term = (-1)**n * (205*(n**2) + 250*n + 77) * d
249 n += 1
250 return s >> (20 + 6)
252"""
253Euler's constant (gamma) is computed using the Brent-McMillan formula,
254gamma ~= I(n)/J(n) - log(n), where
256 I(n) = sum_{k=0,1,2,...} (n**k / k!)**2 * H(k)
257 J(n) = sum_{k=0,1,2,...} (n**k / k!)**2
258 H(k) = 1 + 1/2 + 1/3 + ... + 1/k
260The error is bounded by O(exp(-4n)). Choosing n to be a power
261of two, 2**p, the logarithm becomes particularly easy to calculate.[1]
263We use the formulation of Algorithm 3.9 in [2] to make the summation
264more efficient.
266Reference:
267[1] Xavier Gourdon & Pascal Sebah, The Euler constant: gamma
268http://numbers.computation.free.fr/Constants/Gamma/gamma.pdf
270[2] Jonathan Borwein & David Bailey, Mathematics by Experiment,
271A K Peters, 2003
272"""
274@constant_memo
275def euler_fixed(prec):
276 extra = 30
277 prec += extra
278 # choose p such that exp(-4*(2**p)) < 2**-n
279 p = int(math.log((prec/4) * math.log(2), 2)) + 1
280 n = 2**p
281 A = U = -p*ln2_fixed(prec)
282 B = V = MPZ_ONE << prec
283 k = 1
284 while 1:
285 B = B*n**2//k**2
286 A = (A*n**2//k + B)//k
287 U += A
288 V += B
289 if max(abs(A), abs(B)) < 100:
290 break
291 k += 1
292 return (U<<(prec-extra))//V
294# Use zeta accelerated formulas for the Mertens and twin
295# prime constants; see
296# http://mathworld.wolfram.com/MertensConstant.html
297# http://mathworld.wolfram.com/TwinPrimesConstant.html
299@constant_memo
300def mertens_fixed(prec):
301 wp = prec + 20
302 m = 2
303 s = mpf_euler(wp)
304 while 1:
305 t = mpf_zeta_int(m, wp)
306 if t == fone:
307 break
308 t = mpf_log(t, wp)
309 t = mpf_mul_int(t, moebius(m), wp)
310 t = mpf_div(t, from_int(m), wp)
311 s = mpf_add(s, t)
312 m += 1
313 return to_fixed(s, prec)
315@constant_memo
316def twinprime_fixed(prec):
317 def I(n):
318 return sum(moebius(d)<<(n//d) for d in xrange(1,n+1) if not n%d)//n
319 wp = 2*prec + 30
320 res = fone
321 primes = [from_rational(1,p,wp) for p in [2,3,5,7]]
322 ppowers = [mpf_mul(p,p,wp) for p in primes]
323 n = 2
324 while 1:
325 a = mpf_zeta_int(n, wp)
326 for i in range(4):
327 a = mpf_mul(a, mpf_sub(fone, ppowers[i]), wp)
328 ppowers[i] = mpf_mul(ppowers[i], primes[i], wp)
329 a = mpf_pow_int(a, -I(n), wp)
330 if mpf_pos(a, prec+10, 'n') == fone:
331 break
332 #from libmpf import to_str
333 #print n, to_str(mpf_sub(fone, a), 6)
334 res = mpf_mul(res, a, wp)
335 n += 1
336 res = mpf_mul(res, from_int(3*15*35), wp)
337 res = mpf_div(res, from_int(4*16*36), wp)
338 return to_fixed(res, prec)
341mpf_euler = def_mpf_constant(euler_fixed)
342mpf_apery = def_mpf_constant(apery_fixed)
343mpf_khinchin = def_mpf_constant(khinchin_fixed)
344mpf_glaisher = def_mpf_constant(glaisher_fixed)
345mpf_catalan = def_mpf_constant(catalan_fixed)
346mpf_mertens = def_mpf_constant(mertens_fixed)
347mpf_twinprime = def_mpf_constant(twinprime_fixed)
350#-----------------------------------------------------------------------#
351# #
352# Bernoulli numbers #
353# #
354#-----------------------------------------------------------------------#
356MAX_BERNOULLI_CACHE = 3000
359r"""
360Small Bernoulli numbers and factorials are used in numerous summations,
361so it is critical for speed that sequential computation is fast and that
362values are cached up to a fairly high threshold.
364On the other hand, we also want to support fast computation of isolated
365large numbers. Currently, no such acceleration is provided for integer
366factorials (though it is for large floating-point factorials, which are
367computed via gamma if the precision is low enough).
369For sequential computation of Bernoulli numbers, we use Ramanujan's formula
371 / n + 3 \
372 B = (A(n) - S(n)) / | |
373 n \ n /
375where A(n) = (n+3)/3 when n = 0 or 2 (mod 6), A(n) = -(n+3)/6
376when n = 4 (mod 6), and
378 [n/6]
379 ___
380 \ / n + 3 \
381 S(n) = ) | | * B
382 /___ \ n - 6*k / n-6*k
383 k = 1
385For isolated large Bernoulli numbers, we use the Riemann zeta function
386to calculate a numerical value for B_n. The von Staudt-Clausen theorem
387can then be used to optionally find the exact value of the
388numerator and denominator.
389"""
391bernoulli_cache = {}
392f3 = from_int(3)
393f6 = from_int(6)
395def bernoulli_size(n):
396 """Accurately estimate the size of B_n (even n > 2 only)"""
397 lgn = math.log(n,2)
398 return int(2.326 + 0.5*lgn + n*(lgn - 4.094))
400BERNOULLI_PREC_CUTOFF = bernoulli_size(MAX_BERNOULLI_CACHE)
402def mpf_bernoulli(n, prec, rnd=None):
403 """Computation of Bernoulli numbers (numerically)"""
404 if n < 2:
405 if n < 0:
406 raise ValueError("Bernoulli numbers only defined for n >= 0")
407 if n == 0:
408 return fone
409 if n == 1:
410 return mpf_neg(fhalf)
411 # For odd n > 1, the Bernoulli numbers are zero
412 if n & 1:
413 return fzero
414 # If precision is extremely high, we can save time by computing
415 # the Bernoulli number at a lower precision that is sufficient to
416 # obtain the exact fraction, round to the exact fraction, and
417 # convert the fraction back to an mpf value at the original precision
418 if prec > BERNOULLI_PREC_CUTOFF and prec > bernoulli_size(n)*1.1 + 1000:
419 p, q = bernfrac(n)
420 return from_rational(p, q, prec, rnd or round_floor)
421 if n > MAX_BERNOULLI_CACHE:
422 return mpf_bernoulli_huge(n, prec, rnd)
423 wp = prec + 30
424 # Reuse nearby precisions
425 wp += 32 - (prec & 31)
426 cached = bernoulli_cache.get(wp)
427 if cached:
428 numbers, state = cached
429 if n in numbers:
430 if not rnd:
431 return numbers[n]
432 return mpf_pos(numbers[n], prec, rnd)
433 m, bin, bin1 = state
434 if n - m > 10:
435 return mpf_bernoulli_huge(n, prec, rnd)
436 else:
437 if n > 10:
438 return mpf_bernoulli_huge(n, prec, rnd)
439 numbers = {0:fone}
440 m, bin, bin1 = state = [2, MPZ(10), MPZ_ONE]
441 bernoulli_cache[wp] = (numbers, state)
442 while m <= n:
443 #print m
444 case = m % 6
445 # Accurately estimate size of B_m so we can use
446 # fixed point math without using too much precision
447 szbm = bernoulli_size(m)
448 s = 0
449 sexp = max(0, szbm) - wp
450 if m < 6:
451 a = MPZ_ZERO
452 else:
453 a = bin1
454 for j in xrange(1, m//6+1):
455 usign, uman, uexp, ubc = u = numbers[m-6*j]
456 if usign:
457 uman = -uman
458 s += lshift(a*uman, uexp-sexp)
459 # Update inner binomial coefficient
460 j6 = 6*j
461 a *= ((m-5-j6)*(m-4-j6)*(m-3-j6)*(m-2-j6)*(m-1-j6)*(m-j6))
462 a //= ((4+j6)*(5+j6)*(6+j6)*(7+j6)*(8+j6)*(9+j6))
463 if case == 0: b = mpf_rdiv_int(m+3, f3, wp)
464 if case == 2: b = mpf_rdiv_int(m+3, f3, wp)
465 if case == 4: b = mpf_rdiv_int(-m-3, f6, wp)
466 s = from_man_exp(s, sexp, wp)
467 b = mpf_div(mpf_sub(b, s, wp), from_int(bin), wp)
468 numbers[m] = b
469 m += 2
470 # Update outer binomial coefficient
471 bin = bin * ((m+2)*(m+3)) // (m*(m-1))
472 if m > 6:
473 bin1 = bin1 * ((2+m)*(3+m)) // ((m-7)*(m-6))
474 state[:] = [m, bin, bin1]
475 return numbers[n]
477def mpf_bernoulli_huge(n, prec, rnd=None):
478 wp = prec + 10
479 piprec = wp + int(math.log(n,2))
480 v = mpf_gamma_int(n+1, wp)
481 v = mpf_mul(v, mpf_zeta_int(n, wp), wp)
482 v = mpf_mul(v, mpf_pow_int(mpf_pi(piprec), -n, wp))
483 v = mpf_shift(v, 1-n)
484 if not n & 3:
485 v = mpf_neg(v)
486 return mpf_pos(v, prec, rnd or round_fast)
488def bernfrac(n):
489 r"""
490 Returns a tuple of integers `(p, q)` such that `p/q = B_n` exactly,
491 where `B_n` denotes the `n`-th Bernoulli number. The fraction is
492 always reduced to lowest terms. Note that for `n > 1` and `n` odd,
493 `B_n = 0`, and `(0, 1)` is returned.
495 **Examples**
497 The first few Bernoulli numbers are exactly::
499 >>> from mpmath import *
500 >>> for n in range(15):
501 ... p, q = bernfrac(n)
502 ... print("%s %s/%s" % (n, p, q))
503 ...
504 0 1/1
505 1 -1/2
506 2 1/6
507 3 0/1
508 4 -1/30
509 5 0/1
510 6 1/42
511 7 0/1
512 8 -1/30
513 9 0/1
514 10 5/66
515 11 0/1
516 12 -691/2730
517 13 0/1
518 14 7/6
520 This function works for arbitrarily large `n`::
522 >>> p, q = bernfrac(10**4)
523 >>> print(q)
524 2338224387510
525 >>> print(len(str(p)))
526 27692
527 >>> mp.dps = 15
528 >>> print(mpf(p) / q)
529 -9.04942396360948e+27677
530 >>> print(bernoulli(10**4))
531 -9.04942396360948e+27677
533 .. note ::
535 :func:`~mpmath.bernoulli` computes a floating-point approximation
536 directly, without computing the exact fraction first.
537 This is much faster for large `n`.
539 **Algorithm**
541 :func:`~mpmath.bernfrac` works by computing the value of `B_n` numerically
542 and then using the von Staudt-Clausen theorem [1] to reconstruct
543 the exact fraction. For large `n`, this is significantly faster than
544 computing `B_1, B_2, \ldots, B_2` recursively with exact arithmetic.
545 The implementation has been tested for `n = 10^m` up to `m = 6`.
547 In practice, :func:`~mpmath.bernfrac` appears to be about three times
548 slower than the specialized program calcbn.exe [2]
550 **References**
552 1. MathWorld, von Staudt-Clausen Theorem:
553 http://mathworld.wolfram.com/vonStaudt-ClausenTheorem.html
555 2. The Bernoulli Number Page:
556 http://www.bernoulli.org/
558 """
559 n = int(n)
560 if n < 3:
561 return [(1, 1), (-1, 2), (1, 6)][n]
562 if n & 1:
563 return (0, 1)
564 q = 1
565 for k in list_primes(n+1):
566 if not (n % (k-1)):
567 q *= k
568 prec = bernoulli_size(n) + int(math.log(q,2)) + 20
569 b = mpf_bernoulli(n, prec)
570 p = mpf_mul(b, from_int(q))
571 pint = to_int(p, round_nearest)
572 return (pint, q)
575#-----------------------------------------------------------------------#
576# #
577# Polygamma functions #
578# #
579#-----------------------------------------------------------------------#
581r"""
582For all polygamma (psi) functions, we use the Euler-Maclaurin summation
583formula. It looks slightly different in the m = 0 and m > 0 cases.
585For m = 0, we have
586 oo
587 ___ B
588 (0) 1 \ 2 k -2 k
589 psi (z) ~ log z + --- - ) ------ z
590 2 z /___ (2 k)!
591 k = 1
593Experiment shows that the minimum term of the asymptotic series
594reaches 2^(-p) when Re(z) > 0.11*p. So we simply use the recurrence
595for psi (equivalent, in fact, to summing to the first few terms
596directly before applying E-M) to obtain z large enough.
598Since, very crudely, log z ~= 1 for Re(z) > 1, we can use
599fixed-point arithmetic (if z is extremely large, log(z) itself
600is a sufficient approximation, so we can stop there already).
602For Re(z) << 0, we could use recurrence, but this is of course
603inefficient for large negative z, so there we use the
604reflection formula instead.
606For m > 0, we have
608 N - 1
609 ___
610 ~~~(m) [ \ 1 ] 1 1
611 psi (z) ~ [ ) -------- ] + ---------- + -------- +
612 [ /___ m+1 ] m+1 m
613 k = 1 (z+k) ] 2 (z+N) m (z+N)
615 oo
616 ___ B
617 \ 2 k (m+1) (m+2) ... (m+2k-1)
618 + ) ------ ------------------------
619 /___ (2 k)! m + 2 k
620 k = 1 (z+N)
622where ~~~ denotes the function rescaled by 1/((-1)^(m+1) m!).
624Here again N is chosen to make z+N large enough for the minimum
625term in the last series to become smaller than eps.
627TODO: the current estimation of N for m > 0 is *very suboptimal*.
629TODO: implement the reflection formula for m > 0, Re(z) << 0.
630It is generally a combination of multiple cotangents. Need to
631figure out a reasonably simple way to generate these formulas
632on the fly.
634TODO: maybe use exact algorithms to compute psi for integral
635and certain rational arguments, as this can be much more
636efficient. (On the other hand, the availability of these
637special values provides a convenient way to test the general
638algorithm.)
639"""
641# Harmonic numbers are just shifted digamma functions
642# We should calculate these exactly when x is an integer
643# and when doing so is faster.
645def mpf_harmonic(x, prec, rnd):
646 if x in (fzero, fnan, finf):
647 return x
648 a = mpf_psi0(mpf_add(fone, x, prec+5), prec)
649 return mpf_add(a, mpf_euler(prec+5, rnd), prec, rnd)
651def mpc_harmonic(z, prec, rnd):
652 if z[1] == fzero:
653 return (mpf_harmonic(z[0], prec, rnd), fzero)
654 a = mpc_psi0(mpc_add_mpf(z, fone, prec+5), prec)
655 return mpc_add_mpf(a, mpf_euler(prec+5, rnd), prec, rnd)
657def mpf_psi0(x, prec, rnd=round_fast):
658 """
659 Computation of the digamma function (psi function of order 0)
660 of a real argument.
661 """
662 sign, man, exp, bc = x
663 wp = prec + 10
664 if not man:
665 if x == finf: return x
666 if x == fninf or x == fnan: return fnan
667 if x == fzero or (exp >= 0 and sign):
668 raise ValueError("polygamma pole")
669 # Near 0 -- fixed-point arithmetic becomes bad
670 if exp+bc < -5:
671 v = mpf_psi0(mpf_add(x, fone, prec, rnd), prec, rnd)
672 return mpf_sub(v, mpf_div(fone, x, wp, rnd), prec, rnd)
673 # Reflection formula
674 if sign and exp+bc > 3:
675 c, s = mpf_cos_sin_pi(x, wp)
676 q = mpf_mul(mpf_div(c, s, wp), mpf_pi(wp), wp)
677 p = mpf_psi0(mpf_sub(fone, x, wp), wp)
678 return mpf_sub(p, q, prec, rnd)
679 # The logarithmic term is accurate enough
680 if (not sign) and bc + exp > wp:
681 return mpf_log(mpf_sub(x, fone, wp), prec, rnd)
682 # Initial recurrence to obtain a large enough x
683 m = to_int(x)
684 n = int(0.11*wp) + 2
685 s = MPZ_ZERO
686 x = to_fixed(x, wp)
687 one = MPZ_ONE << wp
688 if m < n:
689 for k in xrange(m, n):
690 s -= (one << wp) // x
691 x += one
692 x -= one
693 # Logarithmic term
694 s += to_fixed(mpf_log(from_man_exp(x, -wp, wp), wp), wp)
695 # Endpoint term in Euler-Maclaurin expansion
696 s += (one << wp) // (2*x)
697 # Euler-Maclaurin remainder sum
698 x2 = (x*x) >> wp
699 t = one
700 prev = 0
701 k = 1
702 while 1:
703 t = (t*x2) >> wp
704 bsign, bman, bexp, bbc = mpf_bernoulli(2*k, wp)
705 offset = (bexp + 2*wp)
706 if offset >= 0: term = (bman << offset) // (t*(2*k))
707 else: term = (bman >> (-offset)) // (t*(2*k))
708 if k & 1: s -= term
709 else: s += term
710 if k > 2 and term >= prev:
711 break
712 prev = term
713 k += 1
714 return from_man_exp(s, -wp, wp, rnd)
716def mpc_psi0(z, prec, rnd=round_fast):
717 """
718 Computation of the digamma function (psi function of order 0)
719 of a complex argument.
720 """
721 re, im = z
722 # Fall back to the real case
723 if im == fzero:
724 return (mpf_psi0(re, prec, rnd), fzero)
725 wp = prec + 20
726 sign, man, exp, bc = re
727 # Reflection formula
728 if sign and exp+bc > 3:
729 c = mpc_cos_pi(z, wp)
730 s = mpc_sin_pi(z, wp)
731 q = mpc_mul_mpf(mpc_div(c, s, wp), mpf_pi(wp), wp)
732 p = mpc_psi0(mpc_sub(mpc_one, z, wp), wp)
733 return mpc_sub(p, q, prec, rnd)
734 # Just the logarithmic term
735 if (not sign) and bc + exp > wp:
736 return mpc_log(mpc_sub(z, mpc_one, wp), prec, rnd)
737 # Initial recurrence to obtain a large enough z
738 w = to_int(re)
739 n = int(0.11*wp) + 2
740 s = mpc_zero
741 if w < n:
742 for k in xrange(w, n):
743 s = mpc_sub(s, mpc_reciprocal(z, wp), wp)
744 z = mpc_add_mpf(z, fone, wp)
745 z = mpc_sub(z, mpc_one, wp)
746 # Logarithmic and endpoint term
747 s = mpc_add(s, mpc_log(z, wp), wp)
748 s = mpc_add(s, mpc_div(mpc_half, z, wp), wp)
749 # Euler-Maclaurin remainder sum
750 z2 = mpc_square(z, wp)
751 t = mpc_one
752 prev = mpc_zero
753 k = 1
754 eps = mpf_shift(fone, -wp+2)
755 while 1:
756 t = mpc_mul(t, z2, wp)
757 bern = mpf_bernoulli(2*k, wp)
758 term = mpc_mpf_div(bern, mpc_mul_int(t, 2*k, wp), wp)
759 s = mpc_sub(s, term, wp)
760 szterm = mpc_abs(term, 10)
761 if k > 2 and mpf_le(szterm, eps):
762 break
763 prev = term
764 k += 1
765 return s
767# Currently unoptimized
768def mpf_psi(m, x, prec, rnd=round_fast):
769 """
770 Computation of the polygamma function of arbitrary integer order
771 m >= 0, for a real argument x.
772 """
773 if m == 0:
774 return mpf_psi0(x, prec, rnd=round_fast)
775 return mpc_psi(m, (x, fzero), prec, rnd)[0]
777def mpc_psi(m, z, prec, rnd=round_fast):
778 """
779 Computation of the polygamma function of arbitrary integer order
780 m >= 0, for a complex argument z.
781 """
782 if m == 0:
783 return mpc_psi0(z, prec, rnd)
784 re, im = z
785 wp = prec + 20
786 sign, man, exp, bc = re
787 if not im[1]:
788 if im in (finf, fninf, fnan):
789 return (fnan, fnan)
790 if not man:
791 if re == finf and im == fzero:
792 return (fzero, fzero)
793 if re == fnan:
794 return (fnan, fnan)
795 # Recurrence
796 w = to_int(re)
797 n = int(0.4*wp + 4*m)
798 s = mpc_zero
799 if w < n:
800 for k in xrange(w, n):
801 t = mpc_pow_int(z, -m-1, wp)
802 s = mpc_add(s, t, wp)
803 z = mpc_add_mpf(z, fone, wp)
804 zm = mpc_pow_int(z, -m, wp)
805 z2 = mpc_pow_int(z, -2, wp)
806 # 1/m*(z+N)^m
807 integral_term = mpc_div_mpf(zm, from_int(m), wp)
808 s = mpc_add(s, integral_term, wp)
809 # 1/2*(z+N)^(-(m+1))
810 s = mpc_add(s, mpc_mul_mpf(mpc_div(zm, z, wp), fhalf, wp), wp)
811 a = m + 1
812 b = 2
813 k = 1
814 # Important: we want to sum up to the *relative* error,
815 # not the absolute error, because psi^(m)(z) might be tiny
816 magn = mpc_abs(s, 10)
817 magn = magn[2]+magn[3]
818 eps = mpf_shift(fone, magn-wp+2)
819 while 1:
820 zm = mpc_mul(zm, z2, wp)
821 bern = mpf_bernoulli(2*k, wp)
822 scal = mpf_mul_int(bern, a, wp)
823 scal = mpf_div(scal, from_int(b), wp)
824 term = mpc_mul_mpf(zm, scal, wp)
825 s = mpc_add(s, term, wp)
826 szterm = mpc_abs(term, 10)
827 if k > 2 and mpf_le(szterm, eps):
828 break
829 #print k, to_str(szterm, 10), to_str(eps, 10)
830 a *= (m+2*k)*(m+2*k+1)
831 b *= (2*k+1)*(2*k+2)
832 k += 1
833 # Scale and sign factor
834 v = mpc_mul_mpf(s, mpf_gamma(from_int(m+1), wp), prec, rnd)
835 if not (m & 1):
836 v = mpf_neg(v[0]), mpf_neg(v[1])
837 return v
840#-----------------------------------------------------------------------#
841# #
842# Riemann zeta function #
843# #
844#-----------------------------------------------------------------------#
846r"""
847We use zeta(s) = eta(s) / (1 - 2**(1-s)) and Borwein's approximation
849 n-1
850 ___ k
851 -1 \ (-1) (d_k - d_n)
852 eta(s) ~= ---- ) ------------------
853 d_n /___ s
854 k = 0 (k + 1)
855where
856 k
857 ___ i
858 \ (n + i - 1)! 4
859 d_k = n ) ---------------.
860 /___ (n - i)! (2i)!
861 i = 0
863If s = a + b*I, the absolute error for eta(s) is bounded by
865 3 (1 + 2|b|)
866 ------------ * exp(|b| pi/2)
867 n
868 (3+sqrt(8))
870Disregarding the linear term, we have approximately,
872 log(err) ~= log(exp(1.58*|b|)) - log(5.8**n)
873 log(err) ~= 1.58*|b| - log(5.8)*n
874 log(err) ~= 1.58*|b| - 1.76*n
875 log2(err) ~= 2.28*|b| - 2.54*n
877So for p bits, we should choose n > (p + 2.28*|b|) / 2.54.
879References:
880-----------
882Peter Borwein, "An Efficient Algorithm for the Riemann Zeta Function"
883http://www.cecm.sfu.ca/personal/pborwein/PAPERS/P117.ps
885http://en.wikipedia.org/wiki/Dirichlet_eta_function
886"""
888borwein_cache = {}
890def borwein_coefficients(n):
891 if n in borwein_cache:
892 return borwein_cache[n]
893 ds = [MPZ_ZERO] * (n+1)
894 d = MPZ_ONE
895 s = ds[0] = MPZ_ONE
896 for i in range(1, n+1):
897 d = d * 4 * (n+i-1) * (n-i+1)
898 d //= ((2*i) * ((2*i)-1))
899 s += d
900 ds[i] = s
901 borwein_cache[n] = ds
902 return ds
904ZETA_INT_CACHE_MAX_PREC = 1000
905zeta_int_cache = {}
907def mpf_zeta_int(s, prec, rnd=round_fast):
908 """
909 Optimized computation of zeta(s) for an integer s.
910 """
911 wp = prec + 20
912 s = int(s)
913 if s in zeta_int_cache and zeta_int_cache[s][0] >= wp:
914 return mpf_pos(zeta_int_cache[s][1], prec, rnd)
915 if s < 2:
916 if s == 1:
917 raise ValueError("zeta(1) pole")
918 if not s:
919 return mpf_neg(fhalf)
920 return mpf_div(mpf_bernoulli(-s+1, wp), from_int(s-1), prec, rnd)
921 # 2^-s term vanishes?
922 if s >= wp:
923 return mpf_perturb(fone, 0, prec, rnd)
924 # 5^-s term vanishes?
925 elif s >= wp*0.431:
926 t = one = 1 << wp
927 t += 1 << (wp - s)
928 t += one // (MPZ_THREE ** s)
929 t += 1 << max(0, wp - s*2)
930 return from_man_exp(t, -wp, prec, rnd)
931 else:
932 # Fast enough to sum directly?
933 # Even better, we use the Euler product (idea stolen from pari)
934 m = (float(wp)/(s-1) + 1)
935 if m < 30:
936 needed_terms = int(2.0**m + 1)
937 if needed_terms < int(wp/2.54 + 5) / 10:
938 t = fone
939 for k in list_primes(needed_terms):
940 #print k, needed_terms
941 powprec = int(wp - s*math.log(k,2))
942 if powprec < 2:
943 break
944 a = mpf_sub(fone, mpf_pow_int(from_int(k), -s, powprec), wp)
945 t = mpf_mul(t, a, wp)
946 return mpf_div(fone, t, wp)
947 # Use Borwein's algorithm
948 n = int(wp/2.54 + 5)
949 d = borwein_coefficients(n)
950 t = MPZ_ZERO
951 s = MPZ(s)
952 for k in xrange(n):
953 t += (((-1)**k * (d[k] - d[n])) << wp) // (k+1)**s
954 t = (t << wp) // (-d[n])
955 t = (t << wp) // ((1 << wp) - (1 << (wp+1-s)))
956 if (s in zeta_int_cache and zeta_int_cache[s][0] < wp) or (s not in zeta_int_cache):
957 zeta_int_cache[s] = (wp, from_man_exp(t, -wp-wp))
958 return from_man_exp(t, -wp-wp, prec, rnd)
960def mpf_zeta(s, prec, rnd=round_fast, alt=0):
961 sign, man, exp, bc = s
962 if not man:
963 if s == fzero:
964 if alt:
965 return fhalf
966 else:
967 return mpf_neg(fhalf)
968 if s == finf:
969 return fone
970 return fnan
971 wp = prec + 20
972 # First term vanishes?
973 if (not sign) and (exp + bc > (math.log(wp,2) + 2)):
974 return mpf_perturb(fone, alt, prec, rnd)
975 # Optimize for integer arguments
976 elif exp >= 0:
977 if alt:
978 if s == fone:
979 return mpf_ln2(prec, rnd)
980 z = mpf_zeta_int(to_int(s), wp, negative_rnd[rnd])
981 q = mpf_sub(fone, mpf_pow(ftwo, mpf_sub(fone, s, wp), wp), wp)
982 return mpf_mul(z, q, prec, rnd)
983 else:
984 return mpf_zeta_int(to_int(s), prec, rnd)
985 # Negative: use the reflection formula
986 # Borwein only proves the accuracy bound for x >= 1/2. However, based on
987 # tests, the accuracy without reflection is quite good even some distance
988 # to the left of 1/2. XXX: verify this.
989 if sign:
990 # XXX: could use the separate refl. formula for Dirichlet eta
991 if alt:
992 q = mpf_sub(fone, mpf_pow(ftwo, mpf_sub(fone, s, wp), wp), wp)
993 return mpf_mul(mpf_zeta(s, wp), q, prec, rnd)
994 # XXX: -1 should be done exactly
995 y = mpf_sub(fone, s, 10*wp)
996 a = mpf_gamma(y, wp)
997 b = mpf_zeta(y, wp)
998 c = mpf_sin_pi(mpf_shift(s, -1), wp)
999 wp2 = wp + max(0,exp+bc)
1000 pi = mpf_pi(wp+wp2)
1001 d = mpf_div(mpf_pow(mpf_shift(pi, 1), s, wp2), pi, wp2)
1002 return mpf_mul(a,mpf_mul(b,mpf_mul(c,d,wp),wp),prec,rnd)
1004 # Near pole
1005 r = mpf_sub(fone, s, wp)
1006 asign, aman, aexp, abc = mpf_abs(r)
1007 pole_dist = -2*(aexp+abc)
1008 if pole_dist > wp:
1009 if alt:
1010 return mpf_ln2(prec, rnd)
1011 else:
1012 q = mpf_neg(mpf_div(fone, r, wp))
1013 return mpf_add(q, mpf_euler(wp), prec, rnd)
1014 else:
1015 wp += max(0, pole_dist)
1017 t = MPZ_ZERO
1018 #wp += 16 - (prec & 15)
1019 # Use Borwein's algorithm
1020 n = int(wp/2.54 + 5)
1021 d = borwein_coefficients(n)
1022 t = MPZ_ZERO
1023 sf = to_fixed(s, wp)
1024 ln2 = ln2_fixed(wp)
1025 for k in xrange(n):
1026 u = (-sf*log_int_fixed(k+1, wp, ln2)) >> wp
1027 #esign, eman, eexp, ebc = mpf_exp(u, wp)
1028 #offset = eexp + wp
1029 #if offset >= 0:
1030 # w = ((d[k] - d[n]) * eman) << offset
1031 #else:
1032 # w = ((d[k] - d[n]) * eman) >> (-offset)
1033 eman = exp_fixed(u, wp, ln2)
1034 w = (d[k] - d[n]) * eman
1035 if k & 1:
1036 t -= w
1037 else:
1038 t += w
1039 t = t // (-d[n])
1040 t = from_man_exp(t, -wp, wp)
1041 if alt:
1042 return mpf_pos(t, prec, rnd)
1043 else:
1044 q = mpf_sub(fone, mpf_pow(ftwo, mpf_sub(fone, s, wp), wp), wp)
1045 return mpf_div(t, q, prec, rnd)
1047def mpc_zeta(s, prec, rnd=round_fast, alt=0, force=False):
1048 re, im = s
1049 if im == fzero:
1050 return mpf_zeta(re, prec, rnd, alt), fzero
1052 # slow for large s
1053 if (not force) and mpf_gt(mpc_abs(s, 10), from_int(prec)):
1054 raise NotImplementedError
1056 wp = prec + 20
1058 # Near pole
1059 r = mpc_sub(mpc_one, s, wp)
1060 asign, aman, aexp, abc = mpc_abs(r, 10)
1061 pole_dist = -2*(aexp+abc)
1062 if pole_dist > wp:
1063 if alt:
1064 q = mpf_ln2(wp)
1065 y = mpf_mul(q, mpf_euler(wp), wp)
1066 g = mpf_shift(mpf_mul(q, q, wp), -1)
1067 g = mpf_sub(y, g)
1068 z = mpc_mul_mpf(r, mpf_neg(g), wp)
1069 z = mpc_add_mpf(z, q, wp)
1070 return mpc_pos(z, prec, rnd)
1071 else:
1072 q = mpc_neg(mpc_div(mpc_one, r, wp))
1073 q = mpc_add_mpf(q, mpf_euler(wp), wp)
1074 return mpc_pos(q, prec, rnd)
1075 else:
1076 wp += max(0, pole_dist)
1078 # Reflection formula. To be rigorous, we should reflect to the left of
1079 # re = 1/2 (see comments for mpf_zeta), but this leads to unnecessary
1080 # slowdown for interesting values of s
1081 if mpf_lt(re, fzero):
1082 # XXX: could use the separate refl. formula for Dirichlet eta
1083 if alt:
1084 q = mpc_sub(mpc_one, mpc_pow(mpc_two, mpc_sub(mpc_one, s, wp),
1085 wp), wp)
1086 return mpc_mul(mpc_zeta(s, wp), q, prec, rnd)
1087 # XXX: -1 should be done exactly
1088 y = mpc_sub(mpc_one, s, 10*wp)
1089 a = mpc_gamma(y, wp)
1090 b = mpc_zeta(y, wp)
1091 c = mpc_sin_pi(mpc_shift(s, -1), wp)
1092 rsign, rman, rexp, rbc = re
1093 isign, iman, iexp, ibc = im
1094 mag = max(rexp+rbc, iexp+ibc)
1095 wp2 = wp + max(0, mag)
1096 pi = mpf_pi(wp+wp2)
1097 pi2 = (mpf_shift(pi, 1), fzero)
1098 d = mpc_div_mpf(mpc_pow(pi2, s, wp2), pi, wp2)
1099 return mpc_mul(a,mpc_mul(b,mpc_mul(c,d,wp),wp),prec,rnd)
1100 n = int(wp/2.54 + 5)
1101 n += int(0.9*abs(to_int(im)))
1102 d = borwein_coefficients(n)
1103 ref = to_fixed(re, wp)
1104 imf = to_fixed(im, wp)
1105 tre = MPZ_ZERO
1106 tim = MPZ_ZERO
1107 one = MPZ_ONE << wp
1108 one_2wp = MPZ_ONE << (2*wp)
1109 critical_line = re == fhalf
1110 ln2 = ln2_fixed(wp)
1111 pi2 = pi_fixed(wp-1)
1112 wp2 = wp+wp
1113 for k in xrange(n):
1114 log = log_int_fixed(k+1, wp, ln2)
1115 # A square root is much cheaper than an exp
1116 if critical_line:
1117 w = one_2wp // isqrt_fast((k+1) << wp2)
1118 else:
1119 w = exp_fixed((-ref*log) >> wp, wp)
1120 if k & 1:
1121 w *= (d[n] - d[k])
1122 else:
1123 w *= (d[k] - d[n])
1124 wre, wim = cos_sin_fixed((-imf*log)>>wp, wp, pi2)
1125 tre += (w * wre) >> wp
1126 tim += (w * wim) >> wp
1127 tre //= (-d[n])
1128 tim //= (-d[n])
1129 tre = from_man_exp(tre, -wp, wp)
1130 tim = from_man_exp(tim, -wp, wp)
1131 if alt:
1132 return mpc_pos((tre, tim), prec, rnd)
1133 else:
1134 q = mpc_sub(mpc_one, mpc_pow(mpc_two, r, wp), wp)
1135 return mpc_div((tre, tim), q, prec, rnd)
1137def mpf_altzeta(s, prec, rnd=round_fast):
1138 return mpf_zeta(s, prec, rnd, 1)
1140def mpc_altzeta(s, prec, rnd=round_fast):
1141 return mpc_zeta(s, prec, rnd, 1)
1143# Not optimized currently
1144mpf_zetasum = None
1147def pow_fixed(x, n, wp):
1148 if n == 1:
1149 return x
1150 y = MPZ_ONE << wp
1151 while n:
1152 if n & 1:
1153 y = (y*x) >> wp
1154 n -= 1
1155 x = (x*x) >> wp
1156 n //= 2
1157 return y
1159# TODO: optimize / cleanup interface / unify with list_primes
1160sieve_cache = []
1161primes_cache = []
1162mult_cache = []
1164def primesieve(n):
1165 global sieve_cache, primes_cache, mult_cache
1166 if n < len(sieve_cache):
1167 sieve = sieve_cache#[:n+1]
1168 primes = primes_cache[:primes_cache.index(max(sieve))+1]
1169 mult = mult_cache#[:n+1]
1170 return sieve, primes, mult
1171 sieve = [0] * (n+1)
1172 mult = [0] * (n+1)
1173 primes = list_primes(n)
1174 for p in primes:
1175 #sieve[p::p] = p
1176 for k in xrange(p,n+1,p):
1177 sieve[k] = p
1178 for i, p in enumerate(sieve):
1179 if i >= 2:
1180 m = 1
1181 n = i // p
1182 while not n % p:
1183 n //= p
1184 m += 1
1185 mult[i] = m
1186 sieve_cache = sieve
1187 primes_cache = primes
1188 mult_cache = mult
1189 return sieve, primes, mult
1191def zetasum_sieved(critical_line, sre, sim, a, n, wp):
1192 if a < 1:
1193 raise ValueError("a cannot be less than 1")
1194 sieve, primes, mult = primesieve(a+n)
1195 basic_powers = {}
1196 one = MPZ_ONE << wp
1197 one_2wp = MPZ_ONE << (2*wp)
1198 wp2 = wp+wp
1199 ln2 = ln2_fixed(wp)
1200 pi2 = pi_fixed(wp-1)
1201 for p in primes:
1202 if p*2 > a+n:
1203 break
1204 log = log_int_fixed(p, wp, ln2)
1205 cos, sin = cos_sin_fixed((-sim*log)>>wp, wp, pi2)
1206 if critical_line:
1207 u = one_2wp // isqrt_fast(p<<wp2)
1208 else:
1209 u = exp_fixed((-sre*log)>>wp, wp)
1210 pre = (u*cos) >> wp
1211 pim = (u*sin) >> wp
1212 basic_powers[p] = [(pre, pim)]
1213 tre, tim = pre, pim
1214 for m in range(1,int(math.log(a+n,p)+0.01)+1):
1215 tre, tim = ((pre*tre-pim*tim)>>wp), ((pim*tre+pre*tim)>>wp)
1216 basic_powers[p].append((tre,tim))
1217 xre = MPZ_ZERO
1218 xim = MPZ_ZERO
1219 if a == 1:
1220 xre += one
1221 aa = max(a,2)
1222 for k in xrange(aa, a+n+1):
1223 p = sieve[k]
1224 if p in basic_powers:
1225 m = mult[k]
1226 tre, tim = basic_powers[p][m-1]
1227 while 1:
1228 k //= p**m
1229 if k == 1:
1230 break
1231 p = sieve[k]
1232 m = mult[k]
1233 pre, pim = basic_powers[p][m-1]
1234 tre, tim = ((pre*tre-pim*tim)>>wp), ((pim*tre+pre*tim)>>wp)
1235 else:
1236 log = log_int_fixed(k, wp, ln2)
1237 cos, sin = cos_sin_fixed((-sim*log)>>wp, wp, pi2)
1238 if critical_line:
1239 u = one_2wp // isqrt_fast(k<<wp2)
1240 else:
1241 u = exp_fixed((-sre*log)>>wp, wp)
1242 tre = (u*cos) >> wp
1243 tim = (u*sin) >> wp
1244 xre += tre
1245 xim += tim
1246 return xre, xim
1248# Set to something large to disable
1249ZETASUM_SIEVE_CUTOFF = 10
1251def mpc_zetasum(s, a, n, derivatives, reflect, prec):
1252 """
1253 Fast version of mp._zetasum, assuming s = complex, a = integer.
1254 """
1256 wp = prec + 10
1257 derivatives = list(derivatives)
1258 have_derivatives = derivatives != [0]
1259 have_one_derivative = len(derivatives) == 1
1261 # parse s
1262 sre, sim = s
1263 critical_line = (sre == fhalf)
1264 sre = to_fixed(sre, wp)
1265 sim = to_fixed(sim, wp)
1267 if a > 0 and n > ZETASUM_SIEVE_CUTOFF and not have_derivatives \
1268 and not reflect and (n < 4e7 or sys.maxsize > 2**32):
1269 re, im = zetasum_sieved(critical_line, sre, sim, a, n, wp)
1270 xs = [(from_man_exp(re, -wp, prec, 'n'), from_man_exp(im, -wp, prec, 'n'))]
1271 return xs, []
1273 maxd = max(derivatives)
1274 if not have_one_derivative:
1275 derivatives = range(maxd+1)
1277 # x_d = 0, y_d = 0
1278 xre = [MPZ_ZERO for d in derivatives]
1279 xim = [MPZ_ZERO for d in derivatives]
1280 if reflect:
1281 yre = [MPZ_ZERO for d in derivatives]
1282 yim = [MPZ_ZERO for d in derivatives]
1283 else:
1284 yre = yim = []
1286 one = MPZ_ONE << wp
1287 one_2wp = MPZ_ONE << (2*wp)
1289 ln2 = ln2_fixed(wp)
1290 pi2 = pi_fixed(wp-1)
1291 wp2 = wp+wp
1293 for w in xrange(a, a+n+1):
1294 log = log_int_fixed(w, wp, ln2)
1295 cos, sin = cos_sin_fixed((-sim*log)>>wp, wp, pi2)
1296 if critical_line:
1297 u = one_2wp // isqrt_fast(w<<wp2)
1298 else:
1299 u = exp_fixed((-sre*log)>>wp, wp)
1300 xterm_re = (u * cos) >> wp
1301 xterm_im = (u * sin) >> wp
1302 if reflect:
1303 reciprocal = (one_2wp // (u*w))
1304 yterm_re = (reciprocal * cos) >> wp
1305 yterm_im = (reciprocal * sin) >> wp
1307 if have_derivatives:
1308 if have_one_derivative:
1309 log = pow_fixed(log, maxd, wp)
1310 xre[0] += (xterm_re * log) >> wp
1311 xim[0] += (xterm_im * log) >> wp
1312 if reflect:
1313 yre[0] += (yterm_re * log) >> wp
1314 yim[0] += (yterm_im * log) >> wp
1315 else:
1316 t = MPZ_ONE << wp
1317 for d in derivatives:
1318 xre[d] += (xterm_re * t) >> wp
1319 xim[d] += (xterm_im * t) >> wp
1320 if reflect:
1321 yre[d] += (yterm_re * t) >> wp
1322 yim[d] += (yterm_im * t) >> wp
1323 t = (t * log) >> wp
1324 else:
1325 xre[0] += xterm_re
1326 xim[0] += xterm_im
1327 if reflect:
1328 yre[0] += yterm_re
1329 yim[0] += yterm_im
1330 if have_derivatives:
1331 if have_one_derivative:
1332 if maxd % 2:
1333 xre[0] = -xre[0]
1334 xim[0] = -xim[0]
1335 if reflect:
1336 yre[0] = -yre[0]
1337 yim[0] = -yim[0]
1338 else:
1339 xre = [(-1)**d * xre[d] for d in derivatives]
1340 xim = [(-1)**d * xim[d] for d in derivatives]
1341 if reflect:
1342 yre = [(-1)**d * yre[d] for d in derivatives]
1343 yim = [(-1)**d * yim[d] for d in derivatives]
1344 xs = [(from_man_exp(xa, -wp, prec, 'n'), from_man_exp(xb, -wp, prec, 'n'))
1345 for (xa, xb) in zip(xre, xim)]
1346 ys = [(from_man_exp(ya, -wp, prec, 'n'), from_man_exp(yb, -wp, prec, 'n'))
1347 for (ya, yb) in zip(yre, yim)]
1348 return xs, ys
1351#-----------------------------------------------------------------------#
1352# #
1353# The gamma function (NEW IMPLEMENTATION) #
1354# #
1355#-----------------------------------------------------------------------#
1357# Higher means faster, but more precomputation time
1358MAX_GAMMA_TAYLOR_PREC = 5000
1359# Need to derive higher bounds for Taylor series to go higher
1360assert MAX_GAMMA_TAYLOR_PREC < 15000
1362# Use Stirling's series if abs(x) > beta*prec
1363# Important: must be large enough for convergence!
1364GAMMA_STIRLING_BETA = 0.2
1366SMALL_FACTORIAL_CACHE_SIZE = 150
1368gamma_taylor_cache = {}
1369gamma_stirling_cache = {}
1371small_factorial_cache = [from_int(ifac(n)) for \
1372 n in range(SMALL_FACTORIAL_CACHE_SIZE+1)]
1374def zeta_array(N, prec):
1375 """
1376 zeta(n) = A * pi**n / n! + B
1378 where A is a rational number (A = Bernoulli number
1379 for n even) and B is an infinite sum over powers of exp(2*pi).
1380 (B = 0 for n even).
1382 TODO: this is currently only used for gamma, but could
1383 be very useful elsewhere.
1384 """
1385 extra = 30
1386 wp = prec+extra
1387 zeta_values = [MPZ_ZERO] * (N+2)
1388 pi = pi_fixed(wp)
1389 # STEP 1:
1390 one = MPZ_ONE << wp
1391 zeta_values[0] = -one//2
1392 f_2pi = mpf_shift(mpf_pi(wp),1)
1393 exp_2pi_k = exp_2pi = mpf_exp(f_2pi, wp)
1394 # Compute exponential series
1395 # Store values of 1/(exp(2*pi*k)-1),
1396 # exp(2*pi*k)/(exp(2*pi*k)-1)**2, 1/(exp(2*pi*k)-1)**2
1397 # pi*k*exp(2*pi*k)/(exp(2*pi*k)-1)**2
1398 exps3 = []
1399 k = 1
1400 while 1:
1401 tp = wp - 9*k
1402 if tp < 1:
1403 break
1404 # 1/(exp(2*pi*k-1)
1405 q1 = mpf_div(fone, mpf_sub(exp_2pi_k, fone, tp), tp)
1406 # pi*k*exp(2*pi*k)/(exp(2*pi*k)-1)**2
1407 q2 = mpf_mul(exp_2pi_k, mpf_mul(q1,q1,tp), tp)
1408 q1 = to_fixed(q1, wp)
1409 q2 = to_fixed(q2, wp)
1410 q2 = (k * q2 * pi) >> wp
1411 exps3.append((q1, q2))
1412 # Multiply for next round
1413 exp_2pi_k = mpf_mul(exp_2pi_k, exp_2pi, wp)
1414 k += 1
1415 # Exponential sum
1416 for n in xrange(3, N+1, 2):
1417 s = MPZ_ZERO
1418 k = 1
1419 for e1, e2 in exps3:
1420 if n%4 == 3:
1421 t = e1 // k**n
1422 else:
1423 U = (n-1)//4
1424 t = (e1 + e2//U) // k**n
1425 if not t:
1426 break
1427 s += t
1428 k += 1
1429 zeta_values[n] = -2*s
1430 # Even zeta values
1431 B = [mpf_abs(mpf_bernoulli(k,wp)) for k in xrange(N+2)]
1432 pi_pow = fpi = mpf_pow_int(mpf_shift(mpf_pi(wp), 1), 2, wp)
1433 pi_pow = mpf_div(pi_pow, from_int(4), wp)
1434 for n in xrange(2,N+2,2):
1435 z = mpf_mul(B[n], pi_pow, wp)
1436 zeta_values[n] = to_fixed(z, wp)
1437 pi_pow = mpf_mul(pi_pow, fpi, wp)
1438 pi_pow = mpf_div(pi_pow, from_int((n+1)*(n+2)), wp)
1439 # Zeta sum
1440 reciprocal_pi = (one << wp) // pi
1441 for n in xrange(3, N+1, 4):
1442 U = (n-3)//4
1443 s = zeta_values[4*U+4]*(4*U+7)//4
1444 for k in xrange(1, U+1):
1445 s -= (zeta_values[4*k] * zeta_values[4*U+4-4*k]) >> wp
1446 zeta_values[n] += (2*s*reciprocal_pi) >> wp
1447 for n in xrange(5, N+1, 4):
1448 U = (n-1)//4
1449 s = zeta_values[4*U+2]*(2*U+1)
1450 for k in xrange(1, 2*U+1):
1451 s += ((-1)**k*2*k* zeta_values[2*k] * zeta_values[4*U+2-2*k])>>wp
1452 zeta_values[n] += ((s*reciprocal_pi)>>wp)//(2*U)
1453 return [x>>extra for x in zeta_values]
1455def gamma_taylor_coefficients(inprec):
1456 """
1457 Gives the Taylor coefficients of 1/gamma(1+x) as
1458 a list of fixed-point numbers. Enough coefficients are returned
1459 to ensure that the series converges to the given precision
1460 when x is in [0.5, 1.5].
1461 """
1462 # Reuse nearby cache values (small case)
1463 if inprec < 400:
1464 prec = inprec + (10-(inprec%10))
1465 elif inprec < 1000:
1466 prec = inprec + (30-(inprec%30))
1467 else:
1468 prec = inprec
1469 if prec in gamma_taylor_cache:
1470 return gamma_taylor_cache[prec], prec
1472 # Experimentally determined bounds
1473 if prec < 1000:
1474 N = int(prec**0.76 + 2)
1475 else:
1476 # Valid to at least 15000 bits
1477 N = int(prec**0.787 + 2)
1479 # Reuse higher precision values
1480 for cprec in gamma_taylor_cache:
1481 if cprec > prec:
1482 coeffs = [x>>(cprec-prec) for x in gamma_taylor_cache[cprec][-N:]]
1483 if inprec < 1000:
1484 gamma_taylor_cache[prec] = coeffs
1485 return coeffs, prec
1487 # Cache at a higher precision (large case)
1488 if prec > 1000:
1489 prec = int(prec * 1.2)
1491 wp = prec + 20
1492 A = [0] * N
1493 A[0] = MPZ_ZERO
1494 A[1] = MPZ_ONE << wp
1495 A[2] = euler_fixed(wp)
1496 # SLOW, reference implementation
1497 #zeta_values = [0,0]+[to_fixed(mpf_zeta_int(k,wp),wp) for k in xrange(2,N)]
1498 zeta_values = zeta_array(N, wp)
1499 for k in xrange(3, N):
1500 a = (-A[2]*A[k-1])>>wp
1501 for j in xrange(2,k):
1502 a += ((-1)**j * zeta_values[j] * A[k-j]) >> wp
1503 a //= (1-k)
1504 A[k] = a
1505 A = [a>>20 for a in A]
1506 A = A[::-1]
1507 A = A[:-1]
1508 gamma_taylor_cache[prec] = A
1509 #return A, prec
1510 return gamma_taylor_coefficients(inprec)
1512def gamma_fixed_taylor(xmpf, x, wp, prec, rnd, type):
1513 # Determine nearest multiple of N/2
1514 #n = int(x >> (wp-1))
1515 #steps = (n-1)>>1
1516 nearest_int = ((x >> (wp-1)) + MPZ_ONE) >> 1
1517 one = MPZ_ONE << wp
1518 coeffs, cwp = gamma_taylor_coefficients(wp)
1519 if nearest_int > 0:
1520 r = one
1521 for i in xrange(nearest_int-1):
1522 x -= one
1523 r = (r*x) >> wp
1524 x -= one
1525 p = MPZ_ZERO
1526 for c in coeffs:
1527 p = c + ((x*p)>>wp)
1528 p >>= (cwp-wp)
1529 if type == 0:
1530 return from_man_exp((r<<wp)//p, -wp, prec, rnd)
1531 if type == 2:
1532 return mpf_shift(from_rational(p, (r<<wp), prec, rnd), wp)
1533 if type == 3:
1534 return mpf_log(mpf_abs(from_man_exp((r<<wp)//p, -wp)), prec, rnd)
1535 else:
1536 r = one
1537 for i in xrange(-nearest_int):
1538 r = (r*x) >> wp
1539 x += one
1540 p = MPZ_ZERO
1541 for c in coeffs:
1542 p = c + ((x*p)>>wp)
1543 p >>= (cwp-wp)
1544 if wp - bitcount(abs(x)) > 10:
1545 # pass very close to 0, so do floating-point multiply
1546 g = mpf_add(xmpf, from_int(-nearest_int)) # exact
1547 r = from_man_exp(p*r,-wp-wp)
1548 r = mpf_mul(r, g, wp)
1549 if type == 0:
1550 return mpf_div(fone, r, prec, rnd)
1551 if type == 2:
1552 return mpf_pos(r, prec, rnd)
1553 if type == 3:
1554 return mpf_log(mpf_abs(mpf_div(fone, r, wp)), prec, rnd)
1555 else:
1556 r = from_man_exp(x*p*r,-3*wp)
1557 if type == 0: return mpf_div(fone, r, prec, rnd)
1558 if type == 2: return mpf_pos(r, prec, rnd)
1559 if type == 3: return mpf_neg(mpf_log(mpf_abs(r), prec, rnd))
1561def stirling_coefficient(n):
1562 if n in gamma_stirling_cache:
1563 return gamma_stirling_cache[n]
1564 p, q = bernfrac(n)
1565 q *= MPZ(n*(n-1))
1566 gamma_stirling_cache[n] = p, q, bitcount(abs(p)), bitcount(q)
1567 return gamma_stirling_cache[n]
1569def real_stirling_series(x, prec):
1570 """
1571 Sums the rational part of Stirling's expansion,
1573 log(sqrt(2*pi)) - z + 1/(12*z) - 1/(360*z^3) + ...
1575 """
1576 t = (MPZ_ONE<<(prec+prec)) // x # t = 1/x
1577 u = (t*t)>>prec # u = 1/x**2
1578 s = ln_sqrt2pi_fixed(prec) - x
1579 # Add initial terms of Stirling's series
1580 s += t//12; t = (t*u)>>prec
1581 s -= t//360; t = (t*u)>>prec
1582 s += t//1260; t = (t*u)>>prec
1583 s -= t//1680; t = (t*u)>>prec
1584 if not t: return s
1585 s += t//1188; t = (t*u)>>prec
1586 s -= 691*t//360360; t = (t*u)>>prec
1587 s += t//156; t = (t*u)>>prec
1588 if not t: return s
1589 s -= 3617*t//122400; t = (t*u)>>prec
1590 s += 43867*t//244188; t = (t*u)>>prec
1591 s -= 174611*t//125400; t = (t*u)>>prec
1592 if not t: return s
1593 k = 22
1594 # From here on, the coefficients are growing, so we
1595 # have to keep t at a roughly constant size
1596 usize = bitcount(abs(u))
1597 tsize = bitcount(abs(t))
1598 texp = 0
1599 while 1:
1600 p, q, pb, qb = stirling_coefficient(k)
1601 term_mag = tsize + pb + texp
1602 shift = -texp
1603 m = pb - term_mag
1604 if m > 0 and shift < m:
1605 p >>= m
1606 shift -= m
1607 m = tsize - term_mag
1608 if m > 0 and shift < m:
1609 w = t >> m
1610 shift -= m
1611 else:
1612 w = t
1613 term = (t*p//q) >> shift
1614 if not term:
1615 break
1616 s += term
1617 t = (t*u) >> usize
1618 texp -= (prec - usize)
1619 k += 2
1620 return s
1622def complex_stirling_series(x, y, prec):
1623 # t = 1/z
1624 _m = (x*x + y*y) >> prec
1625 tre = (x << prec) // _m
1626 tim = (-y << prec) // _m
1627 # u = 1/z**2
1628 ure = (tre*tre - tim*tim) >> prec
1629 uim = tim*tre >> (prec-1)
1630 # s = log(sqrt(2*pi)) - z
1631 sre = ln_sqrt2pi_fixed(prec) - x
1632 sim = -y
1634 # Add initial terms of Stirling's series
1635 sre += tre//12; sim += tim//12;
1636 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1637 sre -= tre//360; sim -= tim//360;
1638 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1639 sre += tre//1260; sim += tim//1260;
1640 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1641 sre -= tre//1680; sim -= tim//1680;
1642 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1643 if abs(tre) + abs(tim) < 5: return sre, sim
1644 sre += tre//1188; sim += tim//1188;
1645 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1646 sre -= 691*tre//360360; sim -= 691*tim//360360;
1647 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1648 sre += tre//156; sim += tim//156;
1649 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1650 if abs(tre) + abs(tim) < 5: return sre, sim
1651 sre -= 3617*tre//122400; sim -= 3617*tim//122400;
1652 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1653 sre += 43867*tre//244188; sim += 43867*tim//244188;
1654 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1655 sre -= 174611*tre//125400; sim -= 174611*tim//125400;
1656 tre, tim = ((tre*ure-tim*uim)>>prec), ((tre*uim+tim*ure)>>prec)
1657 if abs(tre) + abs(tim) < 5: return sre, sim
1659 k = 22
1660 # From here on, the coefficients are growing, so we
1661 # have to keep t at a roughly constant size
1662 usize = bitcount(max(abs(ure), abs(uim)))
1663 tsize = bitcount(max(abs(tre), abs(tim)))
1664 texp = 0
1665 while 1:
1666 p, q, pb, qb = stirling_coefficient(k)
1667 term_mag = tsize + pb + texp
1668 shift = -texp
1669 m = pb - term_mag
1670 if m > 0 and shift < m:
1671 p >>= m
1672 shift -= m
1673 m = tsize - term_mag
1674 if m > 0 and shift < m:
1675 wre = tre >> m
1676 wim = tim >> m
1677 shift -= m
1678 else:
1679 wre = tre
1680 wim = tim
1681 termre = (tre*p//q) >> shift
1682 termim = (tim*p//q) >> shift
1683 if abs(termre) + abs(termim) < 5:
1684 break
1685 sre += termre
1686 sim += termim
1687 tre, tim = ((tre*ure - tim*uim)>>usize), \
1688 ((tre*uim + tim*ure)>>usize)
1689 texp -= (prec - usize)
1690 k += 2
1691 return sre, sim
1694def mpf_gamma(x, prec, rnd='d', type=0):
1695 """
1696 This function implements multipurpose evaluation of the gamma
1697 function, G(x), as well as the following versions of the same:
1699 type = 0 -- G(x) [standard gamma function]
1700 type = 1 -- G(x+1) = x*G(x+1) = x! [factorial]
1701 type = 2 -- 1/G(x) [reciprocal gamma function]
1702 type = 3 -- log(|G(x)|) [log-gamma function, real part]
1703 """
1705 # Specal values
1706 sign, man, exp, bc = x
1707 if not man:
1708 if x == fzero:
1709 if type == 1: return fone
1710 if type == 2: return fzero
1711 raise ValueError("gamma function pole")
1712 if x == finf:
1713 if type == 2: return fzero
1714 return finf
1715 return fnan
1717 # First of all, for log gamma, numbers can be well beyond the fixed-point
1718 # range, so we must take care of huge numbers before e.g. trying
1719 # to convert x to the nearest integer
1720 if type == 3:
1721 wp = prec+20
1722 if exp+bc > wp and not sign:
1723 return mpf_sub(mpf_mul(x, mpf_log(x, wp), wp), x, prec, rnd)
1725 # We strongly want to special-case small integers
1726 is_integer = exp >= 0
1727 if is_integer:
1728 # Poles
1729 if sign:
1730 if type == 2:
1731 return fzero
1732 raise ValueError("gamma function pole")
1733 # n = x
1734 n = man << exp
1735 if n < SMALL_FACTORIAL_CACHE_SIZE:
1736 if type == 0:
1737 return mpf_pos(small_factorial_cache[n-1], prec, rnd)
1738 if type == 1:
1739 return mpf_pos(small_factorial_cache[n], prec, rnd)
1740 if type == 2:
1741 return mpf_div(fone, small_factorial_cache[n-1], prec, rnd)
1742 if type == 3:
1743 return mpf_log(small_factorial_cache[n-1], prec, rnd)
1744 else:
1745 # floor(abs(x))
1746 n = int(man >> (-exp))
1748 # Estimate size and precision
1749 # Estimate log(gamma(|x|),2) as x*log(x,2)
1750 mag = exp + bc
1751 gamma_size = n*mag
1753 if type == 3:
1754 wp = prec + 20
1755 else:
1756 wp = prec + bitcount(gamma_size) + 20
1758 # Very close to 0, pole
1759 if mag < -wp:
1760 if type == 0:
1761 return mpf_sub(mpf_div(fone,x, wp),mpf_shift(fone,-wp),prec,rnd)
1762 if type == 1: return mpf_sub(fone, x, prec, rnd)
1763 if type == 2: return mpf_add(x, mpf_shift(fone,mag-wp), prec, rnd)
1764 if type == 3: return mpf_neg(mpf_log(mpf_abs(x), prec, rnd))
1766 # From now on, we assume having a gamma function
1767 if type == 1:
1768 return mpf_gamma(mpf_add(x, fone), prec, rnd, 0)
1770 # Special case integers (those not small enough to be caught above,
1771 # but still small enough for an exact factorial to be faster
1772 # than an approximate algorithm), and half-integers
1773 if exp >= -1:
1774 if is_integer:
1775 if gamma_size < 10*wp:
1776 if type == 0:
1777 return from_int(ifac(n-1), prec, rnd)
1778 if type == 2:
1779 return from_rational(MPZ_ONE, ifac(n-1), prec, rnd)
1780 if type == 3:
1781 return mpf_log(from_int(ifac(n-1)), prec, rnd)
1782 # half-integer
1783 if n < 100 or gamma_size < 10*wp:
1784 if sign:
1785 w = sqrtpi_fixed(wp)
1786 if n % 2: f = ifac2(2*n+1)
1787 else: f = -ifac2(2*n+1)
1788 if type == 0:
1789 return mpf_shift(from_rational(w, f, prec, rnd), -wp+n+1)
1790 if type == 2:
1791 return mpf_shift(from_rational(f, w, prec, rnd), wp-n-1)
1792 if type == 3:
1793 return mpf_log(mpf_shift(from_rational(w, abs(f),
1794 prec, rnd), -wp+n+1), prec, rnd)
1795 elif n == 0:
1796 if type == 0: return mpf_sqrtpi(prec, rnd)
1797 if type == 2: return mpf_div(fone, mpf_sqrtpi(wp), prec, rnd)
1798 if type == 3: return mpf_log(mpf_sqrtpi(wp), prec, rnd)
1799 else:
1800 w = sqrtpi_fixed(wp)
1801 w = from_man_exp(w * ifac2(2*n-1), -wp-n)
1802 if type == 0: return mpf_pos(w, prec, rnd)
1803 if type == 2: return mpf_div(fone, w, prec, rnd)
1804 if type == 3: return mpf_log(mpf_abs(w), prec, rnd)
1806 # Convert to fixed point
1807 offset = exp + wp
1808 if offset >= 0: absxman = man << offset
1809 else: absxman = man >> (-offset)
1811 # For log gamma, provide accurate evaluation for x = 1+eps and 2+eps
1812 if type == 3 and not sign:
1813 one = MPZ_ONE << wp
1814 one_dist = abs(absxman-one)
1815 two_dist = abs(absxman-2*one)
1816 cancellation = (wp - bitcount(min(one_dist, two_dist)))
1817 if cancellation > 10:
1818 xsub1 = mpf_sub(fone, x)
1819 xsub2 = mpf_sub(ftwo, x)
1820 xsub1mag = xsub1[2]+xsub1[3]
1821 xsub2mag = xsub2[2]+xsub2[3]
1822 if xsub1mag < -wp:
1823 return mpf_mul(mpf_euler(wp), mpf_sub(fone, x), prec, rnd)
1824 if xsub2mag < -wp:
1825 return mpf_mul(mpf_sub(fone, mpf_euler(wp)),
1826 mpf_sub(x, ftwo), prec, rnd)
1827 # Proceed but increase precision
1828 wp += max(-xsub1mag, -xsub2mag)
1829 offset = exp + wp
1830 if offset >= 0: absxman = man << offset
1831 else: absxman = man >> (-offset)
1833 # Use Taylor series if appropriate
1834 n_for_stirling = int(GAMMA_STIRLING_BETA*wp)
1835 if n < max(100, n_for_stirling) and wp < MAX_GAMMA_TAYLOR_PREC:
1836 if sign:
1837 absxman = -absxman
1838 return gamma_fixed_taylor(x, absxman, wp, prec, rnd, type)
1840 # Use Stirling's series
1841 # First ensure that |x| is large enough for rapid convergence
1842 xorig = x
1844 # Argument reduction
1845 r = 0
1846 if n < n_for_stirling:
1847 r = one = MPZ_ONE << wp
1848 d = n_for_stirling - n
1849 for k in xrange(d):
1850 r = (r * absxman) >> wp
1851 absxman += one
1852 x = xabs = from_man_exp(absxman, -wp)
1853 if sign:
1854 x = mpf_neg(x)
1855 else:
1856 xabs = mpf_abs(x)
1858 # Asymptotic series
1859 y = real_stirling_series(absxman, wp)
1860 u = to_fixed(mpf_log(xabs, wp), wp)
1861 u = ((absxman - (MPZ_ONE<<(wp-1))) * u) >> wp
1862 y += u
1863 w = from_man_exp(y, -wp)
1865 # Compute final value
1866 if sign:
1867 # Reflection formula
1868 A = mpf_mul(mpf_sin_pi(xorig, wp), xorig, wp)
1869 B = mpf_neg(mpf_pi(wp))
1870 if type == 0 or type == 2:
1871 A = mpf_mul(A, mpf_exp(w, wp))
1872 if r:
1873 B = mpf_mul(B, from_man_exp(r, -wp), wp)
1874 if type == 0:
1875 return mpf_div(B, A, prec, rnd)
1876 if type == 2:
1877 return mpf_div(A, B, prec, rnd)
1878 if type == 3:
1879 if r:
1880 B = mpf_mul(B, from_man_exp(r, -wp), wp)
1881 A = mpf_add(mpf_log(mpf_abs(A), wp), w, wp)
1882 return mpf_sub(mpf_log(mpf_abs(B), wp), A, prec, rnd)
1883 else:
1884 if type == 0:
1885 if r:
1886 return mpf_div(mpf_exp(w, wp),
1887 from_man_exp(r, -wp), prec, rnd)
1888 return mpf_exp(w, prec, rnd)
1889 if type == 2:
1890 if r:
1891 return mpf_div(from_man_exp(r, -wp),
1892 mpf_exp(w, wp), prec, rnd)
1893 return mpf_exp(mpf_neg(w), prec, rnd)
1894 if type == 3:
1895 if r:
1896 return mpf_sub(w, mpf_log(from_man_exp(r,-wp), wp), prec, rnd)
1897 return mpf_pos(w, prec, rnd)
1900def mpc_gamma(z, prec, rnd='d', type=0):
1901 a, b = z
1902 asign, aman, aexp, abc = a
1903 bsign, bman, bexp, bbc = b
1905 if b == fzero:
1906 # Imaginary part on negative half-axis for log-gamma function
1907 if type == 3 and asign:
1908 re = mpf_gamma(a, prec, rnd, 3)
1909 n = (-aman) >> (-aexp)
1910 im = mpf_mul_int(mpf_pi(prec+10), n, prec, rnd)
1911 return re, im
1912 return mpf_gamma(a, prec, rnd, type), fzero
1914 # Some kind of complex inf/nan
1915 if (not aman and aexp) or (not bman and bexp):
1916 return (fnan, fnan)
1918 # Initial working precision
1919 wp = prec + 20
1921 amag = aexp+abc
1922 bmag = bexp+bbc
1923 if aman:
1924 mag = max(amag, bmag)
1925 else:
1926 mag = bmag
1928 # Close to 0
1929 if mag < -8:
1930 if mag < -wp:
1931 # 1/gamma(z) = z + euler*z^2 + O(z^3)
1932 v = mpc_add(z, mpc_mul_mpf(mpc_mul(z,z,wp),mpf_euler(wp),wp), wp)
1933 if type == 0: return mpc_reciprocal(v, prec, rnd)
1934 if type == 1: return mpc_div(z, v, prec, rnd)
1935 if type == 2: return mpc_pos(v, prec, rnd)
1936 if type == 3: return mpc_log(mpc_reciprocal(v, prec), prec, rnd)
1937 elif type != 1:
1938 wp += (-mag)
1940 # Handle huge log-gamma values; must do this before converting to
1941 # a fixed-point value. TODO: determine a precise cutoff of validity
1942 # depending on amag and bmag
1943 if type == 3 and mag > wp and ((not asign) or (bmag >= amag)):
1944 return mpc_sub(mpc_mul(z, mpc_log(z, wp), wp), z, prec, rnd)
1946 # From now on, we assume having a gamma function
1947 if type == 1:
1948 return mpc_gamma((mpf_add(a, fone), b), prec, rnd, 0)
1950 an = abs(to_int(a))
1951 bn = abs(to_int(b))
1952 absn = max(an, bn)
1953 gamma_size = absn*mag
1954 if type == 3:
1955 pass
1956 else:
1957 wp += bitcount(gamma_size)
1959 # Reflect to the right half-plane. Note that Stirling's expansion
1960 # is valid in the left half-plane too, as long as we're not too close
1961 # to the real axis, but in order to use this argument reduction
1962 # in the negative direction must be implemented.
1963 #need_reflection = asign and ((bmag < 0) or (amag-bmag > 4))
1964 need_reflection = asign
1965 zorig = z
1966 if need_reflection:
1967 z = mpc_neg(z)
1968 asign, aman, aexp, abc = a = z[0]
1969 bsign, bman, bexp, bbc = b = z[1]
1971 # Imaginary part very small compared to real one?
1972 yfinal = 0
1973 balance_prec = 0
1974 if bmag < -10:
1975 # Check z ~= 1 and z ~= 2 for loggamma
1976 if type == 3:
1977 zsub1 = mpc_sub_mpf(z, fone)
1978 if zsub1[0] == fzero:
1979 cancel1 = -bmag
1980 else:
1981 cancel1 = -max(zsub1[0][2]+zsub1[0][3], bmag)
1982 if cancel1 > wp:
1983 pi = mpf_pi(wp)
1984 x = mpc_mul_mpf(zsub1, pi, wp)
1985 x = mpc_mul(x, x, wp)
1986 x = mpc_div_mpf(x, from_int(12), wp)
1987 y = mpc_mul_mpf(zsub1, mpf_neg(mpf_euler(wp)), wp)
1988 yfinal = mpc_add(x, y, wp)
1989 if not need_reflection:
1990 return mpc_pos(yfinal, prec, rnd)
1991 elif cancel1 > 0:
1992 wp += cancel1
1993 zsub2 = mpc_sub_mpf(z, ftwo)
1994 if zsub2[0] == fzero:
1995 cancel2 = -bmag
1996 else:
1997 cancel2 = -max(zsub2[0][2]+zsub2[0][3], bmag)
1998 if cancel2 > wp:
1999 pi = mpf_pi(wp)
2000 t = mpf_sub(mpf_mul(pi, pi), from_int(6))
2001 x = mpc_mul_mpf(mpc_mul(zsub2, zsub2, wp), t, wp)
2002 x = mpc_div_mpf(x, from_int(12), wp)
2003 y = mpc_mul_mpf(zsub2, mpf_sub(fone, mpf_euler(wp)), wp)
2004 yfinal = mpc_add(x, y, wp)
2005 if not need_reflection:
2006 return mpc_pos(yfinal, prec, rnd)
2007 elif cancel2 > 0:
2008 wp += cancel2
2009 if bmag < -wp:
2010 # Compute directly from the real gamma function.
2011 pp = 2*(wp+10)
2012 aabs = mpf_abs(a)
2013 eps = mpf_shift(fone, amag-wp)
2014 x1 = mpf_gamma(aabs, pp, type=type)
2015 x2 = mpf_gamma(mpf_add(aabs, eps), pp, type=type)
2016 xprime = mpf_div(mpf_sub(x2, x1, pp), eps, pp)
2017 y = mpf_mul(b, xprime, prec, rnd)
2018 yfinal = (x1, y)
2019 # Note: we still need to use the reflection formula for
2020 # near-poles, and the correct branch of the log-gamma function
2021 if not need_reflection:
2022 return mpc_pos(yfinal, prec, rnd)
2023 else:
2024 balance_prec += (-bmag)
2026 wp += balance_prec
2027 n_for_stirling = int(GAMMA_STIRLING_BETA*wp)
2028 need_reduction = absn < n_for_stirling
2030 afix = to_fixed(a, wp)
2031 bfix = to_fixed(b, wp)
2033 r = 0
2034 if not yfinal:
2035 zprered = z
2036 # Argument reduction
2037 if absn < n_for_stirling:
2038 absn = complex(an, bn)
2039 d = int((1 + n_for_stirling**2 - bn**2)**0.5 - an)
2040 rre = one = MPZ_ONE << wp
2041 rim = MPZ_ZERO
2042 for k in xrange(d):
2043 rre, rim = ((afix*rre-bfix*rim)>>wp), ((afix*rim + bfix*rre)>>wp)
2044 afix += one
2045 r = from_man_exp(rre, -wp), from_man_exp(rim, -wp)
2046 a = from_man_exp(afix, -wp)
2047 z = a, b
2049 yre, yim = complex_stirling_series(afix, bfix, wp)
2050 # (z-1/2)*log(z) + S
2051 lre, lim = mpc_log(z, wp)
2052 lre = to_fixed(lre, wp)
2053 lim = to_fixed(lim, wp)
2054 yre = ((lre*afix - lim*bfix)>>wp) - (lre>>1) + yre
2055 yim = ((lre*bfix + lim*afix)>>wp) - (lim>>1) + yim
2056 y = from_man_exp(yre, -wp), from_man_exp(yim, -wp)
2058 if r and type == 3:
2059 # If re(z) > 0 and abs(z) <= 4, the branches of loggamma(z)
2060 # and log(gamma(z)) coincide. Otherwise, use the zeroth order
2061 # Stirling expansion to compute the correct imaginary part.
2062 y = mpc_sub(y, mpc_log(r, wp), wp)
2063 zfa = to_float(zprered[0])
2064 zfb = to_float(zprered[1])
2065 zfabs = math.hypot(zfa,zfb)
2066 #if not (zfa > 0.0 and zfabs <= 4):
2067 yfb = to_float(y[1])
2068 u = math.atan2(zfb, zfa)
2069 if zfabs <= 0.5:
2070 gi = 0.577216*zfb - u
2071 else:
2072 gi = -zfb - 0.5*u + zfa*u + zfb*math.log(zfabs)
2073 n = int(math.floor((gi-yfb)/(2*math.pi)+0.5))
2074 y = (y[0], mpf_add(y[1], mpf_mul_int(mpf_pi(wp), 2*n, wp), wp))
2076 if need_reflection:
2077 if type == 0 or type == 2:
2078 A = mpc_mul(mpc_sin_pi(zorig, wp), zorig, wp)
2079 B = (mpf_neg(mpf_pi(wp)), fzero)
2080 if yfinal:
2081 if type == 2:
2082 A = mpc_div(A, yfinal, wp)
2083 else:
2084 A = mpc_mul(A, yfinal, wp)
2085 else:
2086 A = mpc_mul(A, mpc_exp(y, wp), wp)
2087 if r:
2088 B = mpc_mul(B, r, wp)
2089 if type == 0: return mpc_div(B, A, prec, rnd)
2090 if type == 2: return mpc_div(A, B, prec, rnd)
2092 # Reflection formula for the log-gamma function with correct branch
2093 # http://functions.wolfram.com/GammaBetaErf/LogGamma/16/01/01/0006/
2094 # LogGamma[z] == -LogGamma[-z] - Log[-z] +
2095 # Sign[Im[z]] Floor[Re[z]] Pi I + Log[Pi] -
2096 # Log[Sin[Pi (z - Floor[Re[z]])]] -
2097 # Pi I (1 - Abs[Sign[Im[z]]]) Abs[Floor[Re[z]]]
2098 if type == 3:
2099 if yfinal:
2100 s1 = mpc_neg(yfinal)
2101 else:
2102 s1 = mpc_neg(y)
2103 # s -= log(-z)
2104 s1 = mpc_sub(s1, mpc_log(mpc_neg(zorig), wp), wp)
2105 # floor(re(z))
2106 rezfloor = mpf_floor(zorig[0])
2107 imzsign = mpf_sign(zorig[1])
2108 pi = mpf_pi(wp)
2109 t = mpf_mul(pi, rezfloor)
2110 t = mpf_mul_int(t, imzsign, wp)
2111 s1 = (s1[0], mpf_add(s1[1], t, wp))
2112 s1 = mpc_add_mpf(s1, mpf_log(pi, wp), wp)
2113 t = mpc_sin_pi(mpc_sub_mpf(zorig, rezfloor), wp)
2114 t = mpc_log(t, wp)
2115 s1 = mpc_sub(s1, t, wp)
2116 # Note: may actually be unused, because we fall back
2117 # to the mpf_ function for real arguments
2118 if not imzsign:
2119 t = mpf_mul(pi, mpf_floor(rezfloor), wp)
2120 s1 = (s1[0], mpf_sub(s1[1], t, wp))
2121 return mpc_pos(s1, prec, rnd)
2122 else:
2123 if type == 0:
2124 if r:
2125 return mpc_div(mpc_exp(y, wp), r, prec, rnd)
2126 return mpc_exp(y, prec, rnd)
2127 if type == 2:
2128 if r:
2129 return mpc_div(r, mpc_exp(y, wp), prec, rnd)
2130 return mpc_exp(mpc_neg(y), prec, rnd)
2131 if type == 3:
2132 return mpc_pos(y, prec, rnd)
2134def mpf_factorial(x, prec, rnd='d'):
2135 return mpf_gamma(x, prec, rnd, 1)
2137def mpc_factorial(x, prec, rnd='d'):
2138 return mpc_gamma(x, prec, rnd, 1)
2140def mpf_rgamma(x, prec, rnd='d'):
2141 return mpf_gamma(x, prec, rnd, 2)
2143def mpc_rgamma(x, prec, rnd='d'):
2144 return mpc_gamma(x, prec, rnd, 2)
2146def mpf_loggamma(x, prec, rnd='d'):
2147 sign, man, exp, bc = x
2148 if sign:
2149 raise ComplexResult
2150 return mpf_gamma(x, prec, rnd, 3)
2152def mpc_loggamma(z, prec, rnd='d'):
2153 a, b = z
2154 asign, aman, aexp, abc = a
2155 bsign, bman, bexp, bbc = b
2156 if b == fzero and asign:
2157 re = mpf_gamma(a, prec, rnd, 3)
2158 n = (-aman) >> (-aexp)
2159 im = mpf_mul_int(mpf_pi(prec+10), n, prec, rnd)
2160 return re, im
2161 return mpc_gamma(z, prec, rnd, 3)
2163def mpf_gamma_int(n, prec, rnd=round_fast):
2164 if n < SMALL_FACTORIAL_CACHE_SIZE:
2165 return mpf_pos(small_factorial_cache[n-1], prec, rnd)
2166 return mpf_gamma(from_int(n), prec, rnd)