Najszybszy sposób używania Numpy - podsumowuje i produkty wielowymiarowe

Question

mam tych zmiennych o następujących wymiarach:

A - (3,) 
B - (4,) 
X_r - (3,K,N,nS) 
X_u - (4,K,N,nS) 
k - (K,) 

i chcę obliczyć (A.dot(X_r[:,:,n,s])*B.dot(X_u[:,:,n,s])).dot(k) dla każdej możliwej n i s, sposób w jaki to robię teraz jest następujący:

np.array([[(A.dot(X_r[:,:,n,s])*B.dot(X_u[:,:,n,s])).dot(k) for n in xrange(N)] for s in xrange(nS)]) #nSxN 

ale to jest bardzo powolny i zastanawiałem się, czy istnieje lepszy sposób to zrobić, ale nie jestem pewien.

Jednakże istnieje inny obliczenie, czy robię i jestem pewien, że to może być zoptymalizowane:

np.sum(np.array([(X_r[:,:,n,s]*B.dot(X_u[:,:,n,s])).dot(k) for n in xrange(N)]),axis=0) 

W tym jednym Tworzę numpy tablicę tylko podsumować je w jednej osi i wyrzucić tablicę po . Gdyby to była lista w 1-D, użyłbym reduce i zoptymalizowałbym to, co powinienem użyć dla tablic numpy?

Answer 1

Korzystanie kilka np.einsum połączeń -

# Calculation of A.dot(X_r[:,:,n,s]) 
p1 = np.einsum('i,ijkl->jkl',A,X_r) 

# Calculation of B.dot(X_u[:,:,n,s]) 
p2 = np.einsum('i,ijkl->jkl',B,X_u) 

# Include .dot(k) part to get the final output 
out = np.einsum('ijk,i->kj',p1*p2,k) 

O drugim przykładzie to rozwiązuje go:

p1 = np.einsum('i,ijkl->jkl',B,X_u)#OUT_DIM - (k,N,nS) 
sol = np.einsum('ijkl,j->il',X_r*p1[None,:,:,:],k)#OUT_DIM (3,nS) 

Answer 2

Można użyć dot mnożenia macierzy w wyższych wymiarach, ale działające wskaźniki muszą być ostatnie dwa. Kiedy uporządkować swoje macierze

X_r_t = X_r.transpose(2,3,0,1) 
X_u_t = X_u.transpose(2,3,0,1) 

otrzymujemy dla pierwszego wyrazu

res1_imp = (A.dot(X_r_t)*B.dot(X_u_t)).dot(k).T # shape nS x N 

a dla drugiego wyrazu

res2_imp = np.sum((X_r_t * B.dot(X_u_t)[:,:,None,:]).dot(k),axis=0)[-1] 

Szybkość wczytywania

rozwiązanie

Divakars daje na moim komputer 10000 loops, best of 3: 21.7 µs per loop

moje rozwiązanie daje 10000 loops, best of 3: 101 µs per loop

Edit

Moi górne Timings obejmował obliczenie obu wyrażeń. Kiedy zawieram tylko pierwsze wyrażenie (jako Divakar), uzyskuję 10000 loops, best of 3: 41 µs per loop ... który jest wciąż wolniejszy, ale bliższy jego timingom