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

这是一种在3行中完成此操作的方法，但首先是关于我的方法的一些评论：

为了更好的表示法，我将使用大写字母表示矩阵，使用小写字母表示索引。我们以A[a,b,c,d,e]和B[a,x,b]作为输入。
根据您的代码，C对于所有x都是相同的，因此我们并不需要该轴进行计算（如果需要，您可以将其添加为新维度并在之后复制条目。）
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()对您来说太慢，您可以使用for循环实现Cython的最后一步。

Numpy：将维度从一个数组集成到另一个数组中

1 个答案: