numpy is double transposition necessary in this specific case?-CodePudding

I have an array

xx = np.arange(24).reshape(2, 12)

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

and I would like to reshape it, to obtain

array([[[ 0,  1,  2,  3],
        [12, 13, 14, 15]],

       [[ 4,  5,  6,  7],
        [16, 17, 18, 19]],

       [[ 8,  9, 10, 11],
        [20, 21, 22, 23]]])

I can achieve it via

xx.T.reshape(3, 4, 2).transpose(0, 2, 1)

But it has to be transposed twice, which seems unnecessary to me. So could somebody confirm that this is the only way of doing it or provide more readable solution otherwise? Thanks!

CodePudding user response：

It is possible to do a single transpose:

data = np.arange(24).reshape(2, 12)
data = data.reshape(2, 3, 4).transpose(1, 0, 2)

Edit:

I checked this using itertools.permutations and itertools.product:

import itertools
import numpy as np

data = np.arange(24).reshape(2, 12)
desired_data = np.array([[[ 0,  1,  2,  3],
                          [12, 13, 14, 15]],
                         
                         [[ 4,  5,  6,  7],
                          [16, 17, 18, 19]],
                         
                         [[ 8,  9, 10, 11],
                          [20, 21, 22, 23]]])

shapes = [2, 3, 4]
transpose_dims = [0, 1, 2]

shape_permutations = itertools.permutations(shapes)
transpose_permutations = itertools.permutations(transpose_dims)

for shape, transpose in itertools.product(
    list(shape_permutations),
    list(transpose_permutations),
):
    
    new_data = data.reshape(*shape).transpose(*transpose)

    try:
        np.allclose(new_data, desired_data)
    except ValueError as e:
        pass
    else:
        break

print(f"{shape=}, {transpose=}")

shape=(2, 3, 4), transpose=(1, 0, 2)

CodePudding user response：

I would do it this way: first, generate two arrays (shown separated for the sake of decomposition):

xx.reshape(2, -1, 4)
# Output:
# 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]]])

From here, I would then stack along the second dimension in order to combine them like you want:

np.stack(xx.reshape(2, -1, 4), axis=1)
# Output:
# array([[[ 0,  1,  2,  3],
#         [12, 13, 14, 15]],
# 
#        [[ 4,  5,  6,  7],
#         [16, 17, 18, 19]],
# 
#        [[ 8,  9, 10, 11],
#         [20, 21, 22, 23]]])

You'd avoid the transposition. Hopefully it's more readable, but in the end, that's highly subjective, right? '^^

CodePudding user response：

To add on top of @Paul's answer, there is some speedup from removing one of the transpose. The time gain is of ~15%: