How to compute the moving average over 3D array with a step size?

Question

I need to calculate a moving average over a 3D array with a step size set by me. What I am doing right now is

 img = np.ones(10,10,50)
 img_new = bottleneck.move.move_mean(img, window=5, axis=2)

While the bottleneck.move.move_mean is fast enough to deal with images, unfortunately it does not let me set the step size. This leads to a lot of overhead because the mean of 4 out 5 windows is calculated unnecessarily.

Is there a similar function to bottleneck.move.move_mean where I can set the step size?

Answer 1

If IIUC you can directly use numpy mean after reshaping the data:

shape = (2,2,5)
len_ = np.prod(shape)
img = np.linspace(1,len_,len_).reshape(shape) # init

init array:

array([[[ 1.,  2.,  3.,  4.,  5.],
        [ 6.,  7.,  8.,  9., 10.]],

       [[11., 12., 13., 14., 15.],
        [16., 17., 18., 19., 20.]]])

Code:

step = 5 # same as window size
# shape[2]/step == int(shape[2]/step) # maybe add this check!
img.reshape(shape[0], shape[1], step, int(shape[2]/step)).mean(2)

Output:

array([[[ 3.],
        [ 8.]],

       [[13.],
        [18.]]])

Note: This assumes the window and step-size are the same, which is what I think you're asking for in your particular example as well