將浮點數重載到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])