numpy：将维度从一个数组整合到另一个数组

Question

给定两个数组a （shape = a, b, c, d, e ）和b （shape = a, x, b ），我想将b的维度x包含到a中 ，以便得到一个新数组c形状为a, x, b, c, d, e 。 b的值应均匀分布：

• c.sum(1) == a
• b.sum(1) == a.sum(axis=(2, 3, 4)) == c.sum(axis=(1, 3, 4, 5)

有什么聪明的方法可以做到这一点吗？用numpy行几行，还是有必要手动遍历b[x]所有值？

我当前的解决方案：

for a, x, b in zip(*_b_.nonzero()):
    tot = _a_[a, b].sum()
    for c, d, e in zip(*_a_[a, b].nonzero()):
        val = _b_[a, a, b]
        frac = _a_[a, b, c, d, e] / tot
        _c_[a, x, b, c, d, e] = val * frac

Answer 1

这是一种分三行执行此操作的方法，但首先要介绍一下我的方法：

• 为了获得更好的表示法，我将大写字母用于矩阵，将小写字母用于索引。 因此，我们有A[a,b,c,d,e]和B[a,x,b]作为输入。
• 根据您的代码，对于所有x ， C都是相同的，因此我们实际上并不需要该轴进行计算（如果需要，您可以将其添加为新维度，然后再复制条目）。
• B[a,a,b]可以通过沿前两个轴取对角线来收缩。
• tot是索引c,d,e ，我们可以将其存储在预先计算的数组Tot[a,b]
• 为了能够在最后一步中使用numpy.einsum ，我将首先取反Tot = 1/Tot

这是完整的代码：

import numpy

# generate some example input
a = 2
b = 3
c = 4
d = 5
e = 6
x = 7

A = numpy.arange(a*b*c*d*e).reshape((a,b,c,d,e))
B = numpy.arange(a*x*b).reshape((a,x,b))
C = numpy.zeros((a,x,b,c,d,e))

# solution by orange
for a, x, b in zip(*B.nonzero()):
    tot = A[a, b].sum()
    for c, d, e in zip(*A[a, b].nonzero()):
        val = B[a, a, b]
        frac = A[a, b, c, d, e] / tot
        C[a, x, b, c, d, e] = val * frac

# new solution
B2 = numpy.diagonal(B, axis1=0, axis2=1).transpose() # contract B_aab -> B2_ab
Tot = 1/numpy.sum(A, (2,3,4))   # contract \sum_cde A_abcde -> 1 / Tot_ab
C2 = numpy.einsum('ab,abcde,ab->abcde',B2,A,Tot)

# compare (should print x times True)
for i in range(C.shape[1]):
    C_ = C[:,i,:,:,:]
    print(numpy.all(numpy.isclose(C_,C2)))

编辑：如果numpy.einsum()对您来说太慢，则可以在Cython中使用for循环实现最后一步。