Indexing numpy 2D array that wraps around

Question

How do you index a numpy array that wraps around when its out of bounds?

For example, I have 3x3 array:

import numpy as np

matrix = np.array([[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15]])

## 
[[ 1  2  3  4  5]
 [ 6  7  8  9 10]
 [11 12 13 14 15]]

Say I would like to index the values around index (2,4) where value 15 is located. I would like to get back the array with values:

[[9,  10, 6]
 [14, 15, 11]
 [4,  5,  1]]

Basically all the values around 15 was returned, assuming it wraps around

Answer 1

A fairly standard idiom to find the neighboring elements in a numpy array is arr[x-1:x+2, y-1:y+2]. However, since you want to wrap, you can pad your array using wrap mode, and offset your x and y coordinates to account for this padding.

This answer assumes that you want the neighbors of the first occurence of your desired element.

---

First, find the indices of your element, and offset to account for padding:

x, y = np.unravel_index((m==15).argmax(), m.shape)
x += 1; y += 1

Now pad, and index your array to get your neighbors:

t = np.pad(m, 1, mode='wrap')    
out = t[x-1:x+2, y-1:y+2]  

array([[ 9, 10,  6],
       [14, 15, 11],
       [ 4,  5,  1]])