Function of Numpy Array with if-statement

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

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

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.


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.


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