3D Numpy array to 3D Numpy array

Question

I'm working in NumPy. I have an array of floats U, with shape (n,d,d), and a 2D Boolean array B with shape (k,n). This looks kind of like

U = np.array([
    [[0,1],
    [2,3]
    ],
    [[4,5],
    [6,7]
    ]
    [[1,2],
    [3,4]
    ]
])

B = np.array([
    [True,False,False],
    [True,False,True],
    [True,True,False],
    [False,False,True]
])

I want a vectorized function vector_sum(A,B) that will output a shape (4,2,2) array Z, where Z[0] is U[0]; Z[1] is U[0] + U[2]; Z[2] is U[0]+U[1], and Z[3] is U[2]. How can I do this? I'm guessing there's a way to do this with np.einsum, but I don't really understand how that works and I'm on a bit of a time crunch.

Thanks!

Answer 1

IIUC, you could definitely use np.einsum :

In [70]: np.einsum('ij,jkl->ikl', B, U)
Out[70]: 
array([[[ 0,  1],
        [ 2,  3]],

       [[ 1,  3],
        [ 5,  7]],

       [[ 4,  6],
        [ 8, 10]],

       [[ 1,  2],
        [ 3,  4]]])

which will act over the j coordinate of B (the bools) and the j coordinate of U (the dxd subarrays).

Answer 2

This will do it:

import numpy as np

U = np.array([
    [[0,1],[2,3]],
    [[4,5],[6,7]],
    [[1,2],[3,4]]
    ])

B = np.array([
    [True,False,False],
    [True,False,True],
    [True,True,False],
    [False,False,True]
])


Z = np.array([U[i].sum(axis=0) for i in B])