Using numpy to construct an array with rows extracted from another 2D array as 2x2 blocks-CodePudding

Suppose I have the following 2D array:

x = np.array([[10,20,30,40], [50,60,70,80],[90,100,110,120]])  
print(x)

array([[ 10,  20,  30,  40],
       [ 50,  60,  70,  80],
       [ 90, 100, 110, 120]])

I would like to construct a new array, y, where each row has the values of a 2x2 block from x in clockwise order:

print(y)
array([[ 10,  20,  60,  50],
       [ 20,  30,  70,  60],
       [ 30,  40,  80,  70],
       [ 50,  60,  100, 90],
       [ 60,  70,  110, 100],
       [ 70,  80,  120, 110]])

I could achieve that using Python for loops as follows:

n_rows, n_cols = x.shape
y = []
for i in range(n_rows-1): 
     for j in range(n_cols-1): 
         row = [x[i,j],x[i,j 1],x[i 1, j 1],x[i 1,j]] 
         y.append(row) 
y = np.array(y)

I wonder if there is a faster way that takes advantage of Numpy functions and avoid using Python loops.

CodePudding user response：

First, create a sliding_window_view into x with the 2x2 boxes you want to see:

b = np.lib.stride_tricks.sliding_window_view(x, (2, 2))

Each of the innermost 2x2 arrays contains an unraveled version of what you want, but with the second part of the array reversed. So far we didn't copy any data. Now make a copy by raveling the last dimension. The reshape will always make a copy here because b is highly non-contiguous:

c = b.reshape(*b.shape[:2], 4)

Swap the last two columns:

c[..., 2:] = c[..., -1:1:-1]

Now ravel the leading dimensions:

y = c.reshape(-1, c.shape[-1])

If you have a version of numpy that is older whan 1.20, you can replace the definition of b with

b = np.lib.stride_tricks.as_strided(x, shape=(x.shape[0] - 1, x.shape[1] - 1, 2, 2), strides=x.strides * 2)

CodePudding user response：

You can cache your code since the loop is mainly iterating the same matrice again and again (If you want to keep your same code with loop). I have made a speed comparison for your code before and after caching.

# Before caching
def loop_before_cache():
    n_rows, n_cols = x.shape
    y = []
    for i in range(n_rows-1): 
        for j in range(n_cols-1): 
            row = [x[i,j],x[i,j 1],x[i 1, j 1],x[i 1,j]] 
            y.append(row) 
    return np.array(y)


%timeit loop_before_cache()
11.6 µs ± 318 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

And now with caching

# After caching
from functools import lru_cache

@lru_cache()
def loop_after_cache():
    n_rows, n_cols = x.shape
    y = []
    for i in range(n_rows-1): 
        for j in range(n_cols-1): 
            row = [x[i,j],x[i,j 1],x[i 1, j 1],x[i 1,j]] 
            y.append(row) 
    return np.array(y)

%timeit loop_after_cache()
83.6 ns ± 2.42 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)