CODEWITHSUNDEEP
pythonnumpy
Sato
Sato

Reputation: 1133

Piecewise function in numpy with multiple arguments

I tried to define a function (tent map) as following:

def f(r, x):
    return np.piecewise([r, x], [x < 0.5, x >= 0.5], [lambda r, x: 2*r*x, lambda r, x: 2*r*(1-x)])

And r, x will be numpy arrays:

no_r = 10001
r = np.linspace(0, 4, no_r)
x = np.random.rand(no_r)

I would like the result to be a numpy array matching the shapes of r and x, calculated using each pairs of elements of arrays r and x with the same indicies. For example if r = [0, 1, 2, 3] and x = [0.1, 0.7, 0.3, 1], the result should be [0, 0.6, 1.2, 0]. An error occured: "boolean index did not match indexed array along dimension 0; dimension is 2 but corresponding boolean dimension is 10001" So what should I do to achieve the intended purpose?

Upvotes: 2

Views: 2664

Answers (1)

Ben.T
Ben.T

Reputation: 29635

what you want to get as result can be done with np.select such as:

def f(r, x):
    return np.select([x < 0.5,x >= 0.5], [2*r*x, 2*r*(1-x)])

Then with

r = np.array([0, 1, 2, 3])
x = np.array([0.1, 0.7, 0.3, 1])

print (f(r,x))
[0.  0.6 1.2 0. ]

EDIT: in this case, with only 2 conditions that are exclusive, you can also use np.where:

def f(r,x):
    return np.where(x<0.5,2*r*x, 2*r*(1-x))

will give the same result.

Upvotes: 4

Related Questions