Fastest way to create a tilted/successively shifted matrix in numpy

Viewed 103

I need to create a (w,N)-matrix that looks like this:

w//2............N-1,N-1
.             \     N-1
.              \    N-1
.               \   N-1
1...............N-1,N-1
0...................N-1
00..................N-2
. \                 N-3
.  \                .
.   \               .
000000..............N-w//2

Which is an (w,N) matrix, with an odd w. The middle row is the range from 0 to N. For each row index above the middle row, the row is shifted to the left like with scipy.ndimage.shift(mode='nearest') and for each row below the middle row it is shifted to the right with the same method.

N is usually around 10^4 and w is usually between 10 and 10^2.

I've come up with 2 ways to do this:

from scipy.ndimage import shift
middle = np.arange(0, N)
final = np.vstack(
    [shift(middle, i, mode='nearest') for i in range(-w//2, 0)] + 
    [middle] + 
    [shift(middle, i, mode='nearest') for i in range(1, w//2)] ) 

Which takes 0.035 seconds to run.

np.vstack([
        np.maximum(
            0,
            np.minimum(
                N-1,
                np.arange(-step, N-step)
            )
        )
        for step in range(-w//2, w//2)
    ])

Which takes 0.021 seconds to run.

These numbers were with N=10^3 and w=21.

I'd really like to get these numbers down as low as possible, ideally down to around 1ms.

I tried multiprocessing, but that doesn't really help, the overhead to too big to gain something from the concurrency. Also I know I could store this result somewhere, but that'd require a significant change by the caller of this function, so that'll be done later.

Is there any mathematical relation that can represent a tilt/shift operation like this? I couldn't think of one, but if there is, numpy can probably take advantage of that to beat my results.

So yeah, any ideas to make my code faster?

1 Answers

initialise an array with appropriate shape and horizontal values from 0 to N (inclusive)

w, N = 11, 10
arr = np.empty(shape= [w, N], dtype= int)
 
arr[:] = np.arange(N)
arr

>>>   [[0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]]

subtract from each row an appropriate value

arr += np.arange(w).reshape([-1, 1])[::-1] - (1+w//2)
arr

>>>   [[ 5.,  6.,  7.,  8.,  9., 10., 11., 12., 13., 14.],
       [ 4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13.],
       [ 3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12.],
       [ 2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11.],
       [ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.],
       [ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9.],
       [-1.,  0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.],
       [-2., -1.,  0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.],
       [-3., -2., -1.,  0.,  1.,  2.,  3.,  4.,  5.,  6.],
       [-4., -3., -2., -1.,  0.,  1.,  2.,  3.,  4.,  5.],
       [-5., -4., -3., -2., -1.,  0.,  1.,  2.,  3.,  4.]]

where values cross limiting values reassign them the limit values

arr[arr<0] = 0
arr[arr>N-1] = N-1

arr

>>>   [[5., 6., 7., 8., 9., 9., 9., 9., 9., 9.],
       [4., 5., 6., 7., 8., 9., 9., 9., 9., 9.],
       [3., 4., 5., 6., 7., 8., 9., 9., 9., 9.],
       [2., 3., 4., 5., 6., 7., 8., 9., 9., 9.],
       [1., 2., 3., 4., 5., 6., 7., 8., 9., 9.],
       [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [0., 0., 1., 2., 3., 4., 5., 6., 7., 8.],
       [0., 0., 0., 1., 2., 3., 4., 5., 6., 7.],
       [0., 0., 0., 0., 1., 2., 3., 4., 5., 6.],
       [0., 0., 0., 0., 0., 1., 2., 3., 4., 5.],
       [0., 0., 0., 0., 0., 0., 1., 2., 3., 4.]]

Edit

tried timing the script

import timeit

script = '''
w, N   = 21, 10**3
arr    = np.empty(shape= [w, N], dtype= int)
arr[:] = np.arange(N)
arr   += np.arange(w).reshape([-1, 1])[::-1] - (1+w//2)

arr[arr<0] = 0
arr[arr>N-1] = N-1
'''

time = timeit.timeit(script, number= 100000, setup= 'import numpy as np') / 100000

time

>>> 0.00019059010320999733    # 0.19 ms
Related