Merge branch 'master' into master

1 files changed, 195 insertions, 79 deletions

diff --git a/numpy/core/function_base.py b/numpy/core/function_base.py
index 799b1418d..42604ec3f 100644
--- a/numpy/core/function_base.py
+++ b/numpy/core/function_base.py
@@ -1,29 +1,31 @@
 from __future__ import division, absolute_import, print_function
 
+import functools
 import warnings
 import operator
+import types
 
 from . import numeric as _nx
 from .numeric import (result_type, NaN, shares_memory, MAY_SHARE_BOUNDS,
-                      TooHardError,asanyarray)
+                      TooHardError, asanyarray, ndim)
 from numpy.core.multiarray import add_docstring
+from numpy.core import overrides
 
 __all__ = ['logspace', 'linspace', 'geomspace']
 
 
-def _index_deprecate(i, stacklevel=2):
-    try:
-        i = operator.index(i)
-    except TypeError:
-        msg = ("object of type {} cannot be safely interpreted as "
-               "an integer.".format(type(i)))
-        i = int(i)
-        stacklevel += 1
-        warnings.warn(msg, DeprecationWarning, stacklevel=stacklevel)
-    return i
+array_function_dispatch = functools.partial(
+    overrides.array_function_dispatch, module='numpy')
+
+
+def _linspace_dispatcher(start, stop, num=None, endpoint=None, retstep=None,
+                         dtype=None, axis=None):
+    return (start, stop)
 
 
-def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
+@array_function_dispatch(_linspace_dispatcher)
+def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
+             axis=0):
     """
     Return evenly spaced numbers over a specified interval.
 
@@ -32,11 +34,14 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
 
     The endpoint of the interval can optionally be excluded.
 
+    .. versionchanged:: 1.16.0
+        Non-scalar `start` and `stop` are now supported.
+
     Parameters
     ----------
-    start : scalar
+    start : array_like
         The starting value of the sequence.
-    stop : scalar
+    stop : array_like
         The end value of the sequence, unless `endpoint` is set to False.
         In that case, the sequence consists of all but the last of ``num + 1``
         evenly spaced samples, so that `stop` is excluded.  Note that the step
@@ -55,6 +60,13 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
 
         .. versionadded:: 1.9.0
 
+    axis : int, optional
+        The axis in the result to store the samples.  Relevant only if start
+        or stop are array-like.  By default (0), the samples will be along a
+        new axis inserted at the beginning. Use -1 to get an axis at the end.
+
+        .. versionadded:: 1.16.0
+
     Returns
     -------
     samples : ndarray
@@ -79,11 +91,11 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
     Examples
     --------
     >>> np.linspace(2.0, 3.0, num=5)
-    array([ 2.  ,  2.25,  2.5 ,  2.75,  3.  ])
+    array([2.  , 2.25, 2.5 , 2.75, 3.  ])
     >>> np.linspace(2.0, 3.0, num=5, endpoint=False)
-    array([ 2. ,  2.2,  2.4,  2.6,  2.8])
+    array([2. ,  2.2,  2.4,  2.6,  2.8])
     >>> np.linspace(2.0, 3.0, num=5, retstep=True)
-    (array([ 2.  ,  2.25,  2.5 ,  2.75,  3.  ]), 0.25)
+    (array([2.  ,  2.25,  2.5 ,  2.75,  3.  ]), 0.25)
 
     Graphical illustration:
 
@@ -101,8 +113,13 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
     >>> plt.show()
 
     """
-    # 2016-02-25, 1.12
-    num = _index_deprecate(num)
+    try:
+        num = operator.index(num)
+    except TypeError:
+        raise TypeError(
+            "object of type {} cannot be safely interpreted as an integer."
+                .format(type(num)))
+
     if num < 0:
         raise ValueError("Number of samples, %s, must be non-negative." % num)
     div = (num - 1) if endpoint else num
@@ -116,16 +133,15 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
     if dtype is None:
         dtype = dt
 
-    y = _nx.arange(0, num, dtype=dt)
-
     delta = stop - start
+    y = _nx.arange(0, num, dtype=dt).reshape((-1,) + (1,) * ndim(delta))
     # In-place multiplication y *= delta/div is faster, but prevents the multiplicant
     # from overriding what class is produced, and thus prevents, e.g. use of Quantities,
     # see gh-7142. Hence, we multiply in place only for standard scalar types.
