numpy.vectorize function 签名

Question

我有 2 个 arrays：

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

我想计算c[i, j] = f(a[i, j, :], b[i])的等价物，其中f(a0, b0)是一个 function ，它需要 2 个参数， len(a0) == 11和len(b0) == 9 。 这里， i在range(9)上迭代， j在range(3)上迭代。

有没有办法使用numpy.vectorize对此进行编码？ 还是通过一些巧妙的广播更简单？

我已经尝试了 2 个小时，但我只是不明白如何使它工作......我尝试广播或使用签名但无济于事。

Answer 1

numpy.apply_along_axis是您所需要的。

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的形状是 (9,3)。

Answer 2

最后，我可以让它像这样工作：

>>> 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)

这确保f以正确的形状被调用并正确迭代。 坦率地说，从 Numpy 文档中了解为此目的在签名中使用(1)的技巧并不简单。