相当于 Matlab shiftdim() 的 Python

问题描述

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

B = shiftdim(A,n) 将数组 A 的维度移动 n 个位置。当 n 是正整数时，shiftdim 向左移动维度，当 n 是负整数时向右移动维度。例如，如果 A 是一个 2×3×4 数组，则 shiftdim(A,2) 返回一个 4×2×3 数组。

解决方法

如果您使用 numpy，则可以使用 np.moveaxis。

来自docs：

>>> x = np.zeros((3,4,5))
>>> np.moveaxis(x,-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.

,

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

• 对于输入 shiftdim(A,n)，如果 n 为正，则将轴向左移动 n（即旋转），但如果 n 为负，则将轴向右移动并附加大小为 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,5))
print(shiftdim(a,2).shape)      # prints (3,2)
print(shiftdim(a,-2).shape)     # prints (1,1,5)

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