-    _mult_inplace = _nx.isscalar(delta) 
+    _mult_inplace = _nx.isscalar(delta)
     if num > 1:
         step = delta / div
-        if step == 0:
+        if _nx.any(step == 0):
             # Special handling for denormal numbers, gh-5437
             y /= div
             if _mult_inplace:
@@ -148,13 +164,23 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
     if endpoint and num > 1:
         y[-1] = stop
 
+    if axis != 0:
+        y = _nx.moveaxis(y, 0, axis)
+
     if retstep:
         return y.astype(dtype, copy=False), step
     else:
         return y.astype(dtype, copy=False)
 
 
-def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
+def _logspace_dispatcher(start, stop, num=None, endpoint=None, base=None,
+                         dtype=None, axis=None):
+    return (start, stop)
+
+
+@array_function_dispatch(_logspace_dispatcher)
+def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None,
+             axis=0):
     """
     Return numbers spaced evenly on a log scale.
 
@@ -162,11 +188,14 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
     (`base` to the power of `start`) and ends with ``base ** stop``
     (see `endpoint` below).
 
+    .. versionchanged:: 1.16.0
+        Non-scalar `start` and `stop` are now supported.
+
     Parameters
     ----------
-    start : float
+    start : array_like
         ``base ** start`` is the starting value of the sequence.
-    stop : float
+    stop : array_like
         ``base ** stop`` is the final value of the sequence, unless `endpoint`
         is False.  In that case, ``num + 1`` values are spaced over the
         interval in log-space, of which all but the last (a sequence of
@@ -183,6 +212,13 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
     dtype : dtype
         The type of the output array.  If `dtype` is not given, infer the data
         type from the other input arguments.
+    axis : int, optional
+        The axis in the result to store the samples.  Relevant only if start
+        or stop are array-like.  By default (0), the samples will be along a
+        new axis inserted at the beginning. Use -1 to get an axis at the end.
+
+        .. versionadded:: 1.16.0
+
 
     Returns
     -------
@@ -210,11 +246,11 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
     Examples
     --------
     >>> np.logspace(2.0, 3.0, num=4)
-    array([  100.        ,   215.443469  ,   464.15888336,  1000.        ])
+    array([ 100.        ,  215.443469  ,  464.15888336, 1000.        ])
     >>> np.logspace(2.0, 3.0, num=4, endpoint=False)
-    array([ 100.        ,  177.827941  ,  316.22776602,  562.34132519])
+    array([100.        ,  177.827941  ,  316.22776602,  562.34132519])
     >>> np.logspace(2.0, 3.0, num=4, base=2.0)
-    array([ 4.        ,  5.0396842 ,  6.34960421,  8.        ])
+    array([4.        ,  5.0396842 ,  6.34960421,  8.        ])
 
     Graphical illustration:
 
@@ -232,24 +268,33 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
     >>> plt.show()
 
     """
-    y = linspace(start, stop, num=num, endpoint=endpoint)
+    y = linspace(start, stop, num=num, endpoint=endpoint, axis=axis)
     if dtype is None:
         return _nx.power(base, y)
