Векторизация и оптимизация матричного вычитания

Question

Можно ли векторизовать / оптимизировать следующий цикл?

In [33]: a = np.arange(10000 * 700).reshape([10000, 700])

In [34]: b = np.arange(1000 * 700).reshape([1000, 700])

In [35]: c = np.empty([b.shape[0], a.shape[0]])

In [36]: for i in range(b.shape[0]):
    ...:     c[i] = np.argsort(np.linalg.norm(a - b[i], axis=1))
    ...: 

---

Редактировать:

Я считаю, что должно работать следующее:

d = np.argsort(np.linalg.norm(a[:, None] - b, axis=2), axis=1)

Но я получаю MemoryError за a[:, None] - b.Я в правильном направлении?Что можно сделать относительно MemoryError?

Answer 1

Самый простой способ был бы с cdist -

from scipy.spatial.distance import cdist

cdist(b,a).argsort(axis=1)

Эквивалентный с pairwise_distances -

from sklearn.metrics import pairwise_distances

pairwise_distances(b,a).argsort(1)

Времядля данных образца -

In [201]: %%timeit # original solution
     ...: c = np.empty([b.shape[0], a.shape[0]],dtype=int)
     ...: for i in range(b.shape[0]):
     ...:     c[i] = np.argsort(np.linalg.norm(a - b[i], axis=1))
1 loop, best of 3: 40.6 s per loop

In [202]: %timeit pairwise_distances(b,a).argsort(1)
1 loop, best of 3: 384 ms per loop

100x+ ускорение!