Search code examples
pythonnumpyarray-broadcasting

Numpy "orthogonal" broadcasting


This feels like something for which there should already be a question, but I can't find one.

In numpy, suppose I have two arrays A and B, how can I ensure that they will be broadcasted "orthogonally" to each other, regardless of the dimensionality of either array? For example I can do this:

import numpy as np

A1 = np.zeros((2,2))
A2 = np.zeros((2,2,2))
B = np.ones((2,2))

C1 = A1 + B[...,np.newaxis,np.newaxis]
C2 = A2 + B[...,np.newaxis,np.newaxis,np.newaxis]

print(C1.shape)
print(C2.shape)

which gives

(2, 2, 2, 2)
(2, 2, 2, 2, 2)

as output. But in order to do that I had to know the number of dimensions in at least one of the arrays. Is there a way to do it without knowing either of them? Hope that makes sense.


Solution

  • Reshape B to include trailing singleton dimensions as many as there are dims in A -

    A + B.reshape(B.shape + tuple([1]*A.ndim)) # where A is generic ndarray
    

    Some of the NumPy ufuncs have outer method that takes care of all that work. So, for the additions, we could simply use -

    np.add.outer(B,A)