31

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?

Improve this question
5
  • You have to decide whether you want to treat x as int or as a numpy.array. Or use ininstance() to check what type is being passed Commented Nov 7, 2011 at 13:10
  • What do you mean by x<0? x is an array, so it is not clear what this is supposed to mean. Commented Nov 7, 2011 at 13:10
  • @Bjoern Pollex I know that is why python is getting confused, I want to apply this function to each element of the array individually (i.e. elementwise) Commented Nov 7, 2011 at 13:12
  • @pajton Indeed but is there a way to make the second handle this gracefully like the first example? Commented Nov 7, 2011 at 13:13
  • So what exactly do you want the output to be, for the illustrated input? Commented Nov 7, 2011 at 13:52

6 Answers 6

Reset to default
22

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.

Improve this answer
Sign up to request clarification or add additional context in comments.

1 Comment

Your solution is similar to mine (np.where and np.choose behave similiarly in this case), but better!
20

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

from numpy import vectorize
vfunc = vectorize(func)
y = vfunc(x)
Improve this answer

4 Comments

Neat. I do believe however, that my solution is faster as long as you use operations for which Numpy has it's own implementations (have not tested it though, just a feeling).
Sure. All the built-in ops, such as those you compose, will be faster than pure python. I was just targeting the "is there a nice way to make the function handle arrays element wise?" question.
I wonder if there is a way to make this efficient when using functions written in C (e.g. by using Cython). I have a feeling there might be.
Accepted this answer since I was after readable code more than speed (this is only a little script after all) :)
10

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.

Improve this answer

4 Comments

Thanks, this seems to work. I'm curious as to why x[x<0]=0 etc doesn't seem to work?
@Dan: It does work, but the way I show it, the potentially expensive operation x<0 has to be executed only once (on second thought, invert is probably just as expensive, you'd have to test that).
Well, "x[x<0]=0" works for me. (Python 2.6.1 numpy 2.0.0.dev-3071eab)
I think using np.where is even better.
9

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

Improve this answer

Comments

1 
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

Improve this answer

1 Comment

It would be even nicer (and correct) if you'd make y = zeroes(len(x)) and then y[mask] = x[mask] * 10.
1

not sure why you need a function

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

Comments

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.