# 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:\Scripts\test.py", line 17, in <module>
y = func(x)
File "D:\Scripts\test.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?

1
17
11/7/2011 1:14:44 PM

Use numpy.vectorize to wrap func before applying it to array x:

from numpy import vectorize
vfunc = vectorize(func)
y = vfunc(x)
13
4/8/2017 4:45:11 PM

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.