Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to vectorize a function which contains an if statement?

Tags:

Let's say we have the following function:

def f(x, y):
    if y == 0:
        return 0
    return x/y

This works fine with scalar values. Unfortunately when I try to use numpy arrays for x and y the comparison y == 0 is treated as an array operation which results in an error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-13-9884e2c3d1cd> in <module>()
----> 1 f(np.arange(1,10), np.arange(10,20))

<ipython-input-10-fbd24f17ea07> in f(x, y)
      1 def f(x, y):
----> 2     if y == 0:
      3         return 0
      4     return x/y

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I tried to use np.vectorize but it doesn't make a difference, the code still fails with the same error. np.vectorize is one option which gives the result I expect.

The only solution that I can think of is to use np.where on the y array with something like:

def f(x, y):
    np.where(y == 0, 0, x/y)

which doesn't work for scalars.

Is there a better way to write a function which contains an if statement? It should work with both scalars and arrays.


Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!