how to Vectorize the np.linalg.norm between to matices for each row

I am trying to optimize my code and for the following example arrays,(that can be dynamic in length):

arr1 = np.random.random(300).reshape((-1,3))
arr2 = np.random.random(450).reshape((-1,3))
>>> arr1
array([[0.32725109, 0.39912246, 0.9834273 ], [0.78003681, 0.92057381, 0.06478441], [0.15232456, 0.82258548, 0.1685084 ], ... [0.10011657, 0.05840116, 0.07685251], [0.360477 , 0.09061205, 0.54737966], [0.44562439, 0.90453774, 0.27240101]])

The following code are the critical part:

result_arr = np.zeros((arr1.shape[0],arr2.shape[0]))
for arr1_index,arr1_row in enumerate(arr1): for arr2_index, arr2_row in enumerate(arr2): result_arr[arr1_index,arr2_index] = 1/(1+np.linalg.norm(arr1_row - arr2_row,2) ** (1/3) )

and the performance is too poor, is there any idea of how do I turn this np.linalg.norm function to work on a vectorized way and and to skip this use of loops. there is also arithmetics to take into account here and I can get my head over it?

1 Answer

You can use broadcasting and exploit the vectorized nature of the linalg.norm function to perform the operation in one function call as follow (in my computer this achieves 2 orders of magnitude of improvement in speed):

import numpy as np
# Create dummy arrays
arr1 = np.random.random(300).reshape((-1,3))
arr2 = np.random.random(450).reshape((-1,3))
# Output
r = 1 / (1 + np.linalg.norm(arr1[:,None] - arr2, axis=-1)**(1/3))

I suspect further improvements could be obtained via NumPy's einsum function.