python - Using NumPy argsort and take in 2D arrays -

the aim calculate distance matrix between 2 sets of points (set1 , set2), use argsort() obtain sorted indexes , take() extract sorted array. know sort() directly, need indexes next steps.

i using fancy indexing concepts discussed here. not manage use take() directly obtained matrix of indexes, adding each row corresponding quantity makes work, because take() flattens source array making second row elements index += len(set2), third row index += 2*len(set2) , forth (see below):

dist  = np.subtract.outer( set1[:,0], set2[:,0] )**2 dist += np.subtract.outer( set1[:,1], set2[:,1] )**2 dist += np.subtract.outer( set1[:,2], set2[:,2] )**2 = np.argsort( dist, axis=1 ) += np.array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],                [10, 10, 10, 10, 10, 10, 10, 10, 10, 10],                [20, 20, 20, 20, 20, 20, 20, 20, 20, 20],                [30, 30, 30, 30, 30, 30, 30, 30, 30, 30]]) s1 = np.sort(dist,axis=1) s2 = np.take(dist,a) np.nonzero((s1-s2)) == false #true # meaning works... 

the main question is: there direct way use take() without summing these indexes?

data play with:

set1 = np.array([[ 250., 0.,    0.],                  [ 250., 0.,  510.],                  [-250., 0.,    0.],                  [-250., 0.,    0.]])  set2 = np.array([[  61.0, 243.1, 8.3],                  [ -43.6, 246.8, 8.4],                  [ 102.5, 228.8, 8.4],                  [  69.5, 240.9, 8.4],                  [ 133.4, 212.2, 8.4],                  [ -52.3, 245.1, 8.4],                  [-125.8, 216.8, 8.5],                  [-154.9, 197.1, 8.6],                  [  61.0, 243.1, 8.7],                  [ -26.2, 249.3, 8.7]]) 

other related questions:

- euclidean distance between points in 2 different numpy arrays, not within

i don't think there way use np.take without going flat indices. since dimensions change, better off using np.ravel_multi_index that, doing this:

a = np.argsort(dist, axis=1) = np.ravel_multi_index((np.arange(dist.shape[0])[:, none], a), dims=dist.shape) 

alternatively, can use fancy indexing without using take:

s2 = dist[np.arange(4)[:, none], a]