How do I select a window from a numpy array with periodic boundary conditions?

numpy, python

Solution

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

Problem

Suppose I make a 2d array like this: ``` >>> 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]]) ``` and I want to be able to select a 3x3 window around any given element so that the window wraps around the boundaries how would I do that? I know I can do this if the boundaries of the window don't overlap the boundaries of the original array: ``` >>> A[1:4,0:3] array([[ 4, 5, 6], [ 8, 9, 10], [12, 13, 14]]) ``` but if I use an expression like `A[i-1:i+2,j-1:j+2]` it only returns an empty array for i=0, j=0 for example.

Original source