python - How can I broadcast between 1D and nD arrays to obtain a (1+n)D array output?

Question

I have an n-dimensional ndarray z0, and a 1-dimensional ndarray za. The sizes don't correspond to each other in any way. I'd like to be able to create a new n+1-dimensional array, z, where z[i]=z0+za[i]. Is there some simple way to do this with broadcasting?

This is not equivalent to this question. If z0 is 2D, this can be easily achieved as follows:

z0[np.newaxis]+norm.ppf(alphas)[:,None]

However, I need to be able to do this regardless of z0's dimensionality, and so simply adding the correct number of None or np.newaxis terms won't work.

Answer 1

How about:

z = za.reshape(za.shape + (1,)*z0.ndim) + z0

For example:

import numpy as np
z0 = np.ones((2, 3, 4, 5))
za = np.ones(6)

z = za.reshape(za.shape + (1,)*z0.ndim) + z0

print z.shape
# (6, 2, 3, 4, 5)

Answer 2

Maybe something like

>>> z0 = np.random.random((2,3,4))
>>> za = np.random.random(5)
>>> z = np.rollaxis((z0[...,None] + za), -1)
>>> z.shape
(5, 2, 3, 4)
>>> [np.allclose(z[i], z0 + za[i]) for i in range(len(za))]
[True, True, True, True, True]

where I've used ... to mean any number of dimensions, and rollaxis to put it in the shape I think you want. If you don't mind the new axis being at the end, you could get away with z0[..., None] + za, I think.