-    return _nx.power(base, y).astype(dtype)
+    return _nx.power(base, y).astype(dtype, copy=False)
+
+
+def _geomspace_dispatcher(start, stop, num=None, endpoint=None, dtype=None,
+                          axis=None):
+    return (start, stop)
 
 
-def geomspace(start, stop, num=50, endpoint=True, dtype=None):
+@array_function_dispatch(_geomspace_dispatcher)
+def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
     """
     Return numbers spaced evenly on a log scale (a geometric progression).
 
     This is similar to `logspace`, but with endpoints specified directly.
     Each output sample is a constant multiple of the previous.
 
+    .. versionchanged:: 1.16.0
+        Non-scalar `start` and `stop` are now supported.
+
     Parameters
     ----------
-    start : scalar
+    start : array_like
         The starting value of the sequence.
-    stop : scalar
+    stop : array_like
         The final value of the sequence, unless `endpoint` is False.
         In that case, ``num + 1`` values are spaced over the
         interval in log-space, of which all but the last (a sequence of
@@ -262,6 +307,12 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None):
     dtype : dtype
         The type of the output array.  If `dtype` is not given, infer the data
         type from the other input arguments.
+    axis : int, optional
+        The axis in the result to store the samples.  Relevant only if start
+        or stop are array-like.  By default (0), the samples will be along a
+        new axis inserted at the beginning. Use -1 to get an axis at the end.
+
+        .. versionadded:: 1.16.0
 
     Returns
     -------
@@ -304,15 +355,15 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None):
     Negative, decreasing, and complex inputs are allowed:
 
     >>> np.geomspace(1000, 1, num=4)
-    array([ 1000.,   100.,    10.,     1.])
+    array([1000.,  100.,   10.,    1.])
     >>> np.geomspace(-1000, -1, num=4)
     array([-1000.,  -100.,   -10.,    -1.])
     >>> np.geomspace(1j, 1000j, num=4)  # Straight line
-    array([ 0.   +1.j,  0.  +10.j,  0. +100.j,  0.+1000.j])
+    array([0.   +1.j, 0.  +10.j, 0. +100.j, 0.+1000.j])
     >>> np.geomspace(-1+0j, 1+0j, num=5)  # Circle
-    array([-1.00000000+0.j        , -0.70710678+0.70710678j,
-            0.00000000+1.j        ,  0.70710678+0.70710678j,
-            1.00000000+0.j        ])
+    array([-1.00000000e+00+1.22464680e-16j, -7.07106781e-01+7.07106781e-01j,
+            6.12323400e-17+1.00000000e+00j,  7.07106781e-01+7.07106781e-01j,
+            1.00000000e+00+0.00000000e+00j])
 
     Graphical illustration of ``endpoint`` parameter:
 
@@ -320,78 +371,143 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None):
     >>> N = 10
     >>> y = np.zeros(N)
     >>> plt.semilogx(np.geomspace(1, 1000, N, endpoint=True), y + 1, 'o')
+    [<matplotlib.lines.Line2D object at 0x...>]
     >>> plt.semilogx(np.geomspace(1, 1000, N, endpoint=False), y + 2, 'o')
+    [<matplotlib.lines.Line2D object at 0x...>]
     >>> plt.axis([0.5, 2000, 0, 3])
+    [0.5, 2000, 0, 3]
     >>> plt.grid(True, color='0.7', linestyle='-', which='both', axis='both')
     >>> plt.show()
 
     """
-    if start == 0 or stop == 0:
+    start = asanyarray(start)
+    stop = asanyarray(stop)
+    if _nx.any(start == 0) or _nx.any(stop == 0):
         raise ValueError('Geometric sequence cannot include zero')
 
-    dt = result_type(start, stop, float(num))
+    dt = result_type(start, stop, float(num), _nx.zeros((), dtype))
     if dtype is None:
         dtype = dt
     else:
         # complex to dtype('complex128'), for instance
         dtype = _nx.dtype(dtype)
 
+    # Promote both arguments to the same dtype in case, for instance, one is
+    # complex and another is negative and log would produce NaN otherwise.
+    # Copy since we may change things in-place further down.
+    start = start.astype(dt, copy=True)
+    stop = stop.astype(dt, copy=True)
+
+    out_sign = _nx.ones(_nx.broadcast(start, stop).shape, dt)
     # Avoid negligible real or imaginary parts in output by rotating to
     # positive real, calculating, then undoing rotation
-    out_sign = 1
-    if start.real == stop.real == 0:
-        start, stop = start.imag, stop.imag
-        out_sign = 1j * out_sign
-    if _nx.sign(start) == _nx.sign(stop) == -1:
-        start, stop = -start, -stop
-        out_sign = -out_sign
-
-    # Promote both arguments to the same dtype in case, for instance, one is
-    # complex and another is negative and log would produce NaN otherwise
-    start = start + (stop - stop)
-    stop = stop + (start - start)
-    if _nx.issubdtype(dtype, _nx.complexfloating):
-        start = start + 0j
-        stop = stop + 0j
+    if _nx.issubdtype(dt, _nx.complexfloating):
+        all_imag = (start.real == 0.) & (stop.real == 0.)
+        if _nx.any(all_imag):
+            start[all_imag] = start[all_imag].imag
+            stop[all_imag] = stop[all_imag].imag
+            out_sign[all_imag] = 1j
+
+    both_negative = (_nx.sign(start) == -1) & (_nx.sign(stop) == -1)
+    if _nx.any(both_negative):
+        _nx.negative(start, out=start, where=both_negative)
+        _nx.negative(stop, out=stop, where=both_negative)
+        _nx.negative(out_sign, out=out_sign, where=both_negative)
 
     log_start = _nx.log10(start)
     log_stop = _nx.log10(stop)
     result = out_sign * logspace(log_start, log_stop, num=num,
                                  endpoint=endpoint, base=10.0, dtype=dtype)
