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

Question

在 python 中，在单个数组a1上对标量x的函数f(x)进行矢量化很简单：只需使用f(a1) 。 但是假设我有两个（或原则上是多个）数组a1, a2具有相同的形状Nx3 ，并且我想用x,y标量同时对一个函数g(x,y)进行矢量化。 像g(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)

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

result = []
for i, elem in a1:
    result.append(f(elem, a2[i]))
result = np.array(result)

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

result = f(a1, a2) 

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

Answer 1

使用numpy的vectorize 。

使用简单的lambda函数的np.vectorize一个简单示例：

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

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

在过去的几天里，我也一直在尝试做类似的事情。 我终于设法用np.vectorize<\/code>做到了，使用函数签名<\/a>。 试试下面的代码片段：

fn_vectorized = np.vectorize(interpolate.interp1d,
                                     signature='(n),(n)->()')
interp_fn_array = fn_vectorized(x[np.newaxis, :, :], y)