python - numpy 向数组添加可变数量的维度

Question

有没有办法向 numpy 数组添加可变数量的单例维度？我想要类似 atleast_2d 的东西但是对于任意数量的维度。 expand_dims出现了很多，但只增加了一个维度。

我知道的唯一方法是先计算形状并应用它，即

import numpy as np

def atleast_kd(array, k):
    array = np.asarray(array)
    new_shape = array.shape + (1,) * (k - array.ndim)
    return array.reshape(new_shape)

Answer 1

一个更优雅的方法是np.broadcast_to。例如:

a = np.random.rand(2,2)
k = # number_extra dimensions
b = np.broadcast_to(a, (1,) * k + a.shape)

这会将维度添加到 b 的开头。要获得与给定函数完全相同的行为，您可以使用 np.moveaxis。