向量化 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

下面的代码可以解决问题，让您省去创建 array_positions.

的麻烦

首先，使用 ravel 将 X 和 Y 展平为形状为 (10000,).

的 NumPy 数组

X_flattened = X.ravel()
Y_flattened = Y.ravel()

然后，使用apply_along_axis沿着这些展平数组的长度迭代地实现自定义函数。

float_array = np.apply_along_axis(old_func, 0, (X_flattened, Y_flattened))

最后，将输出数组重塑为 (100, 100) 所需的形状。

np.reshape(float_array, (100, 100))

Answer 2

您自己的代码应该（并且确实）可以正常工作：

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 循环。您自己的代码使用了适当的向量化，因此请改用它。

向量化 python 函数

Vectorize python function

python

numpy

vectorization

apply