Как сделать векторизованные вычисления вместо циклов for для всей сетки?

Question

У меня есть двойной цикл for среди всех сеток, и я хочу, чтобы он работал быстрее. r, vec1, vec2, theta - векторы одинаковой длины N. c является константой.

import numpy as np

N = 30
x_coord, y_coord = 300, 300
m1 = np.zeros((x_coord, y_coord))

vec1, vec2 = np.ones(N), np.ones(N)
theta = np.ones(N)

for x in np.arange(x_coord):
    for y in np.arange(y_coord):
        m1[x,y] = np.sum(np.cos(2.*np.pi*(r*(vec1*x + vec2*y))+theta)) * c

Время для двух циклов было:

1.03 s ± 8.96 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Также я пытался использовать np.meshgrid:

def f1(x, y):
    sum1 = vec1*x + vec2*y
    mltpl1 = r * sum1
    sum2 = 2.*np.pi * mltpl1 + theta
    sum3 = np.sum(np.cos(sum2))
    mltpl2 = sum3 * c
    return mltpl2

msh1, msh2 = np.meshgrid(range(x_coord), range(y_coord))
pairs = np.vstack((np.ravel(msh1), np.ravel(msh2))).T

m1 = np.reshape(list(map(lambda x: f1(x[0], x[1]), pairs)), (m1.shape[0], m1.shape[1])).T

Время попытки сетки было больше:

1.25 s ± 48.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Так что мне нужно решение, как это сделать на уровне векторов и матриц. Есть какие-нибудь идеи? Заранее спасибо.

Answer 1

Мы можем использовать тригонометрический трюк здесь -

cos(A + B) = cos A cos B − sin A sin B 

Это позволяет нам использовать matrix-multiplication для решения, которое будет выглядеть примерно так -

# Get x and y as 1D arrays
x = np.arange(x_coord)
y = np.arange(y_coord)

# Get the common constant for scaling vec1 and vec2 parts
k1 = 2.*np.pi*r

# Use outer multiplications for the two vectors against x,y and also scale them
p1 = k1*vec1*x[:,None] + theta
p2 = (k1*vec2)[:,None]*y

# Finally use trigonometry+matrix-multiplication for sum reductions
out = c*(np.cos(p1).dot(np.cos(p2)) - np.sin(p1).dot(np.sin(p2)))

Сроки -

# Setup
In [151]: np.random.seed(0)
     ...: c = 2.3
     ...: N = 30
     ...: x_coord, y_coord = 300, 300
     ...: vec1 = np.random.rand(N)
     ...: vec2 = np.random.rand(N)
     ...: r = np.random.rand(N)
     ...: theta = np.ones(N)

# Original solution
In [152]: %%timeit
     ...: m1 = np.zeros((x_coord, y_coord))
     ...: for x in np.arange(x_coord):
     ...:     for y in np.arange(y_coord):
     ...:         m1[x,y] = np.sum(np.cos(2.*np.pi*(r*(vec1*x + vec2*y))+theta)) * c
1 loop, best of 3: 960 ms per loop

# Proposed solution
In [153]: %%timeit
     ...: x = np.arange(x_coord)
     ...: y = np.arange(y_coord)
     ...: k1 = 2.*np.pi*r
     ...: p1 = k1*vec1*x[:,None] + theta
     ...: p2 = (k1*vec2)[:,None]*y
     ...: out = c*(np.cos(p1).dot(np.cos(p2)) - np.sin(p1).dot(np.sin(p2)))
100 loops, best of 3: 2.54 ms per loop

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

Answer 2

Являются ли r, theta и c постоянными? Если вы опубликуете минимальный рабочий пример, другим будет легче помочь вам.

В любом случае, вы можете воспользоваться вещанием , чтобы сделать что-то вроде:

X = np.outer(x,vec1)
Y = np.outer(y,vec2)

Z = np.reshape(X[:,np.newaxis] + Y[np.newaxis],(x_coord*y_coord,N))

M = np.sum(np.cos(2.*np.pi*r*Z+theta),axis=1)*c
m = np.reshape(M,(x_coord,y_coord))

Когда я попробовал это с r, theta и c в качестве констант, это дало тот же результат. Я думаю, что это также будет хорошо работать, если они являются точечными векторами.

Как это работает

Ключевое наблюдение заключается в том, что почти все операции выполняются по точкам и одинаковы для всех пар x,y. Поэтому нам нужно только выяснить, как векторизовать сложение всех пар vec1*x и vec2*y.

Сначала мы создаем списки X=vec1*x и Y=vec2*y, используя внешний продукт.

Теперь добавим все пары строк X и Y, используя трансляцию.

Затем мы преобразовываем его в список всех пар, применяем остальную часть функции по точкам и суммируем по правильной оси.

Наконец, мы преобразуем его из ndarray длины x_coord*y_coord в 2d массив формы (x_coord,y_coord).

Это не очень эффективно для памяти, так как мы строим все элементы всех сумм одновременно, но если вы не работаете с большими данными, все будет хорошо. Даже если вы работаете с достаточным количеством данных, чтобы не поместиться в память, возможно, имеет смысл заблокировать этот метод, а не использовать цикл Python.