时间戳

首页 > 解决方案 > 将 3D 数组拆分为固定维度的较小块的 Pythonic 方法

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

问题描述

我想在不使用循环的情况下实现以下例程，只使用 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

我尝试过使用“重塑”方法，以便：

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

示例：这是两个实现的最小工作示例。我们想要的是，从尺寸为 Nx3 MX3 M 的初始“立方体”分解成块 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 ))

标签： pythonpython-3.xnumpy

解决方案

如果我理解正确，np.reshape应该可以工作。虽然参数有不同的选项order，但听起来您想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

推荐阅读