Coverage for /usr/lib/python3/dist-packages/sympy/ntheory/partitions_.py: 8%

1from mpmath.libmp import (fzero, from_int, from_rational, 

2 fone, fhalf, bitcount, to_int, to_str, mpf_mul, mpf_div, mpf_sub, 

3 mpf_add, mpf_sqrt, mpf_pi, mpf_cosh_sinh, mpf_cos, mpf_sin) 

4from sympy.core.numbers import igcd 

5from .residue_ntheory import (_sqrt_mod_prime_power, 

6 legendre_symbol, jacobi_symbol, is_quad_residue) 

7 

8import math 

9 

10def _pre(): 

11 maxn = 10**5 

12 global _factor 

13 global _totient 

14 _factor = [0]*maxn 

15 _totient = [1]*maxn 

16 lim = int(maxn**0.5) + 5 

17 for i in range(2, lim): 

18 if _factor[i] == 0: 

19 for j in range(i*i, maxn, i): 

20 if _factor[j] == 0: 

21 _factor[j] = i 

22 for i in range(2, maxn): 

23 if _factor[i] == 0: 

24 _factor[i] = i 

25 _totient[i] = i-1 

26 continue 

27 x = _factor[i] 

28 y = i//x 

29 if y % x == 0: 

30 _totient[i] = _totient[y]*x 

31 else: 

32 _totient[i] = _totient[y]*(x - 1) 

33 

34def _a(n, k, prec): 

35 """ Compute the inner sum in HRR formula [1]_ 

36 

37 References 

38 ========== 

39 

40 .. [1] https://msp.org/pjm/1956/6-1/pjm-v6-n1-p18-p.pdf 

41 

42 """ 

43 if k == 1: 

44 return fone 

45 

46 k1 = k 

47 e = 0 

48 p = _factor[k] 

49 while k1 % p == 0: 

50 k1 //= p 

51 e += 1 

52 k2 = k//k1 # k2 = p^e 

53 v = 1 - 24*n 

54 pi = mpf_pi(prec) 

55 

56 if k1 == 1: 

57 # k = p^e 

58 if p == 2: 

59 mod = 8*k 

60 v = mod + v % mod 

61 v = (v*pow(9, k - 1, mod)) % mod 

62 m = _sqrt_mod_prime_power(v, 2, e + 3)[0] 

63 arg = mpf_div(mpf_mul( 

64 from_int(4*m), pi, prec), from_int(mod), prec) 

65 return mpf_mul(mpf_mul( 

66 from_int((-1)**e*jacobi_symbol(m - 1, m)), 

67 mpf_sqrt(from_int(k), prec), prec), 

68 mpf_sin(arg, prec), prec) 

69 if p == 3: 

70 mod = 3*k 

71 v = mod + v % mod 

72 if e > 1: 

