Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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:\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?

like image 323
Dan Avatar asked Nov 07 '11 13:11

Dan


People also ask

How can we use conditions in NumPy within an array?

It returns a new numpy array, after filtering based on a condition, which is a numpy-like array of boolean values. For example, if condition is array([[True, True, False]]) , and our array is a = ndarray([[1, 2, 3]]) , on applying a condition to array ( a[:, condition] ), we will get the array ndarray([[1 2]]) .

What is NumPy array function?

A numpy array is a grid of values, all of the same type, and is indexed by a tuple of nonnegative integers. The number of dimensions is the rank of the array; the shape of an array is a tuple of integers giving the size of the array along each dimension.

What does NumPy where return?

where() in Python. The numpy. where() function returns the indices of elements in an input array where the given condition is satisfied.


5 Answers

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

from numpy import vectorize
vfunc = vectorize(func)
y = vfunc(x)
like image 161
Chris Kuklewicz Avatar answered Sep 25 '22 08:09

Chris Kuklewicz


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.

like image 26
hurrdrought Avatar answered Sep 28 '22 08:09

hurrdrought


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.

like image 29
Björn Pollex Avatar answered Sep 25 '22 08:09

Björn Pollex


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

like image 26
shx2 Avatar answered Sep 25 '22 08:09

shx2


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

like image 1
Jdog Avatar answered Sep 26 '22 08:09

Jdog