Nikhil Nikhil - 3 months ago 22
Python Question

Numpy Vectorization of sliding-window operation

I have the following numpy arrays:

arr_1 = [[1,2],[3,4],[5,6]] # 3 X 2
arr_2 = [[0.5,0.6],[0.7,0.8],[0.9,1.0],[1.1,1.2],[1.3,1.4]] # 5 X 2


arr_1
is clearly a
3 X 2
array, whereas
arr_2
is a
5 X 2
array.

Now without looping, I want to element-wise multiply arr_1 and arr_2 so that I apply a sliding window technique (window size 3) to arr_2.

Example:

Multiplication 1: np.multiply(arr_1,arr_2[:3,:])

Multiplication 2: np.multiply(arr_1,arr_2[1:4,:])

Multiplication 3: np.multiply(arr_1,arr_2[2:5,:])


I want to do this in some sort of a matrix multiplication form to make it faster than my current solution which is of the form:

for i in (2):
np.multiply(arr_1,arr_2[i:i+3,:])


So if the number of rows in arr_2 are large (of the order of tens of thousands), this solution doesn't really scale very well.

Any help would be much appreciated.

Answer

We can use NumPy broadcasting to create those sliding windowed indices in a vectorized manner. Then, we can simply index into arr_2 with those to create a 3D array and perform element-wise multiplication with 2D array arr_1, which in turn will bring on broadcasting again.

So, we would have a vectorized implementation like so -

W = arr_1.shape[0] # Window size
idx = np.arange(arr_2.shape[0]-W+1)[:,None] + np.arange(W)
out = arr_1*arr_2[idx]

Runtime test and verify results -

In [143]: # Input arrays
     ...: arr_1 = np.random.rand(3,2)
     ...: arr_2 = np.random.rand(10000,2)
     ...: 
     ...: def org_app(arr_1,arr_2):
     ...:     W = arr_1.shape[0] # Window size
     ...:     L = arr_2.shape[0]-W+1
     ...:     out = np.empty((L,W,arr_1.shape[1]))
     ...:     for i in range(L):
     ...:        out[i] = np.multiply(arr_1,arr_2[i:i+W,:])
     ...:     return out
     ...: 
     ...: def vectorized_app(arr_1,arr_2):
     ...:     W = arr_1.shape[0] # Window size
     ...:     idx = np.arange(arr_2.shape[0]-W+1)[:,None] + np.arange(W)
     ...:     return arr_1*arr_2[idx]
     ...: 

In [144]: np.allclose(org_app(arr_1,arr_2),vectorized_app(arr_1,arr_2))
Out[144]: True

In [145]: %timeit org_app(arr_1,arr_2)
10 loops, best of 3: 47.3 ms per loop

In [146]: %timeit vectorized_app(arr_1,arr_2)
1000 loops, best of 3: 1.21 ms per loop