scimath.py 15 KB

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