Find L3 norm of two arrays efficiently in Python

Question

Suppose I have two arrays. A has size n by d, and B has size t by d. Suppose I want to output an array C, where C[i, j] gives the cubed L3 norm between A[i] and B[j] (both of these have size d). i.e.

C[i, j] = |A[i, 0]-B[j, 0]|**3 + |A[i, 1]-B[j, 1]|**3 + ... + |A[i, d-1]-B[j, d-1]|**3

Can anyone redirect me to a really efficient way to do this? I have tried using a double for loop and something with array operations but the runtime is incredibly slow.

Edit: I know TensorFlow.norm works, but how could I implement this myself?

Answer 1

In accordance with @gph answer, an explicit example of the application of Numpy's np.linalg.norm, using broadcasting to avoid any loops:

import numpy as np

n, m, d = 200, 300, 3
A = np.random.random((n,d))
B = np.random.random((m,d))
    
C = np.linalg.norm(A[:,None,:] - B[None,...], ord=3, axis = -1)
# The 'raw' option would be:
C2 = (np.sum(np.abs(A[:,None,:] - B[None,...])**3, axis = -1))**(1/3)

Both ways seems to take around the same time to execute. I've tried to use np.einsum but let's just say my algebra skills have seen better days. Perhaps you have better luck (a nice resource is this wiki by Divakar).

Answer 2

There may be a few optimizations to speed this up, but the performance isn't going to be anywhere near specialized math packages. Those packages are using blas and lapack to vectorize the operations. They also avoid a lot of type checking overhead by enforcing types at the time you set a value. If performance is important there really is no way you are going to do better than a numerical computing package like numpy.

Answer 3

Use a 3rd-party library written in C or create your own

numpy has a linalg library which should be able to compute your L³ norm for each A[i]-B[j]

If numpy works for you, take a look at numba's JIT, which can compile and cache some (numpy) code to be orders of magnitude faster (successive runs will take advantage of it). I believe you will need to describe the parts to accelerate as functions which will be called many times to take advantage of it.

roughly

import numpy as np
from numpy import linalg as LA
from numba import jit as numbajit

@numbajit
def L3norm(A, B, i, j):
    return LA.norm([A[i] - B[j]], ord=3)

# may be able to apply JIT here too
def operate(A, B, n, t):
    # generate array
    C = np.zeros(shape=(n, t))

    for i in range(n):
        for j in range(t):
            C[i, j] = L3norm(A, B, i, j)

    return C