在 NumPy 中向量化 function

Question

我有以下 function

from numpy.random import default_rng

def foo(args):
    [a, x, y, z, b1] = args
    vals = np.random.uniform(0, 10, a)
    rr= np.random.uniform(1, 2, a)
    u_1 = vals - x
    u_2 = vals * rr - y
    u_3 = vals / rr - z
    Q = sum(sum(u_1[None, :] > np.maximum(u_2[None, :], u_3[None, :])))
    # print(Q)
    if Q > b1:
        Q = 10
    return Q

args = [10, 2, 40, 1, 2]
print(foo(args))  # this works fine
x_ = [*range(5, 13, 1)]
y_ = [*range(2, 50, 5)]
z_ = [*range(4, 8, 1)]
x, y, z = np.meshgrid(x_, y_, z_, indexing="ij")
args = [10, x, y, z, 2]
print(foo(args))  # this does not work

我收到以下错误：

ValueError：操作数无法与形状 (10,) (8,10,4) 一起广播

我想在网格的所有点上评估 function foo(args) 。 我应该如何修改foo(args) ？ 我需要此操作的矢量化实现，因为它需要非常快。

我已经知道以下答案： How do I apply some function to a python meshgrid? . 我只是认为可能有一种方法可以加快对 function 的评估。

Answer 1

您需要了解广播规则。 像这样使用np.expand_dims() 。

    x, y, z = [np.expand_dims(e, axis=-1) for e in (x, y, z)]
    u_1 = vals - x
    u_2 = vals * rr - y
    u_3 = vals / rr - z
    Q = np.sum(u_1 > np.maximum(u_2, u_3), axis=-1)
    Q = np.where(Q > b1, 10, Q)

我重构了不必要的索引，如u_1[None, :]和sum()调用。

Answer 2

您可以使用np.vectorize对您的foo function 进行矢量化。并使用exlucded选项排除 arguments 进行矢量化。 在您的情况下， arguments a和b1不需要矢量化。 所以，你需要先修改你的foo ：

def foo(a, x, y, z, b1):
    # [a, x, y, z, b1] = args -> comment this line

使用foo(*args)而不是foo(args)将列表传递给 function。 然后向量化foo function：

vec_func = np.vectorize(foo, excluded=['a', 'b1'])
args = [10, x, y, z, 2]
print(foo(*args))  # it works now