使用 np.einsum 时忽略维度

Question

我使用np.einsum 来计算图表中的材料流量（本例中为 1 个节点到 4 个节点）。流量由amount 给出（amount.shape == (1, 1, 2) 维度定义了某些标准，我们称它们为a、b、c）。

布尔矩阵route 根据a、b、c 标准确定允许流向y（route.shape == (4, 1, 1, 2)；yabc）。我标记了尺寸y、a、b、c。 abc 等价于amounts 尺寸abc，y 是流的方向（0、1、2 或 3）。为了确定y 中的材料量，我计算了np.einsum('abc,yabc->y', amount, route) 并获得了一个流入y 的y-dim 向量。路线还有一个隐含的优先级。例如，任何route[0, ...] == True 对于任何y=1..3 都是False，对于下一个更高的y-dim 路线，任何route[1, ...] == True 是False 等等。 route[3, ...]（最后一个 y-index）定义了 catch-all 路由，也就是说，当之前的 y-index 值为 False ((route[0] ^ route[1] ^ route[2] ^ route[3]).all() == True) 时，它的值为 True。

这很好用。然而，当我引入另一个标准（维度）x，它只存在于route 中，而不存在于amount 中，这个逻辑似乎被打破了。下面的代码演示了这个问题：

>>> import numpy as np

>>> amount = np.asarray([[[5000.0, 0.0]]]) 

>>> route = np.asarray([[[[[False, True]]], [[[False, True]]], [[[False, True]]]], [[[[True, False]]], [[[False, False]]], [[[False, False]]]], [[[[False, False]]], [[[True, False]]], [[[False, False]]]], [[[[False, False]]], [[[False, False]]], [[[True, False]]]]], dtype=bool)

>>> amount.shape
(1, 1, 2)

>>> Added dimension `x`
>>> # y,x,a,b,c
>>> route.shape
(4, 3, 1, 1, 2)

>>> # Attempt 1: `5000` can flow into y=1, 2 or 3. I expect
>>> # `flows1.sum() == amount.sum()` as it would be without `x`.
>>> # Correct solution would be `[0, 5000, 0, 0]` because material is routed
>>> # to y=1, and is not available for y=2 and y=3 as they are lower
>>> # priority (higher index)
>>> flows1 = np.einsum('abc,yxabc->y', amount, route)
>>> flows1
array([   0., 5000., 5000., 5000.])

>>> # Attempt 2: try to collapse `x` => not much different, duplication
>>> np.einsum('abc,yabc->y', amount, route.any(1))
array([   0., 5000., 5000., 5000.])

>>> # This is the flow by `y` and `x`. I'd only expect a `5000` in the
>>> # 2nd row (`[5000.,    0.,    0.]`) not the others.
>>> np.einsum('abc,yxabc->yx', amount, route)
array([[   0.,    0.,    0.],
       [5000.,    0.,    0.],
       [   0., 5000.,    0.],
       [   0.,    0., 5000.]])

是否有任何可行的操作可以应用于route（.all(1) 也不起作用）以忽略 x 维度？

另一个例子：

>>> amount2 = np.asarray([[[5000.0, 1000.0]]])
>>> np.einsum('abc,yabc->y', amount2, route.any(1))
array([1000., 5000., 5000., 5000.])

可以解释为1000.0 被路由到y=0（并且没有其他y 目标）并且5000.0 与目标y=1、y=2 和y=3 兼容，但理想情况下，我'd only like to show 5000.0 up in y=1 (因为那是最低索引和最高目标优先级)。

解决方案尝试

下面的作品，但不是很 numpy-ish。如果能消除循环就好了。

# Initialise destination
result = np.zeros((route.shape[0]))
# Calculate flow by maintaining all dimensions (this will cause
# double ups because `x` is not part of `amount2`
temp = np.einsum('abc,yxabc->yxabc', amount2, route)
temp_ixs = np.asarray(np.where(temp))

# For each original amount, find the destination (`y`)
for a, b, c in zip(*np.where(amount2)):
    # Find where dimensions `abc` are equal in the destination.
    # Take the first vector which contains `yxabc` (we get `yx` as result)
    ix = np.where((temp_ixs[2:].T == [a, b, c]).all(axis=1))[0][0]
    y_ix = temp_ixs.T[ix][0]
    # ignored
    x_ix = temp_ixs.T[ix][1]
    v = amount2[a, b, c]
    # build resulting destination
    result[y_ix] += v

# result == array([1000., 5000.,    0.,    0.])

换句话说，对于amount2 中的每个值，我正在寻找temp 中的最低索引yx，以便可以将值写入result[y] = value（x 被忽略）。

>>> temp = np.einsum('abc,yxabc->yx', amount2, route)
>>> temp
        #  +--- value=1000 at y=0 => result[0] += 1000
        # /
array([[1000., 1000., 1000.],
        #  +--- value=5000 at y=1 => result[1] += 5000
        # /
       [5000.,    0.,    0.],
       [   0., 5000.,    0.],
       [   0.,    0., 5000.]])
>>> result
array([1000., 5000.,    0.,    0.])
>>> amount2
array([[[5000., 1000.]]])

另一个降低route维度的尝试是：

>>> r = route.any(1)
>>> for x  in xrange(1, route.shape[0]):
        r[x] = r[x] & (r[:x] == False).all(axis=0)

>>> np.einsum('abc,yabc->y', amount2, r)
array([1000., 5000.,    0.,    0.])

这基本上保留了route 的第一个维度赋予的上述优先级。当较高优先级数组在该子索引处已具有 True 值时，任何较低优先级（较高索引）数组都不能包含 True 值。虽然这比我的显式方法好很多，但如果 for x in xrange... 循环可以表示为 numpy 向量操作，那就太好了。

Answer 1

我没有尝试遵循您对乘法问题的“流程”解释。我只关注计算选项。

去掉不必要的维度，你的数组是：

In [194]: amount                                                                                       
Out[194]: array([5000.,    0.])
In [195]: route                                                                                        
Out[195]: 
array([[[0, 1],
        [0, 1],
        [0, 1]],

       [[1, 0],
        [0, 0],
        [0, 0]],

       [[0, 0],
        [1, 0],
        [0, 0]],

       [[0, 0],
        [0, 0],
        [1, 0]]])

而yx 的计算是：

In [197]: np.einsum('a,yxa->yx',amount, route)                                                         
Out[197]: 
array([[   0.,    0.,    0.],
       [5000.,    0.,    0.],
       [   0., 5000.,    0.],
       [   0.,    0., 5000.]])

这只是route 的这一部分乘以 5000。

In [198]: route[:,:,0]                                                                                 
Out[198]: 
array([[0, 0, 0],
       [1, 0, 0],
       [0, 1, 0],
       [0, 0, 1]])

在 einsum 的 RHS 上省略 x 会导致跨维度求和。

等效地，我们可以乘法（通过广播）：

In [200]: (amount*route).sum(axis=2)                                                                   
Out[200]: 
array([[   0.,    0.,    0.],
       [5000.,    0.,    0.],
       [   0., 5000.,    0.],
       [   0.,    0., 5000.]])
In [201]: (amount*route).sum(axis=(1,2))                                                               
Out[201]: array([   0., 5000., 5000., 5000.])

也许查看amount*route 将有助于可视化问题。您还可以使用max、min、argmax 等代替sum，或者在一个或多个轴上使用它。