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

【问题标题】：How to avoid iteration inside an array based function in numpy?如何避免在numpy中基于数组的函数内进行迭代？
【发布时间】：2021-03-26 09:56:49
【问题描述】：

假设我在 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) 返回129.26...。

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

我可以通过在函数中添加一个for 循环来做到这一点，但是在我的实际代码中该过程变得非常慢（这比这里的这个例子对计算的要求更高）。我想寻求支持的是如何仅使用 numpy 函数更快地做到这一点？ 这就是 numpy ufunc 的意思吗？没看过。

【问题讨论】：

你知道如何从x 制作一个 (500,2,2) 形状的数组吗？ @ 可以处理这样的 3d 数组，将维度视为“批次”。
“我从来没有调查过。”。你真的应该在问之前先调查一下。请做你自己的研究！

标签： python numpy optimization numpy-ndarray numpy-ufunc

【解决方案1】：

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

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])

注意标量的相同行为。

【讨论】：