Как выбрать окно из массива numpy с периодическими граничными условиями?
Предположим, что я делаю 2d-массив следующим образом:
>>> A=np.arange(16).reshape((4,4))
>>> A
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[12, 13, 14, 15]])
и я хочу иметь возможность выбрать окно 3x3 вокруг любого элемента, чтобы окно обернулось вокруг границ, как я могу это сделать? Я знаю, что могу сделать это, если границы окна не перекрывают границы исходного массива:
>>> A[1:4,0:3]
array([[ 4, 5, 6],
[ 8, 9, 10],
[12, 13, 14]])
но если я использую выражение типа A[i-1:i+2,j-1:j+2]
, оно возвращает только пустой массив для я = 0, j = 0.
Ответы
Ответ 1
import numpy as np
A=np.arange(16).reshape((4,4))
def neighbors(arr,x,y,n=3):
''' Given a 2D-array, returns an nxn array whose "center" element is arr[x,y]'''
arr=np.roll(np.roll(arr,shift=-x+1,axis=0),shift=-y+1,axis=1)
return arr[:n,:n]
print(A)
# [[ 0 1 2 3]
# [ 4 5 6 7]
# [ 8 9 10 11]
# [12 13 14 15]]
print(neighbors(A,0,0))
# [[15 12 13]
# [ 3 0 1]
# [ 7 4 5]]
print(neighbors(A,1,0))
# [[ 3 0 1]
# [ 7 4 5]
# [11 8 9]]