Python-获取 3d 数组的 "subarrays"

Question

我想获取 3D 数组的多个子数组。我可以使用 Stack 帖子中找到的函数在 2D 情况下拆分数组:

def blockshaped(arr, nrows, ncols):
    h, w = arr.shape
    return (arr.reshape(h//nrows, nrows, -1, ncols)
               .swapaxes(1,2)
               .reshape(-1, nrows, ncols))

所以我想将其扩展到 3D 数组情况，将 block 形成为 2D 数组，但在第一维的每个切片中。我尝试使用“for 循环”但不起作用......

例如:

import numpy as np

#2D case (which works)

test=np.array([[ 2.,  1.,  1., 1.],
        [ 1.,  1.,  1., 1.],
        [ 3.,  1.,  1., 1.],
        [ 1.,  1.,  1., 1.]])

def blockshaped(arr, nrows, ncols): 

    h, w = arr.shape
    return (arr.reshape(h//nrows, nrows, -1, ncols)
               .swapaxes(1,2)
               .reshape(-1, nrows, ncols))


sub = blockshaped(test, 2,2)

我得到 4 个“子数组”:

array([[[ 2.,  1.],
        [ 1.,  1.]],

       [[ 1.,  1.],
        [ 1.,  1.]],

       [[ 3.,  1.],
        [ 1.,  1.]],

       [[ 1.,  1.],
        [ 1.,  1.]]])

但是对于 3D 数组作为输入:

test2=np.array([[[ 2.,  1.,  1., 1.],
        [ 1.,  1.,  1., 1.],
        [ 3.,  1.,  1., 1.],
        [ 1.,  1.,  1., 1.]],

       [[ 5.,  1.,  1., 1.],
        [ 1.,  1.,  1., 1.],
        [ 2.,  1.,  1., 1.],
        [ 1.,  1.,  1., 1.]]])       

所以在这里我想要相同的分解，但是在 2 个“切片”中......

def blockshaped(arr, nrows, ncols): 

    h, w, t = arr.shape 
    return (arr.reshape(h//nrows, nrows, -1, ncols)
               .swapaxes(1,2)
               .reshape(-1, nrows, ncols))

我尝试使用“for 循环”但不起作用:

for i in range(test2.shape[0]):                     
    sub = blockshaped(test[i,:,:], 2, 2)

Answer 1

您的 for 循环解决方案可以执行以下操作:

sub = np.array([blockshaped(a, 2, 2) for a in test2])

但是你可以稍微修改blockshape()，在切片之前和之后 reshape 数据:

def blockshaped(arr, nrows, ncols):
    need_reshape = False
    if arr.ndim > 2:
        need_reshape = True
    if need_reshape:
        orig_shape = arr.shape
        arr = arr.reshape(-1, arr.shape[-1])
    h, w = arr.shape
    out = (arr.reshape(h//nrows, nrows, -1, ncols)
               .swapaxes(1, 2)
               .reshape(-1, nrows, ncols))
    if need_reshape:
        new_shape = list(out.shape)
        new_shape[0] //= orig_shape[0]
        out = out.reshape([-1,] + new_shape)
    return out