How can I control each element of a numpy array according to its relative size?

Question

Say we have an array

a = np.array([10,30,50, 20, 10, 90, 0, 25]) 

Pseudocode for what I want -

 if a[x] > 80 then perform funcA on a[x] if 40 < a[x] <= 80 then perform funcB on a[x] if a[x] <= 40 then perform funcC on a[x] 

What is the cleanest way to accomplish this with numpy functions?

Answer 1

Usually you try to avoid any Python loops over NumPy arrays - why you use NumPy in the first place. As an example, I assume that funcA() adds 1 to all elements, funcB() adds 2 and funcC() adds 3 (please clarify what they really do for a more specific example). To achieve what you want, you can do this:

 subset_a = a > 80 subset_b = (40 < a) & (a <= 80) subset_c = a <= 40 a[subset_a] += 1 a[subset_b] += 2 a[subset_c] += 3 

It uses advanced NumPy indexing . For example, a > 80 evaluates an array of booleans that can be used to select entries in the array that satisfy the condition.

Answer 2

Take a look at numpy.piecewise. I think you want:

 np.piecewise( a, [a > 80, (40 < a) & (a <= 80), a <= 40], [funcA, funcB, funcC] ) 

Answer 3

I like it:

 b = np.empty(a.shape) b[a < 40] = funcA(a[a < 40]) b[(a > 40) & (a <= 80)] = funcB(a[(a > 40) & (a <= 80)]) b[a > 80] = funcC(a[a > 80]) 

This avoids strange behavior when funcA sets the element a , which, for example, was from 39 to 41, thereby bringing it to the range for funcB .

Answer 4

If you need more advanced features, you can use

  newfunc=numpy.vectorize(lambda x: func(x)) result=newfunc(yourArray) 

where func (x) is your function.