numpy 中的快速标量三重积

Nic*_*mer 6 python numpy

我有大量的向量三元组,我想计算它们的标量三元组积。我可以

import numpy

n = 871
a = numpy.random.rand(n, 3)
b = numpy.random.rand(n, 3)
c = numpy.random.rand(n, 3)

# <a, b x c>
omega = numpy.einsum('ij, ij->i', a, numpy.cross(b, c))
Run Code Online (Sandbox Code Playgroud)

numpy.cross速度相当慢。问题的对称性(其 Levi-Civita 表达式为eps_{ijk} a_i b_j c_k)表明可能有更好(更快)的计算方法,但我似乎无法弄清楚。

有什么提示吗?

Nic*_*mer 4

我已经对答案中提到的方法进行了比较。结果:

在此输入图像描述

@Divakar's 一点一点地击败了 einsum-cross。

为了完整起见,我要指出的是,还有另一种方法完全依赖于点积和 sqrt,请参见此处。该方法比 einsum-cross 和 slice-sum 稍慢。


该图是用perfplot创建的,

import numpy as np
import perfplot


def einsum_cross(a, b, c):
    return np.einsum("ij, ij->i", a, np.cross(b, c))


def det(a, b, c):
    return np.linalg.det(np.dstack([a, b, c]))


def slice_sum(a, b, c):
    c0 = b[:, 1] * c[:, 2] - b[:, 2] * c[:, 1]
    c1 = b[:, 2] * c[:, 0] - b[:, 0] * c[:, 2]
    c2 = b[:, 0] * c[:, 1] - b[:, 1] * c[:, 0]
    return a[:, 0] * c0 + a[:, 1] * c1 + a[:, 2] * c2


b = perfplot.bench(
    setup=lambda n: (
        np.random.rand(n, 3),
        np.random.rand(n, 3),
        np.random.rand(n, 3),
    ),
    n_range=[2**k for k in range(1, 20)],
    kernels=[einsum_cross, det, slice_sum],
)
b.save("out.png")
b.show()
Run Code Online (Sandbox Code Playgroud)