Python: how to select contiguity neighbors of matrix values?

Question

I have a matrix like the following:

A = array([[12,  6, 14,  8,  4,  1],
       [18, 13,  8, 10,  9, 19],
       [ 8, 15,  6,  5,  6, 18],
       [ 3,  0,  2, 14, 13, 12],
       [ 4,  4,  5, 19,  0, 14],
       [16,  8,  7,  7, 11,  0],
       [ 3, 11,  2, 19, 11,  5],
       [ 4,  2,  1,  9, 12, 12]])

For each cell I want to select the values in a radius of k=2 closest cells.

For instance if I select the A[3,4] I would like a submatrix like the following

array([[18, 13,  8, 10,  9],
       [ 8, 15,  6,  5,  6],
       [ 3,  0,  2, 14, 13],
       [ 4,  4,  5, 19,  0],
       [16,  8,  7,  7, 11]])

I defined the following function

def queen_neighbourhood(Adj, in_row, in_col, k):
    j=k
    k+=1
    neighbourhood = Adj[in_row-j:in_row+k, in_col-j:in_col+k]
    return neighbourhood

such as queen_neighbourhood(A, 3, 2, 2) returns

array([[18, 13,  8, 10,  9],
       [ 8, 15,  6,  5,  6],
       [ 3,  0,  2, 14, 13],
       [ 4,  4,  5, 19,  0],
       [16,  8,  7,  7, 11]])

However it does not work in borders.

For instance, for the cell [0,0] I would like to have

array([[12, 6,  14],
       [18, 13,  8],
       [ 8, 15, 16])

but it returns queen_neighbourhood(A, 0, 0, 2)

array([], shape=(0, 0), dtype=int64)

Answer 1

You could avoid negative indices:

    neighbourhood = Adj[max(in_row-j, 0) : in_row+k,
                        max(in_col-j, 0) : in_col+k]

Answer 2

Adding to the previous answer; taking into consideration the extreme values

def queen_neighbourhood(Adj, in_row, in_col, k):
j=k
k+=1
neighbourhood = Adj[max(in_row-j, 0) : min(in_row+k,Adj.shape[0]),
                    max(in_col-j, 0) : min(in_col+k,Adj.shape[1])]
return(neighbourhood)

Answer 3

You can use numpy roll to ensure you are always dealing with the middle value,

import numpy as np

def queen_neighbourhood(Adj, in_row, in_col, k):
    j=k
    k+=1
    midrow = int(Adj.shape[0]/2.)+1
    midcol = int(Adj.shape[1]/2.)+1
    Ashift = np.roll(Adj,(in_row-midrow,in_col-midcol),(0,1))
    neighbourhood = Ashift[1:k+1, 1:k+1]
    return neighbourhood


A = np.array([[18, 13,  8, 10,  9],
       [ 8, 15,  6,  5,  6],
       [ 3,  0,  2, 14, 13],
       [ 4,  4,  5, 19,  0],
       [16,  8,  7,  7, 11]])

print(A)
An = queen_neighbourhood(A, 0, 0, 2)
print(An)

which gives,

[[11 16  8]
 [ 9 18 13]
 [ 6  8 15]]