Index multiple dimensions of a multi-dimensional array with another - NumPy/ Python

Question

Lets say I have an tensor of the following form:

import numpy as np
a = np.array([ [[1,2],
                [3,4]],
               [[5,6],
                [7,3]] 
             ])
# a.shape : (2,2,2) is a tensor containing 2x2 matrices
indices = np.argmax(a, axis=2)
#print indices
for mat in a:
    max_i = np.argmax(mat,axis=1)
    # Not really working I would like to
    # change 4 in the first matrix to -1
    #  and 3 in the last to -1
    mat[max_i] = -1

print a

Now what I would like to do is to use indices as a mask on a to replace every max element with say -1. Is there a numpy way of doing this ? so far all I have figured out is using for loops.

Your brackets are mismatching. Can you update as appropriate? — gtlambert
– gtlambert, Commented Dec 31, 2015 at 19:57
Also, can you add your loopy code that gives the expected output? — Divakar
– Divakar, Commented Dec 31, 2015 at 19:58

Divakar · Accepted Answer · 2015-12-31 20:48:26Z

Here's one way using linear indexing in 3D -

m,n,r = a.shape
offset = n*r*np.arange(m)[:,None] + r*np.arange(n)
np.put(a,indices + offset,-1)

Sample run -

In [92]: a
Out[92]: 
array([[[28, 59, 26, 70],
        [57, 28, 71, 49],
        [33,  6, 10, 90]],

       [[24, 16, 83, 67],
        [96, 16, 72, 56],
        [74,  4, 71, 81]]])

In [93]: indices = np.argmax(a, axis=2)

In [94]: m,n,r = a.shape
    ...: offset = n*r*np.arange(m)[:,None] + r*np.arange(n)
    ...: np.put(a,indices + offset,-1)
    ...: 

In [95]: a
Out[95]: 
array([[[28, 59, 26, -1],
        [57, 28, -1, 49],
        [33,  6, 10, -1]],

       [[24, 16, -1, 67],
        [-1, 16, 72, 56],
        [74,  4, 71, -1]]])

Here's another way with linear indexing again, but in 2D -

m,n,r = a.shape
a.reshape(-1,r)[np.arange(m*n),indices.ravel()] = -1

Runtime tests and verify output -

In [156]: def vectorized_app1(a,indices): # 3D linear indexing
     ...:   m,n,r = a.shape
     ...:   offset = n*r*np.arange(m)[:,None] + r*np.arange(n)
     ...:   np.put(a,indices + offset,-1)
     ...: 
     ...: def vectorized_app2(a,indices): # 2D linear indexing
     ...:   m,n,r = a.shape
     ...:   a.reshape(-1,r)[np.arange(m*n),indices.ravel()] = -1
     ...:  

In [157]: # Generate random 3D array and the corresponding indices array
     ...: a = np.random.randint(0,99,(100,100,100))
     ...: indices = np.argmax(a, axis=2)
     ...: 
     ...: # Make copies for feeding into functions
     ...: ac1 = a.copy()
     ...: ac2 = a.copy()
     ...: 

In [158]: vectorized_app1(ac1,indices)

In [159]: vectorized_app2(ac2,indices)

In [160]: np.allclose(ac1,ac2)
Out[160]: True

In [161]: # Make copies for feeding into functions
     ...: ac1 = a.copy()
     ...: ac2 = a.copy()
     ...: 

In [162]: %timeit vectorized_app1(ac1,indices)
1000 loops, best of 3: 311 µs per loop

In [163]: %timeit vectorized_app2(ac2,indices)
10000 loops, best of 3: 145 µs per loop

@FranciscoVargas Shouldn't be! So, you can extend the second approach for a 4D case, like so - p,m,n,r = a.shape; a.reshape(-1,r)[np.arange(p*m*n),indices.ravel()] = -1.

ali_m · Accepted Answer · 2015-12-31 20:40:57Z

1

You can use indices to index into the last dimension of a provided that you also specify index arrays into the first two dimensions as well:

import numpy as np

a = np.array([[[1, 2],
               [3, 4]],

              [[5, 6],
               [7, 3]]])

indices = np.argmax(a, axis=2)

print(repr(a[range(a.shape[0]), range(a.shape[1]), indices]))
# array([[2, 3],
#        [2, 7]])

a[range(a.shape[0]), range(a.shape[1]), indices] = -1

print(repr(a))
# array([[[ 1, -1],
#         [ 3,  4]],

#        [[ 5,  6],
#         [-1, -1]]])

edited Dec 31, 2015 at 20:40

answered Dec 31, 2015 at 20:11

ali_m

74.6k28 gold badges230 silver badges314 bronze badges

Collectives™ on Stack Overflow

Index multiple dimensions of a multi-dimensional array with another - NumPy/ Python

2 Answers 2

1 Comment

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Related