Numpy：向量化一个集成2D数组的函数

Question

我需要对一个二维数组执行以下积分： 也就是说，网格中的每个点都获得值RC，这是在2D空间上整个场与某点（x，y）处场的值U的差异的加权积分，该加权核在1D版本中为：

到目前为止，我所做的是对索引的低效迭代：

def normalized_bimodal_kernel_2D(x,y,a,x0=0.0,y0=0.0):
    """ Gives a kernel that is zero in x=0, and its integral from -infty to 
    +infty is 1.0. The parameter a is a length scale where the peaks of the 
    function are."""
    dist = (x-x0)**2 + (y-y0)**2
    return (dist*np.exp(-(dist/a)))/(np.pi*a**2)


def RC_2D(U,a,dx):
    nx,ny=U.shape
    x,y = np.meshgrid(np.arange(0,nx, dx),np.arange(0,ny,dx), sparse=True)
    UB = np.zeros_like(U)
    for i in xrange(0,nx):
        for j in xrange(0,ny):
            field=(U-U[i,j])*normalized_bimodal_kernel_2D(x,y,a,x0=i*dx,y0=j*dx)
            UB[i,j]=np.sum(field)*dx**2
    return UB

def centerlizing_2D(U,a,dx):
    nx,ny=U.shape
    x,y = np.meshgrid(np.arange(0,nx, dx),np.arange(0,ny,dx), sparse=True)
    UB = np.zeros((nx,ny,nx,ny))
    for i in xrange(0,nx):
        for j in xrange(0,ny):
            UB[i,j]=normalized_bimodal_kernel_2D(x,y,a,x0=i*dx,y0=j*dx)
    return UB

你可以在这里查看“集中”功能的结果：centeralizing。

U=np.eye(20)
plt.imshow(centerlizing(U,10,1)[10,10])

我相信还有其他bug，所以欢迎任何反馈，但我真正感兴趣的是了解如何以向量化的方式更快地执行此操作。

- Ohm

---

U.shape 在实际使用中是什么意思？ - Daniel F

我认为这个问题在https://codereview.stackexchange.com上会得到更好的答案。 - Ofer Sadan

不，这里更多的是“numpy”专家。而且这是一个纯粹的“numpy”向量化问题。 - Daniel F

最大的效率改进是，您正在创建和运行 (nx/dx, ny/dx) 的计算 meshgrid，但仅对其进行 (nx, ny) 切片求和。或者也许这是一个错误，很难说。输出是否符合您的预期？ - Daniel F

Answer 1

normalized_bimodal_kernel_2D在两个嵌套的循环中被调用，每个循环只通过小步长移动它的偏移量。这会导致许多重复计算。

对于centerlizing_2D的优化是为整个范围计算一次核，并将UB定义为该核的移位视图。使用stride_tricks可以实现这一点，但不幸的是，这需要相当高级的numpy知识。

def centerlizing_2D_opt(U,a,dx):
    nx,ny=U.shape    
    x,y = np.meshgrid(np.arange(-nx//2, nx+nx//2, dx),
                      np.arange(-nx//2, ny+ny//2, dx),  # note the increased range
                      sparse=True)
    k = normalized_bimodal_kernel_2D(x, y, a, x0=nx//2, y0=ny//2)
    sx, sy = k.strides    
    UB = as_strided(k, shape=(nx, ny, nx*2, ny*2), strides=(sy, sx, sx, sy))
    return UB[:, :, nx:0:-1, ny:0:-1]

assert np.allclose(centerlizing_2D(U,10,1), centerlizing_2D_opt(U,10,1)) # verify it's correct

没错，它更快：

%timeit centerlizing_2D(U,10,1)      #   100 loops, best of 3:  9.88 ms per loop
%timeit centerlizing_2D_opt(U,10,1)  # 10000 loops, best of 3: 85.9  µs per loop

接下来，我们通过使用优化后的centerlizing_2D程序将其表达，以优化RC_2D：

def RC_2D_opt(U,a,dx):
    UB_tmp = centerlizing_2D_opt(U, a, dx)
    U_tmp = U[:, :, None, None] - U[None, None, :, :]
    UB = np.sum(U_tmp * UB_tmp, axis=(0, 1))
    return UB

assert np.allclose(RC_2D(U,10,1), RC_2D_opt(U,10,1))

---

%timeit RC_2D(U,10, 1) 的性能：

#original:    100 loops, best of 3: 13.8 ms per loop
#@DanielF's:  100 loops, best of 3:  6.98 ms per loop
#mine:       1000 loops, best of 3:  1.83 ms per loop

- MB-F

Answer 2

假设 dx=1，因为我不确定您的离散化目的是什么：

def normalized_bimodal_kernel_2D(x, y, a):  
    #generating a 4-d tensor instead of 1d vector
    dist = (x[:,None,None,None] - x[None,None,:,None])**2 +\
           (y[None,:,None,None] - y[None,None,None,:])**2    
    return (dist * np.exp(-(dist / a))) / (np.pi * a**2)

def RC_2D(U, a):
    nx, ny = U.shape
    x, y = np.arange(nx), np.arange(ny)
    U4 = U[:, :, None, None] - U[None, None, :, :] #Another 4d
    k = normalized_bimodal_kernel_2D(x, y, a)
    return np.einsum('ijkl,ijkl->ij', U4, k)


def centerlizing_2D(U, a):
    nx, ny = U.shape
    x, y = np.arange(nx), np.arange(ny)
    return normalized_bimodal_kernel_2D(x, y, a)

基本上，在numpy中将for循环向量化只需要增加更多的维度。你正在对一个二维的U向量进行两个循环，所以为了向量化，只需将其转换为四维。

- Daniel F

Answer 3

为了符合您的公式，让 U 成为一个函数。

然后，您只需使用 np.ix_ 将 x,y,x',y' 放入四个不同的维度中，并仔细翻译您的公式。Numpy 广播会处理其余部分。

a=20
x,y,xp,yp=np.ix_(*[np.linspace(0,1,a)]*4)

def U(x,y) : return np.float32(x == y)  # function "eye"

def f(x,y,xp,yp,a):
    r2=(x-xp)**2+(y-yp)**2
    return r2*np.exp(-r2/a)*(U(xp,yp) - U(x,y))/np.pi/a/a

#f(x,y,xp,yp,a).shape is (20, 20, 20, 20)

RC=f(x,y,xp,yp,a).sum(axis=(2,3))
#RC.shape is (20, 20)

- B. M.

---

非常有趣的解决方案 - 但如果我拥有的输入是给定的二维数组，我该如何使其工作呢？在我的情况下，我正在执行一个变量的时间积分，该变量是二维场中的一个字段，因此在每次迭代中U都是不同的。 - Ohm