如何在部分参数上应用`numpy.vectorize`？

Question

我搜索了很多，但还是找不到这个特定问题的解决办法。我有一个函数，签名如下：

def my_function(self, number: float, lookup: list[str]) -> float:
    # perform some operation
    return some_float_based_on_operation

我想把它向量化，像这样：

my_ndarray = np.vectorize(self.my_function)(my_ndarray, ["a", "b", "c"])

其中 my_ndarray 是一个包含18个浮点数的一维数组。

但是当我尝试运行上面的代码时，出现了以下错误：

ValueError: operands could not be broadcast together with shapes (18,) (3,)

我的使用场景并不需要两个参数的长度相同，因为第二个参数是一个不相关的字符串列表。我该如何向量化这样的函数呢？

Answer 1

正如评论中提到的，你可以使用 np.vectorize 的 excluded 参数：

import numpy as np

def f(x, y):
    assert np.shape(x) == ()
    assert np.shape(y) != ()
    return x*len(y)

f2 = np.vectorize(f, excluded=[1, 'y'])

x = np.arange(3)
y = ['a', 'b', 'c']
f2(x, y)  # array([0, 3, 6])
f2(x, y=y)  # array([0, 3, 6])