How Do I Select A Window From A Numpy Array With Periodic Boundary Conditions?

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], [1

Solution 1:

import numpy as np

A=np.arange(16).reshape((4,4))

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

Solution 2:

I can't comment yet but wanted to suggest an improvement over unutbu's solution:

Their solution can't handle cases like these

A=np.arange(25).reshape((5,5))

print(A) 
# [[ 0  1  2  3  4]#  [ 5  6  7  8  9]#  [10 11 12 13 14]#  [15 16 17 18 19]#  [20 21 22 23 24]]print(neighbors(A, 0, 0, n=5))
# [[24 20 21 22 23]#  [ 4  0  1  2  3]#  [ 9  5  6  7  8]#  [14 10 11 12 13]#  [19 15 16 17 18]]

0 should have been in the center but is off by one row and col.

A small modificatiion in shift values fixes it

defneighbors_updated(arr, x, y, n_row=3, n_col=3):
    ''' Given a 2D-array, returns an nxn array whose "center" element is arr[x,y]'''
    arr=np.roll(np.roll(arr,shift=-x+int(n_row/2),axis=0),shift=-y+int(n_col/2),axis=1)
    return arr[:n_row,:n_col]

print(neighbors(A, 0, 0, n_row=5, n_col=5))
# [[18 19 15 16 17]#  [23 24 20 21 22]#  [ 3  4  0  1  2]#  [ 8  9  5  6  7]#  [13 14 10 11 12]]