Numpy 将数组减少到给定的形状

问题描述

我正在尝试编写一个函数,将 numpy ndarray 减少到给定的形状,从而有效地“取消广播”该数组。例如,使用 add 作为通用函数,我想要以下结果:

  • A = [1,2,3,4,5],reduced_shape = (1,) -> [1+2+3+4+5] = [15](轴上的和=0)
  • A = [[1,2],[1,2]],reduced_shape = (2,) -> [1+1,2+2] = [2,4](轴上的和=0)
  • A = [[1,2]],reduced_shape = (1,) -> [1+2+1+2] = [6](轴上的和=(0,1) )
  • A = [[[1,2]],[[1,2]]],reduction_shape = (2,2) -> [[1+1,2+2],[1+1,2+2]] = [[2,4],[2,4]](轴上的和=0)

这是我想出的解决方案:

def unbroadcast(A,reduced_shape):
    fill = reduced_shape[-1] if reduced_shape else None
    reduced_axes = tuple(i for i,(a,b) in enumerate(itertools.zip_longest(A,shape,fillvalue=fill)) if a!=b)
    return np.add.reduce(A,axis=reduced_axes).reshape(shape)

但是感觉不必要的复杂,有没有办法实现这个依赖于Numpy的公共API?

解决方法

不清楚这是如何“取消广播”的。

进行计算的直接方法是使用 axissum 参数:

In [124]: np.array([1,2,3,4,5]).sum()
Out[124]: 15
In [125]: np.array([[1,2],[1,2]]).sum(axis=0)
Out[125]: array([2,4])
In [126]: np.array([[1,2]]).sum(axis=(0,1))
Out[126]: 6
In [128]: np.array([[[1,2]],[[1,2]]]).sum(axis=0)
Out[128]: 
array([[2,4],[2,4]])

我不经常使用 reduce,但看起来轴也一样:

In [130]: np.add.reduce(np.array([1,5]))
Out[130]: 15
In [132]: np.add.reduce(np.array([[[1,2]]]),axis=0)
Out[132]: 
array([[2,4]])

但我还没有弄清楚从 reduced_shape 到必要的 axis 值的逻辑。对于 (2,) 和 (2,2) 之类的形状,当您说将形状减少到 (2,2) 时可能会产生歧义。如果您使用像 np.arange(24).reshape(2,4)

这样的样本数组可能会更清楚