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