将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) )

但是这似乎以某种方式改变了元素的顺序，块需要按照例程中的精确存储。有人可以提供一种方法来做到这一点，并简要解释为什么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的“立方体”分解成块NxMxMx3x3，这些块是原始块的分块版本。 通过上面的两个实现，可以检查它们给出了相同的结果；问题在于如何以高效的方式实现这一点（即，没有循环）。

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 ))

- Zarathustra

---

你能给出一个输入和输出的例子吗？ - Kevin

我编写了一个实用程序，将3D体积分成较小的块。您可以在此处找到它：https://github.com/neuronets/nobrainer/blob/be89dd2e6059eaf705177f8f04f51fa1fc971f37/nobrainer/volume.py#L271-L299。 - jkr

Answer 1

如果我理解正确的话，np.reshape 应该可以用。虽然order参数有不同的选项，但你想要将原始数组中长度为3M的最后两个轴重塑为(M,3)形状的数组，并以'C'顺序（默认）排序。这可以通过以下方式实现：

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

通过使用 np.swapaxes，就可以得到所期望的输出形状 (N, M, 3, M, 3)：

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

- gofvonx