索引环绕的 numpy 二维数组

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

Answer 2

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.

Answer 3

The take function works in-place.

>>> a = np.arange(1, 16).reshape(3,5)
array([[ 1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10],
       [11, 12, 13, 14, 15]])
>>> b = np.take(a, [3,4,5], axis=1, mode='wrap')
array([[ 4,  5,  1],
       [ 9, 10,  6],
       [14, 15, 11]])
>>> np.take(b, [1,2,3], mode='wrap', axis=0)
array([[ 9, 10,  6],
       [14, 15, 11],
       [ 4,  5,  1]])