¿Cómo encontrar las diferencias por pares entre filas de dos matrices muy grandes usando numpy?

Dadas dos matrices, quiero calcular las diferencias por pares entre todas las filas. Cada matriz tiene 1000 filas y 100 columnas, por lo que son bastante grandes. Intenté usar un bucle for y transmisión pura, pero el bucle for parece funcionar más rápido. ¿Estoy haciendo algo mal? Aquí está el código:

from numpy import * A = random.randn(1000,100) B = random.randn(1000,100) start = time.time() for a in A: sum((a - B)**2,1) print time.time() - start # pure broadcasting start = time.time() ((A[:,newaxis,:] - B)**2).sum(-1) print time.time() - start 

El método de transmisión demora aproximadamente 1 segundo más y es aún más largo para matrices grandes. ¿Alguna idea de cómo acelerar esto puramente usando numpy?

Aquí hay otra forma de realizar:

(ab) ^ 2 = a ^ 2 + b ^ 2 – 2ab

con np.einsum para los dos primeros términos y dot-product para el tercero –

 import numpy as np np.einsum('ij,ij->i',A,A)[:,None] + np.einsum('ij,ij->i',B,B) - 2*np.dot(A,BT) 

Prueba de tiempo de ejecución

Enfoques –

 def loopy_app(A,B): m,n = A.shape[0], B.shape[0] out = np.empty((m,n)) for i,a in enumerate(A): out[i] = np.sum((a - B)**2,1) return out def broadcasting_app(A,B): return ((A[:,np.newaxis,:] - B)**2).sum(-1) # @Paul Panzer's soln def outer_sum_dot_app(A,B): return np.add.outer((A*A).sum(axis=-1), (B*B).sum(axis=-1)) - 2*np.dot(A,BT) # @Daniel Forsman's soln def einsum_all_app(A,B): return np.einsum('ijk,ijk->ij', A[:,None,:] - B[None,:,:], \ A[:,None,:] - B[None,:,:]) # Proposed in this post def outer_einsum_dot_app(A,B): return np.einsum('ij,ij->i',A,A)[:,None] + np.einsum('ij,ij->i',B,B) - \ 2*np.dot(A,BT) 

Tiempos –

 In [51]: A = np.random.randn(1000,100) ...: B = np.random.randn(1000,100) ...: In [52]: %timeit loopy_app(A,B) ...: %timeit broadcasting_app(A,B) ...: %timeit outer_sum_dot_app(A,B) ...: %timeit einsum_all_app(A,B) ...: %timeit outer_einsum_dot_app(A,B) ...: 10 loops, best of 3: 136 ms per loop 1 loops, best of 3: 302 ms per loop 100 loops, best of 3: 8.51 ms per loop 1 loops, best of 3: 341 ms per loop 100 loops, best of 3: 8.38 ms per loop 

Aquí hay una solución que evita el bucle y los intermedios grandes:

 from numpy import * import time A = random.randn(1000,100) B = random.randn(1000,100) start = time.time() for a in A: sum((a - B)**2,1) print time.time() - start # pure broadcasting start = time.time() ((A[:,newaxis,:] - B)**2).sum(-1) print time.time() - start #matmul start = time.time() add.outer((A*A).sum(axis=-1), (B*B).sum(axis=-1)) - 2*dot(A,BT) print time.time() - start 

Huellas dactilares:

 0.546781778336 0.674743175507 0.10723400116 

Otro trabajo para np.einsum

 np.einsum('ijk,ijk->ij', A[:,None,:] - B[None,:,:], A[:,None,:] - B[None,:,:])