Follow

Keep Up to Date with the Most Important News

By pressing the Subscribe button, you confirm that you have read and are agreeing to our Privacy Policy and Terms of Use
Contact

Expand numpy array to be able to broadcast with second array of variable depth

I have a function which can take an np.ndarray of shape (3,) or (3, N), or (3, N, M), etc.. I want to add to the input array an array of shape (3,). At the moment, I have to manually check the shape of the incoming array and if neccessary, expand the array that is added to it, so that I don’t get a broadcasting error. Is there a function in numpy that can expand my array to allow broadcasting for an input array of arbitrary depth?

def myfunction(input_array):
    array_to_add = np.array([1, 2, 3])
    if len(input_array.shape) == 1:
        return input_array + array_to_add
    elif len(input_array.shape) == 2:
        return input_array + array_to_add[:, None]
    elif len(input_array.shape) == 3:
        return input_array + array_to_add[:, None, None]
    ...

>Solution :

MEDevel.com: Open-source for Healthcare and Education

Collecting and validating open-source software for healthcare, education, enterprise, development, medical imaging, medical records, and digital pathology.

Visit Medevel

One option would be to transpose before and after the addition:

(input_array.T + array_to_add).T

You could also use expand_dims to add the extra dimensions:

(np.expand_dims(array_to_add, tuple(range(1, input_array.ndim)))
 + input_array
)

Alternatively with broadcast_to on the reversed shape of input_array + transpose:

(np.broadcast_to(array_to_add, input_array.shape[::-1]).T
 + input_array
)

Or with a custom reshape:

(array_to_add.reshape(array_to_add.shape+(1,)*(input_array.ndim-1))
 + input_array
)
Add a comment

Leave a Reply

Keep Up to Date with the Most Important News

By pressing the Subscribe button, you confirm that you have read and are agreeing to our Privacy Policy and Terms of Use

Discover more from Dev solutions

Subscribe now to keep reading and get access to the full archive.

Continue reading