How to index multiple blocks from 2d array given multiple offsets?-CodePudding

Suppose I have a 2D array A with a size of 10x10

# A
array([[ 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, 49],
       [50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
       [60, 61, 62, 63, 64, 65, 66, 67, 68, 69],
       [70, 71, 72, 73, 74, 75, 76, 77, 78, 79],
       [80, 81, 82, 83, 84, 85, 86, 87, 88, 89],
       [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]])

and given multiple offset points offsets with a size of 3x2 and block_size=3

# offsets
array([[0, 0],
       [2, 3],
       [4, 5]])

I want to index multiple blocks from A using offsets block_size

block_size = 3
results = []
for x, y in offsets:
    block = A[x:x block_size, y:y block_size]
    results.append(block)
results = np.array(results)

# results.shape = (3, 3, 3)
array([[[ 0,  1,  2],
        [10, 11, 12],
        [20, 21, 22]],

       [[23, 24, 25],
        [33, 34, 35],
        [43, 44, 45]],

       [[45, 46, 47],
        [55, 56, 57],
        [65, 66, 67]]])

Is there any way to do this without explicitly using for loop? Maybe something like this

# this doesn't work
results = A[offsets[:,0]:offsets[:,0] block_size, offsets[:,1]:offsets[:,1] block_size]

CodePudding user response：

Updated with Mechanic Pig suggestion

import numpy as np

a = np.array([[ 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, 49],
       [50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
       [60, 61, 62, 63, 64, 65, 66, 67, 68, 69],
       [70, 71, 72, 73, 74, 75, 76, 77, 78, 79],
       [80, 81, 82, 83, 84, 85, 86, 87, 88, 89],
       [90, 91, 92, 93, 94, 95, 96, 97, 98, 99],])

offsets = np.array([[0, 0],
       [2, 3],
       [4, 5]])

block_size = 3

res = (
    np.lib.stride_tricks
    .sliding_window_view(a, (block_size, block_size))
    [tuple(offsets.T)]
)
res