sympy/ntheory/partitions_.py

   1 from sympy.mpmath.lib import (fzero, fadd, fsub, fmul, fsqrt,
   2     fdiv, fpi, pi_fixed, from_man_exp, from_int, from_rational, cosh_sinh,
   3     fone, fcos, fshift, ftwo, fhalf, bitcount, to_int, to_str)
   4 from sympy.core.numbers import igcd
   5 import math
   6
   7 def A(n, j, prec):
   8     """Compute the inner sum in the HRR formula."""
   9     if j == 1:
  10         return fone
  11     s = fzero
  12     pi = pi_fixed(prec)
  13     for h in xrange(1, j):
  14         if igcd(h,j) != 1:
  15             continue
  16         # & with mask to compute fractional part of fixed-point number
  17         one = 1 << prec
  18         onemask = one - 1
  19         half = one >> 1
  20         g = 0
  21         if j >= 3:
  22             for k in xrange(1, j):
  23                 t = h*k*one//j
  24                 if t > 0: frac = t & onemask
  25                 else:     frac = -((-t) & onemask)
  26                 g += k*(frac - half)
  27         g = ((g - 2*h*n*one)*pi//j) >> prec
  28         s = fadd(s, fcos(from_man_exp(g, -prec), prec), prec)
  29     return s
  30
  31 def D(n, j, prec, sq23pi, sqrt8):
  32     """
  33     Compute the sinh term in the outer sum of the HRR formula.
  34     The constants sqrt(2/3*pi) and sqrt(8) must be precomputed.
  35     """
  36     j = from_int(j)
  37     pi = fpi(prec)
  38     a = fdiv(sq23pi, j, prec)
  39     b = fsub(from_int(n), from_rational(1,24,prec), prec)
  40     c = fsqrt(b, prec)
  41     ch, sh = cosh_sinh(fmul(a,c), prec)
  42     D = fdiv(fsqrt(j,prec), fmul(fmul(sqrt8,b),pi), prec)
  43     E = fsub(fmul(a,ch), fdiv(sh,c,prec), prec)
  44     return fmul(D, E)
  45
  46 def npartitions(n, verbose=False):
  47     """
  48     Calculate the partition function P(n), i.e. the number of ways that
  49     n can be written as a sum of positive integers.
  50
  51     P(n) is computed using the Hardy-Ramanujan-Rademacher formula,
  52     described e.g. at http://mathworld.wolfram.com/PartitionFunctionP.html
  53
  54     The correctness of this implementation has been tested for 10**n
  55     up to n = 8.
  56     """
  57     n = int(n)
  58     if n < 0: return 0
  59     if n <= 5: return [1, 1, 2, 3, 5, 7][n]
  60     # Estimate number of bits in p(n). This formula could be tidied
  61     pbits = int((math.pi*(2*n/3.)**0.5-math.log(4*n))/math.log(10)+1)*\
  62         math.log(10,2)
  63     prec = p = int(pbits*1.1 + 100)
  64     s = fzero
  65     M = max(6, int(0.24*n**0.5+4))
  66     sq23pi = fmul(fsqrt(from_rational(2,3,p), p), fpi(p), p)
  67     sqrt8 = fsqrt(from_int(8), p)
  68     for q in xrange(1, M):
  69         a = A(n,q,p)
  70         d = D(n,q,p, sq23pi, sqrt8)
  71         s = fadd(s, fmul(a, d), prec)
  72         if verbose:
  73             print "step", q, "of", M, to_str(a, 10), to_str(d, 10)
  74         # On average, the terms decrease rapidly in magnitude. Dynamically
  75         # reducing the precision greatly improves performance.
  76         p = bitcount(abs(to_int(d))) + 50
  77     np = to_int(fadd(s, fhalf, prec))
  78     return int(np)
  79
  80 __all__ = ['npartitions']