Merge pull request #12388 from mhvk/linspace-allow-array

ENH: allow arrays for start and stop in {lin,log,geom}space

4 files changed, 139 insertions, 49 deletions

diff --git a/doc/release/1.16.0-notes.rst b/doc/release/1.16.0-notes.rst
index 791cbe731..8d176c3ea 100644
--- a/doc/release/1.16.0-notes.rst
+++ b/doc/release/1.16.0-notes.rst
@@ -20,7 +20,7 @@ New functions
 =============
 
  * New functions in the `numpy.lib.recfunctions` module to ease the structured
-   assignment changes: `assign_fields_by_name`, `structured_to_unstructured`, 
+   assignment changes: `assign_fields_by_name`, `structured_to_unstructured`,
    `unstructured_to_structured`, `apply_along_fields`, and `require_fields`.
    See the user guide at <https://docs.scipy.org/doc/numpy/user/basics.rec.html>
    for more info.
@@ -389,6 +389,13 @@ overlapping fields and padding.
 implementation has also changed, ensuring it uses the same BLAS routines as
 `numpy.dot`, ensuring its performance is similar for large matrices.
 
+Start and stop arrays for ``linspace``, ``logspace`` and ``geomspace``
+----------------------------------------------------------------------
+These functions used to be limited to scalar stop and start values, but can
+now take arrays, which will be properly broadcast and result in an output
+which has one axis prepended.  This can be used, e.g., to obtain linearly
+interpolated points between sets of points.
+
 Changes
 =======
 
diff --git a/numpy/core/function_base.py b/numpy/core/function_base.py
index 0fc56e70e..b68fd4068 100644
--- a/numpy/core/function_base.py
+++ b/numpy/core/function_base.py
@@ -29,13 +29,14 @@ def _index_deprecate(i, stacklevel=2):
     return i
 
 
-def _linspace_dispatcher(
-        start, stop, num=None, endpoint=None, retstep=None, dtype=None):
+def _linspace_dispatcher(start, stop, num=None, endpoint=None, retstep=None,
+                         dtype=None, axis=None):
     return (start, stop)
 
 
 @array_function_dispatch(_linspace_dispatcher)
-def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
+def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
+             axis=0):
     """
     Return evenly spaced numbers over a specified interval.
 
@@ -44,11 +45,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
@@ -67,6 +71,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
@@ -128,16 +139,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,) * delta.ndim)
     # 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:
@@ -160,19 +170,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_dispatcher(
-        start, stop, num=None, endpoint=None, base=None, 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):
+def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None,
+             axis=0):
     """
     Return numbers spaced evenly on a log scale.
 
@@ -180,11 +194,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
@@ -201,6 +218,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
     -------
@@ -250,29 +274,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):
+def _geomspace_dispatcher(start, stop, num=None, endpoint=None, dtype=None,
+                          axis=None):
     return (start, stop)
 
 
 @array_function_dispatch(_geomspace_dispatcher)
-def geomspace(start, stop, num=50, endpoint=True, dtype=None):
+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
@@ -285,6 +313,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
     -------
@@ -349,40 +383,48 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None):
     >>> 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
