较小尺寸的阵列上的广播操作

时间:2018-06-27 16:40:41

标签: python performance numpy-broadcasting

我需要提高对不同形状/大小的数组执行的操作的性能。数组pos的形状为(2, 500),数组xa, xb, ya, yb的形状为(30,)

下面MVCE中显示的操作将pos的两个维度与xa, xbya, yb结合在一起。

可以通过numpy广播来做到这一点吗?

import numpy as np

# Some random data
N = 30
xa, xb = np.random.uniform(0., 1., N), np.random.uniform(0., 1., N)
ya, yb = np.random.uniform(0., 1., N), np.random.uniform(0., 1., N)

# Grid
M = 500
ext = [xa.min(), xa.max(), ya.min(), ya.max()]
x, y = np.mgrid[ext[0]:ext[1]:complex(0, M), ext[2]:ext[3]:complex(0, M)]
pos = np.vstack([x.ravel(), y.ravel()])

# Apply broadcasting on the operation performed by this 'for' block?
vals = []
for p in zip(*pos):
    vals.append(np.sum(np.exp(-0.5 * (
        ((p[0] - xa) / xb)**2 + ((p[1] - ya) / yb)**2)) / (xb * yb)))

1 个答案:

答案 0 :(得分:1)

您可以使用np.tile 并如下修改for循环

AResolver

说明:

  1. 在每次迭代中,您都采用pos [0] [i]和pos [1] [i]并对xa,xb,ya,yb进行运算。
  2. Tile复制这250000次中的全部4次,即pos的shape [1]或迭代次数。
  3. 我们还需要重塑pos [0]和pos [1]并将其设为2D才能使操作有效。

计时细节: 在我的机器上,矢量化代码大约需要0.2秒,而非矢量化代码大约需要3秒。下面是要复制的代码:

xa_tiled = np.tile(xa, (pos.shape[1],1))
xb_tiled = np.tile(xb, (pos.shape[1],1))
ya_tiled = np.tile(ya, (pos.shape[1],1))
yb_tiled = np.tile(yb, (pos.shape[1],1))

vals_ = np.exp(-0.5 * (
         ((pos[0].reshape(pos.shape[1],1) - xa_tiled) / xb_tiled)**2 + ((pos[1].reshape(pos.shape[1],1) - ya_tiled) / yb_tiled)**2)) / (xb_tiled * yb_tiled)
vals_ = vals_.sum(axis=1)