[SOLVED] numpy.vectorize function signature

numpy.vectorize function signature

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.

Solution

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.