numpy store the neigborhood as new (3rd) dimension into matrix

Question

Given a numpy array like:

L = 2
np.random.randint([-1,1],size=(L,L), dtype=int)
array([[1, -1],
       [-1, 1]])

How can I transform it into an array of similar shape (efficiently)

np.random.choice([-1, 1], size=(2,2,4))
array([[[-1, -1,  1,  1],
        [-1, -1,  1, -1]],

       [[-1,  1, -1,  1],
        [ 1, -1,  1,  1]]])

But unlike shown here where the 3rd dimension is random to contain the 4 neighbors in it (0-padded on the corners).

I.e.

[[1, -1], [-1, 1]]

has for the first element a neighborhood of:

• 0, 0, -1,-1,
• for the second 1,0, 0, 1 and so on.

I want to store this neighborhood vector into the 3rd dimension of the matrix.

Is this possible without manually looping the matrix?

edit

For the example of:

[[1, -1], [-1, 1]]

[[[0,0,-1-1], [1,0,0,1]],

...]

Answer 1

You can try the following:

#sample array
a = np.arange(9).reshape(3, 3)
print(a)

It gives:

[[0 1 2]
 [3 4 5]
 [6 7 8]]

Compute the array of neighbors:

p = np.pad(a, 1)
out = np.empty((*a.shape, 4), dtype=a.dtype)

out[..., 0] = p[:-2, 1:-1] #up
out[..., 1] = p[2:, 1:-1]  #down
out[..., 2] = p[1:-1, :-2] #left
out[..., 3] = p[1:-1, 2:]  #right

Then, for example out[2, 1] is [4, 0, 6, 8] i.e. the array of neighbors of a[2, 1] in the [up, down, left, right] order (with 0 padding).