5

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?

Jareth Holt
  • 321
  • 1
  • 10

1 Answers1

6

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
Alex Riley
  • 169,130
  • 45
  • 262
  • 238
  • Indeed quite cool! @ajcr just a minor typo, I think you missed the `z` in the example call. – Imanol Luengo Aug 07 '15 at 14:29
  • Thank you so much! I don't know why this was such a hard function for me to find, it seems *incredibly* useful. – Jareth Holt Aug 07 '15 at 22:45
  • Might as well go with `np.broadcast(*args).shape`. Shorter and faster `:)` –  Aug 08 '15 at 10:30
  • @moarningsun: Good idea, that's much neater. Looking at the question/answer today, I'm not sure why I opted for `broadcast_arrays` over `broadcast`. I've edited the answer to use your suggestion. Thanks! – Alex Riley Aug 08 '15 at 10:42
  • Just a note that `numpy.broadcast` can only take 32 values at once. I had to write this little loop to get around that: `while len(numpy_array_list > 1): numpy_array_list =[numpy.broadcast(*numpy_array_list[:32])] + numpy_array_list[32:]` – Rich Jun 28 '18 at 17:21