numpy.vectorize 函数签名

Question

我有 2 个数组：

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

我想计算 c[i, j] = f(a[i, j, :], b[i]) 的等效项，其中 f(a0, b0) 是一个带有 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 sum(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) 的技巧并不简单。

numpy.vectorize 函数签名

numpy.vectorize function signature

python

numpy

numpy-ufunc

array-broadcasting