从第一维进行numpy广播

Question

在NumPy中，有没有什么简单的方法可以让两个不同维度的数组，比如(x,y)和(x,y,z)，进行广播？NumPy的广播规则通常是从最后一个维度开始匹配，所以一般的广播方法不适用（因为第一个数组需要是(y,z)的维度）。

背景：我在处理图像，有些是RGB格式（形状是(h,w,3)），有些是灰度图（形状是(h,w)）。我生成的alpha遮罩的形状是(h,w)，我想通过mask * im来将遮罩应用到图像上。但因为上面提到的问题，这样做不行，所以我最后不得不写类似下面的代码：

mask = mask.reshape(mask.shape + (1,) * (len(im.shape) - len(mask.shape)))

这看起来很麻烦。代码的其他部分也在处理向量和矩阵，但也遇到了同样的问题：当m的形状是(x,y)而v的形状是(x,)时，执行m + v会失败。虽然可以使用atleast_3d这样的函数，但我还得记住我到底想要多少个维度。

Answer 1

使用 np.newaxis 进行索引时，会在指定的位置创建一个新的轴。也就是说，

xyz = #some 3d array
xy = #some 2d array
xyz_sum = xyz + xy[:,:,np.newaxis]
or
xyz_sum = xyz + xy[:,:,None]

这样索引后，会在这个位置生成一个形状为 1、步长为 0 的轴。

Answer 2

为什么不直接用装饰-处理-去装饰的方式呢：

def flipflop(func):
    def wrapper(a, mask):
        if len(a.shape) == 3:
            mask = mask[..., None]
        b = func(a, mask)
        return np.squeeze(b)
    return wrapper

@flipflop
def f(x, mask):
    return x * mask

然后

>>> N = 12
>>> gs = np.random.random((N, N))
>>> rgb = np.random.random((N, N, 3))
>>> 
>>> mask = np.ones((N, N))
>>> 
>>> f(gs, mask).shape
(12, 12)
>>> f(rgb, mask).shape
(12, 12, 3)

Answer 3

从另一个角度来看：如果你经常遇到这种情况，创建一个工具函数来确保正确的广播方式可能会很有用：

def right_broadcasting(arr, target):
    return arr.reshape(arr.shape + (1,) * (target.ndim - arr.ndim))

不过，如果输入只有两种类型（要么已经有3个维度，要么只有2个维度），我觉得用一个简单的if语句会更好。

Answer 4

numpy 的函数通常会有一些代码块，用来检查数组的维度，调整数组的形状，以便它们能够兼容，然后才开始真正的加法或乘法操作。它们可能会调整输出的形状，以匹配输入的形状。所以，自己写一些类似的代码来处理这些操作是完全没问题的。

不要轻易忽视把变量 3 的维度放到最前面的想法。这样做可以利用 numpy 自动在前面添加维度的特性。

如果你想逐个元素地进行乘法运算，einsum 是一个非常强大的工具。

np.einsum('ij...,ij...->ij...',im,mask)

这个方法可以处理 im 和 mask 是2维或3维的各种组合（假设前两个维度总是兼容的）。不过，不幸的是，这个方法并不能推广到加法或其他操作上。

之前我用纯Python模拟过 einsum。为了实现这一点，我使用了 np.lib.stride_tricks.as_strided 和 np.nditer。如果你想在维度的组合和匹配上获得更多的灵活性，可以看看这两个函数。

Answer 5

可以试试用转置的方法：

(a.T + c.T).T