将 3D 数组拆分为固定维度的较小块的 Pythonic 方法

Question

我想在不使用循环的情况下实现以下例程，只使用 Numpy 函数或 ndarray 方法。这是代码:

def split_array_into_blocks( the_array, block_dim, total_blocks_per_row ):

    n_grid = the_array.shape[0]
    res    = np.empty( (n_grid, total_blocks_per_row, total_blocks_per_row, block_dim, 
    block_dim ) )

    for i in range( total_blocks_per_row  ):
        for j in range( total_blocks_per_row ):
        
            subblock          = the_array[ :, block_dim*i:block_dim*(i+1), block_dim*j:block_dim*(j+1) ]   
            res[ :, i,j,:,: ] = subblock 

    return res

我尝试过使用“ reshape ”方法，以便:

the_array = the_array.reshape( ( n_grid, total_blocks_per_row, total_blocks_per_row, block_dim, block_dim) )

但这似乎以某种方式改变了元素的顺序，并且 block 需要完全按照例程中的方式存储。谁能提供一种方法来做到这一点，并简要解释一下为什么 reshape 方法会在这里给出不同的结果？ (也许我还缺少使用 np.transpose() ？)
编辑:我想出了这个替代实现，但我仍然不确定这是否是最有效的方法(也许有人可以在这里阐明):

def split_array_in_blocks( the_array, block_dim, total_blocks_per_row ):
       
    indx   = [ block_dim*j for j in range( 1, total_blocks_per_row ) ] 

    the_array = np.array( [ np.split(  np.split( the_array, indx, axis=1 )[j], indx, axis=-1 ) for j in range( total_blocks_per_row ) ] )
    the_array = np.transpose( the_array, axes=( 2,0,1,3,4 ) )


    return the_array

示例:这是两个实现的最小工作示例。我们想要的是，从尺寸为 Nx3MX3M 的初始“立方体”分解成 block NxMxMx3x3，它们是原始 block 的分 block 版本。使用上面的两种实现，可以检查它们是否给出相同的结果；问题是如何以有效的方式实现这一点(即没有循环)

import numpy as np


def split_array_in_blocks_2( the_array, block_dim, total_blocks_per_row ):
    
    n_grid = the_array.shape[0]
    res    = np.zeros( (n_grid, total_blocks_per_row, total_blocks_per_row, block_dim, block_dim ), dtype=the_array.dtype )
    
    for i in range( total_blocks_per_row ):
        for j in range( total_blocks_per_row ):            
            subblock          = the_array[ :, block_dim*i:block_dim*(i+1), block_dim*j:block_dim*(j+1) ]   
            res[ :, i,j,:,: ] = subblock 
    
    return res


def split_array_in_blocks( the_array, block_dim, total_blocks_per_row ):
           
    indx   = [ block_dim*j for j in range( 1, total_blocks_per_row ) ] 
    
    the_array = np.array( [ np.split(  np.split( the_array, indx, axis=1 )[j], indx, axis=-1 ) for j in range( total_blocks_per_row ) ] )
    the_array = np.transpose( the_array, axes=( 2,0,1,3,4 ) )
    
    
    return the_array

A = np.random.rand( 1001, 63, 63 )
n = 3
D = 21
from time import time

ts = time()
An = split_array_in_blocks( A, n, D )

t2 = time()

Bn = split_array_in_blocks_2( A, n, D )

t3 = time()
print( t2-ts )
print(t3-t2)
print(np.allclose( An, Bn ))

Answer 1

如果我理解正确，np.reshape应该管用。虽然 order 有不同的选项论点，听起来你想 reshape 长度的最后两个轴 3M在原始数组中都为形状数组 (M,3)在 'C'顺序(默认)。这可以通过

A.reshape(A.shape[0], A.shape[1]//3, 3, A.shape[1]//3, 3, order='C')

获得所需的输出形状 (N, M, 3, M, 3)然后是使用 np.swapaxes 的问题:

reshaped = np.swapaxes(A.reshape(A.shape[0], A.shape[1]//3, 3, A.shape[1]//3, 3, order='C'), axis1=-2, axis2=-3)

np.allclose(reshaped, An)  # This is true