Numpy умножение векторов разного размера, избегая циклов for

Распространенной стратегией устранения for-loop в вычислениях NumPy является работа с многомерными массивами.

Рассмотрим, например, строку

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

Эта строка зависит от индексов m, n, i и j. Значит, 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 для создания массива формы (M, N, X) . Таким образом, выше новые оси используются несколько как заполнители, так что 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 оказывается двумерным массивом, индексированным только 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))

Проверка работоспособности (показывая, что alt и orig возвращают один и тот же результат):

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