Create numpy matrix from output of function on vector vs vector-CodePudding

I have two vectors / one-dimensional numpy arrays and a function I want to apply:

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

I now want to construct a n * m matrix (with n, m being the lengths of arr1, and arr2 respectively) containing the values of the function outputs. The naive approach using for loops would look like this:

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

I was wondering if there is a smarter vectorized approach using the arr1[:, None] syntax to apply my function - please note my actual function is significantly more complicated and can't be broken down to simple numpy operations (arr1[:, None] * arr2[None, :] won't work).

CodePudding user response：

When you have numpy.array, One approach can be numpy.einsum. Because you want to compute this : 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]])

Or you can use numpy.matmul or use @.

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

CodePudding user response：

Here is some comparison between your loop approach and @I'mahdi 's approach:

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)

Output:

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