Search code examples
numpyreshape

How to you concat inner matrixes of tensor along axis?


Let we have a tensor with shape $n\times d\times h\times w\times p\times p$ and we want to concat innet matrix with shape $p\time p$, such that we made a matrix with shape $n\times d\times ph\times pw$. How can I do it?

array([[[[[[ 0,  1,  2],
           [ 3,  4,  5],
           [ 6,  7,  8]],

          [[ 9, 10, 11],
           [12, 13, 14],
           [15, 16, 17]]],


         [[[18, 19, 20],
           [21, 22, 23],
           [24, 25, 26]],

          [[27, 28, 29],
           [30, 31, 32],
           [33, 34, 35]]]]]])

after concat

array([[[[0,   1,  2,  9, 10, 11],
         [3,   4,  5, 12, 13, 14],
         [6,   7,  8, 15, 16, 17],
         [18, 19, 20, 27, 28, 29],
         [21, 22, 23, 30, 31, 32],
         [24, 25, 26, 33, 34, 35]]]])

I have done a lot of experiments using reshape, but without success. One of my experiments

a.reshape(n, d, p*h, p*w)

I can do it using for loop, but I think that without this, too, it is possible. Please help me. Code using for loop

p = 3
arr = np.arange(1*1*2*2*p*p).reshape(1, 1, 2, 2, p, p)
answer = np.zeros(shape=(1, 1, 2*p, 2*p))

for (n, d, h, w) in np.ndindex(*arr.shape[:4]):
    answer[n, d, h:h+p, w:w+p] = arr[n, d, h, w]

Solution

  • In [15]: arr=np.arange(0,36).reshape(2,2,3,3)
    

    reshape cannot reorder the elements of the array. I started with [0,1,...35], and reshape retains that:

    In [18]: arr.reshape(2,3,6)
    Out[18]: 
    array([[[ 0,  1,  2,  3,  4,  5],
            [ 6,  7,  8,  9, 10, 11],
            [12, 13, 14, 15, 16, 17]],
    
           [[18, 19, 20, 21, 22, 23],
            [24, 25, 26, 27, 28, 29],
            [30, 31, 32, 33, 34, 35]]])
    

    We have to somehow reorder the elements, putting the [9,10,11] block adjacent to the [0,1,2]. transpose is one such tool:

    In [19]: arr.transpose(0,2,1,3)
    Out[19]: 
    array([[[[ 0,  1,  2],
             [ 9, 10, 11]],
    
            [[ 3,  4,  5],
             [12, 13, 14]],
    
            [[ 6,  7,  8],
             [15, 16, 17]]],
    
    
           [[[18, 19, 20],
             [27, 28, 29]],
    
            [[21, 22, 23],
             [30, 31, 32]],
    
            [[24, 25, 26],
             [33, 34, 35]]]])
    
    In [20]: arr.transpose(0,2,1,3).reshape(6,6)
    Out[20]: 
    array([[ 0,  1,  2,  9, 10, 11],
           [ 3,  4,  5, 12, 13, 14],
           [ 6,  7,  8, 15, 16, 17],
           [18, 19, 20, 27, 28, 29],
           [21, 22, 23, 30, 31, 32],
           [24, 25, 26, 33, 34, 35]])
    

    To do the same by assignment to a 'blank' we need something like:

    In [32]: res=np.zeros((6,6),int)
    In [33]: res[:,:3] = arr[:,::2,:].reshape(6,3)
    
    In [34]: res
    Out[34]: 
    array([[ 0,  1,  2,  0,  0,  0],
           [ 3,  4,  5,  0,  0,  0],
           [ 6,  7,  8,  0,  0,  0],
           [18, 19, 20,  0,  0,  0],
           [21, 22, 23,  0,  0,  0],
           [24, 25, 26,  0,  0,  0]])
    
    In [35]: res[:,3:] = arr[:,1::2,:].reshape(6,3)
    
    In [36]: res
    Out[36]: 
    array([[ 0,  1,  2,  9, 10, 11],
           [ 3,  4,  5, 12, 13, 14],
           [ 6,  7,  8, 15, 16, 17],
           [18, 19, 20, 27, 28, 29],
           [21, 22, 23, 30, 31, 32],
           [24, 25, 26, 33, 34, 35]])
    

    A concatenate version of that same block join:

    In [41]: np.concatenate((arr[:,::2], arr[:,1::2]), axis=3)
    Out[41]: 
    array([[[[ 0,  1,  2,  9, 10, 11],
             [ 3,  4,  5, 12, 13, 14],
             [ 6,  7,  8, 15, 16, 17]]],
    
    
           [[[18, 19, 20, 27, 28, 29],
             [21, 22, 23, 30, 31, 32],
             [24, 25, 26, 33, 34, 35]]]])