Problem with array[np.argwhere(array==value)]-CodePudding

Taking np.argwhere as a list of coordinates in 2d leads to a wrong result. For instance:

import numpy as np

array = np.array([[1, 0, 1], [2, 0, 0], [2, 3, 0]])
coordinates = np.argwhere(array==1)
array[coordinates] = 3
print(array)

Gives:

[[3 3 3]
 [2 0 0]
 [3 3 3]]

even though only the values 1 should be changed.

CodePudding user response：

In [163]: array = np.array([[1, 0, 1], [2, 0, 0], [2, 3, 0]])   
In [164]: array
Out[164]: 
array([[1, 0, 1],
       [2, 0, 0],
       [2, 3, 0]])   
In [165]: array==1
Out[165]: 
array([[ True, False,  True],
       [False, False, False],
       [False, False, False]])

np.where/nonzero finds the coordinates of the True values, giving a tuple of arrays:

In [166]: np.nonzero(array==1)
Out[166]: (array([0, 0], dtype=int64), array([0, 2], dtype=int64))

That tuple can be used - as is - to index the array:

In [167]: array[np.nonzero(array==1)]
Out[167]: array([1, 1])

argwhere returns the same values, but as 2d array, one column per dimension:

In [168]: np.argwhere(array==1)
Out[168]: 
array([[0, 0],
       [0, 2]], dtype=int64)

It is wrong to use it as an index:

In [169]: array[np.argwhere(array==1),:]
Out[169]: 
array([[[1, 0, 1],
        [1, 0, 1]],

       [[1, 0, 1],
        [2, 3, 0]]])

I added the : to show more clearly that it is just indexing the first dimension of array, not both as done with the tuple of arrays in [167].

To use argwhere to index the array we have to do:

In [170]: x = np.argwhere(array==1)
In [171]: x
Out[171]: 
array([[0, 0],
       [0, 2]], dtype=int64)

Apply each column separately to the dimensions:

In [172]: array[x[:,0], x[:,1]]
Out[172]: array([1, 1])

argwhere is more useful if you want to iterate through the nonzero elements:

In [173]: for i in x:
     ...:     print(array[tuple(i)])
     ...:     
1
1

But as commented, you don't need the where/argwhere step; just use the boolean array as index

In [174]: array[array==1]
Out[174]: array([1, 1])