numpy assignment with masking

Question

I am a newbie to Python. During an exercise I am supposed to use a mask to multiply all values below 100 in the following list by 2:

a = np.array([230, 10, 284, 39, 76])

So I wrote the following code:

import numpy as np
a = np.array([230, 10, 284, 39, 76])
cut = 100
a[a < cut] = a*2    

This results in the following error:
IndexError: index 230 is out of bounds for axis 0 with size 5

This is confusing since to my understanding, the a in [a < cut] actually refers to each value in array a, but the a in a*2 refers to the whole array.

How can I correct this code using the masking method, instead of using a loop?

Answer 1

Not exactly sure what you want, if you want to assign to places where a < cut holds (a < cut = [0, 1, 0, 1, 1] is the boolean index), when you assign to a[a < cut], you assign to the places where the element is 1, meaning on the right side it expects a numpy array of size 3 (or of course one number). You can do this

In [1]: a = np.array([230, 10, 284, 39, 76])

In [2]: a[a < cut] = 999

In [3]: a
Out[3]: array([230, 999, 284, 999, 999])

Or

In [1]: a = np.array([230, 10, 284, 39, 76])

In [2]: a[a < cut] = a[a < cut] * 2

In [3]: a
Out[3]: array([230,  20, 284,  78, 152])

To multiply the selected elements by 2.

Answer 2

Alternatively, once the mask is defined, you can use numpy.where or numpy.putmask

import numpy as np

a = np.array([230, 10, 284, 39, 76])
cut = 100
mask = a < cut # defines the mask

The first does't change the original array:

res = np.where(mask, a*2,a)
a #=> [230  10 284  39  76]
res #=> [230  20 284  78 152]

The second modifies the original array:

np.putmask(a, mask, a*2)
a #=> [230  20 284  78 152]