python - 将 float 重载到 numpy 数组

Question

我有一个函数，处理一维 numpy 数组，如下所示:

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

使用 asarray，我允许调用带有 f([4,5,6]) 列表的函数。现在，我想将参数也“重载”到单个 float ，这样我就可以使用 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

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

Answer 1

我相信您正在寻找np.atleast_1d .

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