function_base.py 19 KB

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