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

Question

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]

Since offsets has arbitrary length, I'm not sure how useful not using a loop will be. If you're after conciseness, maybe try list comprehension? results = [A[x:x+size, y:y+size] for x, y in offsets] — Tony Liu
– Tony Liu, Commented Oct 18, 2022 at 3:09

Nikolay Zakirov · Accepted Answer · 2022-10-18 03:46:15Z

2

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

edited Oct 18, 2022 at 3:46

answered Oct 18, 2022 at 3:27

Nikolay Zakirov

1,58410 silver badges21 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Mechanic Pig Over a year ago

reshape can omit: np.lib.stride_tricks.sliding_window_view(a, (block_size, block_size))[tuple(offsets.T)]

Nikolay Zakirov Over a year ago

oh thanks! Could not make it work. Let me update

Collectives™ on Stack Overflow

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

1 Answer 1

2 Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Related