numpy array multiplication with arrays of arbitrary dimensions

Question

I have a numpy array A, which has shape (10,).

I also have, as of this moment, a numpy array B with shape (10,3,5). I want to do a multiplication between these two to get C such that C[0,:,:]=A[0]*B[0,:,:], C[1]=A[1]*B[1,:,:], etc.

I do not want to work this out with loops, one reason being the aesthetics of the thing, the other being that this code needs to be very generic. I want the user to be able to input pretty much any B of any shape as long as the leading dimension is 10. For instance, I want the user to be able to also put in a B of shape (10,4).

So: How can I implement this multiplication using numpy? Thanks.

ADDENDUM: Have been asked for example. Will go smaller. Let's say A is the numpy array [1,2,3] and B is the numpy array [[1,2],[4,5],[7,8]]. I want the multiplication of the two to result in [[1,2],[8,10],[21,24]]. ...

>>> a
array([1, 2, 3])
>>> b
array([[1, 2],
       [4, 5],
       [7, 8]])
>>> #result
>>> c
array([[ 1,  2],
       [ 8, 10],
       [21, 24]])
>>>

Answer 1

You could use None (or np.newaxis) to expand A to match B:

>>> A = np.arange(10)
>>> B = np.random.random((10,3,5))
>>> C0 = np.array([A[i]*B[i,:,:] for i in range(len(A))])
>>> C1 = A[:,None,None] * B
>>> np.allclose(C0, C1)
True

But this will only work for the 2 case. Borrowing from @ajcr, with enough transposes we can get implicit broadcasting to work for the general case:

>>> C3 = (A * B.T).T
>>> np.allclose(C0, C3)
True

Alternatively, you could use einsum to provide the generality. In retrospect it's probably overkill here compared with the transpose route, but it's handy when the multiplications are more complicated.

>>> C2 = np.einsum('i,i...->i...', A, B)
>>> np.allclose(C0, C2)
True

and

>>> B = np.random.random((10,4))
>>> D0 = np.array([A[i]*B[i,:] for i in range(len(A))])
>>> D2 = np.einsum('i,i...->i...', A, B)
>>> np.allclose(D0, D2)
True

Answer 2

Although I like the einsum notation, I'll add a little variety to the mix ....

You can add enough extra dimensions to a so that it will broadcast across b.

>>> a.shape
(3,)
>>> b.shape
(3,2)

b has more dimensions than a

extra_dims = b.ndim - a.ndim

Add the extra dimension(s) to a

new_shape = a.shape + (1,)*extra_dims    # (3,1)
new_a = a.reshape(new_shape)

Multiply

new_a * b

---

As a function:

def f(a, b):
    '''Product across the first dimension of b.

    Assumes a is 1-dimensional.
    Raises AssertionError if a.ndim > b.ndim or
     - the first dimensions are different
    '''
    assert a.shape[0] == b.shape[0], 'First dimension is different'
    assert b.ndim >= a.ndim, 'a has more dimensions than b'

    # add extra dimensions so that a will broadcast
    extra_dims = b.ndim - a.ndim
    newshape = a.shape + (1,)*extra_dims
    new_a = a.reshape(newshape)

    return new_a * b