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

Question

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?

This looks like a sliding window problem, which has been addressed a number of times. — hpaulj, Apr 28 '15 at 17:46
See scikit image's functions view_as_windows and view_as_blocks. There are some answers here that go into the details of the shape and stride fiddling. eg. http://stackoverflow.com/a/4947453/553404 — YXD, Apr 28 '15 at 19:16
@MrE those functions looks really cool and I think they do what I need. I'll definitely look into skimage more. For this question's purposes -- and I should have mentioned this -- I'd like to avoid additional dependencies. — Paul H, Apr 28 '15 at 19:34
I wouldn't be worried about having scikit-image as a dependency but still you could adapt the source for these functions as they are very concise https://github.com/scikit-image/scikit-image/blob/master/skimage/util/shape.py — YXD, Apr 28 '15 at 19:55
possible duplicate of [Using strides for an efficient moving average filter](http://stackoverflow.com/questions/4936620/using-strides-for-an-efficient-moving-average-filter) — ali_m, Apr 29 '15 at 00:22

Divakar · Accepted Answer · 2015-04-28T17:53:36.173

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

Thanks. This is a very elegant solution. I notice that you use `N1` for both `idx1` and `idx2`. Should one of those use `M1`? — Paul H, Apr 28 '15 at 18:33
@PaulH Well `N1` doesn't have any role to play for the starting block, so that's good there. For the offset indices, we do need `M1`, which comes in as `collen`. — Divakar, Apr 28 '15 at 18:43
thanks for be info. This solution performs well. I'll leave the question open for a bit longer, — Paul H, Apr 28 '15 at 19:27

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

Setup

Prior research

What I have now

Primary question

1 Answers1

Linked