python - 选择一个 ndarray 的所有维度，但一个

Question

我有一个维度为 (2,3,4) 的 ndarray m。我想选择一个特定的维度索引并像这样总结其他所有内容:

m[:, 2, :].sum()

但我需要一个函数来处理元素选择，因为 ndarray 可以有不同的维度，总和应该跨越不同的切片。该函数将采用感兴趣的维度(此处为 0、1 或 2)和索引(也为 0、1 或 2)并返回 m[:, 2, :].sum()如果输入是 (1, 2)。如果它是 (2, 0)，我希望它们具有 m[:, :, 0].sum() 的等价物。

虽然我不知道如何编写这样的函数，也许可以让我将切片“:”放入变量中？到目前为止，我还没有找到解决方案。谢谢!

Answer 1

您可以使用 sl = slice(None, None, None) 将切片 : 放入一个变量中。然而，numpy 有一个简写形式，np.s_[:]

总体而言:

def sum_axis_i(arr, axis, i):
    idx = (np.s_[:],) * axis + (i,)
    return arr[idx].sum()

请注意，尾随 : 是多余的，因此我们不必费心创建它们。

如果你使用与np.s_非常相似的np.index_exp，你可以这样写:

def sum_axis_i(arr, axis, i):
    idx = np.index_exp[:] * axis + np.index_exp[i]
    return arr[idx].sum()