Elementwise if elif function in python using arrays

Question

I have a definition

def myfunc(a, b):
    if a < (b*10):
        result = a*2
    else:
        result = a*(-1)
    return result

Now this obviously works perfectly when I feed in my a and b values one by one using for loops, however it takes forever (I've simplified the definition a wee bit) and I know from experience that passing in the values as an array will speed it up.

So how do I modify this code to accept arrays. I've used the any() and all() commands but I must be using them wrong as my function only spits out one value rather than an array of values.

An example of my desired output would be:

>>>a = np.array([1,5,50,500])
>>>b = 1
>>>print myfunc(a, b)
array([-1, -5, 100, 1000])

Answer 1

You could use np.where:

def myfunc(a, b):
    return np.where(a < b*10, a*2, -a)    

For example,

In [48]: a = np.array([1, 5, 50, 500])

In [49]: b = 1

In [50]: myfunc(a, b)
Out[50]: array([   2,   10,  -50, -500])

Note the output is not the same as your desired output, but is consistent with the code you posted. You can of course get the desired output by reversing the inequality:

def myfunc(a, b):
    return np.where(a > b*10, a*2, -a)

then

In [52]: myfunc(a, b)
Out[52]: array([  -1,   -5,  100, 1000])