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

look at np.take

hpaulj
– hpaulj

2018-11-05 02:21:44 +00:00
Commented Nov 5, 2018 at 2:21 — hpaulj
– hpaulj, Commented Nov 5, 2018 at 2:21
@hpaulj something like this?

Antony Hatchkins
– Antony Hatchkins

2020-12-21 18:23:24 +00:00
Commented Dec 21, 2020 at 18:23 — Antony Hatchkins
– Antony Hatchkins, Commented Dec 21, 2020 at 18:23

user3483203 · Accepted Answer · 2018-11-05 02:30:44Z

6

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

edited Nov 5, 2018 at 2:30

answered Nov 5, 2018 at 2:25

user3483203

51.3k10 gold badges72 silver badges104 bronze badges

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

Comments

Silver · Accepted Answer · 2018-11-05 02:55:25Z

3

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.

edited Nov 5, 2018 at 2:55

answered Nov 5, 2018 at 2:37

Silver

1,4991 gold badge13 silver badges27 bronze badges

4 Comments

user1179317 Over a year ago

I see, this is actually very adaptable where if i want 2 neighbors i can just add 2 elements in the row and col list with a -2 and a +2. The only thing Im a little confuse about is whats the : for matrix[rows][:, cols]. I would think its just for matrix[rows][cols] but obviously that doesnt work.

Silver Over a year ago

So in numpy you can index into matrices using matrix[row, col] instead of pythons list[row][col]. So what that is doing is really two separate indexings a = matrix[rows] and b = a[:, cols]. The : just means all.

user1179317 Over a year ago

Ah right, that makes a lot of sense. Thank you for the explanation

hpaulj Over a year ago

Each '[]` is evaluated (by the Python interpreter) individually. matrix[rows] is done, then [cols] on the result. So you have to pay attention to the intermediate result. In [rows, cols] numpy handles both parts together.

Antony Hatchkins · Accepted Answer · 2020-12-21 18:22:11Z

0

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

answered Dec 21, 2020 at 18:22

Antony Hatchkins

34.4k13 gold badges121 silver badges114 bronze badges

Collectives™ on Stack Overflow

Indexing numpy 2D array that wraps around

3 Answers 3

Comments

4 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

4 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related