Как выбрать окно из массива с периодическими граничными условиями?

Question

Предположим, я сделал 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]])

и я хочу иметь возможность выбрать окно 3х3 вокруг любого заданного элемента, чтобы окно обвивало границы, как бы я это сделал? Я знаю, что могу сделать это, если границы окна не перекрывают границы исходного массива:

>>> A[1:4,0:3]
array([[ 4,  5,  6],
       [ 8,  9, 10],
       [12, 13, 14]])

но если я использую выражение типа A[i-1:i+2,j-1:j+2], оно возвращает только пустой массив для i = 0, например, j = 0.

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