Wsrtjtyk

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

Here's how you can do it without padding. This can generalize easily to when you want more than just one neighbor and without the overhead of padding the array.

def get_wrapped(matrix, i, j):

  m, n = matrix.shape

  rows = [(i-1) % m, i, (i+1) % m]

  cols = [(j-1) % n, j, (j+1) % n]

  return matrix[rows][:, cols]



res = get_wrapped(matrix, 2, 4)

Let me explain what's happening here return matrix[rows][:, cols]. This is really two operations.

The first is matrix[rows] which is short hand for matrix[rows, :] which means give me the selected rows, and all columns for those rows.

Then next we do [:, cols] which means give me all the rows and the selected cols.

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]])