Create 3D array from elements of a 2D array

Question

I'm working on a code and a question just pop up in my head. So basically I have a 2D numpy array with shape L,W = (4, 4) (this is just an example, the array, can be much bigger).

What I need is to create 3D array with elements from the 2D numpy array, where the elements in each cell of the output are: array[i:i+l, j:j+w](the elements of the subarray of dimension (l,w) starting from i, j position): output[i,j,:] = array[i:i+l,j:j+w].reshape(l*w,)

I thought about non-vectorized solution :

import numpy as np 

L = 4
W = 4
array = np.arange(16).reshape(L,W)

l= 2
w = 2

subarrays_elements = []
for i in range(L-(l-1)):
  for j in range(W-(w-1)):
    subarrays_elements.append(array[i:i+l,j:j+w].reshape(l*w,))


output = np.array(subarrays_elements).reshape(L-(l-1),W-(w-1),l*w)

the shape of the output is W-(w-1),L-(l-1),l*w,because we can't get a (l, w) subarray for the last l-1 rows and for w-1 columns.

The expected output would be array with (3,3,4):

expected_output = np.array([[[0 1 4 5],
                             [1 2 5 6],
                             [2 3 6 7]],

                             [[4 5 8 9],
                              [5 6 9 10],
                              [6 7 10 11]],

                             [[8 9 12 13],
                              [9 10 13 14],
                              [10 11 14 15]]])

I need solutions using only numpy and with vectorization, because I have a huge array, so any help will be appreciated, thank you!

Answer 1

This kind of problem is quite similar to getting the input for a convolutional layer in a neural network, so you can use the same tool, which is numpy.lib.stride_tricks.sliding_window_view.

Here's an example:

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

L = 4
W = 4
array = np.arange(L * W).reshape(L,W)

l = 2
w = 2

new_length = L-(l-1)
new_width = W-(w-1)
output = sliding_window_view(array, window_shape=[l, w]).reshape(new_length, new_width, l*w)

In theory, a sliding window should take almost no time at all, because it's just manipulating array strides and not copying data, but in this case the reshape forces it to make a complete copy of the array. On my computer this runs at about 50 million array elements per second on large arrays.