CODEWITHSUNDEEP
pythonnumpyvectorization
Managarm
Managarm

Reputation: 1121

Vectorize an index-based matrix operation in numpy

How can I vectorize the following loop? 

def my_fnc():
   m = np.arange(27.).reshape((3,3,3))
   ret = np.empty_like(m)

   it = np.nditer(m, flags=['multi_index'])
   for x in it:
      i,j,k = it.multi_index
      ret[i,j,k] = x / m[i,j,i]
   return ret

Basically I'm dividing each value in m by something similar to a diagonal. Not all values in m will be different, the arange is just an example.

Thanks in advance! ~

P.S.: here's the output of the function above, don't mind the nans :)

array([[[        nan,         inf,         inf],
        [ 1.        ,  1.33333333,  1.66666667],
        [ 1.        ,  1.16666667,  1.33333333]],

       [[ 0.9       ,  1.        ,  1.1       ],
        [ 0.92307692,  1.        ,  1.07692308],
        [ 0.9375    ,  1.        ,  1.0625    ]],

       [[ 0.9       ,  0.95      ,  1.        ],
        [ 0.91304348,  0.95652174,  1.        ],
        [ 0.92307692,  0.96153846,  1.        ]]])

Upvotes: 2

Views: 438

Answers (1)

Divakar
Divakar

Reputation: 221704

Use advanced-indexing to get the m[i,j,i] equivalent in one go and then simply divide input array by it -

r = np.arange(len(m))
ret = m/m[r,:,r,None] # Add new axis with None to allow for broadcasting

Upvotes: 2

Related Questions