从向量对向量的函数输出创建numpy矩阵

Question

我有两个矢量/一维numpy数组和一个我想应用的函数:

arr1 = np.arange(1, 5)
arr2 = np.arange(2, 6)
func = lambda x, y: x * y

我现在想构造一个包含函数输出值的n * m矩阵(n,m分别是arr1和arr2的长度)。使用for循环的简单方法是这样的:

np.array([[func(x, y) for x in arr1] for y in arr2])

我想知道是否有一种更智能的矢量化方法使用arr1[:, None]语法来应用我的函数-请注意，我的实际函数要复杂得多，不能分解为简单的numpy操作(arr1[:, None] * arr2[None, :]不起作用)。

Answer 1

当您有numpy.array时，一个方法可以是numpy.einsum。因为你想计算这个:arr1_i * arr2_j -> insert to arr_result_ji.

>>> np.einsum('i, j -> ji', arr1, arr2)
array([[ 2,  4,  6,  8],
[ 3,  6,  9, 12],
[ 4,  8, 12, 16],
[ 5, 10, 15, 20]])

或者您可以使用numpy.matmul或@。

>>> np.matmul(arr2[:,None], arr1[None,:])
# OR
>>> arr2[:,None] @ arr1[None,:]
# Or by thanks @hpaulj by elementwise multiplication with broadcasting
>>> arr2[:,None] * arr1[None,:]
array([[ 2,  4,  6,  8],
[ 3,  6,  9, 12],
[ 4,  8, 12, 16],
[ 5, 10, 15, 20]])

Answer 2

下面是你的循环方法和@I'mahdi的方法之间的一些比较:

import time
arr1 = np.arange(1, 10000)
arr2 = np.arange(2, 10001)

start = time.time()
np.array([[func(x, y) for x in arr1] for y in arr2])
print('loop: __time__', time.time()-start)
start = time.time()
(arr1[:, None]*arr2[None, :]).T
print('* __time__', time.time()-start)
start = time.time()
np.einsum('i, j -> ji', arr1, arr2)
print('einsum __time__', time.time()-start)
start = time.time()
np.matmul(arr2[:,None], arr1[None,:])
print('matmul __time__', time.time()-start)

输出:

loop: __time__ 70.3061535358429
* __time__ 0.43536829948425293
einsum __time__ 0.508014440536499
matmul __time__ 0.7149899005889893