function_base.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546
  1. import functools
  2. import warnings
  3. import operator
  4. import types
  5. import numpy as np
  6. from . import numeric as _nx
  7. from .numeric import result_type, nan, asanyarray, ndim
  8. from numpy._core.multiarray import add_docstring
  9. from numpy._core._multiarray_umath import _array_converter
  10. from numpy._core import overrides
  11. __all__ = ['logspace', 'linspace', 'geomspace']
  12. array_function_dispatch = functools.partial(
  13. overrides.array_function_dispatch, module='numpy')
  14. def _linspace_dispatcher(start, stop, num=None, endpoint=None, retstep=None,
  15. dtype=None, axis=None, *, device=None):
  16. return (start, stop)
  17. @array_function_dispatch(_linspace_dispatcher)
  18. def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
  19. axis=0, *, device=None):
  20. """
  21. Return evenly spaced numbers over a specified interval.
  22. Returns `num` evenly spaced samples, calculated over the
  23. interval [`start`, `stop`].
  24. The endpoint of the interval can optionally be excluded.
  25. .. versionchanged:: 1.20.0
  26. Values are rounded towards ``-inf`` instead of ``0`` when an
  27. integer ``dtype`` is specified. The old behavior can
  28. still be obtained with ``np.linspace(start, stop, num).astype(int)``
  29. Parameters
  30. ----------
  31. start : array_like
  32. The starting value of the sequence.
  33. stop : array_like
  34. The end value of the sequence, unless `endpoint` is set to False.
  35. In that case, the sequence consists of all but the last of ``num + 1``
  36. evenly spaced samples, so that `stop` is excluded. Note that the step
  37. size changes when `endpoint` is False.
  38. num : int, optional
  39. Number of samples to generate. Default is 50. Must be non-negative.
  40. endpoint : bool, optional
  41. If True, `stop` is the last sample. Otherwise, it is not included.
  42. Default is True.
  43. retstep : bool, optional
  44. If True, return (`samples`, `step`), where `step` is the spacing
  45. between samples.
  46. dtype : dtype, optional
  47. The type of the output array. If `dtype` is not given, the data type
  48. is inferred from `start` and `stop`. The inferred dtype will never be
  49. an integer; `float` is chosen even if the arguments would produce an
  50. array of integers.
  51. axis : int, optional
  52. The axis in the result to store the samples. Relevant only if start
  53. or stop are array-like. By default (0), the samples will be along a
  54. new axis inserted at the beginning. Use -1 to get an axis at the end.
  55. device : str, optional
  56. The device on which to place the created array. Default: None.
  57. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  58. .. versionadded:: 2.0.0
  59. Returns
  60. -------
  61. samples : ndarray
  62. There are `num` equally spaced samples in the closed interval
  63. ``[start, stop]`` or the half-open interval ``[start, stop)``
  64. (depending on whether `endpoint` is True or False).
  65. step : float, optional
  66. Only returned if `retstep` is True
  67. Size of spacing between samples.
  68. See Also
  69. --------
  70. arange : Similar to `linspace`, but uses a step size (instead of the
  71. number of samples).
  72. geomspace : Similar to `linspace`, but with numbers spaced evenly on a log
  73. scale (a geometric progression).
  74. logspace : Similar to `geomspace`, but with the end points specified as
  75. logarithms.
  76. :ref:`how-to-partition`
  77. Examples
  78. --------
  79. >>> import numpy as np
  80. >>> np.linspace(2.0, 3.0, num=5)
  81. array([2. , 2.25, 2.5 , 2.75, 3. ])
  82. >>> np.linspace(2.0, 3.0, num=5, endpoint=False)
  83. array([2. , 2.2, 2.4, 2.6, 2.8])
  84. >>> np.linspace(2.0, 3.0, num=5, retstep=True)
  85. (array([2. , 2.25, 2.5 , 2.75, 3. ]), 0.25)
  86. Graphical illustration:
  87. >>> import matplotlib.pyplot as plt
  88. >>> N = 8
  89. >>> y = np.zeros(N)
  90. >>> x1 = np.linspace(0, 10, N, endpoint=True)
  91. >>> x2 = np.linspace(0, 10, N, endpoint=False)
  92. >>> plt.plot(x1, y, 'o')
  93. [<matplotlib.lines.Line2D object at 0x...>]
  94. >>> plt.plot(x2, y + 0.5, 'o')
  95. [<matplotlib.lines.Line2D object at 0x...>]
  96. >>> plt.ylim([-0.5, 1])
  97. (-0.5, 1)
  98. >>> plt.show()
  99. """
  100. num = operator.index(num)
  101. if num < 0:
  102. raise ValueError(
  103. "Number of samples, %s, must be non-negative." % num
  104. )
  105. div = (num - 1) if endpoint else num
  106. conv = _array_converter(start, stop)
  107. start, stop = conv.as_arrays()
  108. dt = conv.result_type(ensure_inexact=True)
  109. if dtype is None:
  110. dtype = dt
  111. integer_dtype = False
  112. else:
  113. integer_dtype = _nx.issubdtype(dtype, _nx.integer)
  114. # Use `dtype=type(dt)` to enforce a floating point evaluation:
  115. delta = np.subtract(stop, start, dtype=type(dt))
  116. y = _nx.arange(
  117. 0, num, dtype=dt, device=device
  118. ).reshape((-1,) + (1,) * ndim(delta))
  119. # In-place multiplication y *= delta/div is faster, but prevents
  120. # the multiplicant from overriding what class is produced, and thus
  121. # prevents, e.g. use of Quantities, see gh-7142. Hence, we multiply
  122. # in place only for standard scalar types.
  123. if div > 0:
  124. _mult_inplace = _nx.isscalar(delta)
  125. step = delta / div
  126. any_step_zero = (
  127. step == 0 if _mult_inplace else _nx.asanyarray(step == 0).any())
  128. if any_step_zero:
  129. # Special handling for denormal numbers, gh-5437
  130. y /= div
  131. if _mult_inplace:
  132. y *= delta
  133. else:
  134. y = y * delta
  135. else:
  136. if _mult_inplace:
  137. y *= step
  138. else:
  139. y = y * step
  140. else:
  141. # sequences with 0 items or 1 item with endpoint=True (i.e. div <= 0)
  142. # have an undefined step
  143. step = nan
  144. # Multiply with delta to allow possible override of output class.
  145. y = y * delta
  146. y += start
  147. if endpoint and num > 1:
  148. y[-1, ...] = stop
  149. if axis != 0:
  150. y = _nx.moveaxis(y, 0, axis)
  151. if integer_dtype:
  152. _nx.floor(y, out=y)
  153. y = conv.wrap(y.astype(dtype, copy=False))
  154. if retstep:
  155. return y, step
  156. else:
  157. return y
  158. def _logspace_dispatcher(start, stop, num=None, endpoint=None, base=None,
  159. dtype=None, axis=None):
  160. return (start, stop, base)
  161. @array_function_dispatch(_logspace_dispatcher)
  162. def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None,
  163. axis=0):
  164. """
  165. Return numbers spaced evenly on a log scale.
  166. In linear space, the sequence starts at ``base ** start``
  167. (`base` to the power of `start`) and ends with ``base ** stop``
  168. (see `endpoint` below).
  169. .. versionchanged:: 1.25.0
  170. Non-scalar 'base` is now supported
  171. Parameters
  172. ----------
  173. start : array_like
  174. ``base ** start`` is the starting value of the sequence.
  175. stop : array_like
  176. ``base ** stop`` is the final value of the sequence, unless `endpoint`
  177. is False. In that case, ``num + 1`` values are spaced over the
  178. interval in log-space, of which all but the last (a sequence of
  179. length `num`) are returned.
  180. num : integer, optional
  181. Number of samples to generate. Default is 50.
  182. endpoint : boolean, optional
  183. If true, `stop` is the last sample. Otherwise, it is not included.
  184. Default is True.
  185. base : array_like, optional
  186. The base of the log space. The step size between the elements in
  187. ``ln(samples) / ln(base)`` (or ``log_base(samples)``) is uniform.
  188. Default is 10.0.
  189. dtype : dtype
  190. The type of the output array. If `dtype` is not given, the data type
  191. is inferred from `start` and `stop`. The inferred type will never be
  192. an integer; `float` is chosen even if the arguments would produce an
  193. array of integers.
  194. axis : int, optional
  195. The axis in the result to store the samples. Relevant only if start,
  196. stop, or base are array-like. By default (0), the samples will be
  197. along a new axis inserted at the beginning. Use -1 to get an axis at
  198. the end.
  199. Returns
  200. -------
  201. samples : ndarray
  202. `num` samples, equally spaced on a log scale.
  203. See Also
  204. --------
  205. arange : Similar to linspace, with the step size specified instead of the
  206. number of samples. Note that, when used with a float endpoint, the
  207. endpoint may or may not be included.
  208. linspace : Similar to logspace, but with the samples uniformly distributed
  209. in linear space, instead of log space.
  210. geomspace : Similar to logspace, but with endpoints specified directly.
  211. :ref:`how-to-partition`
  212. Notes
  213. -----
  214. If base is a scalar, logspace is equivalent to the code
  215. >>> y = np.linspace(start, stop, num=num, endpoint=endpoint)
  216. ... # doctest: +SKIP
  217. >>> power(base, y).astype(dtype)
  218. ... # doctest: +SKIP
  219. Examples
  220. --------
  221. >>> import numpy as np
  222. >>> np.logspace(2.0, 3.0, num=4)
  223. array([ 100. , 215.443469 , 464.15888336, 1000. ])
  224. >>> np.logspace(2.0, 3.0, num=4, endpoint=False)
  225. array([100. , 177.827941 , 316.22776602, 562.34132519])
  226. >>> np.logspace(2.0, 3.0, num=4, base=2.0)
  227. array([4. , 5.0396842 , 6.34960421, 8. ])
  228. >>> np.logspace(2.0, 3.0, num=4, base=[2.0, 3.0], axis=-1)
  229. array([[ 4. , 5.0396842 , 6.34960421, 8. ],
  230. [ 9. , 12.98024613, 18.72075441, 27. ]])
  231. Graphical illustration:
  232. >>> import matplotlib.pyplot as plt
  233. >>> N = 10
  234. >>> x1 = np.logspace(0.1, 1, N, endpoint=True)
  235. >>> x2 = np.logspace(0.1, 1, N, endpoint=False)
  236. >>> y = np.zeros(N)
  237. >>> plt.plot(x1, y, 'o')
  238. [<matplotlib.lines.Line2D object at 0x...>]
  239. >>> plt.plot(x2, y + 0.5, 'o')
  240. [<matplotlib.lines.Line2D object at 0x...>]
  241. >>> plt.ylim([-0.5, 1])
  242. (-0.5, 1)
  243. >>> plt.show()
  244. """
  245. if not isinstance(base, (float, int)) and np.ndim(base):
  246. # If base is non-scalar, broadcast it with the others, since it
  247. # may influence how axis is interpreted.
  248. ndmax = np.broadcast(start, stop, base).ndim
  249. start, stop, base = (
  250. np.array(a, copy=None, subok=True, ndmin=ndmax)
  251. for a in (start, stop, base)
  252. )
  253. base = np.expand_dims(base, axis=axis)
  254. y = linspace(start, stop, num=num, endpoint=endpoint, axis=axis)
  255. if dtype is None:
  256. return _nx.power(base, y)
  257. return _nx.power(base, y).astype(dtype, copy=False)
  258. def _geomspace_dispatcher(start, stop, num=None, endpoint=None, dtype=None,
  259. axis=None):
  260. return (start, stop)
  261. @array_function_dispatch(_geomspace_dispatcher)
  262. def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
  263. """
  264. Return numbers spaced evenly on a log scale (a geometric progression).
  265. This is similar to `logspace`, but with endpoints specified directly.
  266. Each output sample is a constant multiple of the previous.
  267. Parameters
  268. ----------
  269. start : array_like
  270. The starting value of the sequence.
  271. stop : array_like
  272. The final value of the sequence, unless `endpoint` is False.
  273. In that case, ``num + 1`` values are spaced over the
  274. interval in log-space, of which all but the last (a sequence of
  275. length `num`) are returned.
  276. num : integer, optional
  277. Number of samples to generate. Default is 50.
  278. endpoint : boolean, optional
  279. If true, `stop` is the last sample. Otherwise, it is not included.
  280. Default is True.
  281. dtype : dtype
  282. The type of the output array. If `dtype` is not given, the data type
  283. is inferred from `start` and `stop`. The inferred dtype will never be
  284. an integer; `float` is chosen even if the arguments would produce an
  285. array of integers.
  286. axis : int, optional
  287. The axis in the result to store the samples. Relevant only if start
  288. or stop are array-like. By default (0), the samples will be along a
  289. new axis inserted at the beginning. Use -1 to get an axis at the end.
  290. Returns
  291. -------
  292. samples : ndarray
  293. `num` samples, equally spaced on a log scale.
  294. See Also
  295. --------
  296. logspace : Similar to geomspace, but with endpoints specified using log
  297. and base.
  298. linspace : Similar to geomspace, but with arithmetic instead of geometric
  299. progression.
  300. arange : Similar to linspace, with the step size specified instead of the
  301. number of samples.
  302. :ref:`how-to-partition`
  303. Notes
  304. -----
  305. If the inputs or dtype are complex, the output will follow a logarithmic
  306. spiral in the complex plane. (There are an infinite number of spirals
  307. passing through two points; the output will follow the shortest such path.)
  308. Examples
  309. --------
  310. >>> import numpy as np
  311. >>> np.geomspace(1, 1000, num=4)
  312. array([ 1., 10., 100., 1000.])
  313. >>> np.geomspace(1, 1000, num=3, endpoint=False)
  314. array([ 1., 10., 100.])
  315. >>> np.geomspace(1, 1000, num=4, endpoint=False)
  316. array([ 1. , 5.62341325, 31.6227766 , 177.827941 ])
  317. >>> np.geomspace(1, 256, num=9)
  318. array([ 1., 2., 4., 8., 16., 32., 64., 128., 256.])
  319. Note that the above may not produce exact integers:
  320. >>> np.geomspace(1, 256, num=9, dtype=int)
  321. array([ 1, 2, 4, 7, 16, 32, 63, 127, 256])
  322. >>> np.around(np.geomspace(1, 256, num=9)).astype(int)
  323. array([ 1, 2, 4, 8, 16, 32, 64, 128, 256])
  324. Negative, decreasing, and complex inputs are allowed:
  325. >>> np.geomspace(1000, 1, num=4)
  326. array([1000., 100., 10., 1.])
  327. >>> np.geomspace(-1000, -1, num=4)
  328. array([-1000., -100., -10., -1.])
  329. >>> np.geomspace(1j, 1000j, num=4) # Straight line
  330. array([0. +1.j, 0. +10.j, 0. +100.j, 0.+1000.j])
  331. >>> np.geomspace(-1+0j, 1+0j, num=5) # Circle
  332. array([-1.00000000e+00+1.22464680e-16j, -7.07106781e-01+7.07106781e-01j,
  333. 6.12323400e-17+1.00000000e+00j, 7.07106781e-01+7.07106781e-01j,
  334. 1.00000000e+00+0.00000000e+00j])
  335. Graphical illustration of `endpoint` parameter:
  336. >>> import matplotlib.pyplot as plt
  337. >>> N = 10
  338. >>> y = np.zeros(N)
  339. >>> plt.semilogx(np.geomspace(1, 1000, N, endpoint=True), y + 1, 'o')
  340. [<matplotlib.lines.Line2D object at 0x...>]
  341. >>> plt.semilogx(np.geomspace(1, 1000, N, endpoint=False), y + 2, 'o')
  342. [<matplotlib.lines.Line2D object at 0x...>]
  343. >>> plt.axis([0.5, 2000, 0, 3])
  344. [0.5, 2000, 0, 3]
  345. >>> plt.grid(True, color='0.7', linestyle='-', which='both', axis='both')
  346. >>> plt.show()
  347. """
  348. start = asanyarray(start)
  349. stop = asanyarray(stop)
  350. if _nx.any(start == 0) or _nx.any(stop == 0):
  351. raise ValueError('Geometric sequence cannot include zero')
  352. dt = result_type(start, stop, float(num), _nx.zeros((), dtype))
  353. if dtype is None:
  354. dtype = dt
  355. else:
  356. # complex to dtype('complex128'), for instance
  357. dtype = _nx.dtype(dtype)
  358. # Promote both arguments to the same dtype in case, for instance, one is
  359. # complex and another is negative and log would produce NaN otherwise.
  360. # Copy since we may change things in-place further down.
  361. start = start.astype(dt, copy=True)
  362. stop = stop.astype(dt, copy=True)
  363. # Allow negative real values and ensure a consistent result for complex
  364. # (including avoiding negligible real or imaginary parts in output) by
  365. # rotating start to positive real, calculating, then undoing rotation.
  366. out_sign = _nx.sign(start)
  367. start /= out_sign
  368. stop = stop / out_sign
  369. log_start = _nx.log10(start)
  370. log_stop = _nx.log10(stop)
  371. result = logspace(log_start, log_stop, num=num,
  372. endpoint=endpoint, base=10.0, dtype=dt)
  373. # Make sure the endpoints match the start and stop arguments. This is
  374. # necessary because np.exp(np.log(x)) is not necessarily equal to x.
  375. if num > 0:
  376. result[0] = start
  377. if num > 1 and endpoint:
  378. result[-1] = stop
  379. result *= out_sign
  380. if axis != 0:
  381. result = _nx.moveaxis(result, 0, axis)
  382. return result.astype(dtype, copy=False)
  383. def _needs_add_docstring(obj):
  384. """
  385. Returns true if the only way to set the docstring of `obj` from python is
  386. via add_docstring.
  387. This function errs on the side of being overly conservative.
  388. """
  389. Py_TPFLAGS_HEAPTYPE = 1 << 9
  390. if isinstance(obj, (types.FunctionType, types.MethodType, property)):
  391. return False
  392. if isinstance(obj, type) and obj.__flags__ & Py_TPFLAGS_HEAPTYPE:
  393. return False
  394. return True
  395. def _add_docstring(obj, doc, warn_on_python):
  396. if warn_on_python and not _needs_add_docstring(obj):
  397. warnings.warn(
  398. "add_newdoc was used on a pure-python object {}. "
  399. "Prefer to attach it directly to the source."
  400. .format(obj),
  401. UserWarning,
  402. stacklevel=3)
  403. try:
  404. add_docstring(obj, doc)
  405. except Exception:
  406. pass
  407. def add_newdoc(place, obj, doc, warn_on_python=True):
  408. """
  409. Add documentation to an existing object, typically one defined in C
  410. The purpose is to allow easier editing of the docstrings without requiring
  411. a re-compile. This exists primarily for internal use within numpy itself.
  412. Parameters
  413. ----------
  414. place : str
  415. The absolute name of the module to import from
  416. obj : str or None
  417. The name of the object to add documentation to, typically a class or
  418. function name.
  419. doc : {str, Tuple[str, str], List[Tuple[str, str]]}
  420. If a string, the documentation to apply to `obj`
  421. If a tuple, then the first element is interpreted as an attribute
  422. of `obj` and the second as the docstring to apply -
  423. ``(method, docstring)``
  424. If a list, then each element of the list should be a tuple of length
  425. two - ``[(method1, docstring1), (method2, docstring2), ...]``
  426. warn_on_python : bool
  427. If True, the default, emit `UserWarning` if this is used to attach
  428. documentation to a pure-python object.
  429. Notes
  430. -----
  431. This routine never raises an error if the docstring can't be written, but
  432. will raise an error if the object being documented does not exist.
  433. This routine cannot modify read-only docstrings, as appear
  434. in new-style classes or built-in functions. Because this
  435. routine never raises an error the caller must check manually
  436. that the docstrings were changed.
  437. Since this function grabs the ``char *`` from a c-level str object and puts
  438. it into the ``tp_doc`` slot of the type of `obj`, it violates a number of
  439. C-API best-practices, by:
  440. - modifying a `PyTypeObject` after calling `PyType_Ready`
  441. - calling `Py_INCREF` on the str and losing the reference, so the str
  442. will never be released
  443. If possible it should be avoided.
  444. """
  445. new = getattr(__import__(place, globals(), {}, [obj]), obj)
  446. if isinstance(doc, str):
  447. if "${ARRAY_FUNCTION_LIKE}" in doc:
  448. doc = overrides.get_array_function_like_doc(new, doc)
  449. _add_docstring(new, doc.strip(), warn_on_python)
  450. elif isinstance(doc, tuple):
  451. attr, docstring = doc
  452. _add_docstring(getattr(new, attr), docstring.strip(), warn_on_python)
  453. elif isinstance(doc, list):
  454. for attr, docstring in doc:
  455. _add_docstring(
  456. getattr(new, attr), docstring.strip(), warn_on_python
  457. )