numpy.vectorize function 签名

Question

I have 2 arrays:

>>> a.shape
(9, 3, 11)
>>> b.shape
(9,)

I would like to compute the equivalent of c[i, j] = f(a[i, j, :], b[i]) where f(a0, b0) is a function that takes 2 parameters, with len(a0) == 11 and len(b0) == 9 . Here, i is iterating on range(9) and j is iterating on range(3) .

Is there a way to code this using numpy.vectorize ? Or is it simpler with some clever broadcasting?

I have been trying for 2 hours and I just don't understand how to make it work... I tried to broadcast or to use signatures but to no avail.

Answer 1

numpy.apply_along_axis is what you need.

import numpy as np

a = np.ones( (9,3,11) )
b = np.ones( 9 )

def f(a0, b0):
    return a0[:9]+b0

c = np.apply_along_axis( f, 2, a, b )
print(c)

c 's shape is (9,3).

Answer 2

In the end, I could make it work like this:

>>> f = np.vectorize(f, signature="(k),(1)->()")
>>> print(a.shape)
(9, 3, 11)
>>> print(b.shape)
(9,)
>>> print(f(a, b[:, None, None]).shape)
(9, 3)

This ensures that f gets called with the correct shapes and iterates properly. It is frankly not straightforward from the Numpy documentation to understand the trick to use a (1) in the signature for this purpose.