In Python, is it possible to return multiple 3D array slices all at once without a for loop?-CodePudding

Here is example code. I want to find the way to run the last two lines at once and return the results in the same array all at once without concatenation, is this possible?

import numpy as np
arr = np.ones((3,3,3))
arr[0:2,0:2,0:2]
arr[1:3,1:3,1:3]

The resulting command should be like

arr[(0:1,2:3),(0:1,2:3),(0:1,2:3)]

And the dimensionality of the results will be (2,2,2,2).

CodePudding user response：

Your array and slices:

In [128]: arr = np.arange(27).reshape(3,3,3)
In [129]: a1=arr[0:2,0:2,0:2]
     ...: a2=arr[1:3,1:3,1:3]
In [130]: a1.shape
Out[130]: (2, 2, 2)
In [131]: a2.shape
Out[131]: (2, 2, 2)

a1 and a2 are views, sharing the databuffer with arr.

Joining them on a new dimension (np.stack will also do this):

In [132]: a3 = np.array((a1,a2))
In [133]: a3.shape
Out[133]: (2, 2, 2, 2)
In [134]: a3
Out[134]: 
array([[[[ 0,  1],
         [ 3,  4]],

        [[ 9, 10],
         [12, 13]]],


       [[[13, 14],
         [16, 17]],

        [[22, 23],
         [25, 26]]]])

Notice how the flattened values are not contiguous (or other regular pattern). So they have to a copy of some sort:

In [135]: a3.ravel()
Out[135]: array([ 0,  1,  3,  4,  9, 10, 12, 13, 13, 14, 16, 17, 22, 23, 25, 26])

An alternative is to construct the indices, join them, and then do one indexing. That times about the same. And in this case I think that would be more complicated.

===

Another way with stride_tricks. I won't promise anything about speed.

In [147]: x = np.lib.stride_tricks.sliding_window_view(arr,(2,2,2))
In [148]: x.shape
Out[148]: (2, 2, 2, 2, 2, 2)
In [149]: x[0,0,0]
Out[149]: 
array([[[ 0,  1],
        [ 3,  4]],

       [[ 9, 10],
        [12, 13]]])
In [150]: x[1,1,1]
Out[150]: 
array([[[13, 14],
        [16, 17]],

       [[22, 23],
        [25, 26]]])
In [151]: x[[0,1],[0,1],[0,1]]
Out[151]: 
array([[[[ 0,  1],
         [ 3,  4]],

        [[ 9, 10],
         [12, 13]]],


       [[[13, 14],
         [16, 17]],

        [[22, 23],
         [25, 26]]]])