martianwars martianwars - 9 days ago 4
Python Question

How to split numpy array keeping a few elements from previous split?

I have a numpy array which I wish to split across a certain dimension. While splitting the array, I need to prepend (to the beginning of each element) a trailing part of the previous element. For instance,

Let my array be

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
. Let my
split_size = 2
and
pad_length = 1
.
split_size
will always be a divisor of array length. My resultant splits would look like,

[random, 0, 1], [1, 2, 3], [3, 4, 5], [5, 6, 7], [7, 8, 9]
. My splits were all prepended by the last value of the previous element.

Needless to say, my arrays are multidimensional and I need an efficent vectorized way to do this along a certain dimension.

Here, I can provide the value of
random
.

Answer

Sounds like a job for as_strided.

as_strided returns a memory efficient view on an array and can be used for retrieving a moving window over an array. The numpy documentation on it is scarce, but there's a number of decent blog posts, online slide decks, and SO issues that you can find that explain it in more detail.

>>> import numpy as np
>>> from numpy.lib.stride_tricks import as_strided
>>> a = np.arange(10)
>>> split_size = 2
>>> pad_length = 1
>>> random = -9
>>> # prepend the desired constant value
>>> b = np.pad(a, (pad_length, 0), mode='constant', constant_values=random)
>>> # return a memory efficient view on the array
>>> as_strided(b,
...     shape=(b.size//split_size, split_size + pad_length),
...     strides=(b.strides[0]*split_size, b.strides[0]))
...
array([[-9,  0,  1],
       [ 1,  2,  3],
       [ 3,  4,  5],
       [ 5,  6,  7],
       [ 7,  8,  9]])

Be aware that if the new strides go out of bounds, you'll see the memory contents of adjacent memory appearing at the end of the array.