python - 将浮点数重载到 numpy 数组

Question

我有一个函数，处理一维 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.sinas sin(array([4,5,6]))、 assin([4,5,6])或 as 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

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

Answer 1

相信你正在寻找np.atleast_1d。

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