python - numpy 切片返回最后两个维度

Question

基本上，我正在寻找一个函数或语法，它允许我获取具有任意维数的 n 维 numpy 数组的最后两个维度的第一个“切片”。

我可以做到这一点，但它太丑陋了，无法忍受，如果有人发送 6d 数组怎么办？必须有一个像椭圆这样的 numpy 函数可以扩展为 0,0,0,... 而不是 :,:,:,...

data_2d = np.ones(5**2).reshape(5,5)
data_3d = np.ones(5**3).reshape(5,5,5)
data_4d = np.ones(5**4).reshape(5,5,5,5)

def get_last2d(data):
    if data.ndim == 2:
        return data[:]
    if data.ndim == 3:
        return data[0, :]
    if data.ndim == 4:
        return data[0, 0, :]

np.array_equal(get_last2d(data_3d), get_last2d(data_4d))

谢谢， 科林

Answer 1

这个怎么样，

def get_last2d(data):
    if data.ndim <= 2:
        return data
    slc = [0] * (data.ndim - 2)
    slc += [slice(None), slice(None)]
    return data[slc]