How to loop back to beginning of the array for out of bounds index in numpy?

Question

I have a 2D numpy array that I want to extract a submatrix from. I get the submatrix by slicing the array as below. Here I want a 3*3 submatrix around an item at the index of (2,3).

>>> import numpy as np
>>> a = np.array([[0, 1, 2, 3],
...        [4, 5, 6, 7],
...        [8, 9, 0, 1],
...        [2, 3, 4, 5]])
>>> a[1:4, 2:5]
array([[6, 7],
   [0, 1],
   [4, 5]])

But what I want is that for indexes that are out of range, it goes back to the beginning of array and continues from there. This is the result I want:

array([[6, 7, 4],
   [0, 1, 8],
   [4, 5, 2]])

I know that I can do things like getting mod of the index to the width of the array; but I'm looking for a numpy function that does that. And also for an one dimensional array this will cause an index out of range error, which is not really useful...

Answer 1

This is one way using np.pad with wraparound mode.

>>> a = np.array([[0, 1, 2, 3],
                  [4, 5, 6, 7],
                  [8, 9, 0, 1],
                  [2, 3, 4, 5]])

>>> pad_width = 1
>>> i, j = 2, 3

>>> startrow, endrow = i-1+pad_width, i+2+pad_width # for 3 x 3 submatrix
>>> startcol, endcol = j-1+pad_width, j+2+pad_width

>>> np.pad(a, (pad_width, pad_width), 'wrap')[startrow:endrow, startcol:endcol]
array([[6, 7, 4],
       [0, 1, 8],
       [4, 5, 2]])

Depending on the shape of your patch (eg. 5 x 5 instead of 3 x 3) you can increase the pad_width and start and end row and column indices accordingly.

Answer 2

np.take does have a mode parameter which can wrap-around out of bound indices. But it's a bit hacky to use np.take for multidimensional arrays since the axis must be a scalar.

However, In your particular case you could do this:

a = np.array([[0, 1, 2, 3],
              [4, 5, 6, 7],
              [8, 9, 0, 1],
              [2, 3, 4, 5]])

np.take(a, np.r_[2:5], axis=1, mode='wrap')[1:4]

Output:

array([[6, 7, 4],
       [0, 1, 8],
       [4, 5, 2]])

EDIT

This function might be what you are looking for (?)

def select3x3(a, idx):
    x,y = idx
    return np.take(np.take(a, np.r_[x-1:x+2], axis=0, mode='wrap'), np.r_[y-1:y+2], axis=1, mode='wrap')

But in retrospect, i recommend using modulo and fancy indexing for this kind of operation (it's basically what the mode='wrap' is doing internally anyways):

def select3x3(a, idx):
    x,y = idx
    return a[np.r_[x-1:x+2][:,None] % a.shape[0], np.r_[y-1:y+2][None,:] % a.shape[1]]

The above solution is also generalized for any 2d shape on a .