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!

DSM · Accepted Answer · 2017-04-30 03:57:32Z

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).

answered Apr 30, 2017 at 3:57

DSM

355k67 gold badges606 silver badges504 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Robbie · Accepted Answer · 2017-04-30 03:28:42Z

0

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])

answered Apr 30, 2017 at 3:28

Robbie

4,8922 gold badges21 silver badges25 bronze badges

1 Comment

Zane Dufour Over a year ago

Sorry, I should have specified that I'm looking for a vectorized operation.

Collectives™ on Stack Overflow

3D Numpy array to 3D Numpy array

2 Answers 2

Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related