将浮点数重载到numpy数组

Question

I have a function, processing a 1D numpy array, like this:

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

With asarray I allow to call the function with a list as f([4,5,6]) . Now, I would like to "overload" the argument also to a single float, so that I can use f(4) instead of f([4]) .

This is a standard numpy feature, since you can call np.sin as sin(array([4,5,6])) , or as sin([4,5,6]) or as sin(4) as well. I came up with this code, that works at least in simple cases:

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

Is this the standard/correct way to do it?

Answer 1

I believe you are looking for np.atleast_1d .

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