How can each element of a numpy array be operated upon according to its relative value?

Question

Let say that we have an array

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

The pseudo code 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 perform this using numpy functions?

Answer 1

Usually, you try to avoid any Python loops over NumPy arrays -- that's why you use NumPy in the first place. For the sake of example, I assume that funcA() adds 1 to all elements, funcB() adds 2 and funcC() adds 3 (please elaborate what they really do for a more tailor-made 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

This uses NumPy advanced indexing. For example a > 80 evaluates to an array of Boolean values which can be used to select the entries in the array fulfilling the condition.

Answer 2

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 this:

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 weird behavior when funcA sets an element of a that had been 39 to 41, for instance, thus bringing it into the range for funcB.