Function of Numpy Array with if-statement
Use numpy.vectorize
to wrap func before applying it to array x
:
from numpy import vectorize
vfunc = vectorize(func)
y = vfunc(x)
I know it is too late for this answer, but I am excited learning NumPy. You can vectorize the function on your own with numpy.where.
def func(x):
import numpy as np
x = np.where(x<0, 0., x*10)
return x
Examples
Using a scalar as data input:
x = 10
y = func(10)
y = array(100.0)
using an array as data input:
x = np.arange(-1,1,0.1)
y = func(x)
y = array([ -1.00000000e+00, -9.00000000e-01, -8.00000000e-01,
-7.00000000e-01, -6.00000000e-01, -5.00000000e-01,
-4.00000000e-01, -3.00000000e-01, -2.00000000e-01,
-1.00000000e-01, -2.22044605e-16, 1.00000000e-01,
2.00000000e-01, 3.00000000e-01, 4.00000000e-01,
5.00000000e-01, 6.00000000e-01, 7.00000000e-01,
8.00000000e-01, 9.00000000e-01])
Caveats:
1) If x
is a masked array, you need to use np.ma.where
instead, since this works for masked arrays.
This should do what you want:
def func(x):
small_indices = x < 10
x[small_indices] = 0
x[invert(small_indices)] *= 10
return x
invert
is a Numpy-function. Note that this modifies the argument. To prevent this, you'd have to modify and return a copy
of x
.
(I realize this is an old question, but ...)
There is one more option which wasn't mentioned here -- using np.choose
.
np.choose(
# the boolean condition
x < 0,
[
# index 0: value if condition is False
10 * x,
# index 1: value if condition is True
0
]
)
Though not terribly readable, this is just a single expression (not a series of statements), and does not compromize numpy's inherent speed (as np.vectorize
does).