Numpy小技巧：将函数应用于两个一维数组的所有组合，以获取一个二维数组。

Question

假设我有两个一维的 numpy 数组，a 和 b，长度分别为 n1 和 n2。我还有一个函数 F(x,y)，它需要两个值作为参数。现在我想对来自我的两个 1D 数组的每一对值应用该函数，结果将是一个形状为 n1, n2 的二维 numpy 数组。这个二维数组的 i, j 元素将是 F(a[i], b[j])。

我没有找到一种不使用大量 for 循环来实现此操作的方法，并且我确信有一种更简单（更快！）的方式在 numpy 中实现。

提前致谢！

- Misconstruction

---

你是否正在寻找一个带有标量的外积类型函数？ - Russia Must Remove Putin

Answer 1

你可以使用 numpy广播 来对两个数组进行计算，使用 newaxis 将 a 转换为一个垂直的二维数组：

In [11]: a = np.array([1, 2, 3]) # n1 = 3
    ...: b = np.array([4, 5]) # n2 = 2
    ...: #if function is c(i, j) = a(i) + b(j)*2:
    ...: c = a[:, None] + b*2

In [12]: c
Out[12]: 
array([[ 9, 11],
       [10, 12],
       [11, 13]])

基准测试：

In [28]: a = arange(100)

In [29]: b = arange(222)

In [30]: timeit r = np.array([[f(i, j) for j in b] for i in a])
10 loops, best of 3: 29.9 ms per loop

In [31]: timeit c = a[:, None] + b*2
10000 loops, best of 3: 71.6 us per loop

- zhangxaochen

---

newaxis 可以是一种正式的方法，用于向 ndarray 添加一个新的轴，即使它们是相同的。 - Kattern

是的，在内部代码中它们是相同的。但是，“newaxis”更容易理解，有点像语法糖。 - Kattern

这看起来就是我要找的，我会立刻尝试一下！ - Misconstruction

这个很好用！能否扩展到比较两个二维数组？这样的结果将是在两个二维数组之间每对列的函数。 - Misconstruction

假设我有一个numpy数组，我想对行成对应用函数并累加结果，例如row0+row1，这个结果再加上row3等等，有什么聪明的方法可以实现这个功能？ - seralouk

Answer 2

如果 F 超出您的控制，您可以使用 numpy.vectorize 自动将其包装为“向量感知”。下面是一个可行的示例，其中我定义了自己的 F，只是为了完整起见。这种方法具有简单性优势，但如果您可以控制 F，则重新编写它以正确向量化可能会带来巨大的速度优势。

import numpy

n1 = 100
n2 = 200

a = numpy.arange(n1)
b = numpy.arange(n2)

def F(x, y):
    return x + y

# Everything above this is setup, the answer to your question lies here:
fv = numpy.vectorize(F)
r = fv(a[:, numpy.newaxis], b)

在我的电脑上，发现以下定时显示了您支付的“自动”矢量化价格：

%timeit fv(a[:, numpy.newaxis], b)
100 loops, best of 3: 3.58 ms per loop

%timeit F(a[:, numpy.newaxis], b)
10000 loops, best of 3: 38.3 µs per loop

- chthonicdaemon

---

我建议使用以下代码：In [5]: %timeit a[:, numpy.newaxis] + b输出结果为：49.9 µs ± 337 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each) - Beatriz Fonseca

太好了！这是一个非常通用的答案，它本身就值得一个清晰的问题。其中很多都涉及到语言 - 广播、外积、约减、带任意函数的相关器等等......你的方法是否存在任何失败的情况？np.vectorize能否处理任意函数/lambda？如果不能呢？ - jtlz2

@jtlz2 numpy.vectorise 可以应用于所有函数，因为它实际上起到了 for 循环的作用（请参阅文档中的注释）。因此，向量化是避免手动编写 for 循环的快速方法，但正如我所展示的，它并不会使代码运行更快。要做到这一点，您需要使用您提到的其他技术之一。 - chthonicdaemon

Answer 3

如果F()可以使用广播参数，请一定使用它，就像其他人所描述的那样。
另一种方法是使用np.fromfunction（function_on_an_int_grid会是一个更好的名称）。以下代码将整数网格映射到您的a-b网格，然后进入F()：

import numpy as np

