Numpyic way to sort a matrix based on another similar matrix

Question

Say I have a matrix Y of random float numbers from 0 to 10 with shape (10, 3):

import numpy as np
np.random.seed(99)
Y = np.random.uniform(0, 10, (10, 3))
print(Y)

Output:

[[6.72278559 4.88078399 8.25495174]
 [0.31446388 8.08049963 5.6561742 ]
 [2.97622499 0.46695721 9.90627399]
 [0.06825733 7.69793028 7.46767101]
 [3.77438936 4.94147452 9.28948392]
 [3.95454044 9.73956297 5.24414715]
 [0.93613093 8.13308413 2.11686786]
 [5.54345785 2.92269116 8.1614236 ]
 [8.28042566 2.21577372 6.44834702]
 [0.95181622 4.11663239 0.96865261]]

I am now given a matrix X with same shape that can be seen as obtained by adding small noises to Y and then shuffling the rows:

X = np.random.normal(Y, scale=0.1)
np.random.shuffle(X)
print(X)

Output:

[[ 4.04067271  9.90959141  5.19126867]
 [ 5.59873104  2.84109306  8.11175891]
 [ 0.10743952  7.74620162  7.51100441]
 [ 3.60396019  4.91708372  9.07551354]
 [ 0.9400948   4.15448712  1.04187208]
 [ 2.91884302  0.47222752 10.12700505]
 [ 0.30995155  8.09263241  5.74876947]
 [ 1.11247872  8.02092335  1.99767444]
 [ 6.68543696  4.8345869   8.17330513]
 [ 8.38904822  2.11830619  6.42013343]]

Now I want to sort the matrix X based on Y by row. I already know each pair of column values in each matching pair of rows are not different from each other more than a tolerance of 0.5. I managed to write the following code and it is working fine.

def sort_X_by_Y(X, Y, tol):
    idxs = [next(i for i in range(len(X)) if all(abs(X[i] - row) <= tol)) for row in Y]
    return X[idxs]

print(sort_X_by_Y(X, Y, tol=0.5))

Output:

[[ 6.68543696  4.8345869   8.17330513]
 [ 0.30995155  8.09263241  5.74876947]
 [ 2.91884302  0.47222752 10.12700505]
 [ 0.10743952  7.74620162  7.51100441]
 [ 3.60396019  4.91708372  9.07551354]
 [ 4.04067271  9.90959141  5.19126867]
 [ 1.11247872  8.02092335  1.99767444]
 [ 5.59873104  2.84109306  8.11175891]
 [ 8.38904822  2.11830619  6.42013343]
 [ 0.9400948   4.15448712  1.04187208]]

However, in reality I am sorting (1000, 3) matrices and my code is way too slow. I feel like there should be more numpyic way to code this. Any suggestions?

Answer 1

This is a vectorized version of your algorithm. It runs ~26.5x faster than your implementation for 1000 samples. But an additional boolean array with shape (1000,1000,3) is created. There is a chance that rows will have similar values within the tolerance and a wrong row is selected.

tol = .5
X[(np.abs(Y[:, np.newaxis] - X) <= tol).all(2).argmax(1)]

Output

array([[ 6.68543696,  4.8345869 ,  8.17330513],
       [ 0.30995155,  8.09263241,  5.74876947],
       [ 2.91884302,  0.47222752, 10.12700505],
       [ 0.10743952,  7.74620162,  7.51100441],
       [ 3.60396019,  4.91708372,  9.07551354],
       [ 4.04067271,  9.90959141,  5.19126867],
       [ 1.11247872,  8.02092335,  1.99767444],
       [ 5.59873104,  2.84109306,  8.11175891],
       [ 8.38904822,  2.11830619,  6.42013343],
       [ 0.9400948 ,  4.15448712,  1.04187208]])

More robust solutions with L1-norm

X[np.abs(Y[:, np.newaxis] - X).sum(2).argmin(1)]

Or L2-norm

X[((Y[:, np.newaxis] - X)**2).sum(2).argmin(1)]