73 v = (v*pow(64, k//3 - 1, mod)) % mod 

74 m = _sqrt_mod_prime_power(v, 3, e + 1)[0] 

75 arg = mpf_div(mpf_mul(from_int(4*m), pi, prec), 

76 from_int(mod), prec) 

77 return mpf_mul(mpf_mul( 

78 from_int(2*(-1)**(e + 1)*legendre_symbol(m, 3)), 

79 mpf_sqrt(from_int(k//3), prec), prec), 

80 mpf_sin(arg, prec), prec) 

81 v = k + v % k 

82 if v % p == 0: 

83 if e == 1: 

84 return mpf_mul( 

85 from_int(jacobi_symbol(3, k)), 

86 mpf_sqrt(from_int(k), prec), prec) 

87 return fzero 

88 if not is_quad_residue(v, p): 

89 return fzero 

90 _phi = p**(e - 1)*(p - 1) 

91 v = (v*pow(576, _phi - 1, k)) 

92 m = _sqrt_mod_prime_power(v, p, e)[0] 

93 arg = mpf_div( 

94 mpf_mul(from_int(4*m), pi, prec), 

95 from_int(k), prec) 

96 return mpf_mul(mpf_mul( 

97 from_int(2*jacobi_symbol(3, k)), 

98 mpf_sqrt(from_int(k), prec), prec), 

99 mpf_cos(arg, prec), prec) 

100 

101 if p != 2 or e >= 3: 

102 d1, d2 = igcd(k1, 24), igcd(k2, 24) 

103 e = 24//(d1*d2) 

104 n1 = ((d2*e*n + (k2**2 - 1)//d1)* 

105 pow(e*k2*k2*d2, _totient[k1] - 1, k1)) % k1 

106 n2 = ((d1*e*n + (k1**2 - 1)//d2)* 

107 pow(e*k1*k1*d1, _totient[k2] - 1, k2)) % k2 

108 return mpf_mul(_a(n1, k1, prec), _a(n2, k2, prec), prec) 

109 if e == 2: 

110 n1 = ((8*n + 5)*pow(128, _totient[k1] - 1, k1)) % k1 

111 n2 = (4 + ((n - 2 - (k1**2 - 1)//8)*(k1**2)) % 4) % 4 

112 return mpf_mul(mpf_mul( 

113 from_int(-1), 

114 _a(n1, k1, prec), prec), 

115 _a(n2, k2, prec)) 

116 n1 = ((8*n + 1)*pow(32, _totient[k1] - 1, k1)) % k1 

117 n2 = (2 + (n - (k1**2 - 1)//8) % 2) % 2 

118 return mpf_mul(_a(n1, k1, prec), _a(n2, k2, prec), prec) 

119 

120def _d(n, j, prec, sq23pi, sqrt8): 

121 """ 

122 Compute the sinh term in the outer sum of the HRR formula. 

123 The constants sqrt(2/3*pi) and sqrt(8) must be precomputed. 

124 """ 

125 j = from_int(j) 

126 pi = mpf_pi(prec) 

127 a = mpf_div(sq23pi, j, prec) 

128 b = mpf_sub(from_int(n), from_rational(1, 24, prec), prec) 

129 c = mpf_sqrt(b, prec) 

130 ch, sh = mpf_cosh_sinh(mpf_mul(a, c), prec) 

131 D = mpf_div( 

132 mpf_sqrt(j, prec), 

133 mpf_mul(mpf_mul(sqrt8, b), pi), prec) 

134 E = mpf_sub(mpf_mul(a, ch), mpf_div(sh, c, prec), prec) 

135 return mpf_mul(D, E) 

136 

137 

138def npartitions(n, verbose=False): 

139 """ 

140 Calculate the partition function P(n), i.e. the number of ways that 

141 n can be written as a sum of positive integers. 

142 

143 P(n) is computed using the Hardy-Ramanujan-Rademacher formula [1]_. 

144 

145 

146 The correctness of this implementation has been tested through $10^{10}$. 

147 

148 Examples 

149 ======== 

150 

151 >>> from sympy.ntheory import npartitions 

152 >>> npartitions(25) 

153 1958 

154 

155 References 

156 ========== 

157 

158 .. [1] https://mathworld.wolfram.com/PartitionFunctionP.html 

159 

160 """ 

161 n = int(n) 

162 if n < 0: 

163 return 0 

164 if n <= 5: 

165 return [1, 1, 2, 3, 5, 7][n] 

166 if '_factor' not in globals(): 

167 _pre() 

168 # Estimate number of bits in p(n). This formula could be tidied 

169 pbits = int(( 

170 math.pi*(2*n/3.)**0.5 - 

171 math.log(4*n))/math.log(10) + 1) * \ 

172 math.log(10, 2) 

173 prec = p = int(pbits*1.1 + 100) 

174 s = fzero 

175 M = max(6, int(0.24*n**0.5 + 4)) 

176 if M > 10**5: 

177 raise ValueError("Input too big") # Corresponds to n > 1.7e11 

178 sq23pi = mpf_mul(mpf_sqrt(from_rational(2, 3, p), p), mpf_pi(p), p) 

179 sqrt8 = mpf_sqrt(from_int(8), p) 

180 for q in range(1, M): 

181 a = _a(n, q, p) 

182 d = _d(n, q, p, sq23pi, sqrt8) 

183 s = mpf_add(s, mpf_mul(a, d), prec) 

184 if verbose: 

185 print("step", q, "of", M, to_str(a, 10), to_str(d, 10)) 

186 # On average, the terms decrease rapidly in magnitude. 

187 # Dynamically reducing the precision greatly improves 

188 # performance. 

189 p = bitcount(abs(to_int(d))) + 50 

190 return int(to_int(mpf_add(s, fhalf, prec))) 

191 

192__all__ = ['npartitions']