multiply numpy ndarray with 1d array along a given axis

Question

It seems I am getting lost in something potentially silly. I have an n-dimensional numpy array, and I want to multiply it with a vector (1d array) along some dimension (whi

Answer 1

I got a similar demand when I was working on some numerical calculation.

Let's assume we have two arrays (A and B) and a user-specified 'axis'. A is a multi-dimensional array. B is a 1-d array.

The basic idea is to expand B so that A and B have the same shape. Here is the solution code

import numpy as np
from numpy.core._internal import AxisError

def multiply_along_axis(A, B, axis):
    A = np.array(A)
    B = np.array(B)
    # shape check
    if axis >= A.ndim:
        raise AxisError(axis, A.ndim)
    if A.shape[axis] != B.size:
        raise ValueError("'A' and 'B' must have the same length along the given axis")
    # Expand the 'B' according to 'axis':
    # 1. Swap the given axis with axis=0 (just need the swapped 'shape' tuple here)
    swapped_shape = A.swapaxes(0, axis).shape
    # 2. Repeat:
    # loop through the number of A's dimensions, at each step:
    # a) repeat 'B':
    #    The number of repetition = the length of 'A' along the 
    #    current looping step; 
    #    The axis along which the values are repeated. This is always axis=0,
    #    because 'B' initially has just 1 dimension
    # b) reshape 'B':
    #    'B' is then reshaped as the shape of 'A'. But this 'shape' only 
    #     contains the dimensions that have been counted by the loop
    for dim_step in range(A.ndim-1):
        B = B.repeat(swapped_shape[dim_step+1], axis=0)\
             .reshape(swapped_shape[:dim_step+2])
    # 3. Swap the axis back to ensure the returned 'B' has exactly the 
    # same shape of 'A'
    B = B.swapaxes(0, axis)
    return A * B

And here is an example

In [33]: A = np.random.rand(3,5)*10; A = A.astype(int); A
Out[33]: 
array([[7, 1, 4, 3, 1],
       [1, 8, 8, 2, 4],
       [7, 4, 8, 0, 2]])

In [34]: B = np.linspace(3,7,5); B
Out[34]: array([3., 4., 5., 6., 7.])

In [35]: multiply_along_axis(A, B, axis=1)
Out[34]: 
array([[21.,  4., 20., 18.,  7.],
       [ 3., 32., 40., 12., 28.],
       [21., 16., 40.,  0., 14.]])