Wrap slice around edges of a 2D array in numpy

Question

Suppose I am working with numpy in Python and I have a two-dimensional array of arbitrary size. For convenience, let\'s say I have a 5 x 5 array. The specific numbers are n

Answer 1

As I mentioned in the comments, there is a good answer at How do I select a window from a numpy array with periodic boundary conditions?

Here is another simple way to do this

# First some setup
import numpy as np
A = np.arange(25).reshape((5, 5))
m, n = A.shape

and then

A[np.arange(i-1, i+2)%m].reshape((3, -1))[:,np.arange(j-1, j+2)%n]

It is somewhat harder to obtain something that you can assign to. Here is a somewhat slower version. In order to get a similar slice of values I would have to do

A.flat[np.array([np.arange(j-1,j+2)%n+a*n for a in xrange(i-1, i+2)]).ravel()].reshape((3,3))

In order to assign to this I would have to avoid the call to reshape and work directly with the flattened version returned by the fancy indexing. Here is an example:

n = 7
A = np.zeros((n, n))
for i in xrange(n-2, 0, -1):
    A.flat[np.array([np.arange(i-1,i+2)%n+a*n for a in xrange(i-1, i+2)]).ravel()] = i+1
print A

which returns

[[ 2.  2.  2.  0.  0.  0.  0.]
 [ 2.  2.  2.  3.  0.  0.  0.]
 [ 2.  2.  2.  3.  4.  0.  0.]
 [ 0.  3.  3.  3.  4.  5.  0.]
 [ 0.  0.  4.  4.  4.  5.  6.]
 [ 0.  0.  0.  5.  5.  5.  6.]
 [ 0.  0.  0.  0.  6.  6.  6.]]