scimath.py 15 KB

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