将浮点数重载为numpy数组

Question

我有一个处理1D numpy数组的函数，就像这样：

def f(arr):
    arr=asarray(arr)
    #process data as numpy array
    #...
    return arr

使用asarray函数，可以将列表作为参数传递给函数，例如f([4,5,6])。现在，我希望还能够使用单个浮点数作为参数，这样就可以使用f(4)代替f([4])。

这是numpy的一个标准功能，因为你可以像这样调用np.sin：sin(array([4,5,6]))，或者sin([4,5,6])，也可以像这样调用sin(4)。下面是代码，至少在简单的情况下有效：

def f(arr):
    arr=asarray(arr)
    if arr.shape is ():
        print 'arr is a single float/int/etc'
        arr = array([arr])
    #process data as numpy array
    #...
    return arr

这是标准/正确的做法吗？

- gg349

Answer 1

我认为您正在寻找np.atleast_1d。

>>> np.atleast_1d(5)
array([5])
>>> np.atleast_1d(np.arange(2))
array([0, 1])

- Daniel