python - 向量化python函数

Question

我有一个名为 old_func 的旧函数，它接受两个位置参数 x 和 y 作为输入。函数的输入是这样写的，使用一个元组作为输入:

def old_func(position):

    x, y = position 
    return x**2 + y**2

我现在想要一种通过值网格调用函数的快速简便的方法:

xx = numpy.linspace(0, 1, 100)
yy = numpy.linspace(0, 1, 100)
X, Y = numpy.meshgrid(xx, yy)

array_positions = (X,Y)
old_fun(array_positions)

目的是函数中对x 的每个操作都在所有X 上完成，对y 也是如此。我尝试使用 numpy.vectorize 对函数进行矢量化但这是行不通的。我不想更改函数以接受 NumPy 数组，因为这会花费太长时间。

Answer 1

您自己的代码应该(并且确实)可以正常工作:

def old_fun(position):

    x, y = position 
    z = x**2 + y**2
    return z

xx = numpy.linspace(0,1,100)
yy = numpy.linspace(0,1,100)
X,Y = numpy.meshgrid(xx, yy)

array_positions = (X,Y)
Z = old_fun(array_positions)

Z.shape 现在是 (100, 100)。

通常，numpy 数组可以与任何标准运算符一起使用，例如 + 和 **。请注意，虽然您的 old_fun 将元组作为输入，并且此元组需要包含两个值，但这两个值的类型可以是任何类型，只要此类型支持数学运算符即可。标准 Python 标量和 numpy 数组都支持这些，因此代码运行良好。

关于@JuanManuel 的回答的注释:虽然它也工作得很好，apply_along_axis 特别不适合矢量化，因为人们通常想要使用它，即从 numpy 中获得良好的性能。 apply_along_axis 将导致缓慢的 Python 循环，而不是像正确矢量化代码那样的快速 C 循环。您自己的代码使用了适当的矢量化，因此请改用它。