Numpy array, multiply and sum over specific axis-CodePudding

Let's say I have an array A such that

A = np.array([
    [[1,1,1,1],[2,2,2,2]],
    [[3,3,3,3],[4,4,4,4]],
    [[5,5,5,5],[6,6,6,6]]
])

where the shape is something like (3,2,4) in this example. Now let's say I have another array B such that

B = np.array([1,2,3,4])

I would like to multiply A and B element wise along the last axis of A and sum, such that

C = np.array([
    [10,20],
    [30,40],
    [50,60]
])

Is there a nice way to do this? I thought about making an equivalent 3D array out of B, doing element wise multiplication, and summing along the last axis of this new array. I was wondering if there is a cleaner way to do this?

Edit: If it makes things easier, A can also be written just as

A = np.array([
    [1,2],
    [3,4],
    [5,6]
])

I made A in the way shown at the top of my post because I thought it was neccessary to do so for the proposed multiplication and summation. If working with this version of A at the bottom of this post is easier/just as easy then that would be preferable.

Cheers

CodePudding user response：

Your example is ambiguous and how you break down the computation is unclear.

It looks to me that you could take any element of the last dimension of A and multiply it with the sum of B:

A[...,0]*B.sum()

Or do you want to sum afterwards?

(A*B[None,:]).sum(2)

output:

array([[10, 20],
       [30, 40],
       [50, 60]])

CodePudding user response：

In [126]: A = np.array([
     ...:     [[1,1,1,1],[2,2,2,2]],
     ...:     [[3,3,3,3],[4,4,4,4]],
     ...:     [[5,5,5,5],[6,6,6,6]]
     ...: ])
In [127]: A.shape
Out[127]: (3, 2, 4)
In [128]: B = np.array([1,2,3,4])

(3,2,4) broadcasts with (4,) (e.g (1,1,4)) to make (3,2,4), then sum on 4:

In [129]: (A*B).sum(axis=-1)
Out[129]: 
array([[10, 20],
       [30, 40],
       [50, 60]])

If the start is (3,2):

In [130]: A1 = np.array([
     ...:     [1,2],
     ...:     [3,4],
     ...:     [5,6]
     ...: ])
     ...:

make it (3,2,1):

In [132]: (A1[:,:,None]*B).shape
Out[132]: (3, 2, 4)
In [133]: (A1[:,:,None]*B).sum(axis=-1)
Out[133]: 
array([[10, 20],
       [30, 40],
       [50, 60]])