Numpy indexing set 1 to max value and zero's to all others

Question

I think I've misunderstood something with indexing in numpy.

I have a 3D-numpy array of shape (dim_x, dim_y, dim_z) and I want to find the maximum along the third axis (dim_z), and set its value to 1 and all the others to zero.

The problem is that I end up with several 1 in the same row, even if values are different.

Here is the code :

>>> test = np.random.rand(2,3,2)
>>> test
array([[[ 0.13110146,  0.07138861],
        [ 0.84444158,  0.35296986],
        [ 0.97414498,  0.63728852]],

       [[ 0.61301975,  0.02313646],
        [ 0.14251848,  0.91090492],
        [ 0.14217992,  0.41549218]]])

>>> result = np.zeros_like(test)
>>> result[:test.shape[0], np.arange(test.shape[1]), np.argmax(test, axis=2)]=1
>>> result
array([[[ 1.,  0.],
        [ 1.,  1.],
        [ 1.,  1.]],

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

I was expecting to end with :

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

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

Probably I'm missing something here. From what I've understood, 0:dim_x, np.arange(dim_y) returns dim_x of dim_y tuples and np.argmax(test, axis=dim_z) has the shape (dim_x, dim_y) so if the indexing is of the form [x, y, z] a couple [x, y] is not supposed to appear twice.

Could someone explain me where I'm wrong ? Thanks in advance.

Answer 1

What we are looking for

We get the argmax indices along the last axis -

idx = np.argmax(test, axis=2)

For the given sample data, we have idx :

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

Now, idx covers the first and second axes, while getting those argmax indices.

To assign the corresponding ones in the output, we need to create range arrays for the first two axes covering the lengths along those and aligned according to the shape of idx. Now, idx is a 2D array of shape (m,n), where m = test.shape[0] and n = test.shape[1].

Thus, the range arrays for assignment into first two axes of output must be -

X = np.arange(test.shape[0])[:,None]
Y = np.arange(test.shape[1])

Notice, the extension of the first range array to 2D is needed to have it aligned against the rows of idx and Y would align against the cols of idx -

In [239]: X
Out[239]: 
array([[0],
       [1]])

In [240]: Y
Out[240]: array([0, 1, 2])

Schematically put -

idx :
    Y array
    --------->
    x x x | X array
    x x x |
          v

The fault in original code

Your code was -

result[:test.shape[0], np.arange(test.shape[1]), ..

This is essentially :

result[:, np.arange(test.shape[1]), ...

So, you are selecting all elements along the first axis, instead of only selecting the corresponding ones that correspond to idx indices. In that process, you were selecting a lot more than required elements for assignment and hence you were seeing many more than required 1s in result array.

The correction

Thus, the only correction needed was indexing into the first axis with the range array and a working solution would be -

result[np.arange(test.shape[0])[:,None], np.arange(test.shape[1]), ...

The alternative(s)

Alternatively, using the range arrays created earlier with X and Y -

result[X,Y,idx] = 1

Another way to get X,Y would be with np.mgrid -

m,n = test.shape[:2]
X,Y = np.ogrid[:m,:n]

Answer 2

Here is an easier way to do it:

>>>  test == test.max(axis=2, keepdims=1)
array([[[ True, False],
        [ True, False],
        [ True, False]],

       [[ True, False],
        [False,  True],
        [False,  True]]], dtype=bool)

...and if you really want that as floating-point 1.0 and 0.0, then convert it:

>>> (test==test.max(axis=2, keepdims=1)).astype(float)
array([[[ 1.,  0.],
        [ 1.,  0.],
        [ 1.,  0.]],

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

Here is a way to do it with only one winner per row-column combo (i.e. no ties, as discussed in comments):

rowmesh, colmesh = np.meshgrid(range(test.shape[0]), range(test.shape[1]), indexing='ij')
maxloc = np.argmax(test, axis=2)
flatind = np.ravel_multi_index( [rowmesh, colmesh, maxloc ], test.shape )
result = np.zeros_like(test)
result.flat[flatind] = 1

UPDATE after reading hpaulj's answer:

rowmesh, colmesh = np.ix_(range(test.shape[0]), range(test.shape[1]))

is a more-efficient, more numpythonic, alternative to my meshgrid call (the rest of the code stays the same)

The issue of why your approach fails is hard to explain, but here's one place where intuition could start: your slicing approach says "all rows, times all columns, times a certain sequence of layers". How many elements is that slice in total? By contrast, how many elements do you actually want to set to 1? It can be instructive to look at the values you get when you view the corresponding test values of the slice you're trying to assign to:

>>> test[:, :, maxloc].shape
(2, 3, 2, 3)   # oops!  it's because maxloc itself is 2x3

>>> test[:, :, maxloc]
array([[[[ 0.13110146,  0.13110146,  0.13110146],
         [ 0.13110146,  0.07138861,  0.07138861]],

        [[ 0.84444158,  0.84444158,  0.84444158],
         [ 0.84444158,  0.35296986,  0.35296986]],

        [[ 0.97414498,  0.97414498,  0.97414498],
         [ 0.97414498,  0.63728852,  0.63728852]]],


       [[[ 0.61301975,  0.61301975,  0.61301975],
         [ 0.61301975,  0.02313646,  0.02313646]],

        [[ 0.14251848,  0.14251848,  0.14251848],
         [ 0.14251848,  0.91090492,  0.91090492]],

        [[ 0.14217992,  0.14217992,  0.14217992],
         [ 0.14217992,  0.41549218,  0.41549218]]]])  # note the repetition, because in maxloc you're repeatedly asking for layer 0 sometimes, and sometimes repeatedly for layer 1

Answer 3

I think there's a problem with mixing basic (slice) and advanced indexing. It's easier to see when selecting value from an array than with this assignment; but it can result in transposed axes. For a problem like this it is better use advanced indexing all around, as provided by ix_

In [24]: test = np.random.rand(2,3,2)
In [25]: idx=np.argmax(test,axis=2)
In [26]: idx
Out[26]: 
array([[1, 0, 1],
       [0, 1, 1]], dtype=int32)

with basic and advanced:

In [31]: res1 = np.zeros_like(test)
In [32]: res1[:, np.arange(test.shape[1]), idx]=1
In [33]: res1
Out[33]: 
array([[[ 1.,  1.],
        [ 1.,  1.],
        [ 0.,  1.]],

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

with advanced:

In [35]: I,J = np.ix_(range(test.shape[0]), range(test.shape[1]))
In [36]: I
Out[36]: 
array([[0],
       [1]])
In [37]: J
Out[37]: array([[0, 1, 2]])
In [38]: res2 = np.zeros_like(test)
In [40]: res2[I, J , idx]=1
In [41]: res2
Out[41]: 
array([[[ 0.,  1.],
        [ 1.,  0.],
        [ 0.,  1.]],

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

On further thought, the use of the slice for the 1st dimension is just wrong , if the goal is to set or find the 6 argmax values

In [54]: test
Out[54]: 
array([[[ 0.15288242,  0.36013289],
        [ 0.90794601,  0.15265616],
        [ 0.34014976,  0.53804266]],

       [[ 0.97979479,  0.15898605],
        [ 0.04933804,  0.89804999],
        [ 0.10199319,  0.76170911]]])
In [55]: test[I, J, idx]
Out[55]: 
array([[ 0.36013289,  0.90794601,  0.53804266],
       [ 0.97979479,  0.89804999,  0.76170911]])

In [56]: test[:, J, idx]
Out[56]: 
array([[[ 0.36013289,  0.90794601,  0.53804266],
        [ 0.15288242,  0.15265616,  0.53804266]],

       [[ 0.15898605,  0.04933804,  0.76170911],
        [ 0.97979479,  0.89804999,  0.76170911]]])

With the slice it selects a (2,3,2) set of values from test (or res), not the intended (2,3). There 2 extra rows.