如何对所有网格进行矢量化计算而不是“ for”循环?

时间:2019-01-10 15:25:28

标签: python numpy vectorization

我在所有网格中都有一个双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)

所以我需要一个解决方案,如何在向量和矩阵级别上执行此操作。有什么想法吗? 预先谢谢你。

2 个答案:

答案 0 :(得分:3)

我们可以在这里使用三角技巧-

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+ 加速!

答案 1 :(得分:2)

rthetac常数吗?如果您发布一个最低限度的工作示例,它将使其他人更容易为您提供帮助。

在任何情况下,您都可以利用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))

当我尝试使用rthetac作为常数时,得到相同的结果。我认为如果它们是逐点应用的向量,也可以正常工作。

工作原理

主要观察到的是,几乎所有操作都是按点进行的,并且在所有x,y对中都是相同的。因此,我们只需要弄清楚如何向量化所有vec1*xvec2*y对。

我们首先使用外部产品列出X=vec1*xY=vec2*y的列表。

我们现在使用广播添加XY的所有成对行。

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

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

由于我们一次构造所有总和的所有元素,因此这不是非常有效的存储,但是除非您使用的是大数据,否则应该没问题。即使您处理的数据不足以容纳到内存中,与使用python循环相比,阻塞此方法也可能是有意义的。