Pythonic way to split 3D array in smaller blocks of fixed dimension

Question

I want to implement the following routine without using loops, just using Numpy functions or ndarray methods. Here is the code:

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

I have tried with the "reshape" method so that:

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

but this seems to change the order of the elements in some way, and the blocks needs to be stored exactly as in the routine. Can anyone provide a way to do this, with a short explanation of why the reshape method gives a different result here? (maybe I am missing using the np.transpose() in addition?)

EDIT: I came up with this alternative implementation, but I am still not sure if this is the most efficient way (maybe someone can shed some light here):

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

EXAMPLE: Here is a minimal working example for the two implementations. What we want is, from an initial "cube" of dimensions Nx3MX3M, decompose into blocks NxMxMx3x3, which are the chunked version of the original block. With the two implementations above, one can check that they give the same result; the question is on how to achieve this in an efficient way (i.e., no loops)

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

Could you given an example of what input and output you are expecting? — Kevin
– Kevin, Commented Mar 29, 2021 at 13:55
I wrote a utility to separate a 3D volume into smaller blocks. You can find it here. — jkr
– jkr, Commented Mar 29, 2021 at 18:26

gofvonx · Accepted Answer · 2021-03-30 08:38:24Z

4

If I understand correctly, np.reshape should work. While there are different options for the order argument, it sounds like you want to reshape the last two axes of length 3M in the original array both to an array of shape (M,3) in the 'C' order (the default). This can be achieved with

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

Getting the desired output shape (N, M, 3, M, 3) is then a matter of using 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

edited Mar 30, 2021 at 8:38

answered Mar 29, 2021 at 18:20

gofvonx

1,4591 gold badge14 silver badges23 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Pythonic way to split 3D array in smaller blocks of fixed dimension

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related