Arman
Arman

Reputation: 907

Get indices of elements that are greater than a threshold in 2D numpy array

I have a 2D numpy array:

x = np.array([
 [  1.92043482e-04,   0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
    0.00000000e+00,   0.00000000e+00,   2.41005634e-03,   0.00000000e+00,
    7.19330120e-04,   0.00000000e+00,   0.00000000e+00,   1.42886875e-04,
    0.00000000e+00,   0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
    0.00000000e+00,   9.79279411e-05,   7.88888657e-04,   0.00000000e+00,
    0.00000000e+00,   1.40425916e-01,   0.00000000e+00,   1.13955893e-02,
    7.36868947e-03,   3.67091988e-04,   0.00000000e+00,   0.00000000e+00,
    0.00000000e+00,   0.00000000e+00,   1.72037105e-03,   1.72377961e-03,
    0.00000000e+00,   0.00000000e+00,   1.19532061e-01,   0.00000000e+00,
    0.00000000e+00,   0.00000000e+00,   0.00000000e+00,   3.37249481e-04,
    0.00000000e+00,   0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
    0.00000000e+00,   0.00000000e+00,   1.75111492e-03,   0.00000000e+00,
    0.00000000e+00,   1.12639313e-02],
 [  0.00000000e+00,   0.00000000e+00,   1.10271735e-04,   5.98736562e-04,
    6.77961628e-04,   7.49569659e-04,   0.00000000e+00,   0.00000000e+00,
    2.91697850e-03,   0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
    0.00000000e+00,   0.00000000e+00,   3.30257021e-04,   2.46629275e-04,
    0.00000000e+00,   1.87586441e-02,   6.49103144e-04,   0.00000000e+00,
    1.19046355e-04,   0.00000000e+00,   0.00000000e+00,   2.69499898e-03,
    1.48525386e-02,   0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
    0.00000000e+00,   0.00000000e+00,   0.00000000e+00,   1.18803119e-03,
    3.93100829e-04,   0.00000000e+00,   3.76245304e-04,   2.79537738e-02,
    0.00000000e+00,   1.20738457e-03,   9.74669064e-06,   7.18680093e-04,
    1.61546793e-02,   3.49360861e-04,   0.00000000e+00,   0.00000000e+00,
    0.00000000e+00,   0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
    0.00000000e+00,   0.00000000e+00]])

How do I get indices of the elements that are greater than 0.01?

Right now, I'm doing t = np.argmax(x, axis=1) to get the index of the maximum value from each and the result of it is: [21 35]. How do I achieve the above?

Upvotes: 51

Views: 97449

Answers (2)

cottontail
cottontail

Reputation: 23021

One can also use np.nonzero() to get a tuple of arrays, one for each dimension of x containing the indices where the condition is True.

x_indices, y_indices = np.nonzero(x > 0.01)
# (array([0, 0, 0, 0, 1, 1, 1, 1], dtype=int64), array([21, 23, 34, 49, 17, 24, 35, 40], dtype=int64))

A nice thing about it is that it can be used to index an array right away. For example, if we want to filter the elements that are greater than 0.01, then

x[np.nonzero(x>0.01)]

nonzero groups indices by dimension while argwhere groups by element (which is just looking at the same thing from a different side), so the following is True:

(np.argwhere(x>0.01).T == np.nonzero(x>0.01)).all()   # True

Upvotes: 2

maxymoo
maxymoo

Reputation: 36545

You can use np.argwhere to return the indices of all the entries in an array matching a boolean condition:

>>> x = np.array([[0,0.2,0.5],[0.05,0.01,0]])

>>> np.argwhere(x > 0.01)
array([[0, 1],
       [0, 2],
       [1, 0]])    

Upvotes: 69

Related Questions