Numpy: partially transpose and stack

user_11384741 0
Viewed 234

Suppse I have a 3 dimension array A

A = np.random.rand(4,4,5)
AT = np.zeros((4,4,5))

# Firstly I want to transpose the first two dimension of A but keep the third dimension fixed
# that is
a, b, c = A.shape

for i in range(c):
    AT[:,:,i] = A[:,:,i].T

# then I want to stack the A and AT into a 4 dimension array
# that is
B = zeros((4,4,2,5))
B[:,:,0,:] = A
B[:,:,1,:] = AT

The first part of code works, but is there any better way to achieve it?

For the second part, how can I fix it?

numpy-ndarray numpy python
1 Answers

Short answer to your question 1) swapaxes and 2) stack:

B = np.stack([A,A.swapaxes(0,1)], axis=2)
user_42384080
Related