numpy：在多个不连续轴上计算均值和标准差（第二次尝试）

Question

之前的版本没有得到任何回复，所以为了让内容更清晰，我重新整理了一下，并增加了一些解释和代码注释。

我想要计算一个numpy的多维数组中某些元素的平均值和标准差，这些元素不对应于单一的轴，而是对应于多个（大于1）不连续的轴，并将结果收集到一个新的数组中，这个新数组的维度是原数组维度减去这些轴的数量再加1。

numpy是否有标准的方法可以高效地完成这个操作呢？

下面的mu_sigma函数是我解决这个问题的最佳尝试，但它有两个明显的缺点：1）它需要复制原始数据；2）它计算平均值的过程重复了两次（因为计算标准差时也需要先计算平均值）。

mu_sigma函数有两个参数：box和axes。box是一个多维的numpy数组（也叫“ndarray”），而axes是一个包含整数的元组，表示box中（不一定是连续的）维度。这个函数返回一个新的数组，维度是原数组维度减去指定的轴的数量再加1，里面包含了在box的“超平面”上计算的平均值和标准差。

下面的代码还包括了mu_sigma的一个示例。在这个示例中，box参数是一个4 x 2 x 4 x 3 x 4的浮点数ndarray，而axes参数是元组(1, 3)。因此，n等于len(box.shape)等于5，k等于len(axes)等于2。对于这个示例输入，返回的结果（我称之为outbox）是一个4 x 4 x 4 x 2的浮点数ndarray。对于每一组索引i、k、j（每个索引的范围是{0, 1, 2, 3}），outbox[i, j, k, 0]是由numpy表达式box[i, 0:2, j, 0:3, k]指定的6个元素的平均值。同样，outbox[i, j, k, 1]是这6个元素的标准差。这意味着结果的前n - k等于3的维度与输入ndarray box中非轴的维度相同，在这个例子中是维度0、2和4。

mu_sigma函数的策略是：

调整维度顺序（使用transpose方法），把第二个参数指定的轴放到最后，剩下的（非轴）维度保持在前面（顺序不变）；
将轴维度合并成一个（使用reshape方法）；新的“合并”维度现在是重塑后的ndarray的最后一个维度；
使用最后的“合并”维度作为轴计算平均值的ndarray；
使用最后的“合并”维度作为轴计算标准差的ndarray；
返回一个通过连接步骤（3）和（4）中生成的ndarray得到的ndarray。

import numpy as np

def mu_sigma(box, axes):
    inshape = box.shape

    # determine the permutation needed to put all the dimensions given in axes
    # at the end (otherwise preserving the relative ordering of the dimensions)
    nonaxes = tuple([i for i in range(len(inshape)) if i not in set(axes)])

    # permute the dimensions
    permuted = box.transpose(nonaxes + axes)

    # determine the shape of the ndarray after permuting the dimensions and
    # collapsing the axes-dimensions; thanks to Bago for the "+ (-1,)"
    newshape = tuple(inshape[i] for i in nonaxes) + (-1,)

    # collapse the axes-dimensions
    # NB: the next line results in copying the input array
    reshaped = permuted.reshape(newshape)

    # determine the shape for the mean and std ndarrays, as required by
    # the subsequent call to np.concatenate (this reshaping is not necessary
    # if the available mean and std methods support the keepdims keyword;
    # instead, just set keepdims to True in both calls).
    outshape = newshape[:-1] + (1,)

    # compute the means and standard deviations
    mean = reshaped.mean(axis=-1).reshape(outshape)
    std = reshaped.std(axis=-1).reshape(outshape)

    # collect the results in a single ndarray, and return it
    return np.concatenate((mean, std), axis=-1)

inshape = 4, 2, 4, 3, 4
inbuf = np.array(map(float, range(np.product(inshape))))
inbox = np.ndarray(inshape, buffer=inbuf)
outbox = mu_sigma(inbox, tuple(range(len(inshape))[1::2]))

# "inline tests"
assert all(outbox[..., 1].ravel() ==
           [inbox[0, :, 0, :, 0].std()] * outbox[..., 1].size)
assert all(outbox[..., 0].ravel() == [float(4*(v + 3*w) + x)
                                      for v in [8*y - 1
                                                for y in [3*z + 1
                                                          for z in range(4)]]
                                      for w in range(4)
                                      for x in range(4)])

numpy 均值计算数组维度多维数组标准差数据重塑轴操作数组合并

numpy：在多个不连续轴上计算均值和标准差（第二次尝试）

1 个回答

撰写回答