From f8cdbbaee1967011d98aa454b14232488df451d3 Mon Sep 17 00:00:00 2001 From: Jonathan Helmus Date: Fri, 28 Jun 2013 12:19:25 -0500 Subject: [PATCH] ENH: Add dtype parameter to linspace and logspace functions. Many NumPy functions such as arange allow users to define the dtype of the returned type with a dtype parameter. This adds this same functionality to the logspace and linspace functions. --- doc/release/1.9.0-notes.rst | 5 ++++ numpy/core/function_base.py | 36 ++++++++++++++++++-------- numpy/core/tests/test_function_base.py | 30 +++++++++++++++++---- 3 files changed, 55 insertions(+), 16 deletions(-) diff --git a/doc/release/1.9.0-notes.rst b/doc/release/1.9.0-notes.rst index 0b1ab77f1834..b78211d24076 100644 --- a/doc/release/1.9.0-notes.rst +++ b/doc/release/1.9.0-notes.rst @@ -31,6 +31,10 @@ functions (ufuncs), ``numpy.core._dotblas.dot``, and ``numpy.core.multiarray.dot`` (the numpy.dot functions). By defining a ``__numpy_ufunc__`` method. +Dtype parameter added to `np.linspace` and `np.logspace` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +The returned data type from the `linspace` and `logspace` functions +can now be specificed using the dtype parameter. Improvements ============ @@ -39,6 +43,7 @@ Improvements Changes ======= + C-API ~~~~~ diff --git a/numpy/core/function_base.py b/numpy/core/function_base.py index f2c895608dcf..3396adf3d6b6 100644 --- a/numpy/core/function_base.py +++ b/numpy/core/function_base.py @@ -3,9 +3,10 @@ __all__ = ['logspace', 'linspace'] from . import numeric as _nx -from .numeric import array +from .numeric import array, result_type -def linspace(start, stop, num=50, endpoint=True, retstep=False): + +def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None): """ Return evenly spaced numbers over a specified interval. @@ -31,6 +32,9 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False): retstep : bool, optional If True, return (`samples`, `step`), where `step` is the spacing between samples. + dtype : dtype + The type of the output array. If `dtype` is not given, infer the data + type from the other input arguments. Returns ------- @@ -74,23 +78,28 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False): """ num = int(num) + + if dtype is None: + dtype = result_type(start, stop, float(num)) + if num <= 0: - return array([], float) + return array([], dtype) if endpoint: if num == 1: - return array([float(start)]) + return array([start], dtype=dtype) step = (stop-start)/float((num-1)) - y = _nx.arange(0, num) * step + start + y = _nx.arange(0, num, dtype=dtype) * step + start y[-1] = stop else: step = (stop-start)/float(num) - y = _nx.arange(0, num) * step + start + y = _nx.arange(0, num, dtype=dtype) * step + start if retstep: - return y, step + return y.astype(dtype), step else: - return y + return y.astype(dtype) + -def logspace(start,stop,num=50,endpoint=True,base=10.0): +def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None): """ Return numbers spaced evenly on a log scale. @@ -116,6 +125,9 @@ def logspace(start,stop,num=50,endpoint=True,base=10.0): The base of the log space. The step size between the elements in ``ln(samples) / ln(base)`` (or ``log_base(samples)``) is uniform. Default is 10.0. + dtype : dtype + The type of the output array. If `dtype` is not given, infer the data + type from the other input arguments. Returns ------- @@ -136,7 +148,7 @@ def logspace(start,stop,num=50,endpoint=True,base=10.0): >>> y = np.linspace(start, stop, num=num, endpoint=endpoint) ... # doctest: +SKIP - >>> power(base, y) + >>> power(base, y).astype(dtype) ... # doctest: +SKIP Examples @@ -165,4 +177,6 @@ def logspace(start,stop,num=50,endpoint=True,base=10.0): """ y = linspace(start, stop, num=num, endpoint=endpoint) - return _nx.power(base, y) + if dtype is None: + return _nx.power(base, y) + return _nx.power(base, y).astype(dtype) diff --git a/numpy/core/tests/test_function_base.py b/numpy/core/tests/test_function_base.py index efca9ef8a26f..b43452e1c9fc 100644 --- a/numpy/core/tests/test_function_base.py +++ b/numpy/core/tests/test_function_base.py @@ -1,23 +1,35 @@ from __future__ import division, absolute_import, print_function from numpy.testing import * -from numpy import logspace, linspace +from numpy import logspace, linspace, dtype + class TestLogspace(TestCase): + def test_basic(self): y = logspace(0, 6) - assert_(len(y)==50) + assert_(len(y) == 50) y = logspace(0, 6, num=100) - assert_(y[-1] == 10**6) + assert_(y[-1] == 10 ** 6) y = logspace(0, 6, endpoint=0) - assert_(y[-1] < 10**6) + assert_(y[-1] < 10 ** 6) y = logspace(0, 6, num=7) assert_array_equal(y, [1, 10, 100, 1e3, 1e4, 1e5, 1e6]) + def test_dtype(self): + y = logspace(0, 6, dtype='float32') + assert_equal(y.dtype, dtype('float32')) + y = logspace(0, 6, dtype='float64') + assert_equal(y.dtype, dtype('float64')) + y = logspace(0, 6, dtype='int32') + assert_equal(y.dtype, dtype('int32')) + + class TestLinspace(TestCase): + def test_basic(self): y = linspace(0, 10) - assert_(len(y)==50) + assert_(len(y) == 50) y = linspace(2, 10, num=100) assert_(y[-1] == 10) y = linspace(2, 10, endpoint=0) @@ -35,3 +47,11 @@ def test_type(self): t3 = linspace(0, 1, 2).dtype assert_equal(t1, t2) assert_equal(t2, t3) + + def test_dtype(self): + y = linspace(0, 6, dtype='float32') + assert_equal(y.dtype, dtype('float32')) + y = linspace(0, 6, dtype='float64') + assert_equal(y.dtype, dtype('float64')) + y = linspace(0, 6, dtype='int32') + assert_equal(y.dtype, dtype('int32'))