Swapping the dimensions of a numpy array

前端 未结 5 1768
不思量自难忘°
不思量自难忘° 2020-12-09 14:37

I would like to do the following:

for i in dimension1:
  for j in dimension2:
    for k in dimension3:
      for l in dimension4:
        B[k,l,i,j] = A[i,j,         


        
5条回答
  •  有刺的猬
    2020-12-09 15:23

    You could rollaxis twice:

    >>> A = np.random.random((2,4,3,5))
    >>> B = np.rollaxis(np.rollaxis(A, 2), 3, 1)
    >>> A.shape
    (2, 4, 3, 5)
    >>> B.shape
    (3, 5, 2, 4)
    >>> from itertools import product
    >>> all(B[k,l,i,j] == A[i,j,k,l] for i,j,k,l in product(*map(range, A.shape)))
    True
    

    or maybe swapaxes twice is easier to follow:

    >>> A = np.random.random((2,4,3,5))
    >>> C = A.swapaxes(0, 2).swapaxes(1,3)
    >>> C.shape
    (3, 5, 2, 4)
    >>> all(C[k,l,i,j] == A[i,j,k,l] for i,j,k,l in product(*map(range, A.shape)))
    True
    

提交回复
热议问题