Merge pull request numpy#5446 from jaimefrio/linspace_denormals_again

BUG: linspace should return the same as arange when possible

numpy/core/function_base.py

@@ -3,7 +3,7 @@
 __all__ = ['logspace', 'linspace']
 
 from . import numeric as _nx
-from .numeric import array, result_type
+from .numeric import array, result_type, NaN
 
 
 def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
@@ -82,6 +82,7 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
 
     """
     num = int(num)
+    div = (num - 1) if endpoint else num
 
     # Convert float/complex array scalars to float, gh-3504
     start = start * 1.
@@ -91,38 +92,31 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
     if dtype is None:
         dtype = dt
 
-    if num <= 0:
-        return array([], dtype)
-    if num == 1:
-        return array([start], dtype=dtype)
     y = _nx.arange(0, num, dtype=dt)
-    if endpoint:
-        num -= 1
-    y /= num
-    y *= stop - start
+
+    if num > 1:
+        delta = stop - start
+        step = delta / div
+        if step == 0:
+            # Special handling for denormal numbers, gh-5437
+            y /= div
+            y *= delta
+        else:
+            y *= step
+    else:
+        # 0 and 1 item long sequences have an undefined step
+        step = NaN
+
     y += start
-    if endpoint:
+
+    if endpoint and num > 1:
         y[-1] = stop
 
     if retstep:
-        return y.astype(dtype, copy=False), (stop - start) / num
+        return y.astype(dtype, copy=False), step
     else:
         return y.astype(dtype, copy=False)
 
-    # if endpoint:
-        # if num == 1:
-            # return array([start], dtype=dtype)
-        # step = (stop-start)/float((num-1))
-        # y = _nx.arange(0, num, dtype=dtype) * step + start
-        # y[-1] = stop
-    # else:
-        # step = (stop-start)/float(num)
-        # y = _nx.arange(0, num, dtype=dtype) * step + start
-    # if retstep:
-        # return y.astype(dtype), step
-    # else:
-        # return y.astype(dtype)
-
 
 def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
     """

numpy/core/tests/test_function_base.py

@@ -1,7 +1,8 @@
 from __future__ import division, absolute_import, print_function
 
 from numpy.testing import *
-from numpy import logspace, linspace, dtype, array, finfo, typecodes
+from numpy import (logspace, linspace, dtype, array, finfo, typecodes, arange,
+                   isnan)
 
 class TestLogspace(TestCase):
 
@@ -117,6 +118,21 @@ def test_denormal_numbers(self):
             stop = finfo(dt).tiny * finfo(dt).resolution
             assert_(any(linspace(0, stop, 10, endpoint=False, dtype=dt)))
 
+    def test_equivalent_to_arange(self):
+        for j in range(1000):
+            assert_equal(linspace(0, j, j+1, dtype=int),
+                         arange(j+1, dtype=int))
+
+    def test_retstep(self):
+        y = linspace(0, 1, 2, retstep=True)
+        assert_(isinstance(y, tuple) and len(y) == 2)
+        for num in (0, 1):
+            for ept in (False, True):
+                y = linspace(0, 1, num, endpoint=ept, retstep=True)
+                assert_(isinstance(y, tuple) and len(y) == 2 and
+                        len(y[0]) == num and isnan(y[1]),
+                        'num={0}, endpoint={1}'.format(num, ept))
+
 
 if __name__ == "__main__":
     run_module_suite()