Home > Net >  Iterating through all subsets of a 3d array
Iterating through all subsets of a 3d array

Time:10-18

Assuming I have created an array by using the code below:

import numpy as np 
array = np.random.randint(0, 255, (3,4,4))
>>> array
array([[[ 72,  11, 158, 252],
        [160,  50, 131, 174],
        [245, 127,  99,   6],
        [152,  25,  58,  96]],

       [[ 29,  37, 211, 215],
        [195,  72, 186,  33],
        [ 12,  68,  44, 241],
        [ 95, 184, 188, 176]],

       [[238,  90, 177,  15],
        [ 48, 221,  41, 236],
        [ 86,  14, 130, 192],
        [ 64,  17,  44, 251]]])

Which created a 3d array of 8x8 matrices I would like to go through all possible adjacent 2x2 sub-matrices. For example, I would like to get the following matrices from the first matrix:

[[72,11],
 [160, 50]]
[[11,158],
 [50, 131]]
[[160,50],
 [245,127]]
etc...

Is there a built-in numpy\pytorch method I can use or do I have to implement this iteration?

CodePudding user response:

Use sliding_window_view to simply generate the entire matrix:

>>> np.lib.stride_tricks.sliding_window_view(array, (1, 2, 2)).reshape(-1, 2, 2)
array([[[ 72,  11],
        [160,  50]],

       [[ 11, 158],
        [ 50, 131]],

       [[158, 252],
        [131, 174]],

       ...

       [[ 14, 130],
        [ 17,  44]],

       [[130, 192],
        [ 44, 251]]])

Note that the data will be copied in the last step of reshaping, and the memory consumption of the results generated by large arrays may be unacceptable. If you just want to iterate over each sub array, you can reshape the original array to generate a view that is easier to iterate:

>>> from itertools import chain
>>> view = np.lib.stride_tricks.sliding_window_view(
...     array.reshape(-1, array.shape[-1]), (2, 2)
... )
>>> for ar in chain.from_iterable(view):
...     print(ar)
...
[[ 72  11]
 [160  50]]
[[ 11 158]
 [ 50 131]]
[[158 252]
 [131 174]]
...
[[ 14 130]
 [ 17  44]]
[[130 192]
 [ 44 251]]

The above method has defects, which will cause the data between two 2d sub arrays to be treated as sub views (such as [[152, 25], [29, 37]]). The possible choice is not to use reshape and iterate through multiple loops.

  • Related