Taozi Taozi - 7 months ago 18
Python Question

A python function that accepts as an argument either a scalar or a numpy array

As the title says, suppose I want to write a sign function (let's forget sign(0) for now), obviously we expect sign(2) = 1 and sign(array([-2,-2,2])) = array([-1,-1,1]). The following function won't work however, because it can't handle numpy arrays.

def sign(x):
if x>0: return 1
else: return -1


The next function won't work either since x doesn't have a shape member if it's just a single number. Even if some trick like y = x*0 + 1 is used, y won't have a [] method.

def sign(x):
y = ones(x.shape)
y[x<0] = -1
return y


Even with the idea from another question(how can I make a numpy function that accepts a numpy array, an iterable, or a scalar?), the next function won't work when x is a single number because in this case x.shape and y.shape are just () and indexing y is illegal.

def sign(x):
x = asarray(x)
y = ones(x.shape)
y[x<0] = -1
return y


The only solution seems to be that first decide if x is an array or a number, but I want to know if there is something better. Writing branchy code would be cumbersome if you have lots of small functions like this.

Answer

i wonder if it's a vectorized function that you want:

>>> import numpy as NP

>>> def fnx(a):
        if a > 0:
            return 1
        else:
            return -1

>>> vfnx = NP.vectorize(fnx)

>>> a = NP.random.randint(1, 10, 5)
array([4, 9, 7, 9, 2])

>>> a0 = 7

>>> vfnx(a)
array([1, 1, 1, 1])

>>> vfnx(a0)
array(1)