How to apply function which returns vector to each numpy array element (and get array with higher dimension)

Question

Let\'s write it directly in code

Note: I edited mapper (original example use x -> (x, 2 * x, 3 * x) just for example), to generic blackbox function, which cause the trou

Answer 1

You can use basic NumPy broadcasting for these kind of "outer products"

np.arange(3)[:, None] * np.arange(2)
# array([[0, 0],
#        [0, 1],
#        [0, 2]])

In your case it would be

def mapper(x):
    return (np.arange(3)[:, None, None] * x).transpose((1, 2, 0))

note the .transpose() is only needed if you specifically need the new axis to be at the end.

And it is almost 3x as fast as stacking 3 separate multiplications:

def mapper(x):
    return (np.arange(3)[:, None, None] * x).transpose((1, 2, 0))


def mapper2(x):
    return np.stack((x, 2 * x, 3 * x), axis = -1)

a = np.arange(30000).reshape(-1, 30)

%timeit mapper(a)   # 48.2 µs ± 417 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit mapper2(a)  # 137 µs ± 3.57 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)