如何对所有网格进行矢量化计算而不是“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

尝试meshgrid时间比较多：

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 是常量吗？如果您发布一个最低限度的工作示例，它将使其他人更容易帮助您。

无论如何，您都可以利用broadcasting 执行以下操作：

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 的所有行对。

然后我们将其重塑为所有对的列表，逐点应用函数的其余部分并沿正确的轴求和。

最后，我们将其从长度为 x_coord*y_coord 的 ndarray 重塑为形状为 (x_coord,y_coord) 的二维数组。

这不是很节省内存，因为我们一次构建所有和的所有元素，但除非您正在处理非常大的数据，否则应该没问题。即使您正在使用足够多的数据而无法放入内存，与使用 python 循环相比，阻止此方法也可能是有意义的。