Numpy function to get shape of added arrays

Question

tl;dr: How do I predict the shape returned by numpy broadcasting across several arrays without having to actually add the arrays?

I have a lot of scripts that make use of numpy (Python) broadcasting rules so that essentially 1D inputs result in a multiple-dimension output. For a basic example, the ideal gas law (pressure = rho * R_d * temperature) might look like

def rhoIdeal(pressure,temperature):
    rho = np.zeros_like(pressure + temperature)
    rho += pressure / (287.05 * temperature)
    return rho

It's not necessary here, but in more complicated functions it's very useful to initialize the array with the right shape. If pressure and temperature have the same shape, then rho also has that shape. If pressure has shape (n,) and temperature has shape (m,), I can call

rhoIdeal(pressure[:,np.newaxis], temperature[np.newaxis,:])

to get rho with shape (n,m). This lets me make plots with multiple values of temperature without having to loop over rhoIdeal, while still allowing the script to accept arrays of the same shape and compute the result element-by-element.

My question is: Is there a built-in function to return the shape compatible with several inputs? Something that behaves like

def returnShape(list_of_arrays):
    return np.zeros_like(sum(list_of_arrays)).shape

without actually having to sum the arrays? If there's no built-in function, what would a good implementation look like?

Answer 1

You could use np.broadcast. This function returns an object encapsulating the result of broadcasting two or more arrays together. No actual operation (e.g. addition) is performed - the object simply has some of the same attributes that an array produced by means of other operations would have (shape, ndim, etc.).

For example:

x = np.array([1,2,3])  # shape (3,)
y = x.reshape(3,1)     # shape (3, 1)
z = np.ones((5,1,1))   # shape (5, 1, 1)

Then you can check what the shape of the array returned by broadcasting x, y and z would be by inspecting the shape attribute:

>>> np.broadcast(x, y, z).shape
(5, 3, 3)

This means that you could implement your function simply as follows:

def returnShape(*args):
    return np.broadcast(*args).shape