How to properly index a big matrix in python

Question

I have a big numpy array with certain entries. Let's say a dummy example is:

    arr = np.array([[[1.0, 2.0, 3.0],[1.5, 1.8, 3.2]],
                    [[1.3, 1.7, 1.9],[1.4, 1.9, 2.1]],
                    [[1.8, 2.2, 2.5],[2.0, 2.2, 2.8]]])

I would like to know all the indexes where the entries of arr fall within some range, say 1.5 and 2.4 . And I would like to fill another matrix of the same shape as arr with 1 at indexes where value of arr falls within the range, otherwise with 0 . That is, I would like to get a matrix like:

mask = np.array([[[0, 1, 0], [1, 1, 0]],
                 [[0, 1, 1], [0, 1, 1]],
                 [[1, 1, 0], [1, 1, 0]]])

Is there any simple numpy trick to do this? I know it is straightforward to do it with a for loop , but since my arr is pretty big in size, I would want it to be reasonably fast. Thanks

Answer 1

You can use masking and np.where : First create a conditional mask combining your two boundary conditions and then pass it to np.where . The matrix will be assigned 1 where this condition holds True else 0 if its False

Minimal working answer

import numpy as np

arr = np.array([[[1.0, 2.0, 3.0],[1.5, 1.8, 3.2]],
                [[1.3, 1.7, 1.9],[1.4, 1.9, 2.1]],
                [[1.8, 2.2, 2.5],[2.0, 2.2, 2.8]]])

mask = ((arr>1.5) & (arr<2.4))
arr = np.where(mask, 1, 0)
print (arr)

Output

array([[[0, 1, 0],
        [0, 1, 0]],

       [[0, 1, 1],
        [0, 1, 1]],

       [[1, 1, 0],
        [1, 1, 0]]])

Answer 2

Create a boolean mask meeting your conditions. Adding 0 to boolean values will convert them to numeric results:

import numpy as np

arr = np.array([[[1.0, 2.0, 3.0],[1.5, 1.8, 3.2]],
                [[1.3, 1.7, 1.9],[1.4, 1.9, 2.1]],
                [[1.8, 2.2, 2.5],[2.0, 2.2, 2.8]]])

arr_out = ((arr>=1.5) & (arr<=2.4)) + 0
print(arr_out)

Alternatively:

arr_out = np.array(((arr>=1.5) & (arr<=2.4)), dtype=np.uint8)
print(arr_out)

Or, as suggested by @hpaulj:

arr_out = ((arr>=1.5) & (arr<=2.4)).astype(int)
print (arr_out)

Output:

[[[0 1 0]
  [1 1 0]]

 [[0 1 1]
  [0 1 1]]

 [[1 1 0]
  [1 1 0]]]