Selecting multiple slices from a numpy array at once

stride_tricks can do that

a = np.arange(10)b = np.lib.stride_tricks.as_strided(a, (3, 5), 2 * a.strides)b# array([[0, 1, 2, 3, 4],#        [1, 2, 3, 4, 5],#        [2, 3, 4, 5, 6]])

Please note that b references the same memory as a, in fact multiple times (for example b[0, 1] and b[1, 0] are the same memory address). It is therefore safest to make a copy before working with the new structure.

nd can be done in a similar fashion, for example 2d -> 4d

a = np.arange(16).reshape(4, 4)b = np.lib.stride_tricks.as_strided(a, (3,3,2,2), 2*a.strides)b.reshape(9,2,2) # this forces a copy# array([[[ 0,  1],#         [ 4,  5]],#        [[ 1,  2],#         [ 5,  6]],#        [[ 2,  3],#         [ 6,  7]],#        [[ 4,  5],#         [ 8,  9]],#        [[ 5,  6],#         [ 9, 10]],#        [[ 6,  7],#         [10, 11]],#        [[ 8,  9],#         [12, 13]],#        [[ 9, 10],#         [13, 14]],#        [[10, 11],#         [14, 15]]])

You can use the indexes to select the rows you want into the appropriate shape.For example:

 data = np.random.normal(size=(100,2,2,2)) # Creating an array of row-indexes indexes = np.array([np.arange(0,5), np.arange(1,6), np.arange(2,7)]) # data[indexes] will return an element of shape (3,5,2,2,2). Converting # to list happens along axis 0 data_extractions = list(data[indexes]) np.all(data_extractions[1] == data[1:6]) True

The final comparison is against the original data.

In this post is an approach with strided-indexing scheme using np.lib.stride_tricks.as_strided that basically creates a view into the input array and as such is pretty efficient for creation and being a view occupies nomore memory space.Also, this works for ndarrays with generic number of dimensions.

Here's the implementation -

def strided_axis0(a, L):    # Store the shape and strides info    shp = a.shape    s  = a.strides    # Compute length of output array along the first axis    nd0 = shp[0]-L+1    # Setup shape and strides for use with np.lib.stride_tricks.as_strided    # and get (n+1) dim output array    shp_in = (nd0,L)+shp[1:]    strd_in = (s[0],) + s    return np.lib.stride_tricks.as_strided(a, shape=shp_in, strides=strd_in)

Sample run for a 4D array case -

In [44]: a = np.random.randint(11,99,(10,4,2,3)) # ArrayIn [45]: L = 5      # Window length along the first axisIn [46]: out = strided_axis0(a, L)In [47]: np.allclose(a[0:L], out[0])  # Verify outputsOut[47]: TrueIn [48]: np.allclose(a[1:L+1], out[1])Out[48]: TrueIn [49]: np.allclose(a[2:L+2], out[2])Out[49]: True