CODEWITHSUNDEEP
pythonarraysnumpy
Luca
Luca

Reputation: 11016

numpy: issues with indexing

I am having some issue where I sort a numpy array and get the indexes for the sorting array but applying the index to the original array does not do what I expected. So, here is a test case for what I am doing:

 import numpy as np

 # Two 3x3 matrices
 x = np.random.rand(2, 3, 3)
 # Perform some decomposition (Never mind the matrices are not hermitian...) 
 evals, evecs = np.linalg.eigh(x)

 # evals has shape (2, 3), evecs has shape (2, 3, 3)
 indices = evals.argsort(axis=1)[..., ::-1] # Do descending sort

 # Now I want to apply the index to evals.
 evals = evals[:, indices]

Instead of getting back a (2, 3) array, I am getting a (2, 3, 3) array back where the rows are getting replicated. Something like:

array([[[ 1.15628047,  0.16853886, -0.28607138],
        [ 1.15628047,  0.16853886, -0.28607138]],

       [[ 2.4311532 , -0.00754817, -0.24086572],
        [ 2.4311532 , -0.00754817, -0.24086572]]])

I am not sure why that is. Would appreciate any help.

Upvotes: 1

Views: 75

Answers (1)

gg349
gg349

Reputation: 22701

This should work:

import numpy as np
idx0 = np.arange(evals.shape[0])[:,np.newaxis]
idx1 = evals.argsort(1)[...,::-1]
evals[idx0,idx1]

This sorts each row, individually, by decreasing order.

EDIT: In this case you need the (idx0,idx1) to further process the eigenvectors evecs. If this wasn't the case, it is straightforward to do

evals.sort()
evals = evals[:,::-1]

Upvotes: 1

Related Questions