+    if axis != 0:
+        result = _nx.moveaxis(result, 0, axis)
 
-    return result.astype(dtype)
+    return result.astype(dtype, copy=False)
 
 
-#always succeed
-def add_newdoc(place, obj, doc):
+def _needs_add_docstring(obj):
     """
-    Adds documentation to obj which is in module place.
+    Returns true if the only way to set the docstring of `obj` from python is
+    via add_docstring.
+
+    This function errs on the side of being overly conservative.
+    """
+    Py_TPFLAGS_HEAPTYPE = 1 << 9
+
+    if isinstance(obj, (types.FunctionType, types.MethodType, property)):
+        return False
+
+    if isinstance(obj, type) and obj.__flags__ & Py_TPFLAGS_HEAPTYPE:
+        return False
 
-    If doc is a string add it to obj as a docstring
+    return True
 
-    If doc is a tuple, then the first element is interpreted as
-       an attribute of obj and the second as the docstring
-          (method, docstring)
 
-    If doc is a list, then each element of the list should be a
-       sequence of length two --> [(method1, docstring1),
-       (method2, docstring2), ...]
+def _add_docstring(obj, doc, warn_on_python):
+    if warn_on_python and not _needs_add_docstring(obj):
+        warnings.warn(
+            "add_newdoc was used on a pure-python object {}. "
+            "Prefer to attach it directly to the source."
+            .format(obj),
+            UserWarning,
+            stacklevel=3)
+    try:
+        add_docstring(obj, doc)
+    except Exception:
+        pass
+
+
+def add_newdoc(place, obj, doc, warn_on_python=True):
+    """
+    Add documentation to an existing object, typically one defined in C
 
-    This routine never raises an error.
+    The purpose is to allow easier editing of the docstrings without requiring
+    a re-compile. This exists primarily for internal use within numpy itself.
+
+    Parameters
+    ----------
+    place : str
+        The absolute name of the module to import from
+    obj : str
+        The name of the object to add documentation to, typically a class or
+        function name
+    doc : {str, Tuple[str, str], List[Tuple[str, str]]}
+        If a string, the documentation to apply to `obj`
+
+        If a tuple, then the first element is interpreted as an attribute of
+        `obj` and the second as the docstring to apply - ``(method, docstring)``
+
+        If a list, then each element of the list should be a tuple of length
+        two - ``[(method1, docstring1), (method2, docstring2), ...]``
+    warn_on_python : bool
+        If True, the default, emit `UserWarning` if this is used to attach
+        documentation to a pure-python object.
+
+    Notes
+    -----
+    This routine never raises an error if the docstring can't be written, but
+    will raise an error if the object being documented does not exist.
 
     This routine cannot modify read-only docstrings, as appear
     in new-style classes or built-in functions. Because this
     routine never raises an error the caller must check manually
     that the docstrings were changed.
+
+    Since this function grabs the ``char *`` from a c-level str object and puts
+    it into the ``tp_doc`` slot of the type of `obj`, it violates a number of
+    C-API best-practices, by:
+
+    - modifying a `PyTypeObject` after calling `PyType_Ready`
+    - calling `Py_INCREF` on the str and losing the reference, so the str
+      will never be released
+
+    If possible it should be avoided.
     """
-    try:
-        new = getattr(__import__(place, globals(), {}, [obj]), obj)
-        if isinstance(doc, str):
-            add_docstring(new, doc.strip())
-        elif isinstance(doc, tuple):
-            add_docstring(getattr(new, doc[0]), doc[1].strip())
-        elif isinstance(doc, list):
-            for val in doc:
-                add_docstring(getattr(new, val[0]), val[1].strip())
-    except Exception:
-        pass
+    new = getattr(__import__(place, globals(), {}, [obj]), obj)
+    if isinstance(doc, str):
+        _add_docstring(new, doc.strip(), warn_on_python)
+    elif isinstance(doc, tuple):
+        attr, docstring = doc
+        _add_docstring(getattr(new, attr), docstring.strip(), warn_on_python)
+    elif isinstance(doc, list):
+        for attr, docstring in doc:
+            _add_docstring(getattr(new, attr), docstring.strip(), warn_on_python)