如何避免在 numpy 中基于数组的函数内进行迭代？

Question

假设我在 numpy 中有一个由 2 个变量 x 和 y:

参数化的矩阵乘法函数

import numpy as np

def func(x, y):
    a = np.array([[1, x],
                  [x, 2]])
    b = np.array([[y,         2*x],
                  [x, np.exp(y+x)]])
    M = np.array([[3.2, 2*1j],
                  [4  ,   93]])

    prod  = a @ M @ b
    final = np.abs(prod[0,0])

    return final

我可以运行这个函数很容易地用于任何两个数值，例如func(1.1, 2.2) returns 129.26....

到目前为止一切顺利，但现在我想运行这对于 x 和 y 的几个值，例如x=np.linspace(0,10,500) 和 y=np.linspace(0,10,500)。我想将这 500 个值一对一对应，即 x 列表中的第一个与 y 列表中的第一个，第二个与第二个，等等

我可以通过在函数内添加一个 for 循环来做到这一点，但在我的实际代码中这个过程变得非常慢（这比这里的这个例子在计算上要求更高）。我想寻求支持的是 如何仅使用 numpy 函数更快地完成此操作？ 这就是 numpy ufunc 的含义吗？我从来没有研究过它。

Answer 1

矢量化非常直接，只要您跟踪维度即可。事实上，只要它们一起广播，您就可以完全不知道 x 和 y 的大小。此处显示的解决方案将 return 具有相同广播形状的内容：

def func(x, y):
    x = np.array(x, copy=False)
    y = np.array(y, copy=False)
    bc = np.broadcast(x, y)

    a = np.stack((np.ones_like(x), x, x, np.full_like(x, 2)), axis=-1).reshape(*x.shape, 2, 2)
    b = np.stack((np.broadcast_to(y, bc.shape),
                  np.broadcast_to(2 * x, bc.shape),
                  np.broadcast_to(x, bc.shape),
                  np.exp(y + x)), axis=-1).reshape(*bc.shape, 2, 2)
    M = np.array([[3.2, 2*1j],
                  [4  ,   93]])

    prod  = a @ M @ b
    final = np.abs(prod[..., 0, 0])

    return final

保留最后两个维度的矩阵可确保您的矩阵乘法正常进行。堆叠是多维的，相当于您使用 np.array 和标量进行的串联。

>>> func(1.1, 1.1)
120.9100165412279
>>> func(1.1, 2.2)
129.26872205
>>> func(2.2, 1.1)
463.34089222
>>> func(1.1, [1.1, 2.2])
array([120.91001654, 129.26872205])
>>> func([1.1, 2.2], 1.1)
array([120.91001654, 463.34089222])
>>> func([1.1, 2.2], [2.2, 1.1])
array([129.26872205, 463.34089222])

注意标量的相同行为。

如何避免在 numpy 中基于数组的函数内进行迭代？

How to avoid iteration inside an array based function in numpy?

python

optimization

numpy

numpy-ufunc

numpy-ndarray