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.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With