Generalize stacking of array elements' neighbors into 3-D array

问题

Setup

Given a 2D array, I would like to create a 3D array where the values along the third dimension at (i.e. stacked[row, col, :]) are the flattened neighbors of the original array at [row, col]. I would like to generalize this process to handle an arbitrary (but reasonable) search radius.

Prior research

This question seemed promising, but I'm not sure I can really utilize its approach without a (couple of) for loops. My current approach, applied with a search radius of 1, for brevity's sake is illustrated with the example below.

Also this question + answer were close, but I'm specifically looking for a solution that purely uses smart indexing to avoid loops.

What I have now

import numpy as np
np.random.seed(0)

x = np.random.random_integers(0, 10, size=(4, 5))
print(x)  # * highlights the neighbors we'll see later

[[ 5   0   3   3   7]
 [ 9  *3  *5  *2   4]
 [ 7  *6  *8  *8  10]
 [ 1  *6  *7  *7   8]]

# padding the edges 
padded = np.pad(x, mode='edge', pad_width=1) # pad_width -> search radius
print(padded)

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

So then we can stack up all of the neighbors. This is the operation that I would like to generalize

blocked = np.dstack([
    padded[0:-2, 0:-2], # upper left
    padded[0:-2, 1:-1], # upper center
    padded[0:-2, 2:],   # upper right
    padded[1:-1, 0:-2], # middle left...
    padded[1:-1, 1:-1],
    padded[1:-1, 2:],
    padded[2:, 0:-2],   # lower left ...
    padded[2:, 1:-1],
    padded[2:, 2:],
])

And accessing the neighbors if a cell looks like this (the call to reshape for illustrative purposes only)

print(blocked[2, 2, :].reshape(3, 3))
[[3 5 2]
 [6 8 8]
 [6 7 7]]

Primary question

For a given search radius, is there an effective way to generalize the call to np.dstack?

回答1:

This could be one approach -

import numpy as np

# Parameters
R = 3  # Radius
M1,N1 = padded.shape
rowlen = N1 - R + 1
collen = M1 - R + 1

# Linear indices for the starting R x R block
idx1 = np.arange(R)[:,None]*N1 + np.arange(R)

# Offset (from the starting block indices) linear indices for all the blocks
idx2 = np.arange(collen)[:,None]*N1 + np.arange(rowlen)

# Finally, get the linear indices for all blocks
all_idx = idx1.ravel()[None,None,:] + idx2[:,:,None]

# Index into padded for the final output
out = padded.ravel()[all_idx]

Here's a sample run for radius, R = 4 -

In [259]: padded
Out[259]: 
array([[ 5,  5,  0,  3,  3,  3],
       [ 5,  5,  0,  3,  3,  3],
       [ 7,  7,  9,  3,  5,  5],
       [ 2,  2,  4,  7,  6,  6],
       [ 8,  8,  8, 10,  1,  1],
       [ 6,  6,  7,  7,  8,  8],
       [ 6,  6,  7,  7,  8,  8]])

In [260]: out
Out[260]: 
array([[[ 5,  5,  0,  3,  5,  5,  0,  3,  7,  7,  9,  3,  2,  2,  4,  7],
        [ 5,  0,  3,  3,  5,  0,  3,  3,  7,  9,  3,  5,  2,  4,  7,  6],
        [ 0,  3,  3,  3,  0,  3,  3,  3,  9,  3,  5,  5,  4,  7,  6,  6]],

       [[ 5,  5,  0,  3,  7,  7,  9,  3,  2,  2,  4,  7,  8,  8,  8, 10],
        [ 5,  0,  3,  3,  7,  9,  3,  5,  2,  4,  7,  6,  8,  8, 10,  1],
        [ 0,  3,  3,  3,  9,  3,  5,  5,  4,  7,  6,  6,  8, 10,  1,  1]],

       [[ 7,  7,  9,  3,  2,  2,  4,  7,  8,  8,  8, 10,  6,  6,  7,  7],
        [ 7,  9,  3,  5,  2,  4,  7,  6,  8,  8, 10,  1,  6,  7,  7,  8],
        [ 9,  3,  5,  5,  4,  7,  6,  6,  8, 10,  1,  1,  7,  7,  8,  8]],

       [[ 2,  2,  4,  7,  8,  8,  8, 10,  6,  6,  7,  7,  6,  6,  7,  7],
        [ 2,  4,  7,  6,  8,  8, 10,  1,  6,  7,  7,  8,  6,  7,  7,  8],
        [ 4,  7,  6,  6,  8, 10,  1,  1,  7,  7,  8,  8,  7,  7,  8,  8]]])

来源：https://stackoverflow.com/questions/29925684/generalize-stacking-of-array-elements-neighbors-into-3-d-array

标签

python

arrays

numpy

vectorization