1"""
2Wrapper functions to more user-friendly calling of certain math functions
3whose output data-type is different than the input data-type in certain
4domains of the input.
5
6For example, for functions like `log` with branch cuts, the versions in this
7module provide the mathematically valid answers in the complex plane::
8
9 >>> import math
10 >>> np.emath.log(-math.exp(1)) == (1+1j*math.pi)
11 True
12
13Similarly, `sqrt`, other base logarithms, `power` and trig functions are
14correctly handled. See their respective docstrings for specific examples.
15
16"""
17import numpy._core.numeric as nx
18import numpy._core.numerictypes as nt
19from numpy._core.numeric import asarray, any
20from numpy._core.overrides import array_function_dispatch, set_module
21from numpy.lib._type_check_impl import isreal
22
23
24__all__ = [
25 'sqrt', 'log', 'log2', 'logn', 'log10', 'power', 'arccos', 'arcsin',
26 'arctanh'
27 ]
28
29
30_ln2 = nx.log(2.0)
31
32
33def _tocomplex(arr):
34 """Convert its input `arr` to a complex array.
35
36 The input is returned as a complex array of the smallest type that will fit
37 the original data: types like single, byte, short, etc. become csingle,
38 while others become cdouble.
39
40 A copy of the input is always made.
41
42 Parameters
43 ----------
44 arr : array
45
46 Returns
47 -------
48 array
49 An array with the same input data as the input but in complex form.
50
51 Examples
52 --------
53 >>> import numpy as np
54
55 First, consider an input of type short:
56
57 >>> a = np.array([1,2,3],np.short)
58
59 >>> ac = np.lib.scimath._tocomplex(a); ac
60 array([1.+0.j, 2.+0.j, 3.+0.j], dtype=complex64)
61
62 >>> ac.dtype
63 dtype('complex64')
64
65 If the input is of type double, the output is correspondingly of the
66 complex double type as well:
67
68 >>> b = np.array([1,2,3],np.double)
69
70 >>> bc = np.lib.scimath._tocomplex(b); bc
71 array([1.+0.j, 2.+0.j, 3.+0.j])
72
73 >>> bc.dtype
74 dtype('complex128')
75
76 Note that even if the input was complex to begin with, a copy is still
77 made, since the astype() method always copies:
78
79 >>> c = np.array([1,2,3],np.csingle)
80
81 >>> cc = np.lib.scimath._tocomplex(c); cc
82 array([1.+0.j, 2.+0.j, 3.+0.j], dtype=complex64)
83
84 >>> c *= 2; c
85 array([2.+0.j, 4.+0.j, 6.+0.j], dtype=complex64)
86
87 >>> cc
88 array([1.+0.j, 2.+0.j, 3.+0.j], dtype=complex64)
89 """
90 if issubclass(arr.dtype.type, (nt.single, nt.byte, nt.short, nt.ubyte,
91 nt.ushort, nt.csingle)):
92 return arr.astype(nt.csingle)
93 else:
94 return arr.astype(nt.cdouble)
95
96
97def _fix_real_lt_zero(x):
98 """Convert `x` to complex if it has real, negative components.
99
100 Otherwise, output is just the array version of the input (via asarray).
101
102 Parameters
103 ----------
104 x : array_like
105
106 Returns
107 -------
108 array
109
110 Examples
111 --------
112 >>> import numpy as np
113 >>> np.lib.scimath._fix_real_lt_zero([1,2])
114 array([1, 2])
115
116 >>> np.lib.scimath._fix_real_lt_zero([-1,2])
117 array([-1.+0.j, 2.+0.j])
118
119 """
120 x = asarray(x)
121 if any(isreal(x) & (x < 0)):
122 x = _tocomplex(x)
123 return x
124
125
126def _fix_int_lt_zero(x):
127 """Convert `x` to double if it has real, negative components.
128
129 Otherwise, output is just the array version of the input (via asarray).
130
131 Parameters
132 ----------
133 x : array_like
134
135 Returns
136 -------
137 array
138
139 Examples
140 --------
141 >>> import numpy as np
142 >>> np.lib.scimath._fix_int_lt_zero([1,2])
143 array([1, 2])
144
145 >>> np.lib.scimath._fix_int_lt_zero([-1,2])
146 array([-1., 2.])
147 """
148 x = asarray(x)
149 if any(isreal(x) & (x < 0)):
150 x = x * 1.0
151 return x
152
153
154def _fix_real_abs_gt_1(x):
155 """Convert `x` to complex if it has real components x_i with abs(x_i)>1.
156
157 Otherwise, output is just the array version of the input (via asarray).
158
159 Parameters
160 ----------
161 x : array_like
162
163 Returns
164 -------
165 array
166
167 Examples
168 --------
169 >>> import numpy as np
170 >>> np.lib.scimath._fix_real_abs_gt_1([0,1])
171 array([0, 1])
172
173 >>> np.lib.scimath._fix_real_abs_gt_1([0,2])
174 array([0.+0.j, 2.+0.j])
175 """
176 x = asarray(x)
177 if any(isreal(x) & (abs(x) > 1)):
178 x = _tocomplex(x)
179 return x
180
181
182def _unary_dispatcher(x):
183 return (x,)
184
185
186@set_module('numpy.lib.scimath')
187@array_function_dispatch(_unary_dispatcher)
188def sqrt(x):
189 """
190 Compute the square root of x.
191
192 For negative input elements, a complex value is returned
193 (unlike `numpy.sqrt` which returns NaN).
194
195 Parameters
196 ----------
197 x : array_like
198 The input value(s).
199
200 Returns
201 -------
202 out : ndarray or scalar
203 The square root of `x`. If `x` was a scalar, so is `out`,
204 otherwise an array is returned.
205
206 See Also
207 --------
208 numpy.sqrt
209
210 Examples
211 --------
212 For real, non-negative inputs this works just like `numpy.sqrt`:
213
214 >>> import numpy as np
215
216 >>> np.emath.sqrt(1)
217 1.0
218 >>> np.emath.sqrt([1, 4])
219 array([1., 2.])
220
221 But it automatically handles negative inputs:
222
223 >>> np.emath.sqrt(-1)
224 1j
225 >>> np.emath.sqrt([-1,4])
226 array([0.+1.j, 2.+0.j])
227
228 Different results are expected because:
229 floating point 0.0 and -0.0 are distinct.
230
231 For more control, explicitly use complex() as follows:
232
233 >>> np.emath.sqrt(complex(-4.0, 0.0))
234 2j
235 >>> np.emath.sqrt(complex(-4.0, -0.0))
236 -2j
237 """
238 x = _fix_real_lt_zero(x)
239 return nx.sqrt(x)
240
241
242@set_module('numpy.lib.scimath')
243@array_function_dispatch(_unary_dispatcher)
244def log(x):
245 """
246 Compute the natural logarithm of `x`.
247
248 Return the "principal value" (for a description of this, see `numpy.log`)
249 of :math:`log_e(x)`. For real `x > 0`, this is a real number (``log(0)``
250 returns ``-inf`` and ``log(np.inf)`` returns ``inf``). Otherwise, the
251 complex principle value is returned.
252
253 Parameters
254 ----------
255 x : array_like
256 The value(s) whose log is (are) required.
257
258 Returns
259 -------
260 out : ndarray or scalar
261 The log of the `x` value(s). If `x` was a scalar, so is `out`,
262 otherwise an array is returned.
263
264 See Also
265 --------
266 numpy.log
267
268 Notes
269 -----
270 For a log() that returns ``NAN`` when real `x < 0`, use `numpy.log`
271 (note, however, that otherwise `numpy.log` and this `log` are identical,
272 i.e., both return ``-inf`` for `x = 0`, ``inf`` for `x = inf`, and,
273 notably, the complex principle value if ``x.imag != 0``).
274
275 Examples
276 --------
277 >>> import numpy as np
278 >>> np.emath.log(np.exp(1))
279 1.0
280
281 Negative arguments are handled "correctly" (recall that
282 ``exp(log(x)) == x`` does *not* hold for real ``x < 0``):
283
284 >>> np.emath.log(-np.exp(1)) == (1 + np.pi * 1j)
285 True
286
287 """
288 x = _fix_real_lt_zero(x)
289 return nx.log(x)
290
291
292@set_module('numpy.lib.scimath')
293@array_function_dispatch(_unary_dispatcher)
294def log10(x):
295 """
296 Compute the logarithm base 10 of `x`.
297
298 Return the "principal value" (for a description of this, see
299 `numpy.log10`) of :math:`log_{10}(x)`. For real `x > 0`, this
300 is a real number (``log10(0)`` returns ``-inf`` and ``log10(np.inf)``
301 returns ``inf``). Otherwise, the complex principle value is returned.
302
303 Parameters
304 ----------
305 x : array_like or scalar
306 The value(s) whose log base 10 is (are) required.
307
308 Returns
309 -------
310 out : ndarray or scalar
311 The log base 10 of the `x` value(s). If `x` was a scalar, so is `out`,
312 otherwise an array object is returned.
313
314 See Also
315 --------
316 numpy.log10
317
318 Notes
319 -----
320 For a log10() that returns ``NAN`` when real `x < 0`, use `numpy.log10`
321 (note, however, that otherwise `numpy.log10` and this `log10` are
322 identical, i.e., both return ``-inf`` for `x = 0`, ``inf`` for `x = inf`,
323 and, notably, the complex principle value if ``x.imag != 0``).
324
325 Examples
326 --------
327 >>> import numpy as np
328
329 (We set the printing precision so the example can be auto-tested)
330
331 >>> np.set_printoptions(precision=4)
332
333 >>> np.emath.log10(10**1)
334 1.0
335
336 >>> np.emath.log10([-10**1, -10**2, 10**2])
337 array([1.+1.3644j, 2.+1.3644j, 2.+0.j ])
338
339 """
340 x = _fix_real_lt_zero(x)
341 return nx.log10(x)
342
343
344def _logn_dispatcher(n, x):
345 return (n, x,)
346
347
348@set_module('numpy.lib.scimath')
349@array_function_dispatch(_logn_dispatcher)
350def logn(n, x):
351 """
352 Take log base n of x.
353
354 If `x` contains negative inputs, the answer is computed and returned in the
355 complex domain.
356
357 Parameters
358 ----------
359 n : array_like
360 The integer base(s) in which the log is taken.
361 x : array_like
362 The value(s) whose log base `n` is (are) required.
363
364 Returns
365 -------
366 out : ndarray or scalar
367 The log base `n` of the `x` value(s). If `x` was a scalar, so is
368 `out`, otherwise an array is returned.
369
370 Examples
371 --------
372 >>> import numpy as np
373 >>> np.set_printoptions(precision=4)
374
375 >>> np.emath.logn(2, [4, 8])
376 array([2., 3.])
377 >>> np.emath.logn(2, [-4, -8, 8])
378 array([2.+4.5324j, 3.+4.5324j, 3.+0.j ])
379
380 """
381 x = _fix_real_lt_zero(x)
382 n = _fix_real_lt_zero(n)
383 return nx.log(x)/nx.log(n)
384
385
386@set_module('numpy.lib.scimath')
387@array_function_dispatch(_unary_dispatcher)
388def log2(x):
389 """
390 Compute the logarithm base 2 of `x`.
391
392 Return the "principal value" (for a description of this, see
393 `numpy.log2`) of :math:`log_2(x)`. For real `x > 0`, this is
394 a real number (``log2(0)`` returns ``-inf`` and ``log2(np.inf)`` returns
395 ``inf``). Otherwise, the complex principle value is returned.
396
397 Parameters
398 ----------
399 x : array_like
400 The value(s) whose log base 2 is (are) required.
401
402 Returns
403 -------
404 out : ndarray or scalar
405 The log base 2 of the `x` value(s). If `x` was a scalar, so is `out`,
406 otherwise an array is returned.
407
408 See Also
409 --------
410 numpy.log2
411
412 Notes
413 -----
414 For a log2() that returns ``NAN`` when real `x < 0`, use `numpy.log2`
415 (note, however, that otherwise `numpy.log2` and this `log2` are
416 identical, i.e., both return ``-inf`` for `x = 0`, ``inf`` for `x = inf`,
417 and, notably, the complex principle value if ``x.imag != 0``).
418
419 Examples
420 --------
421
422 We set the printing precision so the example can be auto-tested:
423
424 >>> np.set_printoptions(precision=4)
425
426 >>> np.emath.log2(8)
427 3.0
428 >>> np.emath.log2([-4, -8, 8])
429 array([2.+4.5324j, 3.+4.5324j, 3.+0.j ])
430
431 """
432 x = _fix_real_lt_zero(x)
433 return nx.log2(x)
434
435
436def _power_dispatcher(x, p):
437 return (x, p)
438
439
440@set_module('numpy.lib.scimath')
441@array_function_dispatch(_power_dispatcher)
442def power(x, p):
443 """
444 Return x to the power p, (x**p).
445
446 If `x` contains negative values, the output is converted to the
447 complex domain.
448
449 Parameters
450 ----------
451 x : array_like
452 The input value(s).
453 p : array_like of ints
454 The power(s) to which `x` is raised. If `x` contains multiple values,
455 `p` has to either be a scalar, or contain the same number of values
456 as `x`. In the latter case, the result is
457 ``x[0]**p[0], x[1]**p[1], ...``.
458
459 Returns
460 -------
461 out : ndarray or scalar
462 The result of ``x**p``. If `x` and `p` are scalars, so is `out`,
463 otherwise an array is returned.
464
465 See Also
466 --------
467 numpy.power
468
469 Examples
470 --------
471 >>> import numpy as np
472 >>> np.set_printoptions(precision=4)
473
474 >>> np.emath.power(2, 2)
475 4
476
477 >>> np.emath.power([2, 4], 2)
478 array([ 4, 16])
479
480 >>> np.emath.power([2, 4], -2)
481 array([0.25 , 0.0625])
482
483 >>> np.emath.power([-2, 4], 2)
484 array([ 4.-0.j, 16.+0.j])
485
486 >>> np.emath.power([2, 4], [2, 4])
487 array([ 4, 256])
488
489 """
490 x = _fix_real_lt_zero(x)
491 p = _fix_int_lt_zero(p)
492 return nx.power(x, p)
493
494
495@set_module('numpy.lib.scimath')
496@array_function_dispatch(_unary_dispatcher)
497def arccos(x):
498 """
499 Compute the inverse cosine of x.
500
501 Return the "principal value" (for a description of this, see
502 `numpy.arccos`) of the inverse cosine of `x`. For real `x` such that
503 `abs(x) <= 1`, this is a real number in the closed interval
504 :math:`[0, \\pi]`. Otherwise, the complex principle value is returned.
505
506 Parameters
507 ----------
508 x : array_like or scalar
509 The value(s) whose arccos is (are) required.
510
511 Returns
512 -------
513 out : ndarray or scalar
514 The inverse cosine(s) of the `x` value(s). If `x` was a scalar, so
515 is `out`, otherwise an array object is returned.
516
517 See Also
518 --------
519 numpy.arccos
520
521 Notes
522 -----
523 For an arccos() that returns ``NAN`` when real `x` is not in the
524 interval ``[-1,1]``, use `numpy.arccos`.
525
526 Examples
527 --------
528 >>> import numpy as np
529 >>> np.set_printoptions(precision=4)
530
531 >>> np.emath.arccos(1) # a scalar is returned
532 0.0
533
534 >>> np.emath.arccos([1,2])
535 array([0.-0.j , 0.-1.317j])
536
537 """
538 x = _fix_real_abs_gt_1(x)
539 return nx.arccos(x)
540
541
542@set_module('numpy.lib.scimath')
543@array_function_dispatch(_unary_dispatcher)
544def arcsin(x):
545 """
546 Compute the inverse sine of x.
547
548 Return the "principal value" (for a description of this, see
549 `numpy.arcsin`) of the inverse sine of `x`. For real `x` such that
550 `abs(x) <= 1`, this is a real number in the closed interval
551 :math:`[-\\pi/2, \\pi/2]`. Otherwise, the complex principle value is
552 returned.
553
554 Parameters
555 ----------
556 x : array_like or scalar
557 The value(s) whose arcsin is (are) required.
558
559 Returns
560 -------
561 out : ndarray or scalar
562 The inverse sine(s) of the `x` value(s). If `x` was a scalar, so
563 is `out`, otherwise an array object is returned.
564
565 See Also
566 --------
567 numpy.arcsin
568
569 Notes
570 -----
571 For an arcsin() that returns ``NAN`` when real `x` is not in the
572 interval ``[-1,1]``, use `numpy.arcsin`.
573
574 Examples
575 --------
576 >>> import numpy as np
577 >>> np.set_printoptions(precision=4)
578
579 >>> np.emath.arcsin(0)
580 0.0
581
582 >>> np.emath.arcsin([0,1])
583 array([0. , 1.5708])
584
585 """
586 x = _fix_real_abs_gt_1(x)
587 return nx.arcsin(x)
588
589
590@set_module('numpy.lib.scimath')
591@array_function_dispatch(_unary_dispatcher)
592def arctanh(x):
593 """
594 Compute the inverse hyperbolic tangent of `x`.
595
596 Return the "principal value" (for a description of this, see
597 `numpy.arctanh`) of ``arctanh(x)``. For real `x` such that
598 ``abs(x) < 1``, this is a real number. If `abs(x) > 1`, or if `x` is
599 complex, the result is complex. Finally, `x = 1` returns``inf`` and
600 ``x=-1`` returns ``-inf``.
601
602 Parameters
603 ----------
604 x : array_like
605 The value(s) whose arctanh is (are) required.
606
607 Returns
608 -------
609 out : ndarray or scalar
610 The inverse hyperbolic tangent(s) of the `x` value(s). If `x` was
611 a scalar so is `out`, otherwise an array is returned.
612
613
614 See Also
615 --------
616 numpy.arctanh
617
618 Notes
619 -----
620 For an arctanh() that returns ``NAN`` when real `x` is not in the
621 interval ``(-1,1)``, use `numpy.arctanh` (this latter, however, does
622 return +/-inf for ``x = +/-1``).
623
624 Examples
625 --------
626 >>> import numpy as np
627 >>> np.set_printoptions(precision=4)
628
629 >>> np.emath.arctanh(0.5)
630 0.5493061443340549
631
632 >>> from numpy.testing import suppress_warnings
633 >>> with suppress_warnings() as sup:
634 ... sup.filter(RuntimeWarning)
635 ... np.emath.arctanh(np.eye(2))
636 array([[inf, 0.],
637 [ 0., inf]])
638 >>> np.emath.arctanh([1j])
639 array([0.+0.7854j])
640
641 """
642 x = _fix_real_abs_gt_1(x)
643 return nx.arctanh(x)