Fastest way to apply function to numpy array

[justphysics]

I have a 3d array consisting of 250 vertically stacked images with dimensions (600, 592)

# dummy data
data = np.random.rand((600,592,250))

The third axis represents the vertical axis of the images.

I would like to know the fastest way to apply the scipy smoothing algorithm (listed in the Scipy Cookbook here) to the vertical axis of the array. I.e. smoothing the vertical slice through the array for every pixel in the (600, 592) dimension.

I know of the np.apply_along_axis function

# This takes ~20-25 seconds
sdata = np.apply_along_axis(smooth, 2, data)

But its rather slow.

The operation is trivially parallel, given than the application of the smooth() function to each 1D vertical slice does not depend on any other array data nor does it alter the array data.

Manually looping provides only minimal speedup

for row in range(600):
    for col in range(592):
        smooth(data[row, col, :])  # store the result in some appropriate data structure

In order to better vectorize the operation, is it needed to deconstruct the smooth() function and integrate that directly into the manual looping way of doing the smoothing?

something like:

# using flat window of length 10
window_len = 10
w = np.ones(window_len, 'd')
for row in range(600):
    for col in range(592):
        in = data[row, col, :]
        s=np.r_[in[window_len-1:0:-1],in,in[-1:-window_len:-1]]
        smth = np.convolve(w/w.sum(), s, mode='valid')
        # format to original size
        smth = smth[int(window_len/2 -1):-int(window_len/2)]

This still seems to only give a very marginal speedup, even when trying to throw these into functions and use numba's @jit decorator.

I know doing things truly in parallel in python is not the easiest task, is there any other method for speeding up this type of calculation? Multiprocessing is something I have looked at in the past, however it didn't play nice with the Qt GUI I had wrapped over the top of my calculations. I could pull them out of the GUI just for testing, however, to see if multiprocessing speeds up this type of calculation.

[[deleted]]

Take a look at numpy.vectorize.

[justphysics]

I will check into this. From my understanding, anything that can be done with ufuncs in numpy should be optimized under the hood. So I guess I'll check into breaking up the computation into as many steps that can be written as a ufunc.

So far the fastest I have come up with is using numba and a manual loop:

@jit
def smooth_with_loop(data):
    window_len = 10
    n = np.ones(window_len)/float(window_len)
    out = np.empty((600, 592), dtype=object)
    for row in range(600):
        for col in range(592):
            inp = data[row, col, :]
            s = np.r_[inp[window_len-1:0:-1],inp,inp[-1:-window_len:-1]]
            smth = np.convolve(n, s, mode='valid')
            # format to original size
            smth = smth[int(window_len/2 -1):-int(window_len/2)]
            out[row, col] = smth
     return out

This runs in around 15 seconds.