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-vectora2 = np.array of size 20x3 # dittodef 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 np

a = 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 np

f = np.vectorize(lambda x: 2*x)

f([[2,3],[3,4],[1,1]])

# output:

array([[4, 6],

       [6, 8],

       [2, 2]])