def func_allpairs( F, a, b ):
    """ -> array len(a) x len(b):
        [[ F( a0 b0 )  F( a0 b1 ) ... ]
         [ F( a1 b0 )  F( a1 b1 ) ... ]
         ...
        ]
    """
    def fab( i, j ):
        return F( a[i], b[j] )  # F scalar or vec, e.g. gradient

    return np.fromfunction( fab, (len(a), len(b)), dtype=int )  # -> fab( all pairs )


#...............................................................................
def F( x, y ):
    return x + 10*y

a = np.arange( 100 )
b = np.arange( 222 )
A = func_allpairs( F, a, b )
# %timeit: 1000 loops, best of 3: 241 µs per loop -- imac i5, np 1.9.3

- denis

---

我真的很喜欢这种方法对于具有两个1D向量的全对矩阵是多么通用。 - user394430

Answer 4

你可以使用列表推导式来创建一个数组的数组：

import numpy as np

# Arrays
a = np.array([1, 2, 3]) # n1 = 3
b = np.array([4, 5]) # n2 = 2

# Your function (just an example)
def f(i, j):
    return i + j

result = np.array([[f(i, j)for j in b ]for i in a])
print result

输出：

[[5 6]
 [6 7]
 [7 8]]

- Christian Tapia

---

列表推导式对于NumPy代码而言，仅比for循环略好。 - user2357112

那是我的。我已经提供了反馈，但我会详细说明。列表推导可以缩短代码，但它们的速度与基于循环的解决方案一样慢。在使用NumPy时，最好养成始终首先寻找使用NumPy的向量化操作的习惯。 - user2357112

这比其他选项慢得多。 - Russia Must Remove Putin

Answer 5

作为另一种比点积更具可扩展性的选择，使用numpy.newaxis可以比嵌套列表推导式节省1/5至1/9的时间（需要进行一些挖掘才能找到：链接）：

>>> import numpy
>>> a = numpy.array([0,1,2])
>>> b = numpy.array([0,1,2,3])

这次，使用幂函数：

>>> pow(a[:,numpy.newaxis], b)
array([[1, 0, 0, 0],
       [1, 1, 1, 1],
       [1, 2, 4, 8]])

与另一种选择相比：

>>> numpy.array([[pow(i,j) for j in b] for i in a])
array([[1, 0, 0, 0],
       [1, 1, 1, 1],
       [1, 2, 4, 8]])

同时比较：

>>> import timeit
>>> timeit.timeit('numpy.array([[pow(i,j) for i in a] for j in b])', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
31.943181037902832
>>> timeit.timeit('pow(a[:, numpy.newaxis], b)', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
5.985810041427612

>>> timeit.timeit('numpy.array([[pow(i,j) for i in a] for j in b])', 'import numpy; a=numpy.arange(10); b=numpy.arange(10)')
109.74687385559082
>>> timeit.timeit('pow(a[:, numpy.newaxis], b)', 'import numpy; a=numpy.arange(10); b=numpy.arange(10)')
11.989138126373291

- Russia Must Remove Putin

Answer 6

如果您的使用场景更多限于产品，我建议您使用外部产品outer-product。

e.g.:

import numpy

a = array([0, 1, 2])
b = array([0, 1, 2, 3])

numpy.outer(a,b)

返回

array([[0, 0, 0, 0],
       [0, 1, 2, 3],
       [0, 2, 4, 6]])

您可以应用其他转换：

numpy.outer(a,b) + 1

返回

array([[1, 1, 1, 1],
       [1, 2, 3, 4],
       [1, 3, 5, 7]])

这个更快：

>>> import timeit
>>> timeit.timeit('numpy.array([[i*j for i in a] for j in b])', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
31.79583477973938

>>> timeit.timeit('numpy.outer(a,b)', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
9.351550102233887
>>> timeit.timeit('numpy.outer(a,b)+1', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
12.308301210403442

- Russia Must Remove Putin

---

有趣，但这不是假定该函数始终取两个输入值的乘积吗？如果你想要更复杂的函数怎么办？ - Misconstruction

这很公平，但相对于其他可能需要的操作，这是一种非常常见的操作。如果您想要一个更复杂的函数，我可以建议您查看我的另一个答案：https://dev59.com/_mEi5IYBdhLWcg3wfsXI#21273998 如果我可以稍微冒昧一点，您是否愿意分享您的神秘函数？ - Russia Must Remove Putin