pythonnumpynumpy-ndarraynumpy-slicing

Arbitrary Stencil Slicing in Numpy

Is there a simple syntax for creating references to an arbitrary number of neighbouring array elements in numpy?

The syntax is relatively straightforward when the number of neighbours is hard-coded. A stencil width of three for example is

import numpy as np

x = np.arange(8)

# Hard-coded stencil width of 3
x_neighbours = (
    x[  :-2]
    x[ 1:-1]
    x[ 2:  ]
)

However, my attempt at arbitrary width stencils is not particularly readable


nStencil = 3
x_neighbours = (
   x[indexStart:indexStop]
   for indexStart, indexStop in zip(
      (None, *range(1,nStencil)),
      (*range(1-nStencil,0), None),
   )
)

Is there a better approach?

Solution

  • I'd recommend using sliding_window_view. Change:

    nStencil = 3
x_neighbours = (
   x[indexStart:indexStop]
   for indexStart, indexStop in zip(
      (None, *range(1,nStencil)),
      (*range(1-nStencil,0), None),
   )
)

    To:

    nStencil = 3
sliding_view = sliding_window_view(x, nStencil) 
x_neighbours = tuple(sliding_view[:, i] for i in range(nStencil))