Coverage for /usr/lib/python3/dist-packages/mpmath/math2.py: 21%

1"""

2This module complements the math and cmath builtin modules by providing

3fast machine precision versions of some additional functions (gamma, ...)

4and wrapping math/cmath functions so that they can be called with either

5real or complex arguments.

6"""

8import operator

9import math

10import cmath

12# Irrational (?) constants

13pi = 3.1415926535897932385

14e = 2.7182818284590452354

15sqrt2 = 1.4142135623730950488

16sqrt5 = 2.2360679774997896964

17phi = 1.6180339887498948482

18ln2 = 0.69314718055994530942

19ln10 = 2.302585092994045684

20euler = 0.57721566490153286061

21catalan = 0.91596559417721901505

22khinchin = 2.6854520010653064453

23apery = 1.2020569031595942854

25logpi = 1.1447298858494001741

27def _mathfun_real(f_real, f_complex):

28 def f(x, **kwargs):

29 if type(x) is float:

30 return f_real(x)

31 if type(x) is complex:

32 return f_complex(x)

33 try:

34 x = float(x)

35 return f_real(x)

36 except (TypeError, ValueError):

37 x = complex(x)

38 return f_complex(x)

39 f.__name__ = f_real.__name__

40 return f

42def _mathfun(f_real, f_complex):

43 def f(x, **kwargs):

44 if type(x) is complex:

45 return f_complex(x)

46 try:

47 return f_real(float(x))

48 except (TypeError, ValueError):

49 return f_complex(complex(x))

50 f.__name__ = f_real.__name__

51 return f

53def _mathfun_n(f_real, f_complex):

54 def f(*args, **kwargs):

55 try:

56 return f_real(*(float(x) for x in args))

57 except (TypeError, ValueError):

58 return f_complex(*(complex(x) for x in args))

59 f.__name__ = f_real.__name__

60 return f

62# Workaround for non-raising log and sqrt in Python 2.5 and 2.4

63# on Unix system

64try:

65 math.log(-2.0)

66 def math_log(x):

67 if x <= 0.0:

68 raise ValueError("math domain error")

69 return math.log(x)

70 def math_sqrt(x):

71 if x < 0.0:

72 raise ValueError("math domain error")

73 return math.sqrt(x)

74except (ValueError, TypeError):

75 math_log = math.log

76 math_sqrt = math.sqrt

78pow = _mathfun_n(operator.pow, lambda x, y: complex(x)**y)

79log = _mathfun_n(math_log, cmath.log)

80sqrt = _mathfun(math_sqrt, cmath.sqrt)

81exp = _mathfun_real(math.exp, cmath.exp)

83cos = _mathfun_real(math.cos, cmath.cos)

84sin = _mathfun_real(math.sin, cmath.sin)

85tan = _mathfun_real(math.tan, cmath.tan)

87acos = _mathfun(math.acos, cmath.acos)

88asin = _mathfun(math.asin, cmath.asin)

89atan = _mathfun_real(math.atan, cmath.atan)

91cosh = _mathfun_real(math.cosh, cmath.cosh)

92sinh = _mathfun_real(math.sinh, cmath.sinh)

93tanh = _mathfun_real(math.tanh, cmath.tanh)

95floor = _mathfun_real(math.floor,

96 lambda z: complex(math.floor(z.real), math.floor(z.imag)))

97ceil = _mathfun_real(math.ceil,

98 lambda z: complex(math.ceil(z.real), math.ceil(z.imag)))

100

101cos_sin = _mathfun_real(lambda x: (math.cos(x), math.sin(x)),

102 lambda z: (cmath.cos(z), cmath.sin(z)))

103

104cbrt = _mathfun(lambda x: x**(1./3), lambda z: z**(1./3))

105

106def nthroot(x, n):

107 r = 1./n

108 try:

109 return float(x) ** r

110 except (ValueError, TypeError):

111 return complex(x) ** r

112

113def _sinpi_real(x):

114 if x < 0:

