python numpy argmax to max in multidimensional array

Question

I have the following code:

import numpy as np
sample = np.random.random((10,10,3))
argmax_indices = np.argmax(sample, axis=2)

i.e. I take the argmax along axis=2 and it gives me a (10,10) matrix. Now, I want to assign these indices value 0. For this, I want to index the sample array. I tried:

max_values = sample[argmax_indices]

but it doesn't work. I want something like

max_values = sample[argmax_indices]
sample[argmax_indices] = 0

I simply validate by checking that max_values - np.max(sample, axis=2) should give a zero matrix of shape (10,10). Any help will be appreciated.

Why so? The value of a matrix at maximizing indices along a plane/axis should be equal to the maximum value of matrix along that plane. No? Or is numpy taking maximum over all values of axis=2 irrespective of what is 1st and 2nd axis? — Rayyan Riaz, Feb 28 '17 at 22:05
Let me ask that differently. My array is 10x10x3. So, for every i(x-axis) and j(y-axis), I want the best k(z-axis). I want its index as well as the value. Is above approach not right for it? — Rayyan Riaz, Feb 28 '17 at 22:07
Sorry I don't get it but it looks like your indexing is wrong anyway, look at the answers. — Stop harming Monica, Feb 28 '17 at 22:17

Divakar · Accepted Answer · 2017-02-28T22:27:17.317

Here's one approach -

m,n = sample.shape[:2]
I,J = np.ogrid[:m,:n]
max_values = sample[I,J, argmax_indices]
sample[I,J, argmax_indices] = 0

Sample step-by-step run

1) Sample input array :

In [261]: a = np.random.randint(0,9,(2,2,3))

In [262]: a
Out[262]: 
array([[[8, 4, 6],
        [7, 6, 2]],

       [[1, 8, 1],
        [4, 6, 4]]])

2) Get the argmax indices along axis=2 :

In [263]: idx = a.argmax(axis=2)

3) Get the shape and arrays for indexing into first two dims :

In [264]: m,n = a.shape[:2]

In [265]: I,J = np.ogrid[:m,:n]

4) Index using I, J and idx for storing the max values using advanced-indexing :

In [267]: max_values = a[I,J,idx]

In [268]: max_values
Out[268]: 
array([[8, 7],
       [8, 6]])

5) Verify that we are getting an all zeros array after subtracting np.max(a,axis=2) from max_values :

In [306]: max_values - np.max(a, axis=2)
Out[306]: 
array([[0, 0],
       [0, 0]])

6) Again using advanced-indexing assign those places as zeros and do one more level of visual verification :

In [269]: a[I,J,idx] = 0

In [270]: a
Out[270]: 
array([[[0, 4, 6], # <=== Compare this against the original version
        [0, 6, 2]],

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

Thanks a lot :) It is perfect – Rayyan Riaz Feb 28 '17 at 22:36 — Rayyan Riaz, Feb 28 '17 at 22:36

score 2 · Answer 2 · answered Jun 17 '17 at 16:14

2

An alternative to np.ogrid is np.indices.

I, J = np.indices(argmax_indices.shape)

sample[I,J,argmax_indices] = 0

answered Jun 17 '17 at 16:14

TianLe Ma

41
3

score 0 · Answer 3 · answered Mar 07 '18 at 17:11

This can also be generalized to handle matrices of any dimensionality. The resulting function will set the largest value in every 1-d vector of the matrix along any dimension d desired (dimension 2 in the case of the original question) to 0 (or to whatever value is desired):

def set_zero(sample, d, val):
    """Set all max value along dimension d in matrix sample to value val."""
    argmax_idxs = sample.argmax(d)
    idxs = [np.indices(argmax_idxs.shape)[j].flatten() for j in range(len(argmax_idxs.shape))]
    idxs.insert(d, argmax_idxs.flatten())
    sample[idxs] = val
    return sample

set_zero(sample, d=2, val=0)

(Tested for numpy 1.14.1 on python 3.6.4 and python 2.7.14)

python numpy argmax to max in multidimensional array

3 Answers3

Linked