Fast summing of subarrays in Python - parallel-processing

Question

I have a data cube a of radius w and for every element of that cube, I would like to add the element and all surrounding values within a cube of radius r, where r < w. The result should be returned in an array of the same shape, b.
As a simple example, suppose:
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)
then b would have the value (2r)(2r)(2r) for all the indices not on the edge.
Currently I am using a loop to do this and it is very slow, especially for larger w and r. I tried scipy convolution but got little speedup over the loop. I am now looking at numba's parallel computation feature but cannot figure out how to rewrite the code to work with numba. I have a Nvidia RTX card so CUDA GPU calculations are also possible.
Suggestions are welcome.
Here is my current code:
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

Here is a very simple version of your code so that it works with numba. I was finding speed-ups of a factor of 10 relative to the pure numpy code. However, you should be able to get even greater speed-ups using a FFT convolution algorithm (e.g. scipy's fftconvolve). Can you share your attempt at getting convolution to work?
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
EDIT: Your original code has a small bug in it. The way numpy indexing works, the final line should be
b[x,y,z] = np.sum(a[x1:x2+1,y1:y2+1,z1:z2+1])
unless you want the cube to be off-centre.
Assuming you do want the cube to be centred, then a much faster way to do this calculation is using scipy's uniform filter:
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
A few quick runtime comparisons for randomly generated data with w = 50, r = 10:
Original raw numpy code - 15.1 sec
Numba'd numpy code - 8.1 sec
uniform_filter - 13.1 ms