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循環實現最后一步。