Piecewise function in numpy with multiple arguments

Question

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?

Ben.T · Accepted Answer · 2018-11-24 02:11:42Z

4

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.

edited Nov 24, 2018 at 2:11

answered Nov 24, 2018 at 1:40

Ben.T

29.7k6 gold badges39 silver badges57 bronze badges

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

Collectives™ on Stack Overflow

Piecewise function in numpy with multiple arguments

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related