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