# [Solved] Function of Numpy Array with if-statement

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?

## Solution #1:

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

``````from numpy import vectorize
vfunc = vectorize(func)
y = vfunc(x)
``````

## Solution #2:

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.

## Solution #3:

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`.

## Solution #4:

(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).

## Solution #5:

``````x = numpy.arrange(-1,1,0.01)
y = numpy.zeros(len(x))
``````

`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

## Solution #6:

not sure why you need a function

``````x = np.arange(-1, 1, 0.01)
y = x * np.where(x < 0, 0, 10)
``````

The answers/resolutions are collected from stackoverflow, are licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0 .