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

【问题标题】：Python: vectorise function over two multi-dimensional arrays simultaneouslyPython：同时对两个多维数组进行矢量化函数
【发布时间】：2020-01-11 04:50:06
【问题描述】：

在 python 中，将标量 x 的函数 f(x) 向量化到单个数组 a1 上很简单：只需使用 f(a1)。但是假设我有两个（或原则上是多个）数组a1, a2 具有相同的形状Nx3，并且我想对一个函数g(x,y) 和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 (https://github.com/google/jax) 的 numpy 包装器，然后在 GPU 上加速。有点天真像

result = f(a1, a2)

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

【问题讨论】：

我想你想要numpy.apply_along_axis docs.scipy.org/doc/numpy/reference/generated/…
操作是什么？也许问题更多是关于广播docs.scipy.org/doc/numpy/user/…？
我只是添加了一个示意图来阐明我的意思
很大程度上取决于f 的样子，正如 Christian K. 所写。 np.vectorize` 只是隐藏循环，请参阅我对答案的评论。最好的方法是只使用纯矢量化 Numpy 函数和广播。显示f()，我们可以提供进一步的帮助。
明确您的函数可以处理哪些维度。标量、1d、2d 等。它处理广播吗？没有明确的例子，很难给出有意义的答案。

标签： python numpy vectorization

【解决方案1】：

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

【讨论】：

如果您的目标是性能，这并不能解决问题。来自doc：提供vectorize 函数主要是为了方便，而不是为了性能。该实现本质上是一个 for 循环。
对，当我回答时，作者没有编辑问题，直到他谈到“加快速度”，即性能......

【解决方案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]]

首先将输入数据转置，以便在广播操作中涉及正确的维度。求和后，输出转回。

【讨论】：