Python中子数组的快速求和

Question

我有一个半径为 w 的数据立方体 a，对于该立方体的每个元素，我想在半径为 r 的立方体内添加元素和所有周围的值，其中 r

举个简单的例子，假设：

a = numpy.ones(shape=(2*w,2*w,2*w),dtype='float32')
kernel = numpy.ones(shape=(2*r,2*r,2*r),dtype='float32')
b = convolve(a,kernel,mode='constant',cval=0)

那么对于不在边缘的所有索引，b 的值将是 (2r)(2r)(2r)。

目前我正在使用循环来执行此操作，它非常慢，尤其是对于较大的 w 和 r。我尝试了 scipy 卷积，但在循环中几乎没有加速。我现在正在研究 numba 的并行计算功能，但无法弄清楚如何重写代码以使用 numba。我有一张 Nvidia RTX 卡，所以也可以进行 CUDA GPU 计算。

欢迎提出建议。

这是我当前的代码：

for x in range(0,w*2):
    print(x)
    for y in range(0,w*2):
        for z in range(0,w*2):
            if x >= r:
                x1 = x - r
            else:
                x1 = 0
            if x < w*2-r:
                x2 = x + r
            else:
                x2 = w*2 - 1
            
            if y >= r:
                y1 = y - r
            else:
                y1 = 0
            if y < w*2-r:
                y2 = y + r
            else:
                y2 = w*2 - 1

            if z >= r:
                z1 = z - r
            else:
                z1 = 0
            if z < w*2-r:
                z2 = z + r
            else:
                z2 = w*2 - 1

            b[x][y][z] = numpy.sum(a[x1:x2,y1:y2,z1:z2])

return b

Answer 1

这是您的代码的一个非常简单的版本，因此它可以与 numba 一起使用。我发现相对于纯 numpy 代码的速度提高了 10 倍。但是，您应该能够使用 FFT 卷积算法（例如 scipy 的 fftconvolve）获得更大的加速。您能分享一下让卷积起作用的尝试吗？

from numba import njit

@njit
def sum_cubes(a,b,w,r):
    for x in range(0,w*2):
        #print(x)
        for y in range(0,w*2):
            for z in range(0,w*2):
                if x >= r:
                    x1 = x - r
                else:
                    x1 = 0
                if x < w*2-r:
                    x2 = x + r
                else:
                    x2 = w*2 - 1

                if y >= r:
                    y1 = y - r
                else:
                    y1 = 0
                if y < w*2-r:
                    y2 = y + r
                else:
                    y2 = w*2 - 1

                if z >= r:
                    z1 = z - r
                else:
                    z1 = 0
                if z < w*2-r:
                    z2 = z + r
                else:
                    z2 = w*2 - 1

                b[x,y,z] = np.sum(a[x1:x2,y1:y2,z1:z2])
    return b

编辑：您的原始代码中有一个小错误。 numpy 索引的工作方式，最后一行应该是

b[x,y,z] = np.sum(a[x1:x2+1,y1:y2+1,z1:z2+1])

除非你想让立方体偏离中心。

假设您确实希望立方体居中，那么执行此计算的更快方法是使用 scipy 的统一过滤器：

from scipy.ndimage import uniform_filter
def sum_cubes_quickly(a,b,w,r):
    b = uniform_filter(a,mode='constant',cval=0,size=2*r+1)*(2*r+1)**3
    return b

对 w = 50, r = 10 的随机生成数据进行一些快速运行时比较：

• 原始 numpy 代码 - 15.1 秒
• Numba'd numpy 代码 - 8.1 秒
• uniform_filter - 13.1 毫秒