Efficient way of computing the cross products between two sets of vectors numpy

Question

I have two sets of 2000 3D vectors each, and I need to compute the cross product between each possible pair. I currently do it like this

for tx in tangents_x:
    for ty in tangents_y:
         cross = np.cross(tx, ty)
         (... do something with the cross variable...)

This works, but it's pretty slow. Is there a way to make it faster?

If I was interested in the element-wise product, I could just do the following

# Define initial vectors
tx = np.array([np.random.randn(3) for i in range(2000)])
ty = np.array([np.random.randn(3) for i in range(2000)])
# Store them into matrices
X = np.array([tx for i in range(2000)])
Y = np.array([ty for i in range(2000)]).T
# Compute the element-wise product
ew = X * Y
# Use the element_wise product as usual
for i,tx in enumerate(tangents_x):
    for j,ty in enumerate(tangents_y):
        (... use the element wise product of tx and ty as ew[i,j])

How can I apply this to the cross product instead of the element-wise one? Or, do you see another alternative?

Thanks much :)

Answer 1

Like many numpy functions cross supports broadcasting, therefore you can simply do:

np.cross(tangents_x[:, None, :], tangents_y)

or - more verbose but maybe easier to read

np.cross(tangents_x[:, None, :], tangents_y[None, :, :])

This reshapes tangents_x and tangents_y to shapes 2000, 1, 3 and 1, 2000, 3. By the rules of broadcasting this will be interpreted like two arrays of shape 2000, 2000, 3 where tangents_x is repeated along axis 1 and tangents_y is repeated along axis 0.

Answer 2

Just write it out and compile it

import numpy as np
import numba as nb

@nb.njit(fastmath=True,parallel=True)
def calc_cros(vec_1,vec_2):
    res=np.empty((vec_1.shape[0],vec_2.shape[0],3),dtype=vec_1.dtype)
    for i in nb.prange(vec_1.shape[0]):
        for j in range(vec_2.shape[0]):
            res[i,j,0]=vec_1[i,1] * vec_2[j,2] - vec_1[i,2] * vec_2[j,1]
            res[i,j,1]=vec_1[i,2] * vec_2[j,0] - vec_1[i,0] * vec_2[j,2]
            res[i,j,2]=vec_1[i,0] * vec_2[j,1] - vec_1[i,1] * vec_2[j,0]
    
    return res

Performance

#create data
tx = np.random.rand(3000,3)
ty = np.random.rand(3000,3)
#don't measure compilation overhead
comb=calc_cros(tx,ty)

t1=time.time()
comb=calc_cros(tx,ty)
print(time.time()-t1)

This gives 0.08s for the two (3000,3) matrices.