Splitting 2D array into small 2D arrays according row/col indexes

Question

Although I made google search for this type splitting, I could not find appropriate solution. For example I have a 6x6 2D array filled with numbers with numpy arange command.

big = np.arange(6*6).reshape([6,6])

[[ 0  1  2  3  4  5]
 [ 6  7  8  9 10 11]
 [12 13 14 15 16 17]
 [18 19 20 21 22 23]
 [24 25 26 27 28 29]
 [30 31 32 33 34 35]]

I want to split this 2D array into 3x3 small 2D array as below.

1th row of big array
 [[ n  n  n]    [[ n  n  n]    [[ n  n  n]    [[ n  n  n]    [[ n  n  n]    [[ n  n  n] 
  [ n  0  1]     [ 0  1  2]     [ 1  2  3]     [ 2  3  4]     [ 3  4  5]     [ 4  5  n] 
  [ n  6  7]]    [ 6  7  8]]    [ 7  8  9]]    [ 8  9 10]]    [ 9 10 11]]    [10 11  n]] 
  
2nd row
 [[ n  0  1]    [[ 0  1  2]    [[ 1  2  3]    [[ 2  3  4]    [[ 3  4  5]    [[ 4  5  n] 
  [ n  6  7]     [ 6  7  8]     [ 7  8  9]     [ 8  9 10]     [ 9 10 11]     [10 11  n] 
  [ n 12 13]]    [12 13 14]]    [13 14 15]]    [14 15 16]]    [15 16 17]]    [16 17  n]] 

and last row
 [[ n 24 25]    [[24 25 26]    [[25 26 27]    [[26 27 28]    [[27 28 29]    [[28 29  n] 
  [ n 30 31]     [30 31 32]     [31 32 33]     [32 33 34]     [33 34 35]     [34 35  n] 
  [ n  n  n]]    [ n  n  n]]    [ n  n  n]]    [ n  n  n]]    [ n  n  n]]    [ n  n  n]] 

As seen, the value in the center of each small array is the cell value in each row of the large array, letter n should be the NaN value.

How can we create these small 2D arrays using numpy? Every little help will be precious to me. Thanks to everyone who has given their valuable help so far.

Answer 1

You can use np.pad to add the NaNs on the edges, and np.lib.stride_tricks.sliding_window_view to split it up:

padded = np.pad(big.astype(float), (1,1), constant_values=np.nan)
s = sliding_window_view(padded, (3,3))

Output:

>>> s.shape
(6, 6, 3, 3)

>>> s[0]
array([[[nan, nan, nan],
        [nan,  0.,  1.],
        [nan,  6.,  7.]],

       [[nan, nan, nan],
        [ 0.,  1.,  2.],
        [ 6.,  7.,  8.]],

       [[nan, nan, nan],
        [ 1.,  2.,  3.],
        [ 7.,  8.,  9.]],

       [[nan, nan, nan],
        [ 2.,  3.,  4.],
        [ 8.,  9., 10.]],

       [[nan, nan, nan],
        [ 3.,  4.,  5.],
        [ 9., 10., 11.]],

       [[nan, nan, nan],
        [ 4.,  5., nan],
        [10., 11., nan]]])

>>> s[1]
array([[[nan,  0.,  1.],
        [nan,  6.,  7.],
        [nan, 12., 13.]],

       [[ 0.,  1.,  2.],
        [ 6.,  7.,  8.],
        [12., 13., 14.]],

       [[ 1.,  2.,  3.],
        [ 7.,  8.,  9.],
        [13., 14., 15.]],

       [[ 2.,  3.,  4.],
        [ 8.,  9., 10.],
        [14., 15., 16.]],

       [[ 3.,  4.,  5.],
        [ 9., 10., 11.],
        [15., 16., 17.]],

       [[ 4.,  5., nan],
        [10., 11., nan],
        [16., 17., nan]]])

>>> s[-1]
array([[[nan, 24., 25.],
        [nan, 30., 31.],
        [nan, nan, nan]],

       [[24., 25., 26.],
        [30., 31., 32.],
        [nan, nan, nan]],

       [[25., 26., 27.],
        [31., 32., 33.],
        [nan, nan, nan]],

       [[26., 27., 28.],
        [32., 33., 34.],
        [nan, nan, nan]],

       [[27., 28., 29.],
        [33., 34., 35.],
        [nan, nan, nan]],

       [[28., 29., nan],
        [34., 35., nan],
        [nan, nan, nan]]])