Coverage for /usr/lib/python3/dist-packages/mpmath/libmp/libmpc.py: 12%
592 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""
2Low-level functions for complex arithmetic.
3"""
5import sys
7from .backend import MPZ, MPZ_ZERO, MPZ_ONE, MPZ_TWO, BACKEND
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)
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)
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)
38_infs = (finf, fninf)
39_infs_nan = (finf, fninf, fnan)
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
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
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"
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))
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)
80def mpc_conjugate(z, prec, rnd=round_fast):
81 re, im = z
82 return re, mpf_neg(im, prec, rnd)
84def mpc_is_nonzero(z):
85 return z != mpc_zero
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)
92def mpc_add_mpf(z, x, prec, rnd=round_fast):
93 a, b = z
94 return mpf_add(a, x, prec, rnd), b
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)
101def mpc_sub_mpf(z, p, prec=0, rnd=round_fast):
102 a, b = z
103 return mpf_sub(a, p, prec, rnd), b
105def mpc_pos(z, prec, rnd=round_fast):
106 a, b = z
107 return mpf_pos(a, prec, rnd), mpf_pos(b, prec, rnd)
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)
113def mpc_shift(z, n):
114 a, b = z
115 return mpf_shift(a, n), mpf_shift(b, n)
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)
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)
128def mpc_floor(z, prec, rnd=round_fast):
129 a, b = z
130 return mpf_floor(a, prec, rnd), mpf_floor(b, prec, rnd)
132def mpc_ceil(z, prec, rnd=round_fast):
133 a, b = z
134 return mpf_ceil(a, prec, rnd), mpf_ceil(b, prec, rnd)
136def mpc_nint(z, prec, rnd=round_fast):
137 a, b = z
138 return mpf_nint(a, prec, rnd), mpf_nint(b, prec, rnd)
140def mpc_frac(z, prec, rnd=round_fast):
141 a, b = z
142 return mpf_frac(a, prec, rnd), mpf_frac(b, prec, rnd)
145def mpc_mul(z, w, prec, rnd=round_fast):
146 """
147 Complex multiplication.
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
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
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
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
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
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)
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
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
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
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
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)
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)
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)
295def mpc_sqrt(z, prec, rnd=round_fast):
296 """Complex square root (principal branch).
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
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
370def mpc_nthroot(z, n, prec, rnd=round_fast):
371 """
372 Complex n-th root.
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
411def mpc_cbrt(z, prec, rnd=round_fast):
412 """
413 Complex cubic root.
414 """
415 return mpc_nthroot(z, 3, prec, rnd)
417def mpc_exp(z, prec, rnd=round_fast):
418 """
419 Complex exponential function.
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.
426 The formula is efficient since mpmath's real exp is quite fast and
427 since we can compute cos and sin simultaneously.
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
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
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.
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)
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
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
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)
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
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)
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)
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)
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
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
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
604beta_crossover = from_float(0.6417)
605alpha_crossover = from_float(1.5)
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
743def mpc_acos(z, prec, rnd=round_fast):
744 return acos_asin(z, prec, rnd, 0)
746def mpc_asin(z, prec, rnd=round_fast):
747 return acos_asin(z, prec, rnd, 1)
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
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)
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
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
793def mpf_expj(x, prec, rnd='f'):
794 raise ComplexResult
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
808def mpf_expjpi(x, prec, rnd='f'):
809 raise ComplexResult
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
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")