Numpy: integrate dimension from one array into another

Question

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 . The values of b should be distributed evenly:

• 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?

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

Answer 1

Here is a way to do this in 3 lines, but first some remarks about my approach:

• For a better notation I will use capital letters for matrices, and small letters for indices. So we have A[a,b,c,d,e] and B[a,x,b] as input.
• According to your code, 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).
• B[a,a,b] can be contracted by taking the diagonal along the first two axes.
• tot is a sum over indices c,d,e , we can store this in a precalculated array Tot[a,b]
• to be able to use numpy.einsum in the final step, I will first take the inverse 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.