Численное умножение векторов разного размера для петель

Question

У меня есть матрица, скажем, P размера (X,Y). Кроме того, у меня есть две матрицы, скажем, Kx и Ky размера (M,N) как матрица pk размера (M,N) и двух векторов u и v из X и Y соответственно. Например, они могут быть определены следующим образом:

import numpy as np 
P = np.zeros((X,Y)); 
pk = np.random.rand(M,N); 
Kx = np.random.rand(M,N); 
Ky = np.random.rand(M,N); 
u = np.random.rand(X); 
v = np.random.rand(Y); 

В реальном коде они не являются случайными, конечно, но это не имеет значения для этого примера. Вопрос в том, если существует чистый Numpy эквивалентно следующему:

for m in range(0, M): 
    for n in range(0, N): 
     for i in range(0,X): 
      for j in range(0,Y): 
       Arg = Kx[m,n]*u[i] + Ky[m,n]*v[j]; 
       P[i,j] += pk[m,n]*np.cos(Arg); 

Все M, N, X, Y разные, но X и Y может быть таким же, если решение не существует иначе.

Answer 1

Общая стратегия устранения for-loop s в вычислениях NumPy заключается в работе с многомерными массивами.

Рассмотрим, например, линия

Arg = Kx[m,n]*u[i] + Ky[m,n]*v[j] 

Эта линия зависит от индексов m, n, i и j. So Arg зависит от m, n, i и j. Это означает, что Arg можно рассматривать как 4-мерную матрицу, индексированную m, n, i и j. Таким образом, мы можем исключить 4 for-loops - насколько Arg обеспокоен - вычисляя

Kxu = Kx[:,:,np.newaxis]*u 
Kyv = Ky[:,:,np.newaxis]*v 
Arg = Kxu[:,:,:,np.newaxis] + Kyv[:,:,np.newaxis,:] 

Kx[:,:,np.newaxis] имеет форму (M, N, 1) и u имеет форму (X,). Умножая их вместе, используется NumPy broadcasting для создания массива формы (M, N, X). Таким образом, выше, new axes are used somewhat like placeholders, так что Arg заканчивается четырьмя осями, индексированными m, n, i, j в указанном порядке.

Аналогично, P может быть определен как

P = (pk[:,:,np.newaxis,np.newaxis]*np.cos(Arg)).sum(axis=0).sum(axis=0) 

В sum(axis=0) (так называемые два раза) суммы вдоль m и n осей, так что P заканчивает тем, что 2-мерный массив индексируется только i и j.

Работая с этими 4-мерными массивами, мы получаем возможность применять операции NumPy для целых массивов NumPy. Напротив, при использовании 4 for-loops нам приходилось выполнять вычисления по значению по скалярам. Рассмотрим, например, что np.cos(Arg) делает, когда Arg представляет собой 4-мерный массив. Это отключает вычисление всех косинусов в одном вызове функции NumPy, который выполняет базовый цикл в скомпилированном C-коде. Это намного быстрее, чем вызов np.cos один раз для каждого скаляра.Именно по этой причине работа с многомерными массивами заканчивается так быстро, как код for-loop.

---

import numpy as np 

def orig(Kx, Ky, u, v, pk): 
    M, N = Kx.shape 
    X = u.size 
    Y = v.size 
    P = np.empty((X, Y), dtype=pk.dtype) 
    for m in range(0, M): 
     for n in range(0, N): 
      for i in range(0,X): 
       for j in range(0,Y): 
        Arg = Kx[m,n]*u[i] + Ky[m,n]*v[j] 
        P[i,j] += pk[m,n]*np.cos(Arg) 
    return P 

def alt(Kx, Ky, u, v, pk): 
    Kxu = Kx[:,:,np.newaxis]*u 
    Kyv = Ky[:,:,np.newaxis]*v 
    Arg = Kxu[:,:,:,np.newaxis] + Kyv[:,:,np.newaxis,:] 
    P = (pk[:,:,np.newaxis,np.newaxis]*np.cos(Arg)).sum(axis=0).sum(axis=0) 
    return P 

M, N = 10, 20 
X, Y = 5, 15 
Kx = np.random.random((M, N)) 
Ky = np.random.random((M, N)) 
u = np.random.random(X) 
v = np.random.random(Y) 
pk = np.random.random((M, N)) 

---

проверки разумности, (показывая альт и ориг возвращают тот же результат):

In [57]: P2 = alt(Kx, Ky, u, v, pk) 

In [58]: P1 = orig(Kx, Ky, u, v, pk) 

In [59]: np.allclose(P1, P2) 
Out[59]: True 

---

Эталонный, показывающий alt значительно быстрее, чем orig:

In [60]: %timeit orig(Kx, Ky, u, v, pk) 
10 loops, best of 3: 33.6 ms per loop 

In [61]: %timeit alt(Kx, Ky, u, v, pk) 
1000 loops, best of 3: 349 µs per loop