[英]Numpy: integrate dimension from one array into another
Given two arrays a (shape= a, b, c, d, e
) and b (shape= a, x, b
), I'd like to include the dimension x
of b into a , so that a new array c results in a shape= a, x, b, c, d, e
. 给定两个数组a (shape =
a, b, c, d, e
)和b (shape = a, x, b
),我想将b的维度x
包含到a中 ,以便得到一个新数组c形状为a, x, b, c, d, e
。 The values of b should be distributed evenly: b的值应均匀分布:
c.sum(1) == a
b.sum(1) == a.sum(axis=(2, 3, 4)) == c.sum(axis=(1, 3, 4, 5)
Is there any smart way of doing this is a few lines with numpy or is it necessary to iterate over all values of b[x]
manually? 有什么聪明的方法可以做到这一点吗?用numpy行几行,还是有必要手动遍历
b[x]
所有值?
My current solution: 我当前的解决方案:
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
Here is a way to do this in 3 lines, but first some remarks about my approach: 这是一种分三行执行此操作的方法,但首先要介绍一下我的方法:
A[a,b,c,d,e]
and B[a,x,b]
as input. A[a,b,c,d,e]
和B[a,x,b]
作为输入。 C
is the same for all x
, so we don't really need that axis for the calculation (if required, you can add it as a new dimension and and duplicate the entries afterwards). x
, C
都是相同的,因此我们实际上并不需要该轴进行计算(如果需要,您可以将其添加为新维度,然后再复制条目)。 B[a,a,b]
can be contracted by taking the diagonal along the first two axes. B[a,a,b]
可以通过沿前两个轴取对角线来收缩。 tot
is a sum over indices c,d,e
, we can store this in a precalculated array Tot[a,b]
tot
是索引c,d,e
,我们可以将其存储在预先计算的数组Tot[a,b]
numpy.einsum
in the final step, I will first take the inverse Tot = 1/Tot
numpy.einsum
,我将首先取反Tot = 1/Tot
Here is the complete code: 这是完整的代码:
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)))
Edit: If numpy.einsum()
is too slow for you, you can implement the last step in Cython with for
loops. 编辑:如果
numpy.einsum()
对您来说太慢,则可以在Cython中使用for
循环实现最后一步。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.