移动Numpy阵列的最快方法

Question

I'm running some simulations that involve repeatedly comparing values in 2D Numpy arrays with their 'neighbours'; eg. the value at indicie location (y,x) is compared to the value at indicie location (y-1,x) from the same array.

At the moment I am using functions like this:

# example of the typical size of the arrays
my_array = np.ndarray((500,500))    

shapey, shapex = my_array.shape
Yshape = (1, shapex)
Yzeros = np.zeros((1, shapex))

def syf(A, E=True):
    if E == True:
        return np.concatenate((A[1:], A[-1].reshape(Yshape)), axis=0)
    else:
        return np.concatenate((A[1:], Yzeros), axis=0)

shifted_array = syf(my_array)

difference_in_y = shifted_array - my_array 

This has the option to use either the edge values or zeros for comparison at the edge of the array. The functions can also do it in either direction in either axis.

Does anybody have any suggestions for a faster way to do this? I've tried np.roll (much slower) and this:

yf = np.ix_(range(shapey)[1:] + [shapey,], range(shapex))
shifted_array = my_array[yf]

which is a little slower.

These functions are called ~200 times a second in a program that takes 10 hours to run so any small speedups are more that welcome!

Thanks.

EDIT:

So if the same differentiation method is required every time the shift function is called then Divakars method below seems to offer a minor speedup, however if just a shifted array is required both that method and the one I use above seem to be equal speed.

Answer 1

You could have the shifting and differentiating both done in a function call like so -

def syf1(A, E=True):
    out = np.empty_like(A)
    out[:-1] = A[1:] - A[:-1] # Or np.diff(my_array,axis=0)
    if E == True:
        out[-1] = 0
    else:
        out[-1] = -A[-1]
    return out

Thus, the equivalent modified version of syf for runtime comparison purposes would be -

def syf(A, E=True):
    if E == True:
        return np.concatenate((A[1:], A[-1].reshape(Yshape)), axis=0) - A
    else:
        return np.concatenate((A[1:], Yzeros), axis=0) - A

Runtime tests

Let's compare the equivalent version of syf with the proposed approach on runtime performance for the inputs listed in the question code -

In [113]: %timeit syf(my_array)
1000 loops, best of 3: 518 µs per loop

In [114]: %timeit syf1(my_array)
1000 loops, best of 3: 494 µs per loop

So, there is some improvement there!