Function of Numpy Array with if-statement
Question:
I am using Matplotlib and Numpy to produce some plots. I wish to define a function which given an array returns another array with values calculated elementwise, for example:
def func(x):
return x*10
x = numpy.arrange(-1,1,0.01)
y = func(x)
This is fine. Now however I wish to have an if-statement inside func
, for example:
def func(x):
if x<0:
return 0
else:
return x*10
x = numpy.arrange(-1,1,0.01)
y = func(x)
This unfortunately throws the following error
Traceback (most recent call last):
File "D:Scriptstest.py", line 17, in <module>
y = func(x)
File "D:Scriptstest.py", line 11, in func
if x<0:
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
I looked at the documentation for all()
and any()
and they do not fit the bill for what I need. So is there a nice way to make the function handle arrays element wise as in the first example?
Answers:
Use numpy.vectorize
to wrap func before applying it to array x
:
from numpy import vectorize
vfunc = vectorize(func)
y = vfunc(x)
x = numpy.arrange(-1,1,0.01)
mask = x>=0
y = numpy.zeros(len(x))
y[mask] = x[mask]*10
mask
is a boolean array that equates to True
are array indices matching the condition and False
elsewhere. The last line replaces all values in the original array with that value mulitplied by 10.
Edited to reflect Bjorn’s pertinent comment
(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).
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.
not sure why you need a function
x = np.arange(-1, 1, 0.01)
y = x * np.where(x < 0, 0, 10)
I am using Matplotlib and Numpy to produce some plots. I wish to define a function which given an array returns another array with values calculated elementwise, for example:
def func(x):
return x*10
x = numpy.arrange(-1,1,0.01)
y = func(x)
This is fine. Now however I wish to have an if-statement inside func
, for example:
def func(x):
if x<0:
return 0
else:
return x*10
x = numpy.arrange(-1,1,0.01)
y = func(x)
This unfortunately throws the following error
Traceback (most recent call last):
File "D:Scriptstest.py", line 17, in <module>
y = func(x)
File "D:Scriptstest.py", line 11, in func
if x<0:
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
I looked at the documentation for all()
and any()
and they do not fit the bill for what I need. So is there a nice way to make the function handle arrays element wise as in the first example?
Use numpy.vectorize
to wrap func before applying it to array x
:
from numpy import vectorize
vfunc = vectorize(func)
y = vfunc(x)
x = numpy.arrange(-1,1,0.01)
mask = x>=0
y = numpy.zeros(len(x))
y[mask] = x[mask]*10
mask
is a boolean array that equates to True
are array indices matching the condition and False
elsewhere. The last line replaces all values in the original array with that value mulitplied by 10.
Edited to reflect Bjorn’s pertinent comment
(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).
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.
not sure why you need a function
x = np.arange(-1, 1, 0.01)
y = x * np.where(x < 0, 0, 10)