Multiply an array (k,) or (k, n) by 1D array (k,)

Viewed 78

I sometimes work with 1D arrays:

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

or 2D arrays (mono or stereo signals read with scipy.io.wavfile):

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

Without having to distinguish these 2 cases with an if A.ndim == 2:..., is there a simple one-line solution to multiply this array A by an 1D array B = np.linspace(0., 1., 4)?

If A is 1D, then it's just A * B, and if A is 2D, what I mean is multiply each line of A by each element of B.


Note: this question arises naturally when working with both mono and stereo sounds read with scipy.io.wavfile.

1 Answers

Approach #1

We can use einsum to cover generic ndarrays -

np.einsum('i...,i->i...',A,B)

The trick that works here is the ellipsis that broadcasts for the trailing dimensions after the first axis as they are in A and keeps them in the output, while keeping the first axes for the two inputs aligned, which is the intended multiplication here. With A as 1D array, there's no broadcasting and that essentially reduces to : np.einsum('i,i->i',A,B) under the hoods.

Schematically put :

A   :  i x ....
B   :  i 
out :  i x ....

Hence, covers for A with any number of dimensions.

More info on the use of ellipsis from the docs :

To enable and control broadcasting, use an ellipsis. Default NumPy-style broadcasting is done by adding an ellipsis to the left of each term, like np.einsum('...ii->...i', a). To take the trace along the first and last axes, you can do np.einsum('i...i', a), or to do a matrix-matrix product with the left-most indices instead of rightmost, you can do np.einsum('ij...,jk...->ik...', a, b).

Approach #2

Using the fact that we are trying to align the first axis of A with the only axis of 1D array B, we can simply transpose A, multiply with B and finally transpose back -

(A.T*B).T
Related