CODEWITHSUNDEEP
pythonarraysnumpy
dipankar
dipankar

Reputation: 79

How to properly index a big matrix in python

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

Upvotes: 1

Views: 85

Answers (2)

fountainhead
fountainhead

Reputation: 3722

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]]]

Upvotes: 0

Sheldore
Sheldore

Reputation: 39062

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]]])

Upvotes: 3

Related Questions