diff --git a/numpy/core/tests/test_function_base.py b/numpy/core/tests/test_function_base.py
index d0ff1c15f..459bacab0 100644
--- a/numpy/core/tests/test_function_base.py
+++ b/numpy/core/tests/test_function_base.py
@@ -2,7 +2,7 @@ from __future__ import division, absolute_import, print_function
 
 from numpy import (
     logspace, linspace, geomspace, dtype, array, sctypes, arange, isnan,
-    ndarray, sqrt, nextafter
+    ndarray, sqrt, nextafter, stack
     )
 from numpy.testing import (
     assert_, assert_equal, assert_raises, assert_array_equal, assert_allclose,
@@ -54,6 +54,20 @@ class TestLogspace(object):
         y = logspace(0, 6, num=7)
         assert_array_equal(y, [1, 10, 100, 1e3, 1e4, 1e5, 1e6])
 
+    def test_start_stop_array(self):
+        start = array([0., 1.])
+        stop = array([6., 7.])
+        t1 = logspace(start, stop, 6)
+        t2 = stack([logspace(_start, _stop, 6)
+                    for _start, _stop in zip(start, stop)], axis=1)
+        assert_equal(t1, t2)
+        t3 = logspace(start, stop[0], 6)
+        t4 = stack([logspace(_start, stop[0], 6)
+                    for _start in start], axis=1)
+        assert_equal(t3, t4)
+        t5 = logspace(start, stop, 6, axis=-1)
+        assert_equal(t5, t2.T)
+
     def test_dtype(self):
         y = logspace(0, 6, dtype='float32')
         assert_equal(y.dtype, dtype('float32'))
@@ -156,7 +170,7 @@ class TestGeomspace(object):
         y = geomspace(1, 1e6, dtype=complex)
         assert_equal(y.dtype, dtype('complex'))
 
-    def test_array_scalar(self):
+    def test_start_stop_array_scalar(self):
         lim1 = array([120, 100], dtype="int8")
         lim2 = array([-120, -100], dtype="int8")
         lim3 = array([1200, 1000], dtype="uint16")
@@ -172,6 +186,21 @@ class TestGeomspace(object):
         assert_allclose(t2, t5, rtol=1e-2)
         assert_allclose(t3, t6, rtol=1e-5)
 
+    def test_start_stop_array(self):
+        # Try to use all special cases.
+        start = array([1.e0, 32., 1j, -4j, 1+1j, -1])
+        stop = array([1.e4, 2., 16j, -324j, 10000+10000j, 1])
+        t1 = geomspace(start, stop, 5)
+        t2 = stack([geomspace(_start, _stop, 5)
+                    for _start, _stop in zip(start, stop)], axis=1)
+        assert_equal(t1, t2)
+        t3 = geomspace(start, stop[0], 5)
+        t4 = stack([geomspace(_start, stop[0], 5)
+                    for _start in start], axis=1)
+        assert_equal(t3, t4)
+        t5 = geomspace(start, stop, 5, axis=-1)
+        assert_equal(t5, t2.T)
+
     def test_physical_quantities(self):
         a = PhysicalQuantity(1.0)
         b = PhysicalQuantity(5.0)
@@ -227,7 +256,7 @@ class TestLinspace(object):
         y = linspace(0, 6, dtype='int32')
         assert_equal(y.dtype, dtype('int32'))
 
-    def test_array_scalar(self):
+    def test_start_stop_array_scalar(self):
         lim1 = array([-120, 100], dtype="int8")
         lim2 = array([120, -100], dtype="int8")
         lim3 = array([1200, 1000], dtype="uint16")
@@ -241,6 +270,20 @@ class TestLinspace(object):
         assert_equal(t2, t5)
         assert_equal(t3, t6)
 
+    def test_start_stop_array(self):
+        start = array([-120, 120], dtype="int8")
+        stop = array([100, -100], dtype="int8")
+        t1 = linspace(start, stop, 5)
+        t2 = stack([linspace(_start, _stop, 5)
+                    for _start, _stop in zip(start, stop)], axis=1)
+        assert_equal(t1, t2)
+        t3 = linspace(start, stop[0], 5)
+        t4 = stack([linspace(_start, stop[0], 5)
+                    for _start in start], axis=1)
+        assert_equal(t3, t4)
+        t5 = linspace(start, stop, 5, axis=-1)
+        assert_equal(t5, t2.T)
+
     def test_complex(self):
         lim1 = linspace(1 + 2j, 3 + 4j, 5)
         t1 = array([1.0+2.j, 1.5+2.5j,  2.0+3j, 2.5+3.5j, 3.0+4j])
@@ -285,9 +328,7 @@ class TestLinspace(object):
 
             @property
             def __array_interface__(self):
-                # Ideally should be `'shape': ()` but the current interface
-                # does not allow that
-                return {'shape': (1,), 'typestr': '<i4', 'data': self._data,
+                return {'shape': (), 'typestr': '<i4', 'data': self._data,
                         'version': 3}
 
             def __mul__(self, other):
diff --git a/numpy/lib/tests/test_histograms.py b/numpy/lib/tests/test_histograms.py
index 5d0212f25..c96b01d42 100644
--- a/numpy/lib/tests/test_histograms.py
+++ b/numpy/lib/tests/test_histograms.py
@@ -289,13 +289,13 @@ class TestHistogram(object):
     def test_object_array_of_0d(self):
         # gh-7864
         assert_raises(ValueError,
-            histogram, [np.array([0.4]) for i in range(10)] + [-np.inf])
+            histogram, [np.array(0.4) for i in range(10)] + [-np.inf])
         assert_raises(ValueError,
-            histogram, [np.array([0.4]) for i in range(10)] + [np.inf])
+            histogram, [np.array(0.4) for i in range(10)] + [np.inf])
 
         # these should not crash
-        np.histogram([np.array([0.5]) for i in range(10)] + [.500000000000001])
-        np.histogram([np.array([0.5]) for i in range(10)] + [.5])
+        np.histogram([np.array(0.5) for i in range(10)] + [.500000000000001])
+        np.histogram([np.array(0.5) for i in range(10)] + [.5])
 
     def test_some_nan_values(self):
         # gh-7503