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

...]

https://stackoverflow.com/questions/56032778/how-to-create-an-array-of-neighbors-out-from-each-element-in-an-numpy-2d-array seems to be partially related — Georg Heiler, Jan 29 '22 at 14:38
Based on your input array, `[[1, 0], [0, 1]]`, will you please show a sample output array that you want to get? — , Jan 29 '22 at 16:58
Basically like the convolution operator https://www.researchgate.net/publication/350400158/figure/fig2/AS:1005540972974080@1616751058064/a-Standard-2D-Convolution-With-a-fixed-filer-size-the-standard-convolution-computes-a.png though I want to consider left, upper, right and bottom (4) neighbors (not the diagonal ones) and store them into an array/matrix which has 4 entries (we start with the original 2x2 matrix) but for each of the entries instead store an array/vector with these 4 neighbors. — Georg Heiler, Jan 29 '22 at 17:21

bb1 · Accepted Answer · 2022-01-30T04:45:30.047

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).

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

edit

1 Answers1