Python 相当于 Matlab shiftdim()

Question

我目前正在将一些 Matlab 代码转换为 Python，我想知道是否有与 Matlab 的 shiftdim(A, n)

类似的功能

B = shiftdim(A,n) shifts the dimensions of an array A by n positions. shiftdim shifts the dimensions to the left when n is a positive integer and to the right when n is a negative integer. For example, if A is a 2-by-3-by-4 array, then shiftdim(A,2) returns a 4-by-2-by-3 array.

Answer 1

如果你使用 numpy 你可以使用 np.moveaxis.

来自docs：

>>> x = np.zeros((3, 4, 5))
>>> np.moveaxis(x, 0, -1).shape
(4, 5, 3)
>>> np.moveaxis(x, -1, 0).shape
(5, 3, 4)

numpy.moveaxis(a, source, destination)[source]

Parameters

a:        np.ndarray 
          The array whose axes should be reordered.

source:   int or sequence of int
          Original positions of the axes to move. These must be unique.

destination: int or sequence of int
             Destination positions for each of the original axes. 
             These must also be unique.

Answer 2

shiftdim的功能比左右移动轴要复杂一些。

对于输入 shiftdim(A, n)，如果 n 为正，则将轴向左移动 n（即旋转），但如果 n 为负，则将轴向右移动并附加 size 的尾部维度1.
对于输入 shiftdim(A)，删除任何大小为 1 的尾随维度。

from collections import deque
import numpy as np

def shiftdim(array, n=None):
    if n is not None:
        if n >= 0:
            axes = tuple(range(len(array.shape)))
            new_axes = deque(axes)
            new_axes.rotate(n)
            return np.moveaxis(array, axes, tuple(new_axes))
        return np.expand_dims(array, axis=tuple(range(-n)))
    else:
        idx = 0
        for dim in array.shape:
            if dim == 1:
                idx += 1
            else:
                break
        axes = tuple(range(idx))
        # Note that this returns a tuple of 2 results
        return np.squeeze(array, axis=axes), len(axes)

与 Matlab 文档相同的示例

a = np.random.uniform(size=(4, 2, 3, 5))
print(shiftdim(a, 2).shape)      # prints (3, 5, 4, 2)
print(shiftdim(a, -2).shape)     # prints (1, 1, 4, 2, 3, 5)

a = np.random.uniform(size=(1, 1, 3, 2, 4))
b, nshifts = shiftdim(a)
print(nshifts)                   # prints 2
print(b.shape)                   # prints (3, 2, 4)

Python 相当于 Matlab shiftdim()

Python equivalent of Matlab shiftdim()

python

matlab

numpy

equivalent