libmpc.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835
  1. """
  2. Low-level functions for complex arithmetic.
  3. """
  4. import sys
  5. from .backend import MPZ, MPZ_ZERO, MPZ_ONE, MPZ_TWO, BACKEND
  6. from .libmpf import (\
  7. round_floor, round_ceiling, round_down, round_up,
  8. round_nearest, round_fast, bitcount,
  9. bctable, normalize, normalize1, reciprocal_rnd, rshift, lshift, giant_steps,
  10. negative_rnd,
  11. to_str, to_fixed, from_man_exp, from_float, to_float, from_int, to_int,
  12. fzero, fone, ftwo, fhalf, finf, fninf, fnan, fnone,
  13. mpf_abs, mpf_pos, mpf_neg, mpf_add, mpf_sub, mpf_mul,
  14. mpf_div, mpf_mul_int, mpf_shift, mpf_sqrt, mpf_hypot,
  15. mpf_rdiv_int, mpf_floor, mpf_ceil, mpf_nint, mpf_frac,
  16. mpf_sign, mpf_hash,
  17. ComplexResult
  18. )
  19. from .libelefun import (\
  20. mpf_pi, mpf_exp, mpf_log, mpf_cos_sin, mpf_cosh_sinh, mpf_tan, mpf_pow_int,
  21. mpf_log_hypot,
  22. mpf_cos_sin_pi, mpf_phi,
  23. mpf_cos, mpf_sin, mpf_cos_pi, mpf_sin_pi,
  24. mpf_atan, mpf_atan2, mpf_cosh, mpf_sinh, mpf_tanh,
  25. mpf_asin, mpf_acos, mpf_acosh, mpf_nthroot, mpf_fibonacci
  26. )
  27. # An mpc value is a (real, imag) tuple
  28. mpc_one = fone, fzero
  29. mpc_zero = fzero, fzero
  30. mpc_two = ftwo, fzero
  31. mpc_half = (fhalf, fzero)
  32. _infs = (finf, fninf)
  33. _infs_nan = (finf, fninf, fnan)
  34. def mpc_is_inf(z):
  35. """Check if either real or imaginary part is infinite"""
  36. re, im = z
  37. if re in _infs: return True
  38. if im in _infs: return True
  39. return False
  40. def mpc_is_infnan(z):
  41. """Check if either real or imaginary part is infinite or nan"""
  42. re, im = z
  43. if re in _infs_nan: return True
  44. if im in _infs_nan: return True
  45. return False
  46. def mpc_to_str(z, dps, **kwargs):
  47. re, im = z
  48. rs = to_str(re, dps)
  49. if im[0]:
  50. return rs + " - " + to_str(mpf_neg(im), dps, **kwargs) + "j"
  51. else:
  52. return rs + " + " + to_str(im, dps, **kwargs) + "j"
  53. def mpc_to_complex(z, strict=False, rnd=round_fast):
  54. re, im = z
  55. return complex(to_float(re, strict, rnd), to_float(im, strict, rnd))
  56. def mpc_hash(z):
  57. if sys.version_info >= (3, 2):
  58. re, im = z
  59. h = mpf_hash(re) + sys.hash_info.imag * mpf_hash(im)
  60. # Need to reduce either module 2^32 or 2^64
  61. h = h % (2**sys.hash_info.width)
  62. return int(h)
  63. else:
  64. try:
  65. return hash(mpc_to_complex(z, strict=True))
  66. except OverflowError:
  67. return hash(z)
  68. def mpc_conjugate(z, prec, rnd=round_fast):
  69. re, im = z
  70. return re, mpf_neg(im, prec, rnd)
  71. def mpc_is_nonzero(z):
  72. return z != mpc_zero
  73. def mpc_add(z, w, prec, rnd=round_fast):
  74. a, b = z
  75. c, d = w
  76. return mpf_add(a, c, prec, rnd), mpf_add(b, d, prec, rnd)
  77. def mpc_add_mpf(z, x, prec, rnd=round_fast):
  78. a, b = z
  79. return mpf_add(a, x, prec, rnd), b
  80. def mpc_sub(z, w, prec=0, rnd=round_fast):
  81. a, b = z
  82. c, d = w
  83. return mpf_sub(a, c, prec, rnd), mpf_sub(b, d, prec, rnd)
  84. def mpc_sub_mpf(z, p, prec=0, rnd=round_fast):
  85. a, b = z
  86. return mpf_sub(a, p, prec, rnd), b
  87. def mpc_pos(z, prec, rnd=round_fast):
  88. a, b = z
  89. return mpf_pos(a, prec, rnd), mpf_pos(b, prec, rnd)
  90. def mpc_neg(z, prec=None, rnd=round_fast):
  91. a, b = z
  92. return mpf_neg(a, prec, rnd), mpf_neg(b, prec, rnd)
  93. def mpc_shift(z, n):
  94. a, b = z
  95. return mpf_shift(a, n), mpf_shift(b, n)
  96. def mpc_abs(z, prec, rnd=round_fast):
  97. """Absolute value of a complex number, |a+bi|.
  98. Returns an mpf value."""
  99. a, b = z
  100. return mpf_hypot(a, b, prec, rnd)
  101. def mpc_arg(z, prec, rnd=round_fast):
  102. """Argument of a complex number. Returns an mpf value."""
  103. a, b = z
  104. return mpf_atan2(b, a, prec, rnd)
  105. def mpc_floor(z, prec, rnd=round_fast):
  106. a, b = z
  107. return mpf_floor(a, prec, rnd), mpf_floor(b, prec, rnd)
  108. def mpc_ceil(z, prec, rnd=round_fast):
  109. a, b = z
  110. return mpf_ceil(a, prec, rnd), mpf_ceil(b, prec, rnd)
  111. def mpc_nint(z, prec, rnd=round_fast):
  112. a, b = z
  113. return mpf_nint(a, prec, rnd), mpf_nint(b, prec, rnd)
  114. def mpc_frac(z, prec, rnd=round_fast):
  115. a, b = z
  116. return mpf_frac(a, prec, rnd), mpf_frac(b, prec, rnd)
  117. def mpc_mul(z, w, prec, rnd=round_fast):
  118. """
  119. Complex multiplication.
  120. Returns the real and imaginary part of (a+bi)*(c+di), rounded to
  121. the specified precision. The rounding mode applies to the real and
  122. imaginary parts separately.
  123. """
  124. a, b = z
  125. c, d = w
  126. p = mpf_mul(a, c)
  127. q = mpf_mul(b, d)
  128. r = mpf_mul(a, d)
  129. s = mpf_mul(b, c)
  130. re = mpf_sub(p, q, prec, rnd)
  131. im = mpf_add(r, s, prec, rnd)
  132. return re, im
  133. def mpc_square(z, prec, rnd=round_fast):
  134. # (a+b*I)**2 == a**2 - b**2 + 2*I*a*b
  135. a, b = z
  136. p = mpf_mul(a,a)
  137. q = mpf_mul(b,b)
  138. r = mpf_mul(a,b, prec, rnd)
  139. re = mpf_sub(p, q, prec, rnd)
  140. im = mpf_shift(r, 1)
  141. return re, im
  142. def mpc_mul_mpf(z, p, prec, rnd=round_fast):
  143. a, b = z
  144. re = mpf_mul(a, p, prec, rnd)
  145. im = mpf_mul(b, p, prec, rnd)
  146. return re, im
  147. def mpc_mul_imag_mpf(z, x, prec, rnd=round_fast):
  148. """
  149. Multiply the mpc value z by I*x where x is an mpf value.
  150. """
  151. a, b = z
  152. re = mpf_neg(mpf_mul(b, x, prec, rnd))
  153. im = mpf_mul(a, x, prec, rnd)
  154. return re, im
  155. def mpc_mul_int(z, n, prec, rnd=round_fast):
  156. a, b = z
  157. re = mpf_mul_int(a, n, prec, rnd)
  158. im = mpf_mul_int(b, n, prec, rnd)
  159. return re, im
  160. def mpc_div(z, w, prec, rnd=round_fast):
  161. a, b = z
  162. c, d = w
  163. wp = prec + 10
  164. # mag = c*c + d*d
  165. mag = mpf_add(mpf_mul(c, c), mpf_mul(d, d), wp)
  166. # (a*c+b*d)/mag, (b*c-a*d)/mag
  167. t = mpf_add(mpf_mul(a,c), mpf_mul(b,d), wp)
  168. u = mpf_sub(mpf_mul(b,c), mpf_mul(a,d), wp)
  169. return mpf_div(t,mag,prec,rnd), mpf_div(u,mag,prec,rnd)
  170. def mpc_div_mpf(z, p, prec, rnd=round_fast):
  171. """Calculate z/p where p is real"""
  172. a, b = z
  173. re = mpf_div(a, p, prec, rnd)
  174. im = mpf_div(b, p, prec, rnd)
  175. return re, im
  176. def mpc_reciprocal(z, prec, rnd=round_fast):
  177. """Calculate 1/z efficiently"""
  178. a, b = z
  179. m = mpf_add(mpf_mul(a,a),mpf_mul(b,b),prec+10)
  180. re = mpf_div(a, m, prec, rnd)
  181. im = mpf_neg(mpf_div(b, m, prec, rnd))
  182. return re, im
  183. def mpc_mpf_div(p, z, prec, rnd=round_fast):
  184. """Calculate p/z where p is real efficiently"""
  185. a, b = z
  186. m = mpf_add(mpf_mul(a,a),mpf_mul(b,b), prec+10)
  187. re = mpf_div(mpf_mul(a,p), m, prec, rnd)
  188. im = mpf_div(mpf_neg(mpf_mul(b,p)), m, prec, rnd)
  189. return re, im
  190. def complex_int_pow(a, b, n):
  191. """Complex integer power: computes (a+b*I)**n exactly for
  192. nonnegative n (a and b must be Python ints)."""
  193. wre = 1
  194. wim = 0
  195. while n:
  196. if n & 1:
  197. wre, wim = wre*a - wim*b, wim*a + wre*b
  198. n -= 1
  199. a, b = a*a - b*b, 2*a*b
  200. n //= 2
  201. return wre, wim
  202. def mpc_pow(z, w, prec, rnd=round_fast):
  203. if w[1] == fzero:
  204. return mpc_pow_mpf(z, w[0], prec, rnd)
  205. return mpc_exp(mpc_mul(mpc_log(z, prec+10), w, prec+10), prec, rnd)
  206. def mpc_pow_mpf(z, p, prec, rnd=round_fast):
  207. psign, pman, pexp, pbc = p
  208. if pexp >= 0:
  209. return mpc_pow_int(z, (-1)**psign * (pman<<pexp), prec, rnd)
  210. if pexp == -1:
  211. sqrtz = mpc_sqrt(z, prec+10)
  212. return mpc_pow_int(sqrtz, (-1)**psign * pman, prec, rnd)
  213. return mpc_exp(mpc_mul_mpf(mpc_log(z, prec+10), p, prec+10), prec, rnd)
  214. def mpc_pow_int(z, n, prec, rnd=round_fast):
  215. a, b = z
  216. if b == fzero:
  217. return mpf_pow_int(a, n, prec, rnd), fzero
  218. if a == fzero:
  219. v = mpf_pow_int(b, n, prec, rnd)
  220. n %= 4
  221. if n == 0:
  222. return v, fzero
  223. elif n == 1:
  224. return fzero, v
  225. elif n == 2:
  226. return mpf_neg(v), fzero
  227. elif n == 3:
  228. return fzero, mpf_neg(v)
  229. if n == 0: return mpc_one
  230. if n == 1: return mpc_pos(z, prec, rnd)
  231. if n == 2: return mpc_square(z, prec, rnd)
  232. if n == -1: return mpc_reciprocal(z, prec, rnd)
  233. if n < 0: return mpc_reciprocal(mpc_pow_int(z, -n, prec+4), prec, rnd)
  234. asign, aman, aexp, abc = a
  235. bsign, bman, bexp, bbc = b
  236. if asign: aman = -aman
  237. if bsign: bman = -bman
  238. de = aexp - bexp
  239. abs_de = abs(de)
  240. exact_size = n*(abs_de + max(abc, bbc))
  241. if exact_size < 10000:
  242. if de > 0:
  243. aman <<= de
  244. aexp = bexp
  245. else:
  246. bman <<= (-de)
  247. bexp = aexp
  248. re, im = complex_int_pow(aman, bman, n)
  249. re = from_man_exp(re, int(n*aexp), prec, rnd)
  250. im = from_man_exp(im, int(n*bexp), prec, rnd)
  251. return re, im
  252. return mpc_exp(mpc_mul_int(mpc_log(z, prec+10), n, prec+10), prec, rnd)
  253. def mpc_sqrt(z, prec, rnd=round_fast):
  254. """Complex square root (principal branch).
  255. We have sqrt(a+bi) = sqrt((r+a)/2) + b/sqrt(2*(r+a))*i where
  256. r = abs(a+bi), when a+bi is not a negative real number."""
  257. a, b = z
  258. if b == fzero:
  259. if a == fzero:
  260. return (a, b)
  261. # When a+bi is a negative real number, we get a real sqrt times i
  262. if a[0]:
  263. im = mpf_sqrt(mpf_neg(a), prec, rnd)
  264. return (fzero, im)
  265. else:
  266. re = mpf_sqrt(a, prec, rnd)
  267. return (re, fzero)
  268. wp = prec+20
  269. if not a[0]: # case a positive
  270. t = mpf_add(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) + a
  271. u = mpf_shift(t, -1) # u = t/2
  272. re = mpf_sqrt(u, prec, rnd) # re = sqrt(u)
  273. v = mpf_shift(t, 1) # v = 2*t
  274. w = mpf_sqrt(v, wp) # w = sqrt(v)
  275. im = mpf_div(b, w, prec, rnd) # im = b / w
  276. else: # case a negative
  277. t = mpf_sub(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) - a
  278. u = mpf_shift(t, -1) # u = t/2
  279. im = mpf_sqrt(u, prec, rnd) # im = sqrt(u)
  280. v = mpf_shift(t, 1) # v = 2*t
  281. w = mpf_sqrt(v, wp) # w = sqrt(v)
  282. re = mpf_div(b, w, prec, rnd) # re = b/w
  283. if b[0]:
  284. re = mpf_neg(re)
  285. im = mpf_neg(im)
  286. return re, im
  287. def mpc_nthroot_fixed(a, b, n, prec):
  288. # a, b signed integers at fixed precision prec
  289. start = 50
  290. a1 = int(rshift(a, prec - n*start))
  291. b1 = int(rshift(b, prec - n*start))
  292. try:
  293. r = (a1 + 1j * b1)**(1.0/n)
  294. re = r.real
  295. im = r.imag
  296. re = MPZ(int(re))
  297. im = MPZ(int(im))
  298. except OverflowError:
  299. a1 = from_int(a1, start)
  300. b1 = from_int(b1, start)
  301. fn = from_int(n)
  302. nth = mpf_rdiv_int(1, fn, start)
  303. re, im = mpc_pow((a1, b1), (nth, fzero), start)
  304. re = to_int(re)
  305. im = to_int(im)
  306. extra = 10
  307. prevp = start
  308. extra1 = n
  309. for p in giant_steps(start, prec+extra):
  310. # this is slow for large n, unlike int_pow_fixed
  311. re2, im2 = complex_int_pow(re, im, n-1)
  312. re2 = rshift(re2, (n-1)*prevp - p - extra1)
  313. im2 = rshift(im2, (n-1)*prevp - p - extra1)
  314. r4 = (re2*re2 + im2*im2) >> (p + extra1)
  315. ap = rshift(a, prec - p)
  316. bp = rshift(b, prec - p)
  317. rec = (ap * re2 + bp * im2) >> p
  318. imc = (-ap * im2 + bp * re2) >> p
  319. reb = (rec << p) // r4
  320. imb = (imc << p) // r4
  321. re = (reb + (n-1)*lshift(re, p-prevp))//n
  322. im = (imb + (n-1)*lshift(im, p-prevp))//n
  323. prevp = p
  324. return re, im
  325. def mpc_nthroot(z, n, prec, rnd=round_fast):
  326. """
  327. Complex n-th root.
  328. Use Newton method as in the real case when it is faster,
  329. otherwise use z**(1/n)
  330. """
  331. a, b = z
  332. if a[0] == 0 and b == fzero:
  333. re = mpf_nthroot(a, n, prec, rnd)
  334. return (re, fzero)
  335. if n < 2:
  336. if n == 0:
  337. return mpc_one
  338. if n == 1:
  339. return mpc_pos((a, b), prec, rnd)
  340. if n == -1:
  341. return mpc_div(mpc_one, (a, b), prec, rnd)
  342. inverse = mpc_nthroot((a, b), -n, prec+5, reciprocal_rnd[rnd])
  343. return mpc_div(mpc_one, inverse, prec, rnd)
  344. if n <= 20:
  345. prec2 = int(1.2 * (prec + 10))
  346. asign, aman, aexp, abc = a
  347. bsign, bman, bexp, bbc = b
  348. pf = mpc_abs((a,b), prec)
  349. if pf[-2] + pf[-1] > -10 and pf[-2] + pf[-1] < prec:
  350. af = to_fixed(a, prec2)
  351. bf = to_fixed(b, prec2)
  352. re, im = mpc_nthroot_fixed(af, bf, n, prec2)
  353. extra = 10
  354. re = from_man_exp(re, -prec2-extra, prec2, rnd)
  355. im = from_man_exp(im, -prec2-extra, prec2, rnd)
  356. return re, im
  357. fn = from_int(n)
  358. prec2 = prec+10 + 10
  359. nth = mpf_rdiv_int(1, fn, prec2)
  360. re, im = mpc_pow((a, b), (nth, fzero), prec2, rnd)
  361. re = normalize(re[0], re[1], re[2], re[3], prec, rnd)
  362. im = normalize(im[0], im[1], im[2], im[3], prec, rnd)
  363. return re, im
  364. def mpc_cbrt(z, prec, rnd=round_fast):
  365. """
  366. Complex cubic root.
  367. """
  368. return mpc_nthroot(z, 3, prec, rnd)
  369. def mpc_exp(z, prec, rnd=round_fast):
  370. """
  371. Complex exponential function.
  372. We use the direct formula exp(a+bi) = exp(a) * (cos(b) + sin(b)*i)
  373. for the computation. This formula is very nice because it is
  374. pefectly stable; since we just do real multiplications, the only
  375. numerical errors that can creep in are single-ulp rounding errors.
  376. The formula is efficient since mpmath's real exp is quite fast and
  377. since we can compute cos and sin simultaneously.
  378. It is no problem if a and b are large; if the implementations of
  379. exp/cos/sin are accurate and efficient for all real numbers, then
  380. so is this function for all complex numbers.
  381. """
  382. a, b = z
  383. if a == fzero:
  384. return mpf_cos_sin(b, prec, rnd)
  385. if b == fzero:
  386. return mpf_exp(a, prec, rnd), fzero
  387. mag = mpf_exp(a, prec+4, rnd)
  388. c, s = mpf_cos_sin(b, prec+4, rnd)
  389. re = mpf_mul(mag, c, prec, rnd)
  390. im = mpf_mul(mag, s, prec, rnd)
  391. return re, im
  392. def mpc_log(z, prec, rnd=round_fast):
  393. re = mpf_log_hypot(z[0], z[1], prec, rnd)
  394. im = mpc_arg(z, prec, rnd)
  395. return re, im
  396. def mpc_cos(z, prec, rnd=round_fast):
  397. """Complex cosine. The formula used is cos(a+bi) = cos(a)*cosh(b) -
  398. sin(a)*sinh(b)*i.
  399. The same comments apply as for the complex exp: only real
  400. multiplications are pewrormed, so no cancellation errors are
  401. possible. The formula is also efficient since we can compute both
  402. pairs (cos, sin) and (cosh, sinh) in single stwps."""
  403. a, b = z
  404. if b == fzero:
  405. return mpf_cos(a, prec, rnd), fzero
  406. if a == fzero:
  407. return mpf_cosh(b, prec, rnd), fzero
  408. wp = prec + 6
  409. c, s = mpf_cos_sin(a, wp)
  410. ch, sh = mpf_cosh_sinh(b, wp)
  411. re = mpf_mul(c, ch, prec, rnd)
  412. im = mpf_mul(s, sh, prec, rnd)
  413. return re, mpf_neg(im)
  414. def mpc_sin(z, prec, rnd=round_fast):
  415. """Complex sine. We have sin(a+bi) = sin(a)*cosh(b) +
  416. cos(a)*sinh(b)*i. See the docstring for mpc_cos for additional
  417. comments."""
  418. a, b = z
  419. if b == fzero:
  420. return mpf_sin(a, prec, rnd), fzero
  421. if a == fzero:
  422. return fzero, mpf_sinh(b, prec, rnd)
  423. wp = prec + 6
  424. c, s = mpf_cos_sin(a, wp)
  425. ch, sh = mpf_cosh_sinh(b, wp)
  426. re = mpf_mul(s, ch, prec, rnd)
  427. im = mpf_mul(c, sh, prec, rnd)
  428. return re, im
  429. def mpc_tan(z, prec, rnd=round_fast):
  430. """Complex tangent. Computed as tan(a+bi) = sin(2a)/M + sinh(2b)/M*i
  431. where M = cos(2a) + cosh(2b)."""
  432. a, b = z
  433. asign, aman, aexp, abc = a
  434. bsign, bman, bexp, bbc = b
  435. if b == fzero: return mpf_tan(a, prec, rnd), fzero
  436. if a == fzero: return fzero, mpf_tanh(b, prec, rnd)
  437. wp = prec + 15
  438. a = mpf_shift(a, 1)
  439. b = mpf_shift(b, 1)
  440. c, s = mpf_cos_sin(a, wp)
  441. ch, sh = mpf_cosh_sinh(b, wp)
  442. # TODO: handle cancellation when c ~= -1 and ch ~= 1
  443. mag = mpf_add(c, ch, wp)
  444. re = mpf_div(s, mag, prec, rnd)
  445. im = mpf_div(sh, mag, prec, rnd)
  446. return re, im
  447. def mpc_cos_pi(z, prec, rnd=round_fast):
  448. a, b = z
  449. if b == fzero:
  450. return mpf_cos_pi(a, prec, rnd), fzero
  451. b = mpf_mul(b, mpf_pi(prec+5), prec+5)
  452. if a == fzero:
  453. return mpf_cosh(b, prec, rnd), fzero
  454. wp = prec + 6
  455. c, s = mpf_cos_sin_pi(a, wp)
  456. ch, sh = mpf_cosh_sinh(b, wp)
  457. re = mpf_mul(c, ch, prec, rnd)
  458. im = mpf_mul(s, sh, prec, rnd)
  459. return re, mpf_neg(im)
  460. def mpc_sin_pi(z, prec, rnd=round_fast):
  461. a, b = z
  462. if b == fzero:
  463. return mpf_sin_pi(a, prec, rnd), fzero
  464. b = mpf_mul(b, mpf_pi(prec+5), prec+5)
  465. if a == fzero:
  466. return fzero, mpf_sinh(b, prec, rnd)
  467. wp = prec + 6
  468. c, s = mpf_cos_sin_pi(a, wp)
  469. ch, sh = mpf_cosh_sinh(b, wp)
  470. re = mpf_mul(s, ch, prec, rnd)
  471. im = mpf_mul(c, sh, prec, rnd)
  472. return re, im
  473. def mpc_cos_sin(z, prec, rnd=round_fast):
  474. a, b = z
  475. if a == fzero:
  476. ch, sh = mpf_cosh_sinh(b, prec, rnd)
  477. return (ch, fzero), (fzero, sh)
  478. if b == fzero:
  479. c, s = mpf_cos_sin(a, prec, rnd)
  480. return (c, fzero), (s, fzero)
  481. wp = prec + 6
  482. c, s = mpf_cos_sin(a, wp)
  483. ch, sh = mpf_cosh_sinh(b, wp)
  484. cre = mpf_mul(c, ch, prec, rnd)
  485. cim = mpf_mul(s, sh, prec, rnd)
  486. sre = mpf_mul(s, ch, prec, rnd)
  487. sim = mpf_mul(c, sh, prec, rnd)
  488. return (cre, mpf_neg(cim)), (sre, sim)
  489. def mpc_cos_sin_pi(z, prec, rnd=round_fast):
  490. a, b = z
  491. if b == fzero:
  492. c, s = mpf_cos_sin_pi(a, prec, rnd)
  493. return (c, fzero), (s, fzero)
  494. b = mpf_mul(b, mpf_pi(prec+5), prec+5)
  495. if a == fzero:
  496. ch, sh = mpf_cosh_sinh(b, prec, rnd)
  497. return (ch, fzero), (fzero, sh)
  498. wp = prec + 6
  499. c, s = mpf_cos_sin_pi(a, wp)
  500. ch, sh = mpf_cosh_sinh(b, wp)
  501. cre = mpf_mul(c, ch, prec, rnd)
  502. cim = mpf_mul(s, sh, prec, rnd)
  503. sre = mpf_mul(s, ch, prec, rnd)
  504. sim = mpf_mul(c, sh, prec, rnd)
  505. return (cre, mpf_neg(cim)), (sre, sim)
  506. def mpc_cosh(z, prec, rnd=round_fast):
  507. """Complex hyperbolic cosine. Computed as cosh(z) = cos(z*i)."""
  508. a, b = z
  509. return mpc_cos((b, mpf_neg(a)), prec, rnd)
  510. def mpc_sinh(z, prec, rnd=round_fast):
  511. """Complex hyperbolic sine. Computed as sinh(z) = -i*sin(z*i)."""
  512. a, b = z
  513. b, a = mpc_sin((b, a), prec, rnd)
  514. return a, b
  515. def mpc_tanh(z, prec, rnd=round_fast):
  516. """Complex hyperbolic tangent. Computed as tanh(z) = -i*tan(z*i)."""
  517. a, b = z
  518. b, a = mpc_tan((b, a), prec, rnd)
  519. return a, b
  520. # TODO: avoid loss of accuracy
  521. def mpc_atan(z, prec, rnd=round_fast):
  522. a, b = z
  523. # atan(z) = (I/2)*(log(1-I*z) - log(1+I*z))
  524. # x = 1-I*z = 1 + b - I*a
  525. # y = 1+I*z = 1 - b + I*a
  526. wp = prec + 15
  527. x = mpf_add(fone, b, wp), mpf_neg(a)
  528. y = mpf_sub(fone, b, wp), a
  529. l1 = mpc_log(x, wp)
  530. l2 = mpc_log(y, wp)
  531. a, b = mpc_sub(l1, l2, prec, rnd)
  532. # (I/2) * (a+b*I) = (-b/2 + a/2*I)
  533. v = mpf_neg(mpf_shift(b,-1)), mpf_shift(a,-1)
  534. # Subtraction at infinity gives correct real part but
  535. # wrong imaginary part (should be zero)
  536. if v[1] == fnan and mpc_is_inf(z):
  537. v = (v[0], fzero)
  538. return v
  539. beta_crossover = from_float(0.6417)
  540. alpha_crossover = from_float(1.5)
  541. def acos_asin(z, prec, rnd, n):
  542. """ complex acos for n = 0, asin for n = 1
  543. The algorithm is described in
  544. T.E. Hull, T.F. Fairgrieve and P.T.P. Tang
  545. 'Implementing the Complex Arcsine and Arcosine Functions
  546. using Exception Handling',
  547. ACM Trans. on Math. Software Vol. 23 (1997), p299
  548. The complex acos and asin can be defined as
  549. acos(z) = acos(beta) - I*sign(a)* log(alpha + sqrt(alpha**2 -1))
  550. asin(z) = asin(beta) + I*sign(a)* log(alpha + sqrt(alpha**2 -1))
  551. where z = a + I*b
  552. alpha = (1/2)*(r + s); beta = (1/2)*(r - s) = a/alpha
  553. r = sqrt((a+1)**2 + y**2); s = sqrt((a-1)**2 + y**2)
  554. These expressions are rewritten in different ways in different
  555. regions, delimited by two crossovers alpha_crossover and beta_crossover,
  556. and by abs(a) <= 1, in order to improve the numerical accuracy.
  557. """
  558. a, b = z
  559. wp = prec + 10
  560. # special cases with real argument
  561. if b == fzero:
  562. am = mpf_sub(fone, mpf_abs(a), wp)
  563. # case abs(a) <= 1
  564. if not am[0]:
  565. if n == 0:
  566. return mpf_acos(a, prec, rnd), fzero
  567. else:
  568. return mpf_asin(a, prec, rnd), fzero
  569. # cases abs(a) > 1
  570. else:
  571. # case a < -1
  572. if a[0]:
  573. pi = mpf_pi(prec, rnd)
  574. c = mpf_acosh(mpf_neg(a), prec, rnd)
  575. if n == 0:
  576. return pi, mpf_neg(c)
  577. else:
  578. return mpf_neg(mpf_shift(pi, -1)), c
  579. # case a > 1
  580. else:
  581. c = mpf_acosh(a, prec, rnd)
  582. if n == 0:
  583. return fzero, c
  584. else:
  585. pi = mpf_pi(prec, rnd)
  586. return mpf_shift(pi, -1), mpf_neg(c)
  587. asign = bsign = 0
  588. if a[0]:
  589. a = mpf_neg(a)
  590. asign = 1
  591. if b[0]:
  592. b = mpf_neg(b)
  593. bsign = 1
  594. am = mpf_sub(fone, a, wp)
  595. ap = mpf_add(fone, a, wp)
  596. r = mpf_hypot(ap, b, wp)
  597. s = mpf_hypot(am, b, wp)
  598. alpha = mpf_shift(mpf_add(r, s, wp), -1)
  599. beta = mpf_div(a, alpha, wp)
  600. b2 = mpf_mul(b,b, wp)
  601. # case beta <= beta_crossover
  602. if not mpf_sub(beta_crossover, beta, wp)[0]:
  603. if n == 0:
  604. re = mpf_acos(beta, wp)
  605. else:
  606. re = mpf_asin(beta, wp)
  607. else:
  608. # to compute the real part in this region use the identity
  609. # asin(beta) = atan(beta/sqrt(1-beta**2))
  610. # beta/sqrt(1-beta**2) = (alpha + a) * (alpha - a)
  611. # alpha + a is numerically accurate; alpha - a can have
  612. # cancellations leading to numerical inaccuracies, so rewrite
  613. # it in differente ways according to the region
  614. Ax = mpf_add(alpha, a, wp)
  615. # case a <= 1
  616. if not am[0]:
  617. # c = b*b/(r + (a+1)); d = (s + (1-a))
  618. # alpha - a = (1/2)*(c + d)
  619. # case n=0: re = atan(sqrt((1/2) * Ax * (c + d))/a)
  620. # case n=1: re = atan(a/sqrt((1/2) * Ax * (c + d)))
  621. c = mpf_div(b2, mpf_add(r, ap, wp), wp)
  622. d = mpf_add(s, am, wp)
  623. re = mpf_shift(mpf_mul(Ax, mpf_add(c, d, wp), wp), -1)
  624. if n == 0:
  625. re = mpf_atan(mpf_div(mpf_sqrt(re, wp), a, wp), wp)
  626. else:
  627. re = mpf_atan(mpf_div(a, mpf_sqrt(re, wp), wp), wp)
  628. else:
  629. # c = Ax/(r + (a+1)); d = Ax/(s - (1-a))
  630. # alpha - a = (1/2)*(c + d)
  631. # case n = 0: re = atan(b*sqrt(c + d)/2/a)
  632. # case n = 1: re = atan(a/(b*sqrt(c + d)/2)
  633. c = mpf_div(Ax, mpf_add(r, ap, wp), wp)
  634. d = mpf_div(Ax, mpf_sub(s, am, wp), wp)
  635. re = mpf_shift(mpf_add(c, d, wp), -1)
  636. re = mpf_mul(b, mpf_sqrt(re, wp), wp)
  637. if n == 0:
  638. re = mpf_atan(mpf_div(re, a, wp), wp)
  639. else:
  640. re = mpf_atan(mpf_div(a, re, wp), wp)
  641. # to compute alpha + sqrt(alpha**2 - 1), if alpha <= alpha_crossover
  642. # replace it with 1 + Am1 + sqrt(Am1*(alpha+1)))
  643. # where Am1 = alpha -1
  644. # if alpha <= alpha_crossover:
  645. if not mpf_sub(alpha_crossover, alpha, wp)[0]:
  646. c1 = mpf_div(b2, mpf_add(r, ap, wp), wp)
  647. # case a < 1
  648. if mpf_neg(am)[0]:
  649. # Am1 = (1/2) * (b*b/(r + (a+1)) + b*b/(s + (1-a))
  650. c2 = mpf_add(s, am, wp)
  651. c2 = mpf_div(b2, c2, wp)
  652. Am1 = mpf_shift(mpf_add(c1, c2, wp), -1)
  653. else:
  654. # Am1 = (1/2) * (b*b/(r + (a+1)) + (s - (1-a)))
  655. c2 = mpf_sub(s, am, wp)
  656. Am1 = mpf_shift(mpf_add(c1, c2, wp), -1)
  657. # im = log(1 + Am1 + sqrt(Am1*(alpha+1)))
  658. im = mpf_mul(Am1, mpf_add(alpha, fone, wp), wp)
  659. im = mpf_log(mpf_add(fone, mpf_add(Am1, mpf_sqrt(im, wp), wp), wp), wp)
  660. else:
  661. # im = log(alpha + sqrt(alpha*alpha - 1))
  662. im = mpf_sqrt(mpf_sub(mpf_mul(alpha, alpha, wp), fone, wp), wp)
  663. im = mpf_log(mpf_add(alpha, im, wp), wp)
  664. if asign:
  665. if n == 0:
  666. re = mpf_sub(mpf_pi(wp), re, wp)
  667. else:
  668. re = mpf_neg(re)
  669. if not bsign and n == 0:
  670. im = mpf_neg(im)
  671. if bsign and n == 1:
  672. im = mpf_neg(im)
  673. re = normalize(re[0], re[1], re[2], re[3], prec, rnd)
  674. im = normalize(im[0], im[1], im[2], im[3], prec, rnd)
  675. return re, im
  676. def mpc_acos(z, prec, rnd=round_fast):
  677. return acos_asin(z, prec, rnd, 0)
  678. def mpc_asin(z, prec, rnd=round_fast):
  679. return acos_asin(z, prec, rnd, 1)
  680. def mpc_asinh(z, prec, rnd=round_fast):
  681. # asinh(z) = I * asin(-I z)
  682. a, b = z
  683. a, b = mpc_asin((b, mpf_neg(a)), prec, rnd)
  684. return mpf_neg(b), a
  685. def mpc_acosh(z, prec, rnd=round_fast):
  686. # acosh(z) = -I * acos(z) for Im(acos(z)) <= 0
  687. # +I * acos(z) otherwise
  688. a, b = mpc_acos(z, prec, rnd)
  689. if b[0] or b == fzero:
  690. return mpf_neg(b), a
  691. else:
  692. return b, mpf_neg(a)
  693. def mpc_atanh(z, prec, rnd=round_fast):
  694. # atanh(z) = (log(1+z)-log(1-z))/2
  695. wp = prec + 15
  696. a = mpc_add(z, mpc_one, wp)
  697. b = mpc_sub(mpc_one, z, wp)
  698. a = mpc_log(a, wp)
  699. b = mpc_log(b, wp)
  700. v = mpc_shift(mpc_sub(a, b, wp), -1)
  701. # Subtraction at infinity gives correct imaginary part but
  702. # wrong real part (should be zero)
  703. if v[0] == fnan and mpc_is_inf(z):
  704. v = (fzero, v[1])
  705. return v
  706. def mpc_fibonacci(z, prec, rnd=round_fast):
  707. re, im = z
  708. if im == fzero:
  709. return (mpf_fibonacci(re, prec, rnd), fzero)
  710. size = max(abs(re[2]+re[3]), abs(re[2]+re[3]))
  711. wp = prec + size + 20
  712. a = mpf_phi(wp)
  713. b = mpf_add(mpf_shift(a, 1), fnone, wp)
  714. u = mpc_pow((a, fzero), z, wp)
  715. v = mpc_cos_pi(z, wp)
  716. v = mpc_div(v, u, wp)
  717. u = mpc_sub(u, v, wp)
  718. u = mpc_div_mpf(u, b, prec, rnd)
  719. return u
  720. def mpf_expj(x, prec, rnd='f'):
  721. raise ComplexResult
  722. def mpc_expj(z, prec, rnd='f'):
  723. re, im = z
  724. if im == fzero:
  725. return mpf_cos_sin(re, prec, rnd)
  726. if re == fzero:
  727. return mpf_exp(mpf_neg(im), prec, rnd), fzero
  728. ey = mpf_exp(mpf_neg(im), prec+10)
  729. c, s = mpf_cos_sin(re, prec+10)
  730. re = mpf_mul(ey, c, prec, rnd)
  731. im = mpf_mul(ey, s, prec, rnd)
  732. return re, im
  733. def mpf_expjpi(x, prec, rnd='f'):
  734. raise ComplexResult
  735. def mpc_expjpi(z, prec, rnd='f'):
  736. re, im = z
  737. if im == fzero:
  738. return mpf_cos_sin_pi(re, prec, rnd)
  739. sign, man, exp, bc = im
  740. wp = prec+10
  741. if man:
  742. wp += max(0, exp+bc)
  743. im = mpf_neg(mpf_mul(mpf_pi(wp), im, wp))
  744. if re == fzero:
  745. return mpf_exp(im, prec, rnd), fzero
  746. ey = mpf_exp(im, prec+10)
  747. c, s = mpf_cos_sin_pi(re, prec+10)
  748. re = mpf_mul(ey, c, prec, rnd)
  749. im = mpf_mul(ey, s, prec, rnd)
  750. return re, im
  751. if BACKEND == 'sage':
  752. try:
  753. import sage.libs.mpmath.ext_libmp as _lbmp
  754. mpc_exp = _lbmp.mpc_exp
  755. mpc_sqrt = _lbmp.mpc_sqrt
  756. except (ImportError, AttributeError):
  757. print("Warning: Sage imports in libmpc failed")