将浮点数重载到 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.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])