Python：同时对两个多维数组进行矢量化函数

Question

在 python 中，对单个数组f(x)上的标量函数进行矢量化很简单：只需使用. 但是假设我有两个（或原则上是多个）具有相同形状的数组，并且我想对一个带有标量的函数同时在两者上进行矢量化。类似的东西，它将再次返回一个具有共同维度的对象。xa1f(a1)a1, a2Nx3g(x,y)x,yg(a1,a2)N

编辑：如果a1,a2都是一维的，这变得微不足道。我们使用如下所述的简单广播。但是，对于多维数组，答案对我来说并不明显。那么，我该怎么做，最好使用numpy？

示例（已编辑）：

a1 = np.array of size 20x3 # so that each row is a 3-vector

a2 = np.array of size 20x3 # ditto

def f(x, y): # acts on each element 

   ... complicated function, using other global variables ...

   return ... (scalar)

如果没有矢量化，我需要f单独循环所有 20 行，并获得输出长度为 20 的矢量：

result = []

for i, elem in a1:

    result.append(f(elem, a2[i]))

result = np.array(result)

但是，我想消除for循环，并使用 numpy 矢量化有一个语句。原因是能够使用jax( https://github.com/google/jax ) 的 numpy 包装器，然后在 GPU 上加速它。像天真的东西

result = f(a1, a2) 

不起作用。那么正确的语法是什么？

Answer 1

在过去的几天里，我也一直在尝试做类似的事情。我终于设法做到了np.vectorize，使用函数签名。试试下面的代码片段：fn_vectorized = np.vectorize(interpolate.interp1d,                                     signature='(n),(n)->()')interp_fn_array = fn_vectorized(x[np.newaxis, :, :], y)在这里，我正在对 interp1d 函数进行矢量化。x 和 y 是形状 (mxn) 的数组。目标是为x 的第 i 行和y的第i行生成一个插值函数数组。数组“interp_fn_array”包含插值函数（形状为（1 xm）。

Answer 2

它可能取决于您需要执行的操作，如果它是一个简单的总和，那么以下将起作用：import numpy as npa = np.arange(3*2*20).reshape((20,3,2))b = np.arange(2*20).reshape((20,2))res = (a.transpose((1,2,0))+b.transpose((1,0))).transpose((2,0,1))print(a[0],b[0])[[0 1] [2 3] [4 5]] [0 1]print(res[0])[[0 2] [2 4] [4 6]]首先，输入数据被转置，以便在广播操作中涉及正确的维度。求和后，输出转回。

Answer 3

使用 numpy 的vectorize。np.vectorize使用简单 lambda 函数的简单示例：import numpy as npf = np.vectorize(lambda x: 2*x)f([[2,3],[3,4],[1,1]])# output:array([[4, 6],       [6, 8],       [2, 2]])