115 return -_sinpi_real(-x)

116 n, r = divmod(x, 0.5)

117 r *= pi

118 n %= 4

119 if n == 0: return math.sin(r)

120 if n == 1: return math.cos(r)

121 if n == 2: return -math.sin(r)

122 if n == 3: return -math.cos(r)

123

124def _cospi_real(x):

125 if x < 0:

126 x = -x

127 n, r = divmod(x, 0.5)

128 r *= pi

129 n %= 4

130 if n == 0: return math.cos(r)

131 if n == 1: return -math.sin(r)

132 if n == 2: return -math.cos(r)

133 if n == 3: return math.sin(r)

134

135def _sinpi_complex(z):

136 if z.real < 0:

137 return -_sinpi_complex(-z)

138 n, r = divmod(z.real, 0.5)

139 z = pi*complex(r, z.imag)

140 n %= 4

141 if n == 0: return cmath.sin(z)

142 if n == 1: return cmath.cos(z)

143 if n == 2: return -cmath.sin(z)

144 if n == 3: return -cmath.cos(z)

145

146def _cospi_complex(z):

147 if z.real < 0:

148 z = -z

149 n, r = divmod(z.real, 0.5)

150 z = pi*complex(r, z.imag)

151 n %= 4

152 if n == 0: return cmath.cos(z)

153 if n == 1: return -cmath.sin(z)

154 if n == 2: return -cmath.cos(z)

155 if n == 3: return cmath.sin(z)

156

157cospi = _mathfun_real(_cospi_real, _cospi_complex)

158sinpi = _mathfun_real(_sinpi_real, _sinpi_complex)

159

160def tanpi(x):

161 try:

162 return sinpi(x) / cospi(x)

163 except OverflowError:

164 if complex(x).imag > 10:

165 return 1j

166 if complex(x).imag < 10:

167 return -1j

168 raise

169

170def cotpi(x):

171 try:

172 return cospi(x) / sinpi(x)

173 except OverflowError:

174 if complex(x).imag > 10:

175 return -1j

176 if complex(x).imag < 10:

177 return 1j

178 raise

179

180INF = 1e300*1e300

181NINF = -INF

182NAN = INF-INF

183EPS = 2.2204460492503131e-16

184

185_exact_gamma = (INF, 1.0, 1.0, 2.0, 6.0, 24.0, 120.0, 720.0, 5040.0, 40320.0,

186 362880.0, 3628800.0, 39916800.0, 479001600.0, 6227020800.0, 87178291200.0,

187 1307674368000.0, 20922789888000.0, 355687428096000.0, 6402373705728000.0,

188 121645100408832000.0, 2432902008176640000.0)

189

190_max_exact_gamma = len(_exact_gamma)-1

191

192# Lanczos coefficients used by the GNU Scientific Library

193_lanczos_g = 7

194_lanczos_p = (0.99999999999980993, 676.5203681218851, -1259.1392167224028,

195 771.32342877765313, -176.61502916214059, 12.507343278686905,

196 -0.13857109526572012, 9.9843695780195716e-6, 1.5056327351493116e-7)

197

198def _gamma_real(x):

199 _intx = int(x)

200 if _intx == x:

201 if _intx <= 0:

202 #return (-1)**_intx * INF

203 raise ZeroDivisionError("gamma function pole")

204 if _intx <= _max_exact_gamma:

205 return _exact_gamma[_intx]

206 if x < 0.5:

207 # TODO: sinpi

208 return pi / (_sinpi_real(x)*_gamma_real(1-x))

209 else:

210 x -= 1.0

211 r = _lanczos_p[0]

212 for i in range(1, _lanczos_g+2):

213 r += _lanczos_p[i]/(x+i)

214 t = x + _lanczos_g + 0.5

215 return 2.506628274631000502417 * t**(x+0.5) * math.exp(-t) * r

216

217def _gamma_complex(x):

218 if not x.imag:

219 return complex(_gamma_real(x.real))

220 if x.real < 0.5:

221 # TODO: sinpi

