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