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])

unutbu · Accepted Answer · 2012-11-08 14:24:32Z

18

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])

edited Nov 8, 2012 at 14:24

answered Nov 8, 2012 at 14:14

unutbu

886k197 gold badges1.9k silver badges1.7k bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Rapid Over a year ago

Thank you very much! Although this does work, there was another comment that simply used np.all(a < (b*10)) with my existing definition; why should I use np.where instead of np.all?...edit-ooops it seems I did put the inequality the wrong way around, fortunately this didn't create confusion to my question :)

unutbu Over a year ago

np.all returns a single boolean value, True or False. It simply tests if each and every element of a (and corresponding element of b) satisfies a < b*10. (It returns True if the condition is true element-wise, otherwise it returns False). I don't think that helps you in this situation.

KenHBS Over a year ago

Great answer! I wasn't aware you can use np.where with the additional two arguments, I will definitely be using this much more often

Eric · Accepted Answer · 2012-11-08 14:09:41Z

0

Use a list comprehension:

myarray = [1, 5, 50, 500]
b = 1
[myfunc(a, b) for a in myarray]

edited Nov 8, 2012 at 14:09

Eric

98.1k54 gold badges257 silver badges389 bronze badges

answered Nov 8, 2012 at 14:09

gefei

20k9 gold badges54 silver badges71 bronze badges

Comments

Eric · Accepted Answer · 2012-11-08 14:11:29Z

0

Your function is simple enough that it can be dropped entirely:

arr = [1, 5, 50, 500]
arr = [a * 2 if a < b * 10 else -a for a in arr]

answered Nov 8, 2012 at 14:11

Eric

98.1k54 gold badges257 silver badges389 bronze badges

2 Comments

mgilson Over a year ago

But I expect the list comp to be significantly slower than the numpy array operations.

Eric Over a year ago

@mgilson: Absolutely. I just have no numpy experience

Collectives™ on Stack Overflow

Elementwise if elif function in python using arrays

3 Answers 3

3 Comments

Comments

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

3 Comments

Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related