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)的技巧並不簡單。