222 return pi / (_sinpi_complex(x)*_gamma_complex(1-x))

223 else:

224 x -= 1.0

225 r = _lanczos_p[0]

226 for i in range(1, _lanczos_g+2):

227 r += _lanczos_p[i]/(x+i)

228 t = x + _lanczos_g + 0.5

229 return 2.506628274631000502417 * t**(x+0.5) * cmath.exp(-t) * r

230

231gamma = _mathfun_real(_gamma_real, _gamma_complex)

232

233def rgamma(x):

234 try:

235 return 1./gamma(x)

236 except ZeroDivisionError:

237 return x*0.0

238

239def factorial(x):

240 return gamma(x+1.0)

241

242def arg(x):

243 if type(x) is float:

244 return math.atan2(0.0,x)

245 return math.atan2(x.imag,x.real)

246

247# XXX: broken for negatives

248def loggamma(x):

249 if type(x) not in (float, complex):

250 try:

251 x = float(x)

252 except (ValueError, TypeError):

253 x = complex(x)

254 try:

255 xreal = x.real

256 ximag = x.imag

257 except AttributeError: # py2.5

258 xreal = x

259 ximag = 0.0

260 # Reflection formula

261 # http://functions.wolfram.com/GammaBetaErf/LogGamma/16/01/01/0003/

262 if xreal < 0.0:

263 if abs(x) < 0.5:

264 v = log(gamma(x))

265 if ximag == 0:

266 v = v.conjugate()

267 return v

268 z = 1-x

269 try:

270 re = z.real

271 im = z.imag

272 except AttributeError: # py2.5

273 re = z

274 im = 0.0

275 refloor = floor(re)

276 if im == 0.0:

277 imsign = 0

278 elif im < 0.0:

279 imsign = -1

280 else:

281 imsign = 1

282 return (-pi*1j)*abs(refloor)*(1-abs(imsign)) + logpi - \

283 log(sinpi(z-refloor)) - loggamma(z) + 1j*pi*refloor*imsign

284 if x == 1.0 or x == 2.0:

285 return x*0

286 p = 0.

287 while abs(x) < 11:

288 p -= log(x)

289 x += 1.0

290 s = 0.918938533204672742 + (x-0.5)*log(x) - x

291 r = 1./x

292 r2 = r*r

293 s += 0.083333333333333333333*r; r *= r2

294 s += -0.0027777777777777777778*r; r *= r2

295 s += 0.00079365079365079365079*r; r *= r2

296 s += -0.0005952380952380952381*r; r *= r2

297 s += 0.00084175084175084175084*r; r *= r2

298 s += -0.0019175269175269175269*r; r *= r2

299 s += 0.0064102564102564102564*r; r *= r2

300 s += -0.02955065359477124183*r

301 return s + p

302

303_psi_coeff = [

3040.083333333333333333333,

305-0.0083333333333333333333,

3060.003968253968253968254,

307-0.0041666666666666666667,

3080.0075757575757575757576,

309-0.021092796092796092796,

3100.083333333333333333333,

311-0.44325980392156862745,

3123.0539543302701197438,

313-26.456212121212121212]

314

315def _digamma_real(x):

316 _intx = int(x)

317 if _intx == x:

318 if _intx <= 0:

319 raise ZeroDivisionError("polygamma pole")

320 if x < 0.5:

321 x = 1.0-x

322 s = pi*cotpi(x)

323 else:

324 s = 0.0

325 while x < 10.0:

326 s -= 1.0/x

327 x += 1.0

328 x2 = x**-2

329 t = x2

330 for c in _psi_coeff:

331 s -= c*t

332 if t < 1e-20:

333 break

334 t *= x2

335 return s + math_log(x) - 0.5/x

336

337def _digamma_complex(x):

338 if not x.imag:

339 return complex(_digamma_real(x.real))

340 if x.real < 0.5:

341 x = 1.0-x

342 s = pi*cotpi(x)

343 else:

344 s = 0.0

345 while abs(x) < 10.0:

346 s -= 1.0/x

347 x += 1.0

348 x2 = x**-2

349 t = x2

350 for c in _psi_coeff:

351 s -= c*t

352 if abs(t) < 1e-20:

353 break

354 t *= x2

355 return s + cmath.log(x) - 0.5/x

356

357digamma = _mathfun_real(_digamma_real, _digamma_complex)

358

359# TODO: could implement complex erf and erfc here. Need

360# to find an accurate method (avoiding cancellation)

361# for approx. 1 < abs(x) < 9.

362

363_erfc_coeff_P = [

364 1.0000000161203922312,

365 2.1275306946297962644,

366 2.2280433377390253297,

367 1.4695509105618423961,

368 0.66275911699770787537,

369 0.20924776504163751585,

370 0.045459713768411264339,

371 0.0063065951710717791934,

372 0.00044560259661560421715][::-1]

373

374_erfc_coeff_Q = [

375 1.0000000000000000000,

376 3.2559100272784894318,

377 4.9019435608903239131,

378 4.4971472894498014205,

379 2.7845640601891186528,

380 1.2146026030046904138,

381 0.37647108453729465912,

382 0.080970149639040548613,

383 0.011178148899483545902,

384 0.00078981003831980423513][::-1]

385

386def _polyval(coeffs, x):

387 p = coeffs[0]

388 for c in coeffs[1:]:

389 p = c + x*p

390 return p

391

392def _erf_taylor(x):

393 # Taylor series assuming 0 <= x <= 1

394 x2 = x*x

395 s = t = x

396 n = 1

397 while abs(t) > 1e-17:

398 t *= x2/n

399 s -= t/(n+n+1)

400 n += 1

401 t *= x2/n

402 s += t/(n+n+1)

403 n += 1

404 return 1.1283791670955125739*s

405

406def _erfc_mid(x):

407 # Rational approximation assuming 0 <= x <= 9

408 return exp(-x*x)*_polyval(_erfc_coeff_P,x)/_polyval(_erfc_coeff_Q,x)

409

410def _erfc_asymp(x):

411 # Asymptotic expansion assuming x >= 9

412 x2 = x*x

413 v = exp(-x2)/x*0.56418958354775628695

414 r = t = 0.5 / x2

415 s = 1.0

416 for n in range(1,22,4):

417 s -= t

418 t *= r * (n+2)

419 s += t

420 t *= r * (n+4)

421 if abs(t) < 1e-17:

422 break

423 return s * v

424

425def erf(x):

426 """

427 erf of a real number.

428 """

429 x = float(x)

430 if x != x:

431 return x

432 if x < 0.0:

433 return -erf(-x)

434 if x >= 1.0:

435 if x >= 6.0:

436 return 1.0

437 return 1.0 - _erfc_mid(x)

438 return _erf_taylor(x)

439

440def erfc(x):

441 """

442 erfc of a real number.

443 """

444 x = float(x)

445 if x != x:

446 return x

447 if x < 0.0:

448 if x < -6.0:

449 return 2.0

450 return 2.0-erfc(-x)

451 if x > 9.0:

452 return _erfc_asymp(x)

453 if x >= 1.0:

454 return _erfc_mid(x)

455 return 1.0 - _erf_taylor(x)

456

457gauss42 = [\

458(0.99839961899006235, 0.0041059986046490839),

459(-0.99839961899006235, 0.0041059986046490839),

460(0.9915772883408609, 0.009536220301748501),

461(-0.9915772883408609,0.009536220301748501),

462(0.97934250806374812, 0.014922443697357493),

463(-0.97934250806374812, 0.014922443697357493),

464(0.96175936533820439,0.020227869569052644),

465(-0.96175936533820439, 0.020227869569052644),

466(0.93892355735498811, 0.025422959526113047),

467(-0.93892355735498811,0.025422959526113047),

468(0.91095972490412735, 0.030479240699603467),

469(-0.91095972490412735, 0.030479240699603467),

470(0.87802056981217269,0.03536907109759211),

471(-0.87802056981217269, 0.03536907109759211),

472(0.8402859832618168, 0.040065735180692258),

473(-0.8402859832618168,0.040065735180692258),

474(0.7979620532554873, 0.044543577771965874),

475(-0.7979620532554873, 0.044543577771965874),

476(0.75127993568948048,0.048778140792803244),

477(-0.75127993568948048, 0.048778140792803244),

478(0.70049459055617114, 0.052746295699174064),

479(-0.70049459055617114,0.052746295699174064),

480(0.64588338886924779, 0.056426369358018376),

481(-0.64588338886924779, 0.056426369358018376),

482(0.58774459748510932, 0.059798262227586649),

483(-0.58774459748510932, 0.059798262227586649),

484(0.5263957499311922, 0.062843558045002565),

485(-0.5263957499311922, 0.062843558045002565),

486(0.46217191207042191, 0.065545624364908975),

487(-0.46217191207042191, 0.065545624364908975),

488(0.39542385204297503, 0.067889703376521934),

489(-0.39542385204297503, 0.067889703376521934),

490(0.32651612446541151, 0.069862992492594159),

491(-0.32651612446541151, 0.069862992492594159),

492(0.25582507934287907, 0.071454714265170971),

493(-0.25582507934287907, 0.071454714265170971),

494(0.18373680656485453, 0.072656175243804091),

495(-0.18373680656485453, 0.072656175243804091),

496(0.11064502720851986, 0.073460813453467527),

497(-0.11064502720851986, 0.073460813453467527),

498(0.036948943165351772, 0.073864234232172879),

499(-0.036948943165351772, 0.073864234232172879)]

500

501EI_ASYMP_CONVERGENCE_RADIUS = 40.0

502

503def ei_asymp(z, _e1=False):

504 r = 1./z

505 s = t = 1.0

506 k = 1

507 while 1:

508 t *= k*r

509 s += t

510 if abs(t) < 1e-16:

511 break

512 k += 1

513 v = s*exp(z)/z

514 if _e1:

515 if type(z) is complex:

516 zreal = z.real

517 zimag = z.imag

518 else:

519 zreal = z

520 zimag = 0.0

521 if zimag == 0.0 and zreal > 0.0:

522 v += pi*1j

523 else:

524 if type(z) is complex:

525 if z.imag > 0:

526 v += pi*1j

527 if z.imag < 0:

528 v -= pi*1j

529 return v

530

531def ei_taylor(z, _e1=False):

532 s = t = z

533 k = 2

534 while 1:

535 t = t*z/k

536 term = t/k

537 if abs(term) < 1e-17:

538 break

539 s += term

540 k += 1

541 s += euler

542 if _e1:

543 s += log(-z)

544 else:

545 if type(z) is float or z.imag == 0.0:

546 s += math_log(abs(z))

547 else:

548 s += cmath.log(z)

549 return s

550

551def ei(z, _e1=False):

552 typez = type(z)

553 if typez not in (float, complex):

554 try:

555 z = float(z)

556 typez = float

557 except (TypeError, ValueError):

558 z = complex(z)

559 typez = complex

560 if not z:

561 return -INF

562 absz = abs(z)

563 if absz > EI_ASYMP_CONVERGENCE_RADIUS:

564 return ei_asymp(z, _e1)

565 elif absz <= 2.0 or (typez is float and z > 0.0):

566 return ei_taylor(z, _e1)

567 # Integrate, starting from whichever is smaller of a Taylor

568 # series value or an asymptotic series value

569 if typez is complex and z.real > 0.0:

570 zref = z / absz

571 ref = ei_taylor(zref, _e1)

572 else:

573 zref = EI_ASYMP_CONVERGENCE_RADIUS * z / absz

574 ref = ei_asymp(zref, _e1)

575 C = (zref-z)*0.5

576 D = (zref+z)*0.5

577 s = 0.0

578 if type(z) is complex:

579 _exp = cmath.exp

580 else:

581 _exp = math.exp

582 for x,w in gauss42:

583 t = C*x+D

584 s += w*_exp(t)/t

585 ref -= C*s

586 return ref

587

588def e1(z):

589 # hack to get consistent signs if the imaginary part if 0

590 # and signed

591 typez = type(z)

592 if type(z) not in (float, complex):

593 try:

594 z = float(z)

595 typez = float

596 except (TypeError, ValueError):

597 z = complex(z)

598 typez = complex

599 if typez is complex and not z.imag:

600 z = complex(z.real, 0.0)

601 # end hack

602 return -ei(-z, _e1=True)

603

604_zeta_int = [\

605-0.5,

6060.0,

6071.6449340668482264365,1.2020569031595942854,1.0823232337111381915,

6081.0369277551433699263,1.0173430619844491397,1.0083492773819228268,

6091.0040773561979443394,1.0020083928260822144,1.0009945751278180853,

6101.0004941886041194646,1.0002460865533080483,1.0001227133475784891,

6111.0000612481350587048,1.0000305882363070205,1.0000152822594086519,

6121.0000076371976378998,1.0000038172932649998,1.0000019082127165539,

6131.0000009539620338728,1.0000004769329867878,1.0000002384505027277,

6141.0000001192199259653,1.0000000596081890513,1.0000000298035035147,

6151.0000000149015548284]

616

617_zeta_P = [-3.50000000087575873, -0.701274355654678147,

618-0.0672313458590012612, -0.00398731457954257841,

619-0.000160948723019303141, -4.67633010038383371e-6,

620-1.02078104417700585e-7, -1.68030037095896287e-9,

621-1.85231868742346722e-11][::-1]

622

623_zeta_Q = [1.00000000000000000, -0.936552848762465319,

624-0.0588835413263763741, -0.00441498861482948666,

625-0.000143416758067432622, -5.10691659585090782e-6,

626-9.58813053268913799e-8, -1.72963791443181972e-9,

627-1.83527919681474132e-11][::-1]

628

629_zeta_1 = [3.03768838606128127e-10, -1.21924525236601262e-8,

6302.01201845887608893e-7, -1.53917240683468381e-6,

631-5.09890411005967954e-7, 0.000122464707271619326,

632-0.000905721539353130232, -0.00239315326074843037,

6330.084239750013159168, 0.418938517907442414, 0.500000001921884009]

634

635_zeta_0 = [-3.46092485016748794e-10, -6.42610089468292485e-9,

6361.76409071536679773e-7, -1.47141263991560698e-6, -6.38880222546167613e-7,

6370.000122641099800668209, -0.000905894913516772796, -0.00239303348507992713,

6380.0842396947501199816, 0.418938533204660256, 0.500000000000000052]

639

640def zeta(s):

641 """

642 Riemann zeta function, real argument

643 """

644 if not isinstance(s, (float, int)):

645 try:

646 s = float(s)

647 except (ValueError, TypeError):

648 try:

649 s = complex(s)

650 if not s.imag:

651 return complex(zeta(s.real))

652 except (ValueError, TypeError):

653 pass

654 raise NotImplementedError

655 if s == 1:

656 raise ValueError("zeta(1) pole")

657 if s >= 27:

658 return 1.0 + 2.0**(-s) + 3.0**(-s)

659 n = int(s)

660 if n == s:

661 if n >= 0:

662 return _zeta_int[n]

663 if not (n % 2):

664 return 0.0

665 if s <= 0.0:

666 return 2.**s*pi**(s-1)*_sinpi_real(0.5*s)*_gamma_real(1-s)*zeta(1-s)

667 if s <= 2.0:

668 if s <= 1.0:

669 return _polyval(_zeta_0,s)/(s-1)

670 return _polyval(_zeta_1,s)/(s-1)

671 z = _polyval(_zeta_P,s) / _polyval(_zeta_Q,s)

672 return 1.0 + 2.0**(-s) + 3.0**(-s) + 4.0**(-s)*z