How can I extract a set of 2D slices from a larger 2D numpy array?-CodePudding

If I have a large 2D numpy array and 2 arrays which correspond to the x and y indices I want to extract, It's easy enough:

h = np.arange(49).reshape(7,7)
# h = [[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, 36, 37, 38, 39, 40, 41],
#      [42, 43, 44, 45, 46, 47, 48]]

x_indices = np.array([1,3,4])
y_indices = np.array([2,3,5])

reduced_h = h[x_indices, y_indices]
#reduced_h = [ 9, 24, 33]

However, I would like to, for each x,y pair cut out a square (denoted by 'a' - the number of indices in each direction from the centre) surrounding this 'coordinate' and return an array of these little 2D arrays.

For example, for h, x,y_indices as above and a=1:

reduced_h = [[[1,2,3],[8,9,10],[15,16,17]], [[16,17,18],[23,24,25],[30,31,32]], [[25,26,27],[32,33,34],[39,40,41]]]

i.e one 3x3 array for each x-y index pair corresponding to the 3x3 square of elements centred on the x-y index. In general, this should return a numpy array which has shape (len(x_indices),2a 1, 2a 1)

By analogy to reduced_h[0] = h[x_indices[0]-1:x_indices[0] 1 , y_indices[0]-1:y_indices[0] 1] = h[1-1:1 1 , 2-1:2 1] = h[0:2, 1:3] my first try was the following:

h[x_indices-a : x_indices a, y_indices-a : y_indices a]

However, perhaps unsurprisingly, slicing between the arrays fails. So the obvious next thing to try is to create this slice manually. np.arange seems to struggle with this but linspace works:

a=1
xrange = np.linspace(x_indices-a, x_indices a, 2*a 1, dtype=int)
# xrange = [ [0, 2, 3], [1, 3, 4], [2, 4, 5] ]
yrange = np.linspace(y_indices-a, y_indices a, 2*a 1, dtype=int)

Now can try h[xrange,yrange] but this unsurprisingly does this element-wise meaning I get only one (2a 1)x(2a 1) array (the same dimensions as xrange and yrange). It there a way to, for every index, take the right slices from these ranges (without loops)? Or is there a way to make the broadcast work initially without having to set up linspace explicitly? Thanks

CodePudding user response：

You can index np.lib.stride_tricks.sliding_window_view using your x and y indices:

import numpy as np

h = np.arange(49).reshape(7,7)

x_indices = np.array([1,3,4])
y_indices = np.array([2,3,5])

a = 1
window = (2*a 1, 2*a 1)

out = np.lib.stride_tricks.sliding_window_view(h, window)[x_indices-a, y_indices-a]

out:

array([[[ 1,  2,  3],
        [ 8,  9, 10],
        [15, 16, 17]],

       [[16, 17, 18],
        [23, 24, 25],
        [30, 31, 32]],

       [[25, 26, 27],
        [32, 33, 34],
        [39, 40, 41]]])

Note that you may need to pad h first to handle windows around your coordinates that